| 91 | |
| 92 | |
| 93 | def compile_kernel(): |
| 94 | kernel_dir = os.path.dirname(os.path.abspath(__file__)) |
| 95 | kernel_path = os.path.join(kernel_dir, "matmul-v2.cu") |
| 96 | so_path = "/tmp/gemm_bf16_level2.so" |
| 97 | |
| 98 | cmd = ["nvcc", "--shared", "-Xcompiler", "-fPIC", "-std=c++17", |
| 99 | "-gencode", "arch=compute_90a,code=sm_90a", "-O3", "-lcuda", |
| 100 | kernel_path, "-o", so_path] |
| 101 | print("Compiling:", " ".join(cmd)) |
| 102 | result = subprocess.run(cmd, capture_output=True, text=True) |
| 103 | if result.returncode != 0: |
| 104 | print("Compilation failed:", result.stderr) |
| 105 | raise RuntimeError("nvcc failed") |
| 106 | print("OK") |
| 107 | |
| 108 | lib = ctypes.CDLL(so_path) |
| 109 | lib.gemm_bf16_launch.argtypes = [ |
| 110 | ctypes.POINTER(CUtensorMap), ctypes.POINTER(CUtensorMap), |
| 111 | ctypes.c_void_p, ctypes.c_int, ctypes.c_int, ctypes.c_int |
| 112 | ] |
| 113 | lib.gemm_bf16_launch.restype = None |
| 114 | |
| 115 | lib.benchmark_kernel.argtypes = [ |
| 116 | ctypes.POINTER(CUtensorMap), ctypes.POINTER(CUtensorMap), |
| 117 | ctypes.c_void_p, ctypes.c_int, ctypes.c_int, ctypes.c_int, |
| 118 | ctypes.c_int, ctypes.c_int |
| 119 | ] |
| 120 | lib.benchmark_kernel.restype = ctypes.c_float |
| 121 | |
| 122 | return lib |
| 123 | |
| 124 | |
| 125 | def cosine_similarity(x, y): |