Compiled CUDA kernel that can be launched from Python.
| 32 | |
| 33 | |
| 34 | class CompiledKernel: |
| 35 | """Compiled CUDA kernel that can be launched from Python.""" |
| 36 | |
| 37 | def __init__( |
| 38 | self, |
| 39 | cuda_code: str, |
| 40 | kernel_name: str, |
| 41 | grid: Tuple[int, int, int], |
| 42 | block: Tuple[int, int, int], |
| 43 | shared_mem_bytes: int = 0, |
| 44 | arch: Optional[str] = None, |
| 45 | cache_dir: Optional[str] = None, |
| 46 | verbose: bool = False, |
| 47 | include_paths: Optional[List[Union[str, Path]]] = None, |
| 48 | cluster_dim: Optional[Tuple[int, int, int]] = None, |
| 49 | ): |
| 50 | self.cuda_code = cuda_code |
| 51 | self.kernel_name = kernel_name |
| 52 | self.grid = grid |
| 53 | self.block = block |
| 54 | self.shared_mem_bytes = shared_mem_bytes |
| 55 | self.arch = arch |
| 56 | self.cache_dir = cache_dir |
| 57 | self.verbose = verbose |
| 58 | self.include_paths = include_paths |
| 59 | self.cluster_dim = cluster_dim |
| 60 | |
| 61 | self._runtime = None |
| 62 | self._module = None |
| 63 | self._launcher = None |
| 64 | self._cubin_data = None |
| 65 | |
| 66 | @property |
| 67 | def runtime(self) -> CUDARuntime: |
| 68 | if self._runtime is None: |
| 69 | self._runtime = CUDARuntime() |
| 70 | return self._runtime |
| 71 | |
| 72 | def _get_cache_key(self) -> str: |
| 73 | code_hash = hashlib.sha256(self.cuda_code.encode()).hexdigest()[:16] |
| 74 | arch_str = self.arch or "auto" |
| 75 | return f"{self.kernel_name}_{code_hash}_{arch_str}" |
| 76 | |
| 77 | def _compile(self) -> bytes: |
| 78 | if self.cache_dir: |
| 79 | cache_key = self._get_cache_key() |
| 80 | cache_file = os.path.join(self.cache_dir, f"{cache_key}.cubin") |
| 81 | if os.path.exists(cache_file): |
| 82 | if self.verbose: |
| 83 | print(f"Loading cached binary: {cache_file}") |
| 84 | with open(cache_file, "rb") as f: |
| 85 | return f.read() |
| 86 | |
| 87 | if self.verbose: |
| 88 | print(f"Compiling kernel: {self.kernel_name}") |
| 89 | |
| 90 | cubin = compile_cuda( |
| 91 | self.cuda_code, |
no outgoing calls
no test coverage detected