mdarray_linalg_tblis/
contract.rs

1use mdarray::{Array, Dim, DynRank, Layout, Shape, Slice};
2use mdarray_linalg::contract::{Axes, Contract, ContractBuilder, MatmulBuilder};
3use num_traits::{MulAdd, One, Zero};
4use tblis::{
5    containers::TblisTensor,
6    einsum_impl::tblis_einsum,
7    float_trait::TblisFloatAPI,
8    tensor_ops::{TblisMultCfgBuilder, tblis_tensor_mult},
9};
10
11use crate::Tblis;
12
13struct TblisMatmulBuilder<'a, T, D0, D1, D2, La, Lb>
14where
15    La: Layout,
16    Lb: Layout,
17    D0: Dim,
18    D1: Dim,
19    D2: Dim,
20{
21    alpha: T,
22    a: &'a Slice<T, (D0, D1), La>,
23    b: &'a Slice<T, (D1, D2), Lb>,
24}
25
26struct TblisContractBuilder<'a, T, Sa, Sb, La, Lb>
27where
28    La: Layout,
29    Lb: Layout,
30    Sa: Shape,
31    Sb: Shape,
32{
33    alpha: T,
34    a: &'a Slice<T, Sa, La>,
35    b: &'a Slice<T, Sb, Lb>,
36    mode: ContractMode<'a>,
37}
38
39enum ContractMode<'a> {
40    Structured { axes: Axes<'a> },
41    Einsum {
42        indices_a: &'a [u8],
43        indices_b: &'a [u8],
44        indices_c: &'a [u8],
45    },
46}
47
48impl<'a, T, D0, D1, D2, La, Lb> MatmulBuilder<'a, T, D0, D1, D2, La, Lb>
49    for TblisMatmulBuilder<'a, T, D0, D1, D2, La, Lb>
50where
51    La: Layout,
52    Lb: Layout,
53    T: TblisFloatAPI + Zero + One,
54    D0: Dim,
55    D1: Dim,
56    D2: Dim,
57{
58    fn scale(mut self, factor: T) -> Self {
59        self.alpha = self.alpha * factor;
60        self
61    }
62
63    fn eval(self) -> Array<T, (D0, D2)> {
64        let (m, _) = *self.a.shape();
65        let (_, n) = *self.b.shape();
66        let mut c = Array::from_elem((m, n), T::zero());
67        self.write(&mut c);
68        c
69    }
70
71    fn write<Lc: Layout>(self, c: &mut Slice<T, (D0, D2), Lc>) {
72        self.add_to_scaled(c, T::zero())
73    }
74
75    fn add_to<Lc: Layout>(self, c: &mut Slice<T, (D0, D2), Lc>) {
76        self.add_to_scaled(c, T::one())
77    }
78
79    fn add_to_scaled<Lc: Layout>(self, c: &mut Slice<T, (D0, D2), Lc>, beta: T) {
80        assert_eq!(self.a.dim(1), self.b.dim(0), "matrix inner dimensions must match");
81        assert_eq!(c.dim(0), self.a.dim(0), "output row count mismatch");
82        assert_eq!(c.dim(1), self.b.dim(1), "output column count mismatch");
83
84        let a_t = slice_to_tblis_tensor(self.a);
85        let b_t = slice_to_tblis_tensor(self.b);
86        let mut c_t = slice_to_tblis_tensor_mut(c);
87        let cfg = TblisMultCfgBuilder::default()
88            .alpha(self.alpha)
89            .beta(beta)
90            .build()
91            .unwrap();
92
93        unsafe {
94            tblis_tensor_mult(&a_t, "ik", &b_t, "kj", &mut c_t, "ij", Some(cfg));
95        }
96    }
97}
98
99impl<'a, T, Sa, Sb, La, Lb> ContractBuilder<'a, T, Sa, Sb, La, Lb>
100    for TblisContractBuilder<'a, T, Sa, Sb, La, Lb>
101where
102    La: Layout,
103    Lb: Layout,
104    T: TblisFloatAPI + Zero + One + MulAdd<Output = T>,
105    Sa: Shape,
106    Sb: Shape,
107{
108    fn scale(mut self, factor: T) -> Self {
109        self.alpha = self.alpha * factor;
110        self
111    }
112
113    fn eval(self) -> Array<T, DynRank> {
114        match &self.mode {
115            ContractMode::Structured { axes } => {
116                let (indices_a, indices_b, indices_c, shape_c) =
117                    build_structured_subscripts(self.a, self.b, axes);
118                let mut c = Array::from_elem(shape_c, T::zero());
119                self.run_structured_into(&indices_a, &indices_b, &indices_c, &mut c, T::zero());
120                c
121            }
122            ContractMode::Einsum {
123                indices_a,
124                indices_b,
125                indices_c,
126            } => self.eval_einsum(indices_a, indices_b, indices_c),
127        }
128    }
129
130    fn write<Sc: Shape, Lc: Layout>(self, c: &mut Slice<T, Sc, Lc>) {
131        match &self.mode {
132            ContractMode::Structured { axes } => {
133                let (indices_a, indices_b, indices_c, shape_c) =
134                    build_structured_subscripts(self.a, self.b, axes);
135                assert_output_shape(c, &shape_c);
136                self.run_structured_into(&indices_a, &indices_b, &indices_c, c, T::zero());
137            }
138            ContractMode::Einsum {
139                indices_a,
140                indices_b,
141                indices_c,
142            } => {
143                let result = self.eval_einsum(indices_a, indices_b, indices_c);
144                copy_result(c, &result);
145            }
146        }
147    }
148
149    fn add_to<Sc: Shape, Lc: Layout>(self, c: &mut Slice<T, Sc, Lc>) {
150        self.add_to_scaled(c, T::one())
151    }
152
153    fn add_to_scaled<Sc: Shape, Lc: Layout>(self, c: &mut Slice<T, Sc, Lc>, beta: T) {
154        match &self.mode {
155            ContractMode::Structured { axes } => {
156                let (indices_a, indices_b, indices_c, shape_c) =
157                    build_structured_subscripts(self.a, self.b, axes);
158                assert_output_shape(c, &shape_c);
159                self.run_structured_into(&indices_a, &indices_b, &indices_c, c, beta);
160            }
161            ContractMode::Einsum {
162                indices_a,
163                indices_b,
164                indices_c,
165            } => {
166                let result = self.eval_einsum(indices_a, indices_b, indices_c);
167                add_result(c, &result, beta);
168            }
169        }
170    }
171}
172
173impl<'a, T, Sa, Sb, La, Lb> TblisContractBuilder<'a, T, Sa, Sb, La, Lb>
174where
175    La: Layout,
176    Lb: Layout,
177    T: TblisFloatAPI + Zero + One + MulAdd<Output = T>,
178    Sa: Shape,
179    Sb: Shape,
180{
181    fn run_structured_into<Sc: Shape, Lc: Layout>(
182        &self,
183        indices_a: &str,
184        indices_b: &str,
185        indices_c: &str,
186        c: &mut Slice<T, Sc, Lc>,
187        beta: T,
188    ) {
189        let a_t = slice_to_tblis_tensor(self.a);
190        let b_t = slice_to_tblis_tensor(self.b);
191        let mut c_t = slice_to_tblis_tensor_mut(c);
192        let cfg = TblisMultCfgBuilder::default()
193            .alpha(self.alpha)
194            .beta(beta)
195            .build()
196            .unwrap();
197
198        unsafe {
199            tblis_tensor_mult(&a_t, indices_a, &b_t, indices_b, &mut c_t, indices_c, Some(cfg));
200        }
201    }
202
203    fn eval_einsum(&self, indices_a: &[u8], indices_b: &[u8], indices_c: &[u8]) -> Array<T, DynRank> {
204        assert_eq!(indices_a.len(), self.a.rank(), "einsum indices_a length must match A rank");
205        assert_eq!(indices_b.len(), self.b.rank(), "einsum indices_b length must match B rank");
206
207        let subscripts = build_einsum_subscripts(indices_a, indices_b, indices_c);
208        let a_t = slice_to_tblis_tensor(self.a);
209        let b_t = slice_to_tblis_tensor(self.b);
210        let operands = [&a_t, &b_t];
211
212        let (vec, tsr) = unsafe {
213            tblis_einsum(&subscripts, &operands, "optimal", None, true, None)
214                .expect("tblis_einsum must allocate an output tensor")
215        };
216
217        let shape: Vec<usize> = tsr.shape.iter().map(|&d| d as usize).collect();
218        let mut result = Array::from_elem(shape, T::zero());
219        for (dst, src) in result.iter_mut().zip(vec) {
220            *dst = self.alpha * src;
221        }
222        result
223    }
224}
225
226impl<T> Contract<T> for Tblis
227where
228    T: TblisFloatAPI + Zero + One + MulAdd<Output = T>,
229{
230    fn matmul<'a, D0, D1, D2, La, Lb>(
231        &self,
232        a: &'a Slice<T, (D0, D1), La>,
233        b: &'a Slice<T, (D1, D2), Lb>,
234    ) -> impl MatmulBuilder<'a, T, D0, D1, D2, La, Lb>
235    where
236        La: Layout,
237        Lb: Layout,
238        D0: Dim,
239        D1: Dim,
240        D2: Dim,
241    {
242        TblisMatmulBuilder {
243            alpha: T::one(),
244            a,
245            b,
246        }
247    }
248
249    fn contract_all<'a, Sa, Sb, La, Lb>(
250        &self,
251        a: &'a Slice<T, Sa, La>,
252        b: &'a Slice<T, Sb, Lb>,
253    ) -> T
254    where
255        T: 'a,
256        Sa: Shape,
257        Sb: Shape,
258        La: Layout,
259        Lb: Layout,
260    {
261        assert_eq!(
262            a.rank(),
263            b.rank(),
264            "contract_all requires tensors with the same rank (got ranks {} and {})",
265            a.rank(),
266            b.rank()
267        );
268
269        self.contract_n(a, b, a.rank()).eval().into_scalar()
270    }
271
272    fn contract_n<'a, Sa, Sb, La, Lb>(
273        &self,
274        a: &'a Slice<T, Sa, La>,
275        b: &'a Slice<T, Sb, Lb>,
276        n: usize,
277    ) -> impl ContractBuilder<'a, T, Sa, Sb, La, Lb>
278    where
279        T: 'a,
280        Sa: Shape,
281        Sb: Shape,
282        La: Layout,
283        Lb: Layout,
284    {
285        TblisContractBuilder {
286            alpha: T::one(),
287            a,
288            b,
289            mode: ContractMode::Structured {
290                axes: Axes::LastFirst { k: n },
291            },
292        }
293    }
294
295    fn contract_pairs<'a, Sa, Sb, La, Lb>(
296        &self,
297        a: &'a Slice<T, Sa, La>,
298        b: &'a Slice<T, Sb, Lb>,
299        axes_a: &'a [usize],
300        axes_b: &'a [usize],
301    ) -> impl ContractBuilder<'a, T, Sa, Sb, La, Lb>
302    where
303        T: 'a,
304        Sa: Shape,
305        Sb: Shape,
306        La: Layout,
307        Lb: Layout,
308    {
309        TblisContractBuilder {
310            alpha: T::one(),
311            a,
312            b,
313            mode: ContractMode::Structured {
314                axes: Axes::Specific(axes_a, axes_b),
315            },
316        }
317    }
318
319    fn contract<'a, Sa, Sb, La, Lb>(
320        &self,
321        a: &'a Slice<T, Sa, La>,
322        b: &'a Slice<T, Sb, Lb>,
323        indices_a: &'a [u8],
324        indices_b: &'a [u8],
325        indices_c: &'a [u8],
326    ) -> impl ContractBuilder<'a, T, Sa, Sb, La, Lb>
327    where
328        T: 'a,
329        Sa: Shape,
330        Sb: Shape,
331        La: Layout,
332        Lb: Layout,
333    {
334        TblisContractBuilder {
335            alpha: T::one(),
336            a,
337            b,
338            mode: ContractMode::Einsum {
339                indices_a,
340                indices_b,
341                indices_c,
342            },
343        }
344    }
345}
346
347fn slice_to_tblis_tensor<T, S, L>(slice: &Slice<T, S, L>) -> TblisTensor<T>
348where
349    T: TblisFloatAPI,
350    S: Shape,
351    L: Layout,
352{
353    let shape: Vec<isize> = (0..slice.rank()).map(|i| slice.dim(i) as isize).collect();
354    let stride: Vec<isize> = (0..slice.rank()).map(|i| slice.stride(i)).collect();
355    TblisTensor::new(slice.as_ptr() as *mut T, &shape, &stride)
356}
357
358fn slice_to_tblis_tensor_mut<T, S, L>(slice: &mut Slice<T, S, L>) -> TblisTensor<T>
359where
360    T: TblisFloatAPI,
361    S: Shape,
362    L: Layout,
363{
364    let shape: Vec<isize> = (0..slice.rank()).map(|i| slice.dim(i) as isize).collect();
365    let stride: Vec<isize> = (0..slice.rank()).map(|i| slice.stride(i)).collect();
366    TblisTensor::new(slice.as_mut_ptr(), &shape, &stride)
367}
368
369fn resolve_axes<T, Sa, Sb, La, Lb>(
370    axes: &Axes<'_>,
371    a: &Slice<T, Sa, La>,
372    b: &Slice<T, Sb, Lb>,
373) -> (Vec<usize>, Vec<usize>)
374where
375    Sa: Shape,
376    Sb: Shape,
377    La: Layout,
378    Lb: Layout,
379{
380    let rank_a = a.rank();
381    let rank_b = b.rank();
382    match axes {
383        Axes::All => {
384            assert_eq!(rank_a, rank_b, "full contraction requires equal ranks");
385            ((0..rank_a).collect(), (0..rank_b).collect())
386        }
387        Axes::LastFirst { k } => {
388            assert!(*k <= rank_a, "cannot contract {k} axes on A of rank {rank_a}");
389            assert!(*k <= rank_b, "cannot contract {k} axes on B of rank {rank_b}");
390            (((rank_a - *k)..rank_a).collect(), (0..*k).collect())
391        }
392        Axes::Specific(ax_a, ax_b) => (ax_a.to_vec(), ax_b.to_vec()),
393        Axes::SpecificOwned(ax_a, ax_b) => (ax_a.clone(), ax_b.clone()),
394    }
395}
396
397fn build_structured_subscripts<T, Sa, Sb, La, Lb>(
398    a: &Slice<T, Sa, La>,
399    b: &Slice<T, Sb, Lb>,
400    axes: &Axes<'_>,
401) -> (String, String, String, Vec<usize>)
402where
403    Sa: Shape,
404    Sb: Shape,
405    La: Layout,
406    Lb: Layout,
407{
408    let (axes_a, axes_b) = resolve_axes(axes, a, b);
409    assert_eq!(axes_a.len(), axes_b.len(), "axis count mismatch");
410
411    let rank_a = a.rank();
412    let rank_b = b.rank();
413    let mut labels_a = vec![usize::MAX; rank_a];
414    let mut labels_b = vec![usize::MAX; rank_b];
415    let mut next_label = 0usize;
416
417    for (&ax_a, &ax_b) in axes_a.iter().zip(axes_b.iter()) {
418        assert!(ax_a < rank_a, "axis {ax_a} out of bounds for A rank {rank_a}");
419        assert!(ax_b < rank_b, "axis {ax_b} out of bounds for B rank {rank_b}");
420        assert_eq!(a.dim(ax_a), b.dim(ax_b), "dimension mismatch on contracted axes");
421        assert_eq!(labels_a[ax_a], usize::MAX, "duplicate contracted axis in A");
422        assert_eq!(labels_b[ax_b], usize::MAX, "duplicate contracted axis in B");
423        labels_a[ax_a] = next_label;
424        labels_b[ax_b] = next_label;
425        next_label += 1;
426    }
427
428    let mut output_labels = Vec::new();
429    let mut output_shape = Vec::new();
430
431    for (ax, label) in labels_a.iter_mut().enumerate() {
432        if *label == usize::MAX {
433            *label = next_label;
434            output_labels.push(next_label);
435            output_shape.push(a.dim(ax));
436            next_label += 1;
437        }
438    }
439
440    for (ax, label) in labels_b.iter_mut().enumerate() {
441        if *label == usize::MAX {
442            *label = next_label;
443            output_labels.push(next_label);
444            output_shape.push(b.dim(ax));
445            next_label += 1;
446        }
447    }
448
449    let idx_a = labels_to_subscript(&labels_a);
450    let idx_b = labels_to_subscript(&labels_b);
451    let idx_c = labels_to_subscript(&output_labels);
452
453    (idx_a, idx_b, idx_c, output_shape)
454}
455
456fn build_einsum_subscripts(indices_a: &[u8], indices_b: &[u8], indices_c: &[u8]) -> String {
457    let mut unique = Vec::<u8>::new();
458    for &label in indices_a.iter().chain(indices_b.iter()).chain(indices_c.iter()) {
459        if !unique.contains(&label) {
460            unique.push(label);
461        }
462    }
463    assert!(unique.len() <= 128, "TBLIS backend supports at most 128 distinct einsum labels");
464
465    let idx_a = labels_to_subscript(
466        &indices_a
467            .iter()
468            .map(|label| unique.iter().position(|x| x == label).unwrap())
469            .collect::<Vec<_>>(),
470    );
471    let idx_b = labels_to_subscript(
472        &indices_b
473            .iter()
474            .map(|label| unique.iter().position(|x| x == label).unwrap())
475            .collect::<Vec<_>>(),
476    );
477    let idx_c = labels_to_subscript(
478        &indices_c
479            .iter()
480            .map(|label| unique.iter().position(|x| x == label).unwrap())
481            .collect::<Vec<_>>(),
482    );
483
484    format!("{idx_a},{idx_b}->{idx_c}")
485}
486
487fn labels_to_subscript(labels: &[usize]) -> String {
488    labels
489        .iter()
490        .map(|&label| {
491            let code = 0x0100_u32 + label as u32;
492            char::from_u32(code).expect("invalid generated label")
493        })
494        .collect()
495}
496
497fn assert_output_shape<T, S, L>(c: &Slice<T, S, L>, expected: &[usize])
498where
499    S: Shape,
500    L: Layout,
501{
502    assert_eq!(c.rank(), expected.len(), "output rank mismatch");
503    for (i, &dim) in expected.iter().enumerate() {
504        assert_eq!(c.dim(i), dim, "output shape mismatch on axis {i}");
505    }
506}
507
508fn copy_result<T, Sc, Lc>(c: &mut Slice<T, Sc, Lc>, result: &Array<T, DynRank>)
509where
510    T: Copy,
511    Sc: Shape,
512    Lc: Layout,
513{
514    assert_output_shape(c, result.shape().dims());
515    for (dst, src) in c.iter_mut().zip(result.iter()) {
516        *dst = *src;
517    }
518}
519
520fn add_result<T, Sc, Lc>(c: &mut Slice<T, Sc, Lc>, result: &Array<T, DynRank>, beta: T)
521where
522    T: Copy + MulAdd<Output = T>,
523    Sc: Shape,
524    Lc: Layout,
525{
526    assert_output_shape(c, result.shape().dims());
527    for (dst, src) in c.iter_mut().zip(result.iter()) {
528        *dst = beta.mul_add(*dst, *src);
529    }
530}
mdarray_linalg_tblis/contract.rs

mdarray_linalg_tblis/
contract.rs
