run_mykernel.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
一键编译并基准测试自定义 CUDA kernel（默认 my_kernel.py）。

用法示例：
    conda run -n CudaForge python scripts/run_mykernel.py

可选参数：
    --ref            参考模型路径（默认 KernelBench/level1/1_Square_matrix_multiplication_.py）
    --test           待测文件路径（默认 my_kernel.py）
    --device-idx     GPU 编号，默认 0
    --warmup         预热次数，默认 5
    --repeat         计时次数，默认 20
    --tol            误差容忍，默认 1e-3
"""

from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path

from utils.compile_and_run import compare_and_bench


def _build_ap() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser("Compile + benchmark a single CUDA kernel file")
    ap.add_argument(
        "--ref",
        type=Path,
        default=Path("KernelBench/level1/1_Square_matrix_multiplication_.py"),
        help="参考 PyTorch 模型脚本路径",
    )
    ap.add_argument(
        "--test",
        type=Path,
        default=Path("my_kernel.py"),
        help="待测 CUDA 扩展脚本路径（需定义 ModelNew）",
    )
    ap.add_argument("--device-idx", type=int, default=0, help="GPU 编号")
    ap.add_argument("--warmup", type=int, default=5, help="预热次数")
    ap.add_argument("--repeat", type=int, default=20, help="计时次数")
    ap.add_argument("--tol", type=float, default=1e-3, help="最大误差容忍")
    return ap


def main():
    args = _build_ap().parse_args()

    try:
        res = compare_and_bench(
            ref_py=args.ref,
            test_py=args.test,
            device_idx=args.device_idx,
            warmup=args.warmup,
            repeat=args.repeat,
            tol=args.tol,
        )
    except Exception as exc:
        print(f"[run_mykernel] failed: {exc}", file=sys.stderr)
        sys.exit(1)

    # 便于脚本化解析：先打印 JSON，再打印可读版
    print(json.dumps(res, ensure_ascii=False, indent=2))
    print("\n[run_mykernel] summary:")
    print(f"  runnable: {res.get('runnable', False)}")
    if "ref_latency_ms" in res and "test_latency_ms" in res:
        ref_avg = res['ref_latency_ms']['avg']
        tst_avg = res['test_latency_ms']['avg']
        print(f"  ref avg (ms):  {ref_avg:.4f}")
        print(f"  test avg (ms): {tst_avg:.4f}")
    if "score" in res:
        print(f"  speedup (ref/test): {res['score']:.4f}")


if __name__ == "__main__":
    main()