()
| 16 | |
| 17 | |
| 18 | def compile_kernel(): |
| 19 | cmd = ["nvcc", "--shared", "-Xcompiler", "-fPIC", |
| 20 | "-gencode", "arch=compute_100a,code=sm_100a", "-O3", |
| 21 | KERNEL_PATH, "-o", SO_PATH, "-lcuda"] |
| 22 | print(f"Compiling: {' '.join(cmd)}") |
| 23 | r = subprocess.run(cmd, capture_output=True, text=True) |
| 24 | if r.returncode != 0: |
| 25 | print("STDOUT:", r.stdout) |
| 26 | print("STDERR:", r.stderr) |
| 27 | sys.exit(1) |
| 28 | print("OK") |
| 29 | |
| 30 | |
| 31 | def cosine_similarity(x: torch.Tensor, y: torch.Tensor) -> float: |