mdarray_linalg/naive/contract/
context.rs

1use std::iter::Sum;
2use std::ops::AddAssign;
3
4use mdarray::{Array, Dim, DynRank, Layout, Shape, Slice};
5use num_complex::ComplexFloat;
6use num_traits::{MulAdd, One, Zero};
7
8use super::simple::naive_matmul;
9use crate::{
10    Naive,
11    contract::{
12        _contract, _empty_einsum_shape, _hypercontract, Axes, Contract, ContractBuilder,
13        MatmulBuilder, einsum_to_contract_axes,
14    },
15};
16
17struct NaiveMatmulBuilder<'a, T, D0, D1, D2, La, Lb>
18where
19    La: Layout,
20    Lb: Layout,
21    D0: Dim,
22    D1: Dim,
23    D2: Dim,
24{
25    alpha: T,
26    a: &'a Slice<T, (D0, D1), La>,
27    b: &'a Slice<T, (D1, D2), Lb>,
28}
29
30struct NaiveContractBuilder<'a, T, Sa, Sb, La, Lb>
31where
32    La: Layout,
33    Lb: Layout,
34    Sa: Shape,
35    Sb: Shape,
36{
37    alpha: T,
38    a: &'a Slice<T, Sa, La>,
39    b: &'a Slice<T, Sb, Lb>,
40    axes: Axes<'a>,
41    einsum: bool,
42    empty_output_shape: Option<Vec<usize>>,
43    einsum_axes_a: Option<Vec<usize>>,
44    einsum_axes_b: Option<Vec<usize>>,
45    current_output_labels: Option<Vec<u8>>,
46    requested_output_labels: Option<Vec<u8>>,
47}
48
49impl<'a, T, D0, D1, D2, La, Lb> MatmulBuilder<'a, T, D0, D1, D2, La, Lb>
50    for NaiveMatmulBuilder<'a, T, D0, D1, D2, La, Lb>
51where
52    La: Layout,
53    Lb: Layout,
54    T: ComplexFloat + Zero + One + MulAdd<Output = T>,
55    D0: Dim,
56    D1: Dim,
57    D2: Dim,
58{
59    fn scale(mut self, factor: T) -> Self {
60        self.alpha = self.alpha * factor;
61        self
62    }
63
64    fn eval(self) -> Array<T, (D0, D2)> {
65        let (m, _) = *self.a.shape();
66        let (_, n) = *self.b.shape();
67        let mut c = Array::from_elem((m, n), T::zero());
68        naive_matmul(self.alpha, self.a, self.b, T::zero(), &mut c);
69        c
70    }
71
72    fn write<Lc: Layout>(self, c: &mut Slice<T, (D0, D2), Lc>) {
73        naive_matmul(self.alpha, self.a, self.b, T::zero(), c);
74    }
75
76    fn add_to<Lc: Layout>(self, c: &mut Slice<T, (D0, D2), Lc>) {
77        naive_matmul(self.alpha, self.a, self.b, T::one(), c);
78    }
79
80    fn add_to_scaled<Lc: Layout>(self, c: &mut Slice<T, (D0, D2), Lc>, beta: T) {
81        naive_matmul(self.alpha, self.a, self.b, beta, c);
82    }
83}
84
85impl<'a, T, Sa, Sb, La, Lb> ContractBuilder<'a, T, Sa, Sb, La, Lb>
86    for NaiveContractBuilder<'a, T, Sa, Sb, La, Lb>
87where
88    La: Layout,
89    Lb: Layout,
90    T: ComplexFloat + Zero + One + MulAdd<Output = T> + AddAssign + Sum,
91    Sa: Shape,
92    Sb: Shape,
93{
94    fn scale(mut self, factor: T) -> Self {
95        self.alpha = self.alpha * factor;
96        self
97    }
98
99    fn eval(self) -> Array<T, DynRank> {
100        if let Some(shape) = self.empty_output_shape {
101            return Array::from_elem(shape, T::zero());
102        }
103        if self.einsum {
104            let a = self.a.to_array().into_dyn();
105            let b = self.b.to_array().into_dyn();
106
107            // TODO: it is very likely that this copy is useless. It
108            // should be removed in a near future (9th april
109            // 2026). However due to tricky typing error I did not
110            // manage to give the view to _hypercontract without
111            // removing the genericity toward Layout and shape :(
112
113            let axes_a = self
114                .einsum_axes_a
115                .as_deref()
116                .expect("missing einsum axis labels for A");
117            let axes_b = self
118                .einsum_axes_b
119                .as_deref()
120                .expect("missing einsum axis labels for B");
121
122            let mut result = _hypercontract(Naive, a.expr(), b.expr(), axes_a, axes_b);
123
124            if let (Some(current), Some(requested)) = (
125                self.current_output_labels.as_deref(),
126                self.requested_output_labels.as_deref(),
127            )
128                && current != requested
129            {
130                    let perm: Vec<usize> = requested
131                        .iter()
132                        .map(|label| {
133                            current
134                                .iter()
135                                .position(|cur| cur == label)
136                                .expect("output label not present in contraction result")
137                        })
138                        .collect();
139
140                    result = result.permute(perm).to_tensor().into_dyn();
141            }
142
143            if self.alpha != T::one() {
144                result = result.map(|x| x * self.alpha).into_dyn();
145            }
146
147            result
148        } else {
149            _contract(Naive, self.a, self.b, self.axes, self.alpha)
150        }
151    }
152
153    fn write<Sc: Shape, Lc: Layout>(self, c: &mut Slice<T, Sc, Lc>) {
154        let result = self.eval();
155        assert_eq!(c.rank(), result.rank(), "output rank mismatch");
156        for i in 0..c.rank() {
157            assert_eq!(
158                c.dim(i),
159                result.dim(i),
160                "output shape mismatch on axis {i}: expected {}, got {}",
161                result.dim(i),
162                c.dim(i)
163            );
164        }
165        if c.is_empty() {
166            return;
167        }
168        for (dst, src) in c.iter_mut().zip(result.iter()) {
169            *dst = *src;
170        }
171    }
172
173    fn add_to<Sc: Shape, Lc: Layout>(self, c: &mut Slice<T, Sc, Lc>) {
174        self.add_to_scaled(c, T::one())
175    }
176
177    fn add_to_scaled<Sc: Shape, Lc: Layout>(self, c: &mut Slice<T, Sc, Lc>, beta: T) {
178        let result = self.eval();
179        assert_eq!(c.rank(), result.rank(), "output rank mismatch");
180        for i in 0..c.rank() {
181            assert_eq!(
182                c.dim(i),
183                result.dim(i),
184                "output shape mismatch on axis {i}: expected {}, got {}",
185                result.dim(i),
186                c.dim(i)
187            );
188        }
189        if c.is_empty() {
190            return;
191        }
192        for (dst, src) in c.iter_mut().zip(result.iter()) {
193            *dst = beta * *dst + *src;
194        }
195    }
196}
197
198impl<T> Contract<T> for Naive
199where
200    T: ComplexFloat + Zero + One + MulAdd<Output = T> + AddAssign + Sum,
201{
202    fn matmul<'a, D0, D1, D2, La, Lb>(
203        &self,
204        a: &'a Slice<T, (D0, D1), La>,
205        b: &'a Slice<T, (D1, D2), Lb>,
206    ) -> impl MatmulBuilder<'a, T, D0, D1, D2, La, Lb>
207    where
208        La: Layout,
209        Lb: Layout,
210        D0: Dim,
211        D1: Dim,
212        D2: Dim,
213    {
214        NaiveMatmulBuilder {
215            alpha: T::one(),
216            a,
217            b,
218        }
219    }
220
221    fn contract_all<'a, Sa, Sb, La, Lb>(
222        &self,
223        a: &'a Slice<T, Sa, La>,
224        b: &'a Slice<T, Sb, Lb>,
225    ) -> T
226    where
227        T: 'a,
228        Sa: Shape,
229        Sb: Shape,
230        La: Layout,
231        Lb: Layout,
232    {
233        _contract(Naive, a, b, Axes::All, T::one()).into_scalar()
234    }
235
236    fn contract_n<'a, Sa, Sb, La, Lb>(
237        &self,
238        a: &'a Slice<T, Sa, La>,
239        b: &'a Slice<T, Sb, Lb>,
240        n: usize,
241    ) -> impl ContractBuilder<'a, T, Sa, Sb, La, Lb>
242    where
243        T: 'a,
244        Sa: Shape,
245        Sb: Shape,
246        La: Layout,
247        Lb: Layout,
248    {
249        NaiveContractBuilder {
250            alpha: T::one(),
251            a,
252            b,
253            axes: Axes::LastFirst { k: n },
254            einsum: false,
255            empty_output_shape: None,
256            einsum_axes_a: None,
257            einsum_axes_b: None,
258            current_output_labels: None,
259            requested_output_labels: None,
260        }
261    }
262
263    fn contract_pairs<'a, Sa, Sb, La, Lb>(
264        &self,
265        a: &'a Slice<T, Sa, La>,
266        b: &'a Slice<T, Sb, Lb>,
267        axes_a: &'a [usize],
268        axes_b: &'a [usize],
269    ) -> impl ContractBuilder<'a, T, Sa, Sb, La, Lb>
270    where
271        T: 'a,
272        Sa: Shape,
273        Sb: Shape,
274        La: Layout,
275        Lb: Layout,
276    {
277        NaiveContractBuilder {
278            alpha: T::one(),
279            a,
280            b,
281            axes: Axes::Specific(axes_a, axes_b),
282            einsum: false,
283            empty_output_shape: None,
284            einsum_axes_a: None,
285            einsum_axes_b: None,
286            current_output_labels: None,
287            requested_output_labels: None,
288        }
289    }
290
291    fn contract<'a, Sa, Sb, La, Lb>(
292        &self,
293        a: &'a Slice<T, Sa, La>,
294        b: &'a Slice<T, Sb, Lb>,
295        indices_a: &'a [u8],
296        indices_b: &'a [u8],
297        indices_c: &'a [u8],
298    ) -> impl ContractBuilder<'a, T, Sa, Sb, La, Lb>
299    where
300        T: 'a,
301        Sa: Shape,
302        Sb: Shape,
303        La: Layout,
304        Lb: Layout,
305    {
306        assert_eq!(
307            indices_a.len(),
308            a.rank(),
309            "einsum indices_a length ({}) must match A rank ({})",
310            indices_a.len(),
311            a.rank()
312        );
313        assert_eq!(
314            indices_b.len(),
315            b.rank(),
316            "einsum indices_b length ({}) must match B rank ({})",
317            indices_b.len(),
318            b.rank()
319        );
320
321        let free: std::collections::HashSet<u8> = indices_c.iter().copied().collect();
322        let current_output_labels: Vec<u8> = indices_a
323            .iter()
324            .chain(indices_b.iter())
325            .copied()
326            .filter(|label| free.contains(label))
327            .collect();
328        let (einsum_axes_a, einsum_axes_b) =
329            einsum_to_contract_axes(indices_a, indices_b, indices_c);
330
331        NaiveContractBuilder {
332            alpha: T::one(),
333            a,
334            b,
335            axes: Axes::SpecificOwned(Vec::new(), Vec::new()),
336            einsum: true,
337            empty_output_shape: _empty_einsum_shape(a, b, indices_a, indices_b, indices_c),
338            einsum_axes_a: Some(einsum_axes_a),
339            einsum_axes_b: Some(einsum_axes_b),
340            current_output_labels: Some(current_output_labels),
341            requested_output_labels: Some(indices_c.to_vec()),
342        }
343    }
344}
mdarray_linalg/naive/contract/context.rs

mdarray_linalg/naive/contract/
context.rs
