Prepare kernel code with necessary includes and helpers.
(self, kernel_code: str)
| 214 | raise RuntimeError(f"Kernel compilation failed:\n{error_info}") |
| 215 | |
| 216 | def _prepare_kernel_code(self, kernel_code: str) -> str: |
| 217 | """Prepare kernel code with necessary includes and helpers.""" |
| 218 | includes = """ |
| 219 | #include <cuda_runtime.h> |
| 220 | #include <device_launch_parameters.h> |
| 221 | #include <cuda_fp16.h> |
| 222 | #include <mma.h> |
| 223 | #include <cooperative_groups.h> |
| 224 | #include <stdio.h> |
| 225 | |
| 226 | #define CUDA_CHECK(call) do { \ |
| 227 | cudaError_t error = call; \ |
| 228 | if (error != cudaSuccess) { \ |
| 229 | printf("CUDA error at %s:%d - %s\\n", __FILE__, __LINE__, \ |
| 230 | cudaGetErrorString(error)); \ |
| 231 | return; \ |
| 232 | } \ |
| 233 | } while(0) |
| 234 | |
| 235 | #define FULL_MASK 0xffffffff |
| 236 | |
| 237 | __device__ inline float warp_reduce_sum(float val) { |
| 238 | for (int offset = 16; offset > 0; offset >>= 1) { |
| 239 | val += __shfl_down_sync(FULL_MASK, val, offset); |
| 240 | } |
| 241 | return val; |
| 242 | } |
| 243 | |
| 244 | __device__ inline float warp_reduce_max(float val) { |
| 245 | for (int offset = 16; offset > 0; offset >>= 1) { |
| 246 | val = fmaxf(val, __shfl_down_sync(FULL_MASK, val, offset)); |
| 247 | } |
| 248 | return val; |
| 249 | } |
| 250 | |
| 251 | """ |
| 252 | |
| 253 | if "#include" not in kernel_code: |
| 254 | return includes + kernel_code |
| 255 | else: |
| 256 | return kernel_code |
| 257 | |
| 258 | def _analyze_compilation_output(self, stderr: str) -> Dict[str, Any]: |
| 259 | """Analyze nvcc compilation output for useful information.""" |