Aluode/PerceptionLabPortable
0
1from cython cimport floating
2
3from scipy.linalg.cython_blas cimport sdot, ddot
4from scipy.linalg.cython_blas cimport sasum, dasum
5from scipy.linalg.cython_blas cimport saxpy, daxpy
6from scipy.linalg.cython_blas cimport snrm2, dnrm2
7from scipy.linalg.cython_blas cimport scopy, dcopy
8from scipy.linalg.cython_blas cimport sscal, dscal
9from scipy.linalg.cython_blas cimport srotg, drotg
10from scipy.linalg.cython_blas cimport srot, drot
11from scipy.linalg.cython_blas cimport sgemv, dgemv
12from scipy.linalg.cython_blas cimport sger, dger
13from scipy.linalg.cython_blas cimport sgemm, dgemm
14
15
16################
17# BLAS Level 1 #
18################
19
20cdef floating _dot(int n, const floating *x, int incx,
21 const floating *y, int incy) noexcept nogil:
22 """x.T.y"""
23 if floating is float:
24 return sdot(&n, <float *> x, &incx, <float *> y, &incy)
25 else:
26 return ddot(&n, <double *> x, &incx, <double *> y, &incy)
27
28
29cpdef _dot_memview(const floating[::1] x, const floating[::1] y):
30 return _dot(x.shape[0], &x[0], 1, &y[0], 1)
31
32
33cdef floating _asum(int n, const floating *x, int incx) noexcept nogil:
34 """sum(|x_i|)"""
35 if floating is float:
36 return sasum(&n, <float *> x, &incx)
37 else:
38 return dasum(&n, <double *> x, &incx)
39
40
41cpdef _asum_memview(const floating[::1] x):
42 return _asum(x.shape[0], &x[0], 1)
43
44
45cdef void _axpy(int n, floating alpha, const floating *x, int incx,
46 floating *y, int incy) noexcept nogil:
47 """y := alpha * x + y"""
48 if floating is float:
49 saxpy(&n, &alpha, <float *> x, &incx, y, &incy)
50 else:
51 daxpy(&n, &alpha, <double *> x, &incx, y, &incy)
52
53
54cpdef _axpy_memview(floating alpha, const floating[::1] x, floating[::1] y):
55 _axpy(x.shape[0], alpha, &x[0], 1, &y[0], 1)
56
57
58cdef floating _nrm2(int n, const floating *x, int incx) noexcept nogil:
59 """sqrt(sum((x_i)^2))"""
60 if floating is float:
61 return snrm2(&n, <float *> x, &incx)
62 else:
63 return dnrm2(&n, <double *> x, &incx)
64
65
66cpdef _nrm2_memview(const floating[::1] x):
67 return _nrm2(x.shape[0], &x[0], 1)
68
69
70cdef void _copy(int n, const floating *x, int incx, const floating *y, int incy) noexcept nogil:
71 """y := x"""
72 if floating is float:
73 scopy(&n, <float *> x, &incx, <float *> y, &incy)
74 else:
75 dcopy(&n, <double *> x, &incx, <double *> y, &incy)
76
77
78cpdef _copy_memview(const floating[::1] x, const floating[::1] y):
79 _copy(x.shape[0], &x[0], 1, &y[0], 1)
80
81
82cdef void _scal(int n, floating alpha, const floating *x, int incx) noexcept nogil:
83 """x := alpha * x"""
84 if floating is float:
85 sscal(&n, &alpha, <float *> x, &incx)
86 else:
87 dscal(&n, &alpha, <double *> x, &incx)
88
89
90cpdef _scal_memview(floating alpha, const floating[::1] x):
91 _scal(x.shape[0], alpha, &x[0], 1)
92
93
94cdef void _rotg(floating *a, floating *b, floating *c, floating *s) noexcept nogil:
95 """Generate plane rotation"""
96 if floating is float:
97 srotg(a, b, c, s)
98 else:
99 drotg(a, b, c, s)
100
101
102cpdef _rotg_memview(floating a, floating b, floating c, floating s):
103 _rotg(&a, &b, &c, &s)
104 return a, b, c, s
105
106
107cdef void _rot(int n, floating *x, int incx, floating *y, int incy,
108 floating c, floating s) noexcept nogil:
109 """Apply plane rotation"""
110 if floating is float:
111 srot(&n, x, &incx, y, &incy, &c, &s)
112 else:
113 drot(&n, x, &incx, y, &incy, &c, &s)
114
115
116cpdef _rot_memview(floating[::1] x, floating[::1] y, floating c, floating s):
117 _rot(x.shape[0], &x[0], 1, &y[0], 1, c, s)
118
119
120################
121# BLAS Level 2 #
122################
123
124cdef void _gemv(BLAS_Order order, BLAS_Trans ta, int m, int n, floating alpha,
125 const floating *A, int lda, const floating *x, int incx,
126 floating beta, floating *y, int incy) noexcept nogil:
127 """y := alpha * op(A).x + beta * y"""
128 cdef char ta_ = ta
129 if order == BLAS_Order.RowMajor:
130 ta_ = BLAS_Trans.NoTrans if ta == BLAS_Trans.Trans else BLAS_Trans.Trans
131 if floating is float:
132 sgemv(&ta_, &n, &m, &alpha, <float *> A, &lda, <float *> x,
133 &incx, &beta, y, &incy)
134 else:
135 dgemv(&ta_, &n, &m, &alpha, <double *> A, &lda, <double *> x,
136 &incx, &beta, y, &incy)
137 else:
138 if floating is float:
139 sgemv(&ta_, &m, &n, &alpha, <float *> A, &lda, <float *> x,
140 &incx, &beta, y, &incy)
141 else:
142 dgemv(&ta_, &m, &n, &alpha, <double *> A, &lda, <double *> x,
143 &incx, &beta, y, &incy)
144
145
146cpdef _gemv_memview(BLAS_Trans ta, floating alpha, const floating[:, :] A,
147 const floating[::1] x, floating beta, floating[::1] y):
148 cdef:
149 int m = A.shape[0]
150 int n = A.shape[1]
151 BLAS_Order order = (
152 BLAS_Order.ColMajor if A.strides[0] == A.itemsize else BLAS_Order.RowMajor
153 )
154 int lda = m if order == BLAS_Order.ColMajor else n
155
156 _gemv(order, ta, m, n, alpha, &A[0, 0], lda, &x[0], 1, beta, &y[0], 1)
157
158
159cdef void _ger(BLAS_Order order, int m, int n, floating alpha,
160 const floating *x, int incx, const floating *y,
161 int incy, floating *A, int lda) noexcept nogil:
162 """A := alpha * x.y.T + A"""
163 if order == BLAS_Order.RowMajor:
164 if floating is float:
165 sger(&n, &m, &alpha, <float *> y, &incy, <float *> x, &incx, A, &lda)
166 else:
167 dger(&n, &m, &alpha, <double *> y, &incy, <double *> x, &incx, A, &lda)
168 else:
169 if floating is float:
170 sger(&m, &n, &alpha, <float *> x, &incx, <float *> y, &incy, A, &lda)
171 else:
172 dger(&m, &n, &alpha, <double *> x, &incx, <double *> y, &incy, A, &lda)
173
174
175cpdef _ger_memview(floating alpha, const floating[::1] x,
176 const floating[::1] y, floating[:, :] A):
177 cdef:
178 int m = A.shape[0]
179 int n = A.shape[1]
180 BLAS_Order order = (
181 BLAS_Order.ColMajor if A.strides[0] == A.itemsize else BLAS_Order.RowMajor
182 )
183 int lda = m if order == BLAS_Order.ColMajor else n
184
185 _ger(order, m, n, alpha, &x[0], 1, &y[0], 1, &A[0, 0], lda)
186
187
188################
189# BLAS Level 3 #
190################
191
192cdef void _gemm(BLAS_Order order, BLAS_Trans ta, BLAS_Trans tb, int m, int n,
193 int k, floating alpha, const floating *A, int lda, const floating *B,
194 int ldb, floating beta, floating *C, int ldc) noexcept nogil:
195 """C := alpha * op(A).op(B) + beta * C"""
196 # TODO: Remove the pointer casts below once SciPy uses const-qualification.
197 # See: https://github.com/scipy/scipy/issues/14262
198 cdef:
199 char ta_ = ta
200 char tb_ = tb
201 if order == BLAS_Order.RowMajor:
202 if floating is float:
203 sgemm(&tb_, &ta_, &n, &m, &k, &alpha, <float*>B,
204 &ldb, <float*>A, &lda, &beta, C, &ldc)
205 else:
206 dgemm(&tb_, &ta_, &n, &m, &k, &alpha, <double*>B,
207 &ldb, <double*>A, &lda, &beta, C, &ldc)
208 else:
209 if floating is float:
210 sgemm(&ta_, &tb_, &m, &n, &k, &alpha, <float*>A,
211 &lda, <float*>B, &ldb, &beta, C, &ldc)
212 else:
213 dgemm(&ta_, &tb_, &m, &n, &k, &alpha, <double*>A,
214 &lda, <double*>B, &ldb, &beta, C, &ldc)
215
216
217cpdef _gemm_memview(BLAS_Trans ta, BLAS_Trans tb, floating alpha,
218 const floating[:, :] A, const floating[:, :] B, floating beta,
219 floating[:, :] C):
220 cdef:
221 int m = A.shape[0] if ta == BLAS_Trans.NoTrans else A.shape[1]
222 int n = B.shape[1] if tb == BLAS_Trans.NoTrans else B.shape[0]
223 int k = A.shape[1] if ta == BLAS_Trans.NoTrans else A.shape[0]
224 int lda, ldb, ldc
225 BLAS_Order order = (
226 BLAS_Order.ColMajor if A.strides[0] == A.itemsize else BLAS_Order.RowMajor
227 )
228
229 if order == BLAS_Order.RowMajor:
230 lda = k if ta == BLAS_Trans.NoTrans else m
231 ldb = n if tb == BLAS_Trans.NoTrans else k
232 ldc = n
233 else:
234 lda = m if ta == BLAS_Trans.NoTrans else k
235 ldb = k if tb == BLAS_Trans.NoTrans else n
236 ldc = m
237
238 _gemm(order, ta, tb, m, n, k, alpha, &A[0, 0],
239 lda, &B[0, 0], ldb, beta, &C[0, 0], ldc)
240