Unified CUDA code emitter. Generates standalone CUDA code using raw PTX intrinsics. No CuTe/CUTLASS dependency. Uses CUtensorMap directly for TMA descriptors.
| 31 | |
| 32 | |
| 33 | class CudaEmitter(CppEmitter): |
| 34 | """Unified CUDA code emitter. |
| 35 | |
| 36 | Generates standalone CUDA code using raw PTX intrinsics. |
| 37 | No CuTe/CUTLASS dependency. Uses CUtensorMap directly for TMA descriptors. |
| 38 | """ |
| 39 | |
| 40 | def __init__(self, ctx, all_types, num_threads=None): |
| 41 | super().__init__(ctx, all_types) |
| 42 | self.num_threads = num_threads |
| 43 | headers = [ |
| 44 | "<cuda_runtime.h>", |
| 45 | "<cuda_bf16.h>", |
| 46 | "<cuda_fp16.h>", # __half (float16) |
| 47 | "<cuda.h>", |
| 48 | ] |
| 49 | for hdr in headers: |
| 50 | self.header_cache.add(hdr) |
| 51 | self.header_buffer.write(f"#include {hdr}\n") |
| 52 | # TmaDescriptor is CUtensorMap from the CUDA Driver API |
| 53 | self.header_buffer.write("\nusing TmaDescriptor = CUtensorMap;\n\n") |
| 54 | |
| 55 | def _lltype_to_cpp(self, ll_type: "ll.LLType") -> str: |
| 56 | """Convert an LLType instance to a valid C++ type string.""" |
| 57 | # ------------------------------ Scalar Types ------------------------------ |
| 58 | if hasattr(ll_type, "kind") and ll_type.kind == "int": # IntType |
| 59 | if ll_type.special == "bool": |
| 60 | return "bool" |
| 61 | if ll_type.special == "binary" or ll_type.bits == 4: |
| 62 | # TODO: Add proper 4-bit (int4/uint4) support. CUDA has native 4-bit types. |
| 63 | raise NotImplementedError("4-bit integer types (int4/uint4, binary) are not yet supported. " |
| 64 | "CUDA has native 4-bit support; this will be added in a future release.") |
| 65 | # Map to stdint types (e.g., uint32 -> uint32_t) |
| 66 | sign_prefix = "u" if not ll_type.signed else "" |
| 67 | bit_map = { |
| 68 | 8: f"{sign_prefix}int8_t", 16: f"{sign_prefix}int16_t", 32: f"{sign_prefix}int32_t", 64: |
| 69 | f"{sign_prefix}int64_t", 128: "unsigned __int128" if not ll_type.signed else "__int128" |
| 70 | } |
| 71 | return bit_map[ll_type.bits] |
| 72 | |
| 73 | elif hasattr(ll_type, "kind") and ll_type.kind == "float": # FloatType |
| 74 | float_map = { |
| 75 | "fp4_e2m1": "__nv_fp4_e2m1", "fp8_e5m2": "__nv_fp8_e5m2", "fp8_e4m3": "__nv_fp8_e4m3", "bfloat16": |
| 76 | "__nv_bfloat16", "float16": "__half", |
| 77 | # "tfloat32": "tfloat32_t", |
| 78 | "float32": "float", "float64": "double" |
| 79 | } |
| 80 | return float_map[ll_type.fmt] |
| 81 | |
| 82 | elif hasattr(ll_type, "kind") and ll_type.kind == "void": # VoidType |
| 83 | return "void" |
| 84 | |
| 85 | elif hasattr(ll_type, "kind") and ll_type.kind == "str": # StringType |
| 86 | return "std::string" |
| 87 | |
| 88 | # ------------------------------ Annotated Types ------------------------------ |
| 89 | elif ll_type.is_const(): # Const/GridConstant |
| 90 | inner_cpp = self._lltype_to_cpp(ll_type.inner_type) |
no outgoing calls
no test coverage detected