burn_cubecl/kernel/
contiguous.rs

1use burn_backend::cubecl::dtype_to_storage_type;
2use burn_backend::{DType, TensorMetadata};
3use cubecl::quant::scheme::{QuantStore, QuantValue};
4use cubecl::server::MemoryLayoutStrategy;
5
6use crate::{ops::empty_qtensor, tensor::CubeTensor};
7
8/// A storage-tiled tensor laid out in rows again, through cubek's untile; a plain tensor as it
9/// is. A tensor is tiled for one matmul and read there through its binding; every other kernel
10/// and every layout rewrite reads rows, so this is their head.
11///
12/// # Panics
13///
14/// A quantized storage-tiled tensor: it is read only by the kernel it was tiled for, and
15/// nothing lays its packed values and their scales back as rows.
16pub fn untile(tensor: CubeTensor) -> CubeTensor {
17    if !tensor.meta.is_tiled() {
18        return tensor;
19    }
20    assert!(
21        tensor.qparams.is_none(),
22        "untile: a storage-tiled quantized tensor is read only by the kernel it was tiled for; \
23         nothing lays its packed values back as rows"
24    );
25    let (client, device, dtype) = (tensor.client.clone(), tensor.device.clone(), tensor.dtype);
26    let output = cubek::matmul::tiled::storage::untile(
27        &client,
28        tensor.binding(),
29        dtype_to_storage_type(dtype),
30    )
31    .expect("a storage-tiled binding describes its own tiles");
32    CubeTensor::new(client, output.handle, *output.metadata, device, dtype)
33}
34
35/// Make a jit tensor contiguous.
36pub fn into_contiguous(tensor: CubeTensor) -> CubeTensor {
37    // A tiled buffer's strides order its tiles, not a matrix's rows: contiguous says nothing.
38    let tensor = untile(tensor);
39    if tensor.is_contiguous() {
40        return tensor;
41    }
42
43    if tensor.qparams.is_some() {
44        return into_contiguous_quantized(tensor, MemoryLayoutStrategy::Contiguous);
45    }
46
47    let (client, device, dtype) = (tensor.client.clone(), tensor.device.clone(), tensor.dtype);
48
49    let output = cubecl::std::tensor::into_contiguous(
50        &client,
51        tensor.binding(),
52        dtype_to_storage_type(dtype),
53    );
54
55    CubeTensor::new(
56        client.clone(),
57        output.handle,
58        *output.metadata,
59        device,
60        dtype,
61    )
62}
63
64/// Make a jit tensor contiguous with an aligned last stride. Tensor is considered already contiguous
65/// if runtime can read it as is. This is equivalent in practice.
66#[cfg_attr(
67    feature = "tracing",
68    tracing::instrument(level = "trace", skip(tensor))
69)]
70pub fn into_contiguous_aligned(tensor: CubeTensor) -> CubeTensor {
71    let tensor = untile(tensor);
72    if tensor
73        .device
74        .can_read_tensor(tensor.meta.shape(), tensor.meta.strides())
75    {
76        return tensor;
77    }
78
79    if tensor.qparams.is_some() {
80        return into_contiguous_quantized(tensor, MemoryLayoutStrategy::Optimized);
81    }
82
83    let (client, device, dtype) = (tensor.client.clone(), tensor.device.clone(), tensor.dtype);
84
85    let output = cubecl::std::tensor::into_contiguous_pitched(
86        &client,
87        tensor.binding(),
88        dtype_to_storage_type(dtype),
89    );
90
91    CubeTensor::new(
92        client.clone(),
93        output.handle,
94        *output.metadata,
95        device,
96        dtype,
97    )
98}
99
100#[cfg_attr(
101    feature = "tracing",
102    tracing::instrument(level = "trace", skip(tensor))
103)]
104fn into_contiguous_quantized(tensor: CubeTensor, strategy: MemoryLayoutStrategy) -> CubeTensor {
105    let scheme = tensor.scheme();
106    let output = empty_qtensor(tensor.shape(), tensor.scheme(), &tensor.device, strategy);
107    let (values, scales) = tensor.quantized_handles().unwrap();
108    let (out_values, out_scales) = output.quantized_handles().unwrap();
109
110    let (client, dtype_scales, dtype_value) = (scales.client.clone(), scales.dtype, values.dtype);
111
112    match scheme.store {
113        QuantStore::PackedU32(packed_dim) => {
114            cubecl::std::tensor::into_contiguous_packed_ref(
115                &client,
116                values.binding(),
117                out_values.binding(),
118                packed_dim,
119                tensor.meta.shape(),
120                scheme.num_quants(),
121                dtype_to_storage_type(DType::U32),
122            );
123        }
124        // e2m1 is special because it has a native packed representation, `e2m1x2`.
125        // It's internally stored as `u8` with a packing factor of 2.
126        QuantStore::PackedNative(packed_dim) if scheme.value == QuantValue::E2M1 => {
127            cubecl::std::tensor::into_contiguous_packed_ref(
128                &client,
129                values.binding(),
130                out_values.binding(),
131                packed_dim,
132                tensor.meta.shape(),
133                scheme.num_quants(),
134                dtype_to_storage_type(DType::U8),
135            );
136        }
137        _ => {
138            cubecl::std::tensor::copy_into(
139                &client,
140                values.binding(),
141                out_values.binding(),
142                dtype_to_storage_type(dtype_value),
143            );
144        }
145    }
146
147    cubecl::std::tensor::copy_into(
148        &client,
149        scales.binding(),
150        out_scales.binding(),
151        dtype_to_storage_type(dtype_scales),
152    );
153
154    if let (Some(global), Some(out_global)) = (tensor.global(), output.global()) {
155        let dtype_global = global.dtype;
156        cubecl::std::tensor::copy_into(
157            &client,
158            global.binding(),
159            out_global.binding(),
160            dtype_to_storage_type(dtype_global),
161        );
162    }
163
164    output
165}
166
167/// A tensor stored in storage tiles through cubek, and what burn does with one: the matmul it
168/// was tiled for reads it through its binding, every layout rewrite lays it back in rows first,
169/// and a row kernel refuses it.
170#[cfg(all(
171    test,
172    any(feature = "wgpu", feature = "cpu", feature = "cuda", feature = "hip")
173))]
174mod storage_tiled {
175    use burn_backend::{DType, cubecl::dtype_to_storage_type};
176    use burn_std::{Shape, TensorData};
177
178    use crate::{
179        CubeDevice,
180        kernel::{
181            into_contiguous,
182            matmul::{MatmulStrategy, matmul},
183            slice, untile,
184        },
185        ops::{from_data, into_data_sync, permute, reshape, swap_dims},
186        tensor::CubeTensor,
187    };
188
189    fn tensor(shape: &[usize], device: &CubeDevice, seed: u32) -> CubeTensor {
190        let n: usize = shape.iter().product();
191        let data: Vec<f32> = (0..n as u32)
192            .map(|i| {
193                ((i.wrapping_mul(2654435761).wrapping_add(seed) >> 8) % 97) as f32 / 97.0 - 0.5
194            })
195            .collect();
196        from_data(TensorData::new(data, shape.to_vec()), device)
197    }
198
199    /// `tensor` stored in `(rows, cols)` tiles, a tile's rows one after another and the tiles
200    /// row-major.
201    fn tiled(tensor: &CubeTensor, (rows, cols): (usize, usize)) -> CubeTensor {
202        use cubek::matmul::tiled::storage::{Axis, StorageLevels, tile};
203        const ROW: Axis = Axis(0);
204        const COL: Axis = Axis(1);
205        let client = tensor.client.clone();
206        let storage = StorageLevels::new(&[(COL, cols), (ROW, rows)]).grid(&[COL, ROW]);
207        let out = tile(
208            &client,
209            tensor.clone().binding(),
210            &[ROW, COL],
211            dtype_to_storage_type(tensor.dtype),
212            None,
213            storage,
214        )
215        .expect("the tile divides the matrix");
216        CubeTensor::new(
217            client,
218            out.handle,
219            *out.metadata,
220            tensor.device.clone(),
221            tensor.dtype,
222        )
223    }
224
225    fn values(tensor: CubeTensor) -> Vec<f32> {
226        into_data_sync(tensor).as_slice::<f32>().unwrap().to_vec()
227    }
228
229    fn assert_close(have: &[f32], want: &[f32], what: &str) {
230        assert_eq!(have.len(), want.len(), "{what}: lengths differ");
231        for (i, (h, w)) in have.iter().zip(want).enumerate() {
232            assert!((h - w).abs() < 1e-3, "{what}: at {i}, got {h}, want {w}");
233        }
234    }
235
236    /// The tile is one the matmul's plan can stage to; the selector honours it for this `m`.
237    #[test]
238    fn a_tiled_weight_computes_the_same_product() {
239        let device = CubeDevice::default();
240        let (m, k, n) = (64, 256, 512);
241        let lhs = tensor(&[m, k], &device, 1);
242        let rhs = tensor(&[k, n], &device, 2);
243        let weight = tiled(&rhs, (32, 64));
244        assert!(weight.meta.is_tiled());
245
246        let plain = matmul(lhs.clone(), rhs, None, MatmulStrategy::Cube, DType::F32).unwrap();
247        let from_tiles = matmul(lhs, weight, None, MatmulStrategy::Cube, DType::F32).unwrap();
248        assert_eq!(from_tiles.meta.shape().as_slice(), &[m, n]);
249        assert_close(&values(from_tiles), &values(plain), "tiled weight");
250    }
251
252    #[test]
253    fn untile_lays_the_rows_back() {
254        let device = CubeDevice::default();
255        let rhs = tensor(&[2, 64, 96], &device, 3);
256        let weight = tiled(&rhs, (16, 32));
257        let back = untile(weight);
258        assert!(!back.meta.is_tiled());
259        assert_eq!(back.meta.shape().as_slice(), &[2, 64, 96]);
260        assert_eq!(values(back), values(rhs));
261    }
262
263    /// The layout rewrites read rows: on a tiled tensor they lay it back first, and agree
264    /// with the same rewrite of the plain one.
265    #[test]
266    fn layout_rewrites_untile_first() {
267        let device = CubeDevice::default();
268        let rhs = tensor(&[64, 96], &device, 4);
269        let weight = tiled(&rhs, (16, 32));
270
271        let reshaped = reshape(weight.clone(), Shape::new([32, 192]));
272        assert!(!reshaped.meta.is_tiled());
273        assert_eq!(
274            values(reshaped),
275            values(reshape(rhs.clone(), Shape::new([32, 192])))
276        );
277
278        let ranges = [8..40, 32..80];
279        let sliced = slice(weight.clone(), &ranges);
280        assert!(!sliced.meta.is_tiled());
281        assert_eq!(values(sliced), values(slice(rhs.clone(), &ranges)));
282
283        let swapped = swap_dims(weight.clone(), 0, 1);
284        assert!(!swapped.meta.is_tiled());
285        assert_eq!(
286            values(into_contiguous(swapped)),
287            values(into_contiguous(swap_dims(rhs.clone(), 0, 1)))
288        );
289
290        let permuted = permute(weight.clone(), &[1, 0]);
291        assert!(!permuted.meta.is_tiled());
292        assert_eq!(
293            values(into_contiguous(permuted)),
294            values(into_contiguous(permute(rhs.clone(), &[1, 0])))
295        );
296
297        let contiguous = into_contiguous(weight);
298        assert!(!contiguous.meta.is_tiled());
299        assert_eq!(values(contiguous), values(rhs));
300    }
301
302    #[test]
303    #[should_panic(expected = "storage-tiled")]
304    fn a_row_kernel_refuses_a_tiled_tensor() {
305        let device = CubeDevice::default();
306        let rhs = tensor(&[64, 96], &device, 5);
307        let weight = tiled(&rhs, (16, 32));
308        let _ = weight.into_tensor_arg();
309    }
310}
burn_cubecl/kernel/contiguous.rs

burn_cubecl/kernel/
contiguous.rs
