mdarray_linalg_blas/matvec/
context.rs

1use std::ops::{Add, Mul};
2
3use mdarray::{Array, Dim, Layout, Shape, Slice};
4use mdarray_linalg::{
5    matvec::{Argmax, MatVec, MatVecBuilder, Outer, OuterBuilder, VecOps},
6    utils::unravel_index,
7};
8use num_complex::ComplexFloat;
9use num_traits::Zero;
10
11use super::{
12    scalar::BlasScalar,
13    simple::{amax, asum, axpy, dotc, dotu, gemv, ger, nrm2},
14};
15use crate::Blas;
16
17struct BlasMatVecBuilder<'a, T, D0, D1, La, Lx>
18where
19    D0: Dim,
20    D1: Dim,
21    La: Layout,
22    Lx: Layout,
23{
24    alpha: T,
25    a: &'a Slice<T, (D0, D1), La>,
26    x: &'a Slice<T, (D1,), Lx>,
27}
28
29impl<'a, T, La, Lx, D0: Dim, D1: Dim> MatVecBuilder<'a, T, La, Lx, D0, D1>
30    for BlasMatVecBuilder<'a, T, D0, D1, La, Lx>
31where
32    La: Layout,
33    Lx: Layout,
34    T: BlasScalar + ComplexFloat,
35    T::Real: Into<T>,
36    D0: Dim,
37    D1: Dim,
38{
39    fn scale(mut self, alpha: T) -> Self {
40        self.alpha = alpha * self.alpha;
41        self
42    }
43
44    fn eval(self) -> Array<T, (D0,)> {
45        let mut y = Array::<T, (D0,)>::from_elem(
46            <(D0,) as Shape>::from_dims(&[self.a.shape().dim(0)]),
47            T::zero(),
48        );
49        gemv(self.alpha, self.a, self.x, T::zero(), &mut y);
50        y
51    }
52
53    fn write<Ly: Layout>(self, y: &mut Slice<T, (D0,), Ly>) {
54        gemv(self.alpha, self.a, self.x, T::zero(), y);
55    }
56
57    fn add_to_vec<Ly: Layout>(self, y: &mut Slice<T, (D0,), Ly>) {
58        gemv(self.alpha, self.a, self.x, T::one(), y);
59    }
60
61    fn add_to_scaled_vec<Ly: Layout>(self, y: &mut Slice<T, (D0,), Ly>, beta: T) {
62        gemv(self.alpha, self.a, self.x, beta, y);
63    }
64}
65
66impl<T, D0: Dim, D1: Dim> MatVec<T, D0, D1> for Blas
67where
68    T: BlasScalar + ComplexFloat,
69    T::Real: Into<T>,
70{
71    fn matvec<'a, La, Lx>(
72        &self,
73        a: &'a Slice<T, (D0, D1), La>,
74        x: &'a Slice<T, (D1,), Lx>,
75    ) -> impl MatVecBuilder<'a, T, La, Lx, D0, D1>
76    where
77        La: Layout,
78        Lx: Layout,
79    {
80        BlasMatVecBuilder {
81            alpha: T::one(),
82            a,
83            x,
84        }
85    }
86}
87
88impl<T: ComplexFloat + BlasScalar + Add<Output = T> + Mul<Output = T> + Zero + Copy, D1: Dim>
89    VecOps<T, D1> for Blas
90{
91    type Real = T::Real;
92
93    fn add_to_scaled<Lx: Layout, Ly: Layout>(
94        &self,
95        alpha: T,
96        x: &Slice<T, (D1,), Lx>,
97        y: &mut Slice<T, (D1,), Ly>,
98    ) {
99        axpy(alpha, x, y);
100    }
101
102    fn dot<Lx: Layout, Ly: Layout>(&self, x: &Slice<T, (D1,), Lx>, y: &Slice<T, (D1,), Ly>) -> T {
103        dotu(x, y)
104    }
105
106    fn dotc<Lx: Layout, Ly: Layout>(&self, x: &Slice<T, (D1,), Lx>, y: &Slice<T, (D1,), Ly>) -> T {
107        dotc(x, y)
108    }
109
110    fn norm2<Lx: Layout>(&self, x: &Slice<T, (D1,), Lx>) -> Self::Real {
111        nrm2(x)
112    }
113
114    fn norm1<Lx: Layout>(&self, x: &Slice<T, (D1,), Lx>) -> Self::Real {
115        asum(x)
116    }
117
118    fn rot<Lx: Layout, Ly: Layout>(
119        &self,
120        _x: &mut Slice<T, (D1,), Lx>,
121        _y: &mut Slice<T, (D1,), Ly>,
122        _c: Self::Real,
123        _s: T,
124    ) {
125        todo!()
126    }
127}
128
129impl<
130    T: ComplexFloat
131        + std::cmp::PartialOrd
132        + BlasScalar
133        + Add<Output = T>
134        + Mul<Output = T>
135        + Zero
136        + Copy,
137> Argmax<T> for Blas
138where
139    T::Real: PartialOrd,
140{
141    fn argmax_write<Lx: Layout, S: Shape>(
142        &self,
143        _x: &Slice<T, S, Lx>,
144        _output: &mut Vec<usize>,
145    ) -> bool {
146        unimplemented!();
147    }
148
149    fn argmax<Lx: Layout, S: Shape>(&self, _x: &Slice<T, S, Lx>) -> Option<Vec<usize>> {
150        unimplemented!();
151    }
152
153    fn argmax_abs_write<Lx: Layout, S: Shape>(
154        &self,
155        x: &Slice<T, S, Lx>,
156        output: &mut Vec<usize>,
157    ) -> bool {
158        output.clear();
159        if x.is_empty() {
160            return false;
161        }
162        if x.rank() == 0 {
163            return true;
164        }
165        let max_flat_idx = amax(x);
166        let indices = unravel_index(x, max_flat_idx);
167        output.extend_from_slice(&indices);
168        true
169    }
170
171    fn argmax_abs<Lx: Layout, S: Shape>(&self, x: &Slice<T, S, Lx>) -> Option<Vec<usize>> {
172        let mut result = Vec::new();
173        if self.argmax_abs_write(x, &mut result) {
174            Some(result)
175        } else {
176            None
177        }
178    }
179}
180
181struct BlasOuterBuilder<'a, T, Dx, Dy, Lx, Ly>
182where
183    Lx: Layout,
184    Ly: Layout,
185    Dx: Dim,
186    Dy: Dim,
187{
188    alpha: T,
189    x: &'a Slice<T, (Dx,), Lx>,
190    y: &'a Slice<T, (Dy,), Ly>,
191}
192
193impl<'a, T, Dx, Dy, Lx, Ly> OuterBuilder<'a, T, Lx, Ly, Dx, Dy>
194    for BlasOuterBuilder<'a, T, Dx, Dy, Lx, Ly>
195where
196    Lx: Layout,
197    Ly: Layout,
198    T: BlasScalar + ComplexFloat,
199    T::Real: Into<T>,
200    Dx: Dim,
201    Dy: Dim,
202{
203    fn scale(mut self, alpha: T) -> Self {
204        self.alpha = alpha * self.alpha;
205        self
206    }
207
208    fn eval(self) -> Array<T, (Dx, Dy)> {
209        let shape = <(Dx, Dy) as Shape>::from_dims(&[self.x.len(), self.y.len()]);
210        let mut a = Array::<T, (Dx, Dy)>::from_elem(shape, T::zero());
211        ger(self.alpha, self.x, self.y, &mut a);
212        a
213    }
214
215    fn write<La: Layout>(self, a: &mut Slice<T, (Dx, Dy), La>) {
216        let zero = T::zero();
217        a.fill(zero);
218        ger(self.alpha, self.x, self.y, a);
219    }
220
221    fn add_to<La: Layout>(self, a: &mut Slice<T, (Dx, Dy), La>) {
222        ger(self.alpha, self.x, self.y, a);
223    }
224
225}
226
227impl<T, Dx, Dy> Outer<T, Dx, Dy> for Blas
228where
229    T: BlasScalar + ComplexFloat,
230    T::Real: Into<T>,
231    Dx: Dim,
232    Dy: Dim,
233{
234    fn outer<'a, Lx, Ly>(
235        &self,
236        x: &'a Slice<T, (Dx,), Lx>,
237        y: &'a Slice<T, (Dy,), Ly>,
238    ) -> impl OuterBuilder<'a, T, Lx, Ly, Dx, Dy>
239    where
240        Lx: Layout,
241        Ly: Layout,
242    {
243        BlasOuterBuilder {
244            alpha: T::one(),
245            x,
246            y,
247        }
248    }
249}
mdarray_linalg_blas/matvec/context.rs

mdarray_linalg_blas/matvec/
context.rs
