| 75 | |
| 76 | |
| 77 | def compile_kernel(): |
| 78 | kernel_dir = os.path.dirname(os.path.abspath(__file__)) |
| 79 | kernel_path = os.path.join(kernel_dir, "matmul-v3.cu") |
| 80 | so_path = "/tmp/gemm_bf16_level3.so" |
| 81 | |
| 82 | cmd = ["nvcc", "--shared", "-Xcompiler", "-fPIC", "-std=c++17", |
| 83 | "-gencode", "arch=compute_90a,code=sm_90a", "-O3", "-lcuda", |
| 84 | kernel_path, "-o", so_path] |
| 85 | print("Compiling:", " ".join(cmd)) |
| 86 | result = subprocess.run(cmd, capture_output=True, text=True) |
| 87 | if result.returncode != 0: |
| 88 | print("Compilation failed:", result.stderr) |
| 89 | raise RuntimeError("nvcc failed") |
| 90 | print("OK") |
| 91 | |
| 92 | lib = ctypes.CDLL(so_path) |
| 93 | lib.gemm_bf16_launch.argtypes = [ |
| 94 | ctypes.POINTER(CUtensorMap), ctypes.POINTER(CUtensorMap), |
| 95 | ctypes.c_void_p, ctypes.c_int, ctypes.c_int, ctypes.c_int |
| 96 | ] |
| 97 | lib.gemm_bf16_launch.restype = None |
| 98 | |
| 99 | lib.gemm_bf16_launch_cooperative.argtypes = [ |
| 100 | ctypes.POINTER(CUtensorMap), ctypes.POINTER(CUtensorMap), |
| 101 | ctypes.c_void_p, ctypes.c_int, ctypes.c_int, ctypes.c_int |
| 102 | ] |
| 103 | lib.gemm_bf16_launch_cooperative.restype = None |
| 104 | |
| 105 | lib.benchmark_kernel.argtypes = [ |
| 106 | ctypes.POINTER(CUtensorMap), ctypes.POINTER(CUtensorMap), |
| 107 | ctypes.c_void_p, ctypes.c_int, ctypes.c_int, ctypes.c_int, |
| 108 | ctypes.c_int, ctypes.c_int |
| 109 | ] |
| 110 | lib.benchmark_kernel.restype = ctypes.c_float |
| 111 | |
| 112 | return lib |
| 113 | |
| 114 | |
| 115 | def cosine_similarity(x, y): |