tract_linalg/frame/
strided_panel.rs

1//! K-major f32 panel that the SME kernel reads in place.
2//!
3//! A packed f32 panel is `r` contiguous values per K, then the next K `r * 4`
4//! bytes later. NCHW (and any layout whose mn axis already has element stride 1)
5//! is that same picture with a different K step: one channel plane apart. The
6//! kernel loads it with that stride and the pack copy does not run.
7//!
8//! A last panel shorter than `r` is predicated. The packed fallback still
9//! zero-fills that tail and reports `valid == r`, so the kernel's full-width
10//! loop stays legal. Layouts that are not mn-contiguous take that fallback.
11
12use std::fmt::Display;
13use std::sync::Arc;
14
15use tract_data::internal::*;
16
17use crate::WeightType;
18use crate::mmm::{EagerPackedInput, MMMInputFormat, MMMInputValue, PackedMatrixStorage};
19use crate::pack::PackedFormat;
20
21/// What one SME `AddMatMul` passes as `pa` when the packing index is the
22/// strided one. Three machine words, little-endian: data pointer, signed K
23/// byte stride, valid lanes in `1..=r` (`r` means the tail is already zero).
24#[repr(C)]
25#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
26struct PanelDesc {
27    ptr: usize,
28    k_byte_stride: usize,
29    valid: usize,
30}
31
32#[derive(Clone, Debug, Hash, PartialEq, Eq)]
33enum KeepAlive {
34    /// Bytes stay in the activation. `PanelDesc::ptr` points into this tensor.
35    Tensor(Arc<Tensor>),
36    /// mn was not contiguous. Same descriptor ABI, aimed at a packed buffer.
37    Packed(Arc<Blob>),
38}
39
40/// f32 operand whose panels the 32×32 SME kernel can address without a pack.
41#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
42pub struct StridedKMajor {
43    r: usize,
44}
45
46impl StridedKMajor {
47    pub const fn new(r: usize) -> Self {
48        Self { r }
49    }
50
51    /// Build the operand from a tensor the caller already owns.
52    ///
53    /// `byte_offset` selects one batch item. No activation bytes are copied
54    /// when the mn axis is contiguous and every addressed element sits in the
55    /// allocation; anything else is packed.
56    pub fn panel_from_arc(
57        &self,
58        tensor: Arc<Tensor>,
59        byte_offset: isize,
60        k_axis: usize,
61        mn_axis: usize,
62    ) -> TractResult<Box<dyn MMMInputValue>> {
63        ensure!(k_axis != mn_axis, "k and mn axes are both {k_axis}");
64        ensure!(k_axis < tensor.rank() && mn_axis < tensor.rank());
65        ensure!(tensor.datum_type() == f32::datum_type());
66        let k = tensor.shape()[k_axis];
67        let mn = tensor.shape()[mn_axis];
68        let k_stride = tensor.strides()[k_axis];
69        let mn_stride = tensor.strides()[mn_axis];
70        if mn_stride == 1 && spans_allocation(&tensor, byte_offset, k, mn, k_stride) {
71            Ok(Box::new(StridedPanelInput::view(*self, tensor, byte_offset, k, mn, k_stride)))
72        } else {
73            let view = unsafe {
74                TensorView::from_bytes(&tensor, byte_offset, tensor.shape(), tensor.strides())
75            };
76            let packed = packed_format(self.r).pack_tensor_view(&view, k_axis, mn_axis)?;
77            self.wrap_packed(packed)
78        }
79    }
80
81    fn wrap_packed(&self, packed: Box<dyn MMMInputValue>) -> TractResult<Box<dyn MMMInputValue>> {
82        let eager = packed
83            .downcast::<EagerPackedInput>()
84            .map_err(|_| format_err!("strided fallback did not pack to EagerPackedInput"))?;
85        let panel_bytes = eager.panel_bytes;
86        let k = eager.k();
87        let mn = eager.mn();
88        let blob = eager.packed.clone();
89        let base = blob.as_bytes().as_ptr() as usize;
90        let mut panels = Vec::with_capacity(mn.div_ceil(self.r).max(1));
91        if mn == 0 {
92            panels.push(PanelDesc { ptr: base, k_byte_stride: (self.r * 4), valid: 0 });
93        } else {
94            // The packer zero-fills a short tail, so every panel is a full `r`.
95            for p in 0..mn.div_ceil(self.r) {
96                panels.push(PanelDesc {
97                    ptr: base + p * panel_bytes,
98                    k_byte_stride: (self.r * 4),
99                    valid: self.r,
100                });
101            }
102        }
103        Ok(Box::new(StridedPanelInput {
104            format: *self,
105            keep: KeepAlive::Packed(blob),
106            fact: NoFact,
107            panels,
108            k,
109            mn,
110        }))
111    }
112}
113
114fn packed_format(r: usize) -> PackedFormat {
115    PackedFormat::new(f32::datum_type(), r, 128)
116}
117
118/// True when both ends of the K range and both ends of the mn range land
119/// inside the allocation. A hostile stride falls back to packing.
120fn spans_allocation(
121    tensor: &Tensor,
122    byte_offset: isize,
123    k: usize,
124    mn: usize,
125    k_stride: isize,
126) -> bool {
127    if k == 0 || mn == 0 {
128        return true;
129    }
130    let bytes = tensor.as_bytes().len() as isize;
131    let elem = 4isize;
132    let last_k = (k as isize) - 1;
133    let last_mn = (mn as isize) - 1;
134    for ik in [0, last_k] {
135        for im in [0, last_mn] {
136            // `byte_offset` is already bytes. The strides are elements.
137            let elem_off = ik.saturating_mul(k_stride).saturating_add(im);
138            let off = byte_offset.saturating_add(elem_off.saturating_mul(elem));
139            if off < 0 || off.saturating_add(elem) > bytes {
140                return false;
141            }
142        }
143    }
144    true
145}
146
147impl MMMInputFormat for StridedKMajor {
148    fn prepare_tensor(&self, t: &Tensor, k_axis: usize, mn_axis: usize) -> TractResult<Tensor> {
149        let packed = self.prepare_one(t, k_axis, mn_axis)?;
150        Ok(PackedMatrixStorage::new(packed).into_tensor(t.datum_type()))
151    }
152
153    fn prepare_one(
154        &self,
155        t: &Tensor,
156        k_axis: usize,
157        mn_axis: usize,
158    ) -> TractResult<Box<dyn MMMInputValue>> {
159        // Deep copy: this is the kernel-test entry, which hands us a borrow.
160        // The graph path uses `panel_from_arc` and does not copy.
161        self.panel_from_arc(Arc::new(t.clone()), 0, k_axis, mn_axis)
162    }
163
164    fn prepare_one_view(
165        &self,
166        t: &TensorView,
167        k_axis: usize,
168        mn_axis: usize,
169    ) -> TractResult<Box<dyn MMMInputValue>> {
170        let packed = packed_format(self.r).pack_tensor_view(t, k_axis, mn_axis)?;
171        self.wrap_packed(packed)
172    }
173
174    fn precursor(&self) -> WeightType {
175        WeightType::Plain(f32::datum_type())
176    }
177
178    fn r(&self) -> usize {
179        self.r
180    }
181
182    fn k_alignment(&self) -> usize {
183        1
184    }
185
186    fn mem_size(&self, _k: TDim, _mn: TDim) -> TDim {
187        0.to_dim()
188    }
189
190    fn extract_at_mn_f16(&self, _: &EagerPackedInput, _: usize, _: &mut [f16]) -> TractResult<()> {
191        bail!("strided f32 panel has no f16 extract")
192    }
193
194    fn extract_at_mn_f32(&self, _: &EagerPackedInput, _: usize, _: &mut [f32]) -> TractResult<()> {
195        bail!("strided f32 panel is not an EagerPackedInput")
196    }
197}
198
199impl Display for StridedKMajor {
200    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
201        write!(f, "StridedF32[{}]", self.r)
202    }
203}
204
205#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
206struct NoFact;
207
208impl ExoticFact for NoFact {
209    fn buffer_sizes(&self) -> TVec<TDim> {
210        tvec!(0.to_dim())
211    }
212}
213
214#[derive(Clone, Debug, PartialEq, Eq, Hash)]
215pub struct StridedPanelInput {
216    format: StridedKMajor,
217    keep: KeepAlive,
218    fact: NoFact,
219    panels: Vec<PanelDesc>,
220    k: usize,
221    mn: usize,
222}
223
224impl StridedPanelInput {
225    fn view(
226        format: StridedKMajor,
227        tensor: Arc<Tensor>,
228        byte_offset: isize,
229        k: usize,
230        mn: usize,
231        k_stride: isize,
232    ) -> Self {
233        let base = unsafe { tensor.as_ptr_unchecked::<u8>() } as isize + byte_offset;
234        let k_byte_stride = k_stride * 4;
235        let mut panels = Vec::with_capacity(mn.div_ceil(format.r).max(1));
236        if mn == 0 {
237            panels.push(PanelDesc {
238                ptr: base as usize,
239                k_byte_stride: k_byte_stride as usize,
240                valid: 0,
241            });
242        } else {
243            let mut left = mn;
244            let mut off = 0isize;
245            while left > 0 {
246                let valid = left.min(format.r);
247                panels.push(PanelDesc {
248                    ptr: (base + off) as usize,
249                    k_byte_stride: k_byte_stride as usize,
250                    valid,
251                });
252                left -= valid;
253                off += format.r as isize * 4;
254            }
255        }
256        Self { format, keep: KeepAlive::Tensor(tensor), fact: NoFact, panels, k, mn }
257    }
258
259    pub fn is_view(&self) -> bool {
260        matches!(self.keep, KeepAlive::Tensor(_))
261    }
262}
263
264impl MMMInputValue for StridedPanelInput {
265    fn format(&self) -> &dyn MMMInputFormat {
266        &self.format
267    }
268
269    fn scratch_panel_buffer_layout(&self) -> Option<std::alloc::Layout> {
270        None
271    }
272
273    fn panel_bytes(&self, i: usize, _buffer: Option<*mut u8>) -> TractResult<*const u8> {
274        self.panels
275            .get(i)
276            .map(|p| p as *const PanelDesc as *const u8)
277            .with_context(|| format!("strided panel {i} of {}", self.panels.len()))
278    }
279
280    fn mn(&self) -> usize {
281        self.mn
282    }
283
284    fn k(&self) -> usize {
285        self.k
286    }
287
288    fn exotic_fact(&self) -> &dyn ExoticFact {
289        &self.fact
290    }
291
292    fn extract_at_mn_f16(&self, _mn: usize, _slice: &mut [f16]) -> TractResult<()> {
293        bail!("strided f32 panel has no f16 extract")
294    }
295
296    fn extract_at_mn_f32(&self, mn: usize, slice: &mut [f32]) -> TractResult<()> {
297        ensure!(slice.len() == self.k);
298        ensure!(mn < self.mn);
299        let panel = mn / self.format.r;
300        let lane = mn % self.format.r;
301        let desc = &self.panels[panel];
302        ensure!(lane < desc.valid);
303        let stride = desc.k_byte_stride as isize;
304        unsafe {
305            let ptr = desc.ptr as *const u8;
306            for (i, slot) in slice.iter_mut().enumerate() {
307                *slot = *(ptr.offset(stride * i as isize).add(lane * 4) as *const f32);
308            }
309        }
310        Ok(())
311    }
312}
313
314impl Display for StridedPanelInput {
315    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
316        write!(f, "{} k={} mn={} {} panels", self.format, self.k, self.mn, self.panels.len())
317    }
318}
319
320#[cfg(test)]
321mod tests {
322    use super::*;
323    use crate::mmm::MMMInputValue;
324
325    fn words(input: &dyn MMMInputValue, panel: usize) -> (usize, usize, usize) {
326        let p = input.panel_bytes(panel, None).unwrap() as *const usize;
327        unsafe { (*p, *p.add(1), *p.add(2)) }
328    }
329
330    #[test]
331    fn nchw_row_is_a_view() {
332        let data: Vec<f32> = (0..3 * 40).map(|i| i as f32).collect();
333        let tensor = Tensor::from_shape(&[3, 40], &data).unwrap();
334        let start = tensor.as_bytes().as_ptr() as usize;
335        let input = StridedKMajor::new(32).panel_from_arc(Arc::new(tensor), 0, 0, 1).unwrap();
336        let view = input.downcast_ref::<StridedPanelInput>().unwrap();
337        assert!(view.is_view());
338        assert_eq!(input.panels_count(), 2);
339        let (ptr, stride, valid) = words(&*input, 0);
340        assert_eq!(stride, 40 * 4);
341        assert_eq!(valid, 32);
342        assert!(ptr >= start);
343        assert_eq!(words(&*input, 1).2, 8);
344        let mut got = [0f32; 3];
345        input.extract_at_mn_f32(33, &mut got).unwrap();
346        assert_eq!(got, [33.0, 73.0, 113.0]);
347    }
348
349    #[test]
350    fn k_inner_falls_back_to_a_pack() {
351        let data = vec![1f32; 32 * 4];
352        let tensor = Tensor::from_shape(&[32, 4], &data).unwrap();
353        let input = StridedKMajor::new(32).panel_from_arc(Arc::new(tensor), 0, 1, 0).unwrap();
354        let view = input.downcast_ref::<StridedPanelInput>().unwrap();
355        assert!(!view.is_view());
356        let (_, stride, valid) = words(&*input, 0);
357        assert_eq!(stride, 32 * 4);
358        assert_eq!(valid, 32);
359    }
360}
tract_linalg/frame/strided_panel.rs

tract_linalg/frame/
strided_panel.rs
