1	# bend-blas: the shaped layer. A Mat<rows, cols> is one row-major
2	# Array<F32> with its shape in the type, a Vec<n> likewise, and every
3	# routine below states the shape rule in its signature, so a wrong
4	# operand is a checker error naming both shapes, not a wrong answer:
5	#   Mat.gemm(rows, k, cols, alpha, a, b, beta, c) needs a Mat<rows, k>, a
6	#   Mat<k, cols> and a Mat<rows, cols>; a transpose is a separate def
7	#   (gemm_t, gemm_tn, gemv_t), since it changes the shape.
8	# The sizes are Nat arguments that double as the type indices. The raw
9	# routines of blas.bend and lapack.bend run underneath, unchanged; the
10	# shaped forms answer the same Result, with .try twins, and consume
11	# their operands the same way (Mat.clone keeps a copy). This is the
12	# scipy.linalg to blas.bend's scipy.linalg.blas, plus the shapes.
13	import Base
14	import ./blas.bend as B
15	import ./lapack.bend as L
16	
17	type Mat<-rows: Nat, -cols: Nat> is Type:
18	  MatA{a: Array<F32>}
19	
20	type Vec<-n: Nat> is Type:
21	  VecA{a: Array<F32>}
22	
23	# Construction
24	# ------------
25	
26	# the depth of the smallest power-of-two block holding count floats
27	def depth.go(fuel: Nat, +n: U32, +d: Nat, +cap: U32) -> Nat:
28	  match fuel:
29	    case 0n:
30	      d
31	    case 1n+p:
32	      Bool.pick(Nat, U32.is_lt(cap, n), depth.go(p, n, 1n+d, (cap * 2 : U32)), d)
33	
34	def depth(+count: Nat) -> Nat:
35	  depth.go(40n, U32.from_nat(count), 0n, 1)
36	
37	def fill.go(xs: List<&2, F32>, +i: U32, a: Array<F32>) -> Array<F32>:
38	  match xs:
39	    case []:
40	      a
41	    case x <> rest:
42	      fill.go(rest, (i + 1 : U32), Array.set(F32, a, i, x))
43	
44	# an array of the smallest power-of-two size holding count floats,
45	# the first ones from xs (cut or padded with 0.0), the rest 0.0
46	def array.of(+count: Nat, xs: List<&2, F32>) -> Array<F32>:
47	  fill.go(List.take(&2, F32, xs, count), 0, Array.new(F32, depth(count), 0.0))
48	
49	def take.go(n: Nat, +i: U32, r: Array<F32> & F32, acc: List<&2, F32>) -> List<&2, F32>:
50	  match n:
51	    case 0n:
52	      (a, x) = r
53	      x <> acc
54	    case 1n+p:
55	      (a, x) = r
56	      take.go(p, (i - 1 : U32), Array.get(F32, a, (i - 1 : U32)), x <> acc)
57	
58	# the first count floats of an array, as a list
59	def list.of(+count: Nat, a: Array<F32>) -> List<&2, F32>:
60	  match count:
61	    case 0n:
62	      []
63	    case 1n+p:
64	      take.go(p, (U32.from_nat(p) : U32), Array.get(F32, a, U32.from_nat(p)), [])
65	
66	# xs cut into rows of cols
67	def rows.of(rows: Nat, +cols: Nat, +xs: List<&2, F32>) -> List<&2, List<&2, F32>>:
68	  match rows:
69	    case 0n:
70	      []
71	    case 1n+p:
72	      List.take(&2, F32, xs, cols) <> rows.of(p, cols, List.drop(&2, F32, xs, cols))
73	
74	def F32.show.of(x: F32) -> String:
75	  F32.show(x)
76	
77	def row.show(xs: List<&2, F32>) -> String:
78	  List.show(~&2, ~F32, ~F32.show.of, xs)
79	
80	# rows * cols floats, row-major, cut or padded with 0.0
81	def Mat.from_list(+rows: Nat, +cols: Nat, xs: List<&2, F32>) -> Mat<rows, cols>:
82	  MatA{array.of(Nat.mul(rows, cols), xs)}
83	
84	def Mat.zeros(+rows: Nat, +cols: Nat) -> Mat<rows, cols>:
85	  MatA{Array.new(F32, depth(Nat.mul(rows, cols)), 0.0)}
86	
87	def Mat.to_list(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> List<&2, F32>:
88	  MatA{a} = m
89	  list.of(Nat.mul(rows, cols), a)
90	
91	def Mat.to_rows(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> List<&2, List<&2, F32>>:
92	  rows.of(rows, cols, Mat.to_list(rows, cols, m))
93	
94	def Mat.show(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> String:
95	  List.show(~&2, ~List<&2, F32>, ~row.show, Mat.to_rows(rows, cols, m))
96	
97	def Mat.clone.fin(-rows: Nat, -cols: Nat, r: Array<F32> & Array<F32>) -> Mat<rows, cols> & Mat<rows, cols>:
98	  (x, y) = r
99	  (MatA{x}, MatA{y})
100	
101	# a Mat is affine: two copies, to be opened by matching in a def of its own
102	def Mat.clone(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Mat<rows, cols> & Mat<rows, cols>:
103	  MatA{a} = m
104	  Mat.clone.fin(rows, cols, Array.clone(F32, a))
105	
106	def Vec.from_list(+n: Nat, xs: List<&2, F32>) -> Vec<n>:
107	  VecA{array.of(n, xs)}
108	
109	def Vec.zeros(+n: Nat) -> Vec<n>:
110	  VecA{Array.new(F32, depth(n), 0.0)}
111	
112	def Vec.to_list(+n: Nat, v: Vec<n>) -> List<&2, F32>:
113	  VecA{a} = v
114	  list.of(n, a)
115	
116	def Vec.show(+n: Nat, v: Vec<n>) -> String:
117	  row.show(Vec.to_list(n, v))
118	
119	def Vec.clone.fin(-n: Nat, r: Array<F32> & Array<F32>) -> Vec<n> & Vec<n>:
120	  (x, y) = r
121	  (VecA{x}, VecA{y})
122	
123	def Vec.clone(+n: Nat, v: Vec<n>) -> Vec<n> & Vec<n>:
124	  VecA{a} = v
125	  Vec.clone.fin(n, Array.clone(F32, a))
126	
127	# the raw answers, shaped
128	def Mat.wrap(-rows: Nat, -cols: Nat, r: Result<&1, &1, U32 & String, Array<F32>>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>):
129	  IO.pure(Result<&1, &1, U32 & String, Mat<rows, cols>>, Result.map(&1, &1, U32 & String, Array<F32>, Mat<rows, cols>, x => MatA{x}, r))
130	
131	def Vec.wrap(-n: Nat, r: Result<&1, &1, U32 & String, Array<F32>>) -> IO(Result<&1, &1, U32 & String, Vec<n>>):
132	  IO.pure(Result<&1, &1, U32 & String, Vec<n>>, Result.map(&1, &1, U32 & String, Array<F32>, Vec<n>, x => VecA{x}, r))
133	
134	# BLAS
135	# ----
136	
137	# C = alpha A B + beta C
138	def Mat.gemm(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<rows, k>, b: Mat<k, cols>, beta: F32, c: Mat<rows, cols>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>):
139	  MatA{x} = a
140	  MatA{y} = b
141	  MatA{z} = c
142	  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<rows, cols>>, B.Blas.sgemm(0, 0, U32.from_nat(rows), U32.from_nat(cols), U32.from_nat(k), alpha, x, y, beta, z), r => Mat.wrap(rows, cols, r))
143	
144	def Mat.gemm.try(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<rows, k>, b: Mat<k, cols>, beta: F32, c: Mat<rows, cols>) -> IO(Mat<rows, cols>):
145	  IO.try(Mat<rows, cols>, Mat.gemm(rows, k, cols, alpha, a, b, beta, c))
146	
147	# C = alpha A B^T + beta C, B stored cols x k
148	def Mat.gemm_t(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<rows, k>, b: Mat<cols, k>, beta: F32, c: Mat<rows, cols>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>):
149	  MatA{x} = a
150	  MatA{y} = b
151	  MatA{z} = c
152	  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<rows, cols>>, B.Blas.sgemm(0, 1, U32.from_nat(rows), U32.from_nat(cols), U32.from_nat(k), alpha, x, y, beta, z), r => Mat.wrap(rows, cols, r))
153	
154	def Mat.gemm_t.try(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<rows, k>, b: Mat<cols, k>, beta: F32, c: Mat<rows, cols>) -> IO(Mat<rows, cols>):
155	  IO.try(Mat<rows, cols>, Mat.gemm_t(rows, k, cols, alpha, a, b, beta, c))
156	
157	# C = alpha A^T B + beta C, A stored k x rows
158	def Mat.gemm_tn(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<k, rows>, b: Mat<k, cols>, beta: F32, c: Mat<rows, cols>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>):
159	  MatA{x} = a
160	  MatA{y} = b
161	  MatA{z} = c
162	  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<rows, cols>>, B.Blas.sgemm(1, 0, U32.from_nat(rows), U32.from_nat(cols), U32.from_nat(k), alpha, x, y, beta, z), r => Mat.wrap(rows, cols, r))
163	
164	def Mat.gemm_tn.try(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<k, rows>, b: Mat<k, cols>, beta: F32, c: Mat<rows, cols>) -> IO(Mat<rows, cols>):
165	  IO.try(Mat<rows, cols>, Mat.gemm_tn(rows, k, cols, alpha, a, b, beta, c))
166	
167	# y = alpha A x + beta y
168	def Mat.gemv(+rows: Nat, +cols: Nat, alpha: F32, a: Mat<rows, cols>, x: Vec<cols>, beta: F32, y: Vec<rows>) -> IO(Result<&1, &1, U32 & String, Vec<rows>>):
169	  MatA{m} = a
170	  VecA{xs} = x
171	  VecA{ys} = y
172	  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<rows>>, B.Blas.sgemv(0, U32.from_nat(rows), U32.from_nat(cols), alpha, m, xs, beta, ys), r => Vec.wrap(rows, r))
173	
174	def Mat.gemv.try(+rows: Nat, +cols: Nat, alpha: F32, a: Mat<rows, cols>, x: Vec<cols>, beta: F32, y: Vec<rows>) -> IO(Vec<rows>):
175	  IO.try(Vec<rows>, Mat.gemv(rows, cols, alpha, a, x, beta, y))
176	
177	# y = alpha A^T x + beta y
178	def Mat.gemv_t(+rows: Nat, +cols: Nat, alpha: F32, a: Mat<rows, cols>, x: Vec<rows>, beta: F32, y: Vec<cols>) -> IO(Result<&1, &1, U32 & String, Vec<cols>>):
179	  MatA{m} = a
180	  VecA{xs} = x
181	  VecA{ys} = y
182	  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<cols>>, B.Blas.sgemv(1, U32.from_nat(rows), U32.from_nat(cols), alpha, m, xs, beta, ys), r => Vec.wrap(cols, r))
183	
184	def Mat.gemv_t.try(+rows: Nat, +cols: Nat, alpha: F32, a: Mat<rows, cols>, x: Vec<rows>, beta: F32, y: Vec<cols>) -> IO(Vec<cols>):
185	  IO.try(Vec<cols>, Mat.gemv_t(rows, cols, alpha, a, x, beta, y))
186	
187	# A = alpha x y^T + A
188	def Mat.ger(+rows: Nat, +cols: Nat, alpha: F32, x: Vec<rows>, y: Vec<cols>, a: Mat<rows, cols>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>):
189	  VecA{xs} = x
190	  VecA{ys} = y
191	  MatA{m} = a
192	  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<rows, cols>>, B.Blas.sger(U32.from_nat(rows), U32.from_nat(cols), alpha, xs, ys, m), r => Mat.wrap(rows, cols, r))
193	
194	def Mat.ger.try(+rows: Nat, +cols: Nat, alpha: F32, x: Vec<rows>, y: Vec<cols>, a: Mat<rows, cols>) -> IO(Mat<rows, cols>):
195	  IO.try(Mat<rows, cols>, Mat.ger(rows, cols, alpha, x, y, a))
196	
197	def Vec.dot(+n: Nat, x: Vec<n>, y: Vec<n>) -> IO(Result<&1, &1, U32 & String, F32>):
198	  VecA{xs} = x
199	  VecA{ys} = y
200	  B.Blas.sdot(U32.from_nat(n), xs, ys)
201	
202	def Vec.dot.try(+n: Nat, x: Vec<n>, y: Vec<n>) -> IO(F32):
203	  IO.try(F32, Vec.dot(n, x, y))
204	
205	def Vec.nrm2(+n: Nat, x: Vec<n>) -> IO(Result<&1, &1, U32 & String, F32>):
206	  VecA{xs} = x
207	  B.Blas.snrm2(U32.from_nat(n), xs)
208	
209	def Vec.nrm2.try(+n: Nat, x: Vec<n>) -> IO(F32):
210	  IO.try(F32, Vec.nrm2(n, x))
211	
212	def Vec.asum(+n: Nat, x: Vec<n>) -> IO(Result<&1, &1, U32 & String, F32>):
213	  VecA{xs} = x
214	  B.Blas.sasum(U32.from_nat(n), xs)
215	
216	def Vec.asum.try(+n: Nat, x: Vec<n>) -> IO(F32):
217	  IO.try(F32, Vec.asum(n, x))
218	
219	def Vec.iamax(+n: Nat, x: Vec<n>) -> IO(Result<&1, &1, U32 & String, U32>):
220	  VecA{xs} = x
221	  B.Blas.isamax(U32.from_nat(n), xs)
222	
223	def Vec.iamax.try(+n: Nat, x: Vec<n>) -> IO(U32):
224	  IO.try(U32, Vec.iamax(n, x))
225	
226	# y = alpha x + y
227	def Vec.axpy(+n: Nat, alpha: F32, x: Vec<n>, y: Vec<n>) -> IO(Result<&1, &1, U32 & String, Vec<n>>):
228	  VecA{xs} = x
229	  VecA{ys} = y
230	  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<n>>, B.Blas.saxpy(U32.from_nat(n), alpha, xs, ys), r => Vec.wrap(n, r))
231	
232	def Vec.axpy.try(+n: Nat, alpha: F32, x: Vec<n>, y: Vec<n>) -> IO(Vec<n>):
233	  IO.try(Vec<n>, Vec.axpy(n, alpha, x, y))
234	
235	def Vec.scal(+n: Nat, alpha: F32, x: Vec<n>) -> IO(Result<&1, &1, U32 & String, Vec<n>>):
236	  VecA{xs} = x
237	  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<n>>, B.Blas.sscal(U32.from_nat(n), alpha, xs), r => Vec.wrap(n, r))
238	
239	def Vec.scal.try(+n: Nat, alpha: F32, x: Vec<n>) -> IO(Vec<n>):
240	  IO.try(Vec<n>, Vec.scal(n, alpha, x))
241	
242	# LAPACK
243	# ------
244	
245	# X = A^-1 B; Fail with the zero pivot as the code when A is singular
246	def Mat.solve(+n: Nat, +nrhs: Nat, a: Mat<n, n>, b: Mat<n, nrhs>) -> IO(Result<&1, &1, U32 & String, Mat<n, nrhs>>):
247	  MatA{x} = a
248	  MatA{y} = b
249	  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<n, nrhs>>, L.Lapack.sgesv(U32.from_nat(n), U32.from_nat(nrhs), x, y), r => Mat.wrap(n, nrhs, r))
250	
251	def Mat.solve.try(+n: Nat, +nrhs: Nat, a: Mat<n, n>, b: Mat<n, nrhs>) -> IO(Mat<n, nrhs>):
252	  IO.try(Mat<n, nrhs>, Mat.solve(n, nrhs, a, b))
253	
254	# the lower Cholesky factor L of a symmetric positive definite A
255	def Mat.cholesky(+n: Nat, a: Mat<n, n>) -> IO(Result<&1, &1, U32 & String, Mat<n, n>>):
256	  MatA{x} = a
257	  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<n, n>>, L.Lapack.spotrf(U32.from_nat(n), x), r => Mat.wrap(n, n, r))
258	
259	def Mat.cholesky.try(+n: Nat, a: Mat<n, n>) -> IO(Mat<n, n>):
260	  IO.try(Mat<n, n>, Mat.cholesky(n, a))
261	
262	# X = A^-1 B given L from Mat.cholesky
263	def Mat.cholesky_solve(+n: Nat, +nrhs: Nat, l: Mat<n, n>, b: Mat<n, nrhs>) -> IO(Result<&1, &1, U32 & String, Mat<n, nrhs>>):
264	  MatA{x} = l
265	  MatA{y} = b
266	  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<n, nrhs>>, L.Lapack.spotrs(U32.from_nat(n), U32.from_nat(nrhs), x, y), r => Mat.wrap(n, nrhs, r))
267	
268	def Mat.cholesky_solve.try(+n: Nat, +nrhs: Nat, l: Mat<n, n>, b: Mat<n, nrhs>) -> IO(Mat<n, nrhs>):
269	  IO.try(Mat<n, nrhs>, Mat.cholesky_solve(n, nrhs, l, b))
270	
271	# the singular values, largest first
272	def Mat.svd_values(+m: Nat, +n: Nat, a: Mat<m, n>) -> IO(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)>>):
273	  MatA{x} = a
274	  IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<Nat.min(m, n)>>, L.Lapack.sgesvd_s(U32.from_nat(m), U32.from_nat(n), x), r => Vec.wrap(Nat.min(m, n), r))
275	
276	def Mat.svd_values.try(+m: Nat, +n: Nat, a: Mat<m, n>) -> IO(Vec<Nat.min(m, n)>):
277	  IO.try(Vec<Nat.min(m, n)>, Mat.svd_values(m, n, a))
278	
279	def Mat.svd.fin(-m: Nat, -n: Nat, s: Array<F32>, uv: Array<F32> & Array<F32>) -> IO(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>):
280	  match uv:
281	    case (u, vt):
282	      IO.pure(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>, Done{(VecA{s}, (MatA{u}, MatA{vt}))})
283	
284	def Mat.svd.wrap(-m: Nat, -n: Nat, r: Result<&1, &1, U32 & String, Array<F32> & (Array<F32> & Array<F32>)>) -> IO(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>):
285	  match r:
286	    case Fail{e}:
287	      IO.pure(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>, Fail{e})
288	    case Done{v}:
289	      match v:
290	        case (s, uv):
291	          Mat.svd.fin(m, n, s, uv)
292	
293	# the thin SVD: (s, (u, vt)) with A = u diag(s) vt
294	def Mat.svd(+m: Nat, +n: Nat, a: Mat<m, n>) -> IO(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>):
295	  MatA{x} = a
296	  IO.bind(Result<&1, &1, U32 & String, Array<F32> & (Array<F32> & Array<F32>)>, Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>, L.Lapack.sgesvd(U32.from_nat(m), U32.from_nat(n), x), r => Mat.svd.wrap(m, n, r))
297	
298	def Mat.svd.try(+m: Nat, +n: Nat, a: Mat<m, n>) -> IO(Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)):
299	  IO.try(Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>), Mat.svd(m, n, a))
300	
301	def Mat.eigh.wrap(-n: Nat, r: Result<&1, &1, U32 & String, Array<F32> & Array<F32>>) -> IO(Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>):
302	  match r:
303	    case Fail{e}:
304	      IO.pure(Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>, Fail{e})
305	    case Done{v}:
306	      match v:
307	        case (w, vs):
308	          IO.pure(Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>, Done{(VecA{w}, MatA{vs})})
309	
310	# the eigenvalues ascending and the eigenvectors as the rows of the Mat
311	def Mat.eigh(+n: Nat, a: Mat<n, n>) -> IO(Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>):
312	  MatA{x} = a
313	  IO.bind(Result<&1, &1, U32 & String, Array<F32> & Array<F32>>, Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>, L.Lapack.ssyev(U32.from_nat(n), x), r => Mat.eigh.wrap(n, r))
314	
315	def Mat.eigh.try(+n: Nat, a: Mat<n, n>) -> IO(Vec<n> & Mat<n, n>):
316	  IO.try(Vec<n> & Mat<n, n>, Mat.eigh(n, a))