Function main

profile.py:939–1121

()

# ---------------------------------------------------------------------------

def main() -> int:
    global WARMUP_ITERS, PROFILE_ITERS

    args = parse_args()
    WARMUP_ITERS = args.warmup_iters
    PROFILE_ITERS = args.profile_iters

    # Discover supported kernel types from kernels/ directory
    global _SUPPORTED_OP_TYPES
    _SUPPORTED_OP_TYPES = _discover_supported_op_types()

    # Parse input shape
    try:
        input_shape = [int(x.strip()) for x in args.input_shape.split(",")]
    except ValueError:
        print(
            f"ERROR: Invalid --input-shape '{args.input_shape}'. "
            "Expected comma-separated integers."
        )
        return 1

    # Resolve dtype
    try:
        dtype = _resolve_dtype(args.dtype)
    except ValueError as e:
        print(f"ERROR: {e}")
        return 1

    # Check GPU availability
    if not torch.cuda.is_available():
        print("ERROR: No CUDA GPU detected. The profiler requires a GPU.")
        return 1

    device = "cuda"

    # Detect GPU
    gpu = detect_gpu()

    print()
    print("=== AutoKernel Profiler ===")

    # Load model
    print("Loading model...")
    try:
        model, model_desc = load_model(args)
    except Exception as e:
        print(f"ERROR loading model: {e}")
        traceback.print_exc()
        return 1

    print(f"  Model: {model_desc}")
    print(f"  Input: shape={input_shape}, dtype={args.dtype}")
    print(f"  GPU:   {gpu.name}")
    print()

    # Prepare model and input
    print("Preparing model and input...")
    try: