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