metrale_gpu_runtime/cuda_backend/
tensormap.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: TMA descriptors (`CUtensorMap`) for `cp.async.bulk.tensor` loads,
4//! encoded with the driver's `cuTensorMapEncodeTiled`.
5//!
6//! Owner: gpu-runtime.
7//! Invariants:
8//! - `TensorMap` is 128 bytes aligned to 64, the size of `CUtensorMap` and the
9//!   alignment `cuTensorMapEncodeTiled` requires (`cuda.h`); a test pins both.
10//! - `tiled_2d_bf16` rejects a global address or row stride that is not 16-byte
11//!   aligned, a row stride narrower than `cols`, and a zero box dimension, before
12//!   it calls the driver.
13//!
14//! The driver takes dimensions fastest-varying first and `rank - 1` strides in
15//! bytes, the innermost stride being the element size. `tiled_2d_bf16` takes
16//! row-major `[rows][cols]` in elements and does that conversion. A descriptor is
17//! passed by value to a `__grid_constant__ const CUtensorMap` kernel parameter
18//! through `KernelLaunch::arg_tensormap`.
19
20use anyhow::{Result, bail};
21
22use crate::gpu::DevicePtr;
23
24/// 2026-09-25: Enum values from `CUtensorMapDataType` and its neighbours in
25/// `/usr/local/cuda/include/cuda.h` (CUDA 13.0): `BFLOAT16` is 9, after
26/// `FLOAT64` (8) and before `FLOAT32_FTZ` (10).
27const CU_TENSOR_MAP_DATA_TYPE_BFLOAT16: u32 = 9;
28const CU_TENSOR_MAP_INTERLEAVE_NONE: u32 = 0;
29const CU_TENSOR_MAP_SWIZZLE_NONE: u32 = 0;
30const CU_TENSOR_MAP_L2_PROMOTION_L2_128B: u32 = 2;
31const CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE: u32 = 0;
32
33/// 2026-09-25: The prototype of `cuTensorMapEncodeTiled` in `cuda.h`. The symbol is
34/// found with `dlsym(RTLD_DEFAULT)` in the libcuda the process has already loaded,
35/// once, and cached (`tensor_map_encode_tiled`).
36type EncodeTiledFn = unsafe extern "C" fn(
37    *mut u8,
38    u32,
39    u32,
40    u64,
41    *const u64,
42    *const u64,
43    *const u32,
44    *const u32,
45    u32,
46    u32,
47    u32,
48    u32,
49) -> i32;
50
51fn tensor_map_encode_tiled() -> Option<EncodeTiledFn> {
52    static CACHED: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
53    let addr = (*CACHED.get_or_init(|| {
54        // 2026-09-25: `dlsym` is unix-only. Elsewhere the symbol is reported
55        // missing and `tiled_2d_bf16` returns an error.
56        #[cfg(unix)]
57        {
58            unsafe extern "C" {
59                fn dlsym(handle: *mut std::ffi::c_void, symbol: *const i8)
60                -> *mut std::ffi::c_void;
61            }
62            let name = c"cuTensorMapEncodeTiled";
63            let p = unsafe { dlsym(std::ptr::null_mut(), name.as_ptr().cast::<i8>()) };
64            if p.is_null() { None } else { Some(p as usize) }
65        }
66        #[cfg(not(unix))]
67        {
68            None
69        }
70    }))?;
71    // 2026-09-25: SAFETY: `addr` is libcuda's `cuTensorMapEncodeTiled`, whose
72    // prototype in `cuda.h` is `EncodeTiledFn`.
73    Some(unsafe { std::mem::transmute::<usize, EncodeTiledFn>(addr) })
74}
75
76/// 2026-09-25: A 128-byte TMA descriptor, aligned to 64 bytes.
77#[repr(C, align(64))]
78#[derive(Clone, Copy)]
79pub struct TensorMap([u8; 128]);
80
81impl std::fmt::Debug for TensorMap {
82    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83        f.write_str("TensorMap(<128 bytes>)")
84    }
85}
86
87impl TensorMap {
88    /// 2026-09-25: The raw bytes, for `KernelLaunch::arg_tensormap`.
89    pub fn bytes(&self) -> &[u8; 128] {
90        &self.0
91    }
92
93    /// 2026-09-25: Describe a row-major bf16 matrix as tiles of `box_rows x box_cols`.
94    ///
95    /// `rows`, `cols` and `row_stride_elems` are in elements. `row_stride_elems`
96    /// is the distance between consecutive rows and may exceed `cols` (a view into
97    /// a wider tensor). Out-of-bounds elements of a tile are filled with zero
98    /// (`CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE`).
99    ///
100    /// Errors: a misaligned address or stride, a stride below `cols`, a zero box
101    /// dimension, a libcuda without `cuTensorMapEncodeTiled`, or a non-zero
102    /// `CUresult` from the encode.
103    pub fn tiled_2d_bf16(
104        global: DevicePtr,
105        rows: u64,
106        cols: u64,
107        row_stride_elems: u64,
108        box_rows: u32,
109        box_cols: u32,
110    ) -> Result<Self> {
111        if !global.0.is_multiple_of(16) {
112            bail!("TMA global address {:#x} is not 16-byte aligned", global.0);
113        }
114        let row_stride_bytes = row_stride_elems * 2;
115        if !row_stride_bytes.is_multiple_of(16) {
116            bail!(
117                "TMA row stride {row_stride_elems} elems ({row_stride_bytes} B) is not \
118                 16-byte aligned; bf16 needs a stride that is a multiple of 8 elements"
119            );
120        }
121        if row_stride_elems < cols {
122            bail!("TMA row stride {row_stride_elems} is narrower than cols {cols}");
123        }
124        if box_rows == 0 || box_cols == 0 {
125            bail!("TMA box dims must be non-zero, got {box_rows}x{box_cols}");
126        }
127
128        // 2026-09-25: Fastest-varying axis first: for row-major [rows][cols]
129        // that is cols.
130        let global_dim: [u64; 2] = [cols, rows];
131        // 2026-09-25: rank - 1 strides, in bytes; the innermost is implicit.
132        let global_strides: [u64; 1] = [row_stride_bytes];
133        let box_dim: [u32; 2] = [box_cols, box_rows];
134        let element_strides: [u32; 2] = [1, 1];
135
136        // 2026-09-25: Resolved at runtime rather than linked, so the crate links
137        // against a libcuda that lacks the symbol; such a driver gets this error.
138        let f = tensor_map_encode_tiled().ok_or_else(|| {
139            anyhow::anyhow!(
140                "cuTensorMapEncodeTiled not present in libcuda — TMA needs a CUDA 12.0+ driver"
141            )
142        })?;
143        let mut map = TensorMap([0u8; 128]);
144        let rc = unsafe {
145            f(
146                map.0.as_mut_ptr(),
147                CU_TENSOR_MAP_DATA_TYPE_BFLOAT16,
148                2,
149                global.0,
150                global_dim.as_ptr(),
151                global_strides.as_ptr(),
152                box_dim.as_ptr(),
153                element_strides.as_ptr(),
154                CU_TENSOR_MAP_INTERLEAVE_NONE,
155                CU_TENSOR_MAP_SWIZZLE_NONE,
156                CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
157                CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE,
158            )
159        };
160        if rc != 0 {
161            bail!(
162                "cuTensorMapEncodeTiled failed: CUresult {rc} \
163                 (rows={rows} cols={cols} stride={row_stride_elems} box={box_rows}x{box_cols})"
164            );
165        }
166        Ok(map)
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    /// 2026-09-25: The guards reject before the driver is called, so these tests
175    /// need no CUDA context.
176    #[test]
177    fn misaligned_stride_is_rejected_without_touching_the_driver() {
178        // 2026-09-25: A stride of 5 bf16 elements is 10 bytes.
179        let e = TensorMap::tiled_2d_bf16(DevicePtr(0x1000), 4, 5, 5, 2, 5).unwrap_err();
180        assert!(
181            e.to_string().contains("not 16-byte aligned"),
182            "expected a stride-alignment error, got: {e}"
183        );
184    }
185
186    #[test]
187    fn misaligned_address_is_rejected() {
188        let e = TensorMap::tiled_2d_bf16(DevicePtr(0x1004), 4, 8, 8, 2, 8).unwrap_err();
189        assert!(e.to_string().contains("not 16-byte aligned"), "got: {e}");
190    }
191
192    #[test]
193    fn a_stride_narrower_than_the_row_is_rejected() {
194        let e = TensorMap::tiled_2d_bf16(DevicePtr(0x1000), 4, 64, 32, 2, 64).unwrap_err();
195        assert!(e.to_string().contains("narrower than cols"), "got: {e}");
196    }
197
198    #[test]
199    fn the_descriptor_has_the_layout_the_driver_requires() {
200        assert_eq!(std::mem::size_of::<TensorMap>(), 128);
201        assert_eq!(std::mem::align_of::<TensorMap>(), 64);
202    }
203}
metrale_gpu_runtime/cuda_backend/tensormap.rs

metrale_gpu_runtime/cuda_backend/
tensormap.rs
