bend_tensors.bend source
bend_tensors.bend on the hub · documented module
# bend_tensors: dense linear algebra with shapes in the types.## A Vec(n) holds 4 * 2^n floats and a Mat(r, c) 4 * 2^r rows of 4 * 2^c# floats, both as balanced trees of 4x4 register blocks. The tree is the# fork tree: every op below a node is a parallel let, and the matrix# product recurses on the tree, cache-oblivious, down to a 64-multiply# kernel. A shape mismatch is a checker error, not a runtime one.## Nothing here matches on a computed value, recurses mutually, or opens# an Array: those are the rules the compiler makes fast code from. Read# AGENTS.md before editing. This file is assembled from gen/ by# gen/assemble.py; edit the parts, not this file. bend_tensors.bend is# the release (CPU BLAS); bend_tensors_gpu.bend adds the GPU paths.import Base# Base's F32.max and F32.show take + parameters, which a template slot# does not accept; these take plain ones.def F32.max.of(a: F32, b: F32) -> F32: F32.max(a, b)def F32.show.of(x: F32) -> String: F32.show(x)# Blocks# ======## A Blk is a 4x4 block of F32 in row-major order, a Vec4 four F32. Both# are Data records of scalars: they ride in registers, and the trees# built from them are shared through + binders with no copy. A block# product is 64 fused multiply-adds with no memory traffic: the kernel# every larger op bottoms out in.type Blk is Data: Blk{a00: F32, a01: F32, a02: F32, a03: F32, a10: F32, a11: F32, a12: F32, a13: F32, a20: F32, a21: F32, a22: F32, a23: F32, a30: F32, a31: F32, a32: F32, a33: F32}type Vec4 is Data: Vec4{x0: F32, x1: F32, x2: F32, x3: F32}# Vec4# ----def Vec4.full(+x: F32) -> Vec4: Vec4{x, x, x, x}def Vec4.zeros() -> Vec4: Vec4{0.0, 0.0, 0.0, 0.0}# x, x + 1, x + 2, x + 3def Vec4.arange(+x: F32) -> Vec4: Vec4{x, (x + 1.0 : F32), (x + 2.0 : F32), (x + 3.0 : F32)}def Vec4.map(~f: F32 -> F32, v: Vec4) -> Vec4: Vec4{x0, x1, x2, x3} = v Vec4{f(x0), f(x1), f(x2), f(x3)}def Vec4.zip(~f: F32 -> F32 -> F32, a: Vec4, b: Vec4) -> Vec4: Vec4{x0, x1, x2, x3} = a Vec4{y0, y1, y2, y3} = b Vec4{f(x0, y0), f(x1, y1), f(x2, y2), f(x3, y3)}def Vec4.add(a: Vec4, b: Vec4) -> Vec4: Vec4{x0, x1, x2, x3} = a Vec4{y0, y1, y2, y3} = b Vec4{(x0 + y0 : F32), (x1 + y1 : F32), (x2 + y2 : F32), (x3 + y3 : F32)}def Vec4.sub(a: Vec4, b: Vec4) -> Vec4: Vec4{x0, x1, x2, x3} = a Vec4{y0, y1, y2, y3} = b Vec4{(x0 - y0 : F32), (x1 - y1 : F32), (x2 - y2 : F32), (x3 - y3 : F32)}def Vec4.mul(a: Vec4, b: Vec4) -> Vec4: Vec4{x0, x1, x2, x3} = a Vec4{y0, y1, y2, y3} = b Vec4{(x0 * y0 : F32), (x1 * y1 : F32), (x2 * y2 : F32), (x3 * y3 : F32)}def Vec4.scale(+k: F32, v: Vec4) -> Vec4: Vec4{x0, x1, x2, x3} = v Vec4{(k * x0 : F32), (k * x1 : F32), (k * x2 : F32), (k * x3 : F32)}# the sum, as (x0 + x1) + (x2 + x3): every reduction in this file# pairs its terms the same way, so laws can equate themdef Vec4.sum(v: Vec4) -> F32: Vec4{x0, x1, x2, x3} = v ((x0 + x1) + (x2 + x3) : F32)def Vec4.dot(a: Vec4, b: Vec4) -> F32: Vec4{x0, x1, x2, x3} = a Vec4{y0, y1, y2, y3} = b ((x0 * y0 + x1 * y1) + (x2 * y2 + x3 * y3) : F32)def Vec4.max(v: Vec4) -> F32: Vec4{x0, x1, x2, x3} = v F32.max(F32.max(x0, x1), F32.max(x2, x3))def Vec4.min(v: Vec4) -> F32: Vec4{x0, x1, x2, x3} = v F32.min(F32.min(x0, x1), F32.min(x2, x3))# the index (0-3) and value of the largest element; the first wins tiesdef Vec4.argmax.pick(c: Bool, +i: U32, +x: F32, +j: U32, +y: F32) -> U32 & F32: match c: case True{}: (i, x) case False{}: (j, y)def Vec4.argmax.two(ix: U32 & F32, jy: U32 & F32) -> U32 & F32: (i, x) = ix (j, y) = jy +x2 = x +y2 = y Vec4.argmax.pick((x2 >= y2 : F32), i, x2, j, y2)def Vec4.argmax(v: Vec4) -> U32 & F32: Vec4{x0, x1, x2, x3} = v Vec4.argmax.two(Vec4.argmax.two((0, x0), (1, x1)), Vec4.argmax.two((2, x2), (3, x3)))def Vec4.to_list(v: Vec4, acc: List<&2, F32>) -> List<&2, F32>: Vec4{x0, x1, x2, x3} = v x0 <> x1 <> x2 <> x3 <> acc# the first four of xs, padded with 0.0, and the restdef Vec4.from_list(xs: List<&2, F32>) -> Vec4: match xs: case x0 <> x1 <> x2 <> x3 <> rest: Vec4{x0, x1, x2, x3} case x0 <> x1 <> x2 <> []: Vec4{x0, x1, x2, 0.0} case x0 <> x1 <> []: Vec4{x0, x1, 0.0, 0.0} case x0 <> []: Vec4{x0, 0.0, 0.0, 0.0} case []: Vec4{0.0, 0.0, 0.0, 0.0}def Vec4.show(v: Vec4) -> String: Vec4{x0, x1, x2, x3} = v F32.show(x0) ++ " " ++ F32.show(x1) ++ " " ++ F32.show(x2) ++ " " ++ F32.show(x3)# Blk# ---def Blk.full(+x: F32) -> Blk: Blk{x, x, x, x, x, x, x, x, x, x, x, x, x, x, x, x}def Blk.zeros() -> Blk: Blk.full(0.0)def Blk.eye() -> Blk: Blk{1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0}def Blk.map(~f: F32 -> F32, m: Blk) -> Blk: Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m Blk{f(a00), f(a01), f(a02), f(a03), f(a10), f(a11), f(a12), f(a13), f(a20), f(a21), f(a22), f(a23), f(a30), f(a31), f(a32), f(a33)}def Blk.zip(~f: F32 -> F32 -> F32, m: Blk, n: Blk) -> Blk: Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m Blk{b00, b01, b02, b03, b10, b11, b12, b13, b20, b21, b22, b23, b30, b31, b32, b33} = n Blk{f(a00, b00), f(a01, b01), f(a02, b02), f(a03, b03), f(a10, b10), f(a11, b11), f(a12, b12), f(a13, b13), f(a20, b20), f(a21, b21), f(a22, b22), f(a23, b23), f(a30, b30), f(a31, b31), f(a32, b32), f(a33, b33)}def Blk.add(m: Blk, n: Blk) -> Blk: Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m Blk{b00, b01, b02, b03, b10, b11, b12, b13, b20, b21, b22, b23, b30, b31, b32, b33} = n Blk{(a00 + b00 : F32), (a01 + b01 : F32), (a02 + b02 : F32), (a03 + b03 : F32), (a10 + b10 : F32), (a11 + b11 : F32), (a12 + b12 : F32), (a13 + b13 : F32), (a20 + b20 : F32), (a21 + b21 : F32), (a22 + b22 : F32), (a23 + b23 : F32), (a30 + b30 : F32), (a31 + b31 : F32), (a32 + b32 : F32), (a33 + b33 : F32)}def Blk.sub(m: Blk, n: Blk) -> Blk: Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m Blk{b00, b01, b02, b03, b10, b11, b12, b13, b20, b21, b22, b23, b30, b31, b32, b33} = n Blk{(a00 - b00 : F32), (a01 - b01 : F32), (a02 - b02 : F32), (a03 - b03 : F32), (a10 - b10 : F32), (a11 - b11 : F32), (a12 - b12 : F32), (a13 - b13 : F32), (a20 - b20 : F32), (a21 - b21 : F32), (a22 - b22 : F32), (a23 - b23 : F32), (a30 - b30 : F32), (a31 - b31 : F32), (a32 - b32 : F32), (a33 - b33 : F32)}def Blk.mul(m: Blk, n: Blk) -> Blk: Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m Blk{b00, b01, b02, b03, b10, b11, b12, b13, b20, b21, b22, b23, b30, b31, b32, b33} = n Blk{(a00 * b00 : F32), (a01 * b01 : F32), (a02 * b02 : F32), (a03 * b03 : F32), (a10 * b10 : F32), (a11 * b11 : F32), (a12 * b12 : F32), (a13 * b13 : F32), (a20 * b20 : F32), (a21 * b21 : F32), (a22 * b22 : F32), (a23 * b23 : F32), (a30 * b30 : F32), (a31 * b31 : F32), (a32 * b32 : F32), (a33 * b33 : F32)}def Blk.scale(+k: F32, m: Blk) -> Blk: Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m Blk{(k * a00 : F32), (k * a01 : F32), (k * a02 : F32), (k * a03 : F32), (k * a10 : F32), (k * a11 : F32), (k * a12 : F32), (k * a13 : F32), (k * a20 : F32), (k * a21 : F32), (k * a22 : F32), (k * a23 : F32), (k * a30 : F32), (k * a31 : F32), (k * a32 : F32), (k * a33 : F32)}def Blk.transpose(m: Blk) -> Blk: Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m Blk{a00, a10, a20, a30, a01, a11, a21, a31, a02, a12, a22, a32, a03, a13, a23, a33}# m * n: 64 multiply-adds, each cell as (a0 b0 + a1 b1) + (a2 b2 + a3 b3)def Blk.matmul(m: Blk, n: Blk) -> Blk: Blk{+a00, +a01, +a02, +a03, +a10, +a11, +a12, +a13, +a20, +a21, +a22, +a23, +a30, +a31, +a32, +a33} = m Blk{+b00, +b01, +b02, +b03, +b10, +b11, +b12, +b13, +b20, +b21, +b22, +b23, +b30, +b31, +b32, +b33} = n Blk{((a00 * b00 + a01 * b10) + (a02 * b20 + a03 * b30) : F32), ((a00 * b01 + a01 * b11) + (a02 * b21 + a03 * b31) : F32), ((a00 * b02 + a01 * b12) + (a02 * b22 + a03 * b32) : F32), ((a00 * b03 + a01 * b13) + (a02 * b23 + a03 * b33) : F32), ((a10 * b00 + a11 * b10) + (a12 * b20 + a13 * b30) : F32), ((a10 * b01 + a11 * b11) + (a12 * b21 + a13 * b31) : F32), ((a10 * b02 + a11 * b12) + (a12 * b22 + a13 * b32) : F32), ((a10 * b03 + a11 * b13) + (a12 * b23 + a13 * b33) : F32), ((a20 * b00 + a21 * b10) + (a22 * b20 + a23 * b30) : F32), ((a20 * b01 + a21 * b11) + (a22 * b21 + a23 * b31) : F32), ((a20 * b02 + a21 * b12) + (a22 * b22 + a23 * b32) : F32), ((a20 * b03 + a21 * b13) + (a22 * b23 + a23 * b33) : F32), ((a30 * b00 + a31 * b10) + (a32 * b20 + a33 * b30) : F32), ((a30 * b01 + a31 * b11) + (a32 * b21 + a33 * b31) : F32), ((a30 * b02 + a31 * b12) + (a32 * b22 + a33 * b32) : F32), ((a30 * b03 + a31 * b13) + (a32 * b23 + a33 * b33) : F32)}# m * v: each row dotted with vdef Blk.matvec(m: Blk, v: Vec4) -> Vec4: Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m Vec4{+y0, +y1, +y2, +y3} = v Vec4{((a00 * y0 + a01 * y1) + (a02 * y2 + a03 * y3) : F32), ((a10 * y0 + a11 * y1) + (a12 * y2 + a13 * y3) : F32), ((a20 * y0 + a21 * y1) + (a22 * y2 + a23 * y3) : F32), ((a30 * y0 + a31 * y1) + (a32 * y2 + a33 * y3) : F32)}# row i's sum, as (a_i0 + a_i1) + (a_i2 + a_i3)def Blk.sum_rows(m: Blk) -> Vec4: Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m Vec4{((a00 + a01) + (a02 + a03) : F32), ((a10 + a11) + (a12 + a13) : F32), ((a20 + a21) + (a22 + a23) : F32), ((a30 + a31) + (a32 + a33) : F32)}# the sum of every cell, as the sum of the row sumsdef Blk.sum(m: Blk) -> F32: Vec4.sum(Blk.sum_rows(m))# column j's sum, as (a_0j + a_1j) + (a_2j + a_3j): the row sums of the transposedef Blk.sum_cols(m: Blk) -> Vec4: Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m Vec4{((a00 + a10) + (a20 + a30) : F32), ((a01 + a11) + (a21 + a31) : F32), ((a02 + a12) + (a22 + a32) : F32), ((a03 + a13) + (a23 + a33) : F32)}def Blk.max_rows(m: Blk) -> Vec4: Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m Vec4{F32.max(F32.max(a00, a01), F32.max(a02, a03)), F32.max(F32.max(a10, a11), F32.max(a12, a13)), F32.max(F32.max(a20, a21), F32.max(a22, a23)), F32.max(F32.max(a30, a31), F32.max(a32, a33))}# v added to every rowdef Blk.add_row(m: Blk, v: Vec4) -> Blk: Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m Vec4{+y0, +y1, +y2, +y3} = v Blk{(a00 + y0 : F32), (a01 + y1 : F32), (a02 + y2 : F32), (a03 + y3 : F32), (a10 + y0 : F32), (a11 + y1 : F32), (a12 + y2 : F32), (a13 + y3 : F32), (a20 + y0 : F32), (a21 + y1 : F32), (a22 + y2 : F32), (a23 + y3 : F32), (a30 + y0 : F32), (a31 + y1 : F32), (a32 + y2 : F32), (a33 + y3 : F32)}# v subtracted from every rowdef Blk.sub_row(m: Blk, v: Vec4) -> Blk: Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m Vec4{+y0, +y1, +y2, +y3} = v Blk{(a00 - y0 : F32), (a01 - y1 : F32), (a02 - y2 : F32), (a03 - y3 : F32), (a10 - y0 : F32), (a11 - y1 : F32), (a12 - y2 : F32), (a13 - y3 : F32), (a20 - y0 : F32), (a21 - y1 : F32), (a22 - y2 : F32), (a23 - y3 : F32), (a30 - y0 : F32), (a31 - y1 : F32), (a32 - y2 : F32), (a33 - y3 : F32)}# row i subtracted by v_i (the column vector v broadcast along rows)def Blk.sub_col(m: Blk, v: Vec4) -> Blk: Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m Vec4{+y0, +y1, +y2, +y3} = v Blk{(a00 - y0 : F32), (a01 - y0 : F32), (a02 - y0 : F32), (a03 - y0 : F32), (a10 - y1 : F32), (a11 - y1 : F32), (a12 - y1 : F32), (a13 - y1 : F32), (a20 - y2 : F32), (a21 - y2 : F32), (a22 - y2 : F32), (a23 - y2 : F32), (a30 - y3 : F32), (a31 - y3 : F32), (a32 - y3 : F32), (a33 - y3 : F32)}# row i scaled by v_idef Blk.scale_rows(m: Blk, v: Vec4) -> Blk: Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m Vec4{+y0, +y1, +y2, +y3} = v Blk{(y0 * a00 : F32), (y0 * a01 : F32), (y0 * a02 : F32), (y0 * a03 : F32), (y1 * a10 : F32), (y1 * a11 : F32), (y1 * a12 : F32), (y1 * a13 : F32), (y2 * a20 : F32), (y2 * a21 : F32), (y2 * a22 : F32), (y2 * a23 : F32), (y3 * a30 : F32), (y3 * a31 : F32), (y3 * a32 : F32), (y3 * a33 : F32)}# the 16 cells in row-major order, before accdef Blk.to_list(m: Blk, acc: List<&2, F32>) -> List<&2, F32>: Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m a00 <> a01 <> a02 <> a03 <> a10 <> a11 <> a12 <> a13 <> a20 <> a21 <> a22 <> a23 <> a30 <> a31 <> a32 <> a33 <> acc# a block from its four rowsdef Blk.from_vec4s(v0: Vec4, v1: Vec4, v2: Vec4, v3: Vec4) -> Blk: Vec4{a00, a01, a02, a03} = v0 Vec4{a10, a11, a12, a13} = v1 Vec4{a20, a21, a22, a23} = v2 Vec4{a30, a31, a32, a33} = v3 Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}# a block from four rows of four; a short row pads with 0.0def Blk.from_rows(r0: List<&2, F32>, r1: List<&2, F32>, r2: List<&2, F32>, r3: List<&2, F32>) -> Blk: Blk.from_vec4s(Vec4.from_list(r0), Vec4.from_list(r1), Vec4.from_list(r2), Vec4.from_list(r3))def Blk.row(m: Blk, i: Nat) -> Vec4: Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = m match i: case 0n: Vec4{a00, a01, a02, a03} case 1n: Vec4{a10, a11, a12, a13} case 2n: Vec4{a20, a21, a22, a23} case _: Vec4{a30, a31, a32, a33}# Masks# -----## k counts the valid elements (or rows, or columns) of a block; the# rest are padding, set to 0.0 or left out of a reduction. A k of 0# only reaches a reduction whose result the caller discards.# elements at index >= k set to 0.0def Vec4.mask(k: Nat, v: Vec4) -> Vec4: match k v: case 0n Vec4{x0, x1, x2, x3}: Vec4{0.0, 0.0, 0.0, 0.0} case 1n Vec4{x0, x1, x2, x3}: Vec4{x0, 0.0, 0.0, 0.0} case 2n Vec4{x0, x1, x2, x3}: Vec4{x0, x1, 0.0, 0.0} case 3n Vec4{x0, x1, x2, x3}: Vec4{x0, x1, x2, 0.0} case _ v2: v2# the largest of the first k elementsdef Vec4.max_upto(k: Nat, v: Vec4) -> F32: match k v: case 1n Vec4{x0, x1, x2, x3}: x0 case 2n Vec4{x0, x1, x2, x3}: F32.max(x0, x1) case 3n Vec4{x0, x1, x2, x3}: F32.max(F32.max(x0, x1), x2) case _ v2: Vec4.max(v2)# the smallest of the first k elementsdef Vec4.min_upto(k: Nat, v: Vec4) -> F32: match k v: case 1n Vec4{x0, x1, x2, x3}: x0 case 2n Vec4{x0, x1, x2, x3}: F32.min(x0, x1) case 3n Vec4{x0, x1, x2, x3}: F32.min(F32.min(x0, x1), x2) case _ v2: Vec4.min(v2)# the index and value of the largest of the first k elementsdef Vec4.argmax_upto(k: Nat, v: Vec4) -> U32 & F32: match k v: case 1n Vec4{x0, x1, x2, x3}: (0, x0) case 2n Vec4{x0, x1, x2, x3}: Vec4.argmax.two((0, x0), (1, x1)) case 3n Vec4{x0, x1, x2, x3}: Vec4.argmax.two(Vec4.argmax.two((0, x0), (1, x1)), (2, x2)) case _ v2: Vec4.argmax(v2)# rows at index >= k set to 0.0def Blk.mask_rows(k: Nat, m: Blk) -> Blk: match k m: case 0n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}: Blk{0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0} case 1n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}: Blk{a00, a01, a02, a03, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0} case 2n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}: Blk{a00, a01, a02, a03, a10, a11, a12, a13, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0} case 3n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}: Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, 0.0, 0.0, 0.0, 0.0} case _ m2: m2# columns at index >= k set to 0.0def Blk.mask_cols(k: Nat, m: Blk) -> Blk: match k m: case 0n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}: Blk{0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0} case 1n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}: Blk{a00, 0.0, 0.0, 0.0, a10, 0.0, 0.0, 0.0, a20, 0.0, 0.0, 0.0, a30, 0.0, 0.0, 0.0} case 2n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}: Blk{a00, a01, 0.0, 0.0, a10, a11, 0.0, 0.0, a20, a21, 0.0, 0.0, a30, a31, 0.0, 0.0} case 3n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}: Blk{a00, a01, a02, 0.0, a10, a11, a12, 0.0, a20, a21, a22, 0.0, a30, a31, a32, 0.0} case _ m2: m2# rows at index >= ki and columns at index >= kj set to 0.0def Blk.mask(ki: Nat, kj: Nat, m: Blk) -> Blk: Blk.mask_cols(kj, Blk.mask_rows(ki, m))# each row's largest among its first k columnsdef Blk.max_rows_upto(k: Nat, m: Blk) -> Vec4: match k m: case 1n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}: Vec4{a00, a10, a20, a30} case 2n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}: Vec4{F32.max(a00, a01), F32.max(a10, a11), F32.max(a20, a21), F32.max(a30, a31)} case 3n Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33}: Vec4{F32.max(F32.max(a00, a01), a02), F32.max(F32.max(a10, a11), a12), F32.max(F32.max(a20, a21), a22), F32.max(F32.max(a30, a31), a32)} case _ m2: Blk.max_rows(m2)# Vectors# =======## VecT(n) is a balanced binary tree of depth n over Vec4 leaves: 4 * 2^n# floats. The depth is a live Nat, so every op matches on it first and# then opens the constructor the depth implies, as Base does for Word.# Every recursion below a node is a parallel let: the tree is the fork# tree, balanced by construction.law VecT: for n: Nat Datatype VecT.Leaf is Data: VecL{v: Vec4}type VecT.Node<-p: Nat> is Data: VecN{l: VecT(p), r: VecT(p)}def VecT(n): match n: case 0n: VecT.Leaf case 1n+p: VecT.Node<p># the element count of a VecT(n): 4 * 2^ndef VecT.size(+n: Nat) -> Nat: match n: case 0n: 4n case 1n+p: +h = VecT.size(p) Nat.add(h, h)def VecT.full(+n: Nat, +x: F32) -> VecT(n): match n: case 0n: VecL{Vec4.full(x)} case 1n+p: l r = VecT.full(p, x) VecT.full(p, x) VecN{l, r}def VecT.zeros(+n: Nat) -> VecT(n): VecT.full(n, 0.0)def VecT.ones(+n: Nat) -> VecT(n): VecT.full(n, 1.0)# x, x + 1, x + 2, ..def VecT.arange.from(+n: Nat, +x: F32) -> VecT(n): match n: case 0n: VecL{Vec4.arange(x)} case 1n+p: +h = F32.from_nat(VecT.size(p)) l r = VecT.arange.from(p, x) VecT.arange.from(p, (x + h : F32)) VecN{l, r}# 0, 1, 2, ..def VecT.arange(+n: Nat) -> VecT(n): VecT.arange.from(n, 0.0)def VecT.map(~f: F32 -> F32, +n: Nat, v: VecT(n)) -> VecT(n): match n: case 0n: VecL{x} = v VecL{Vec4.map(~f, x)} case 1n+p: VecN{l, r} = v a b = VecT.map(~f, p, l) VecT.map(~f, p, r) VecN{a, b}def VecT.zip(~f: F32 -> F32 -> F32, +n: Nat, a: VecT(n), b: VecT(n)) -> VecT(n): match n: case 0n: VecL{x} = a VecL{y} = b VecL{Vec4.zip(~f, x, y)} case 1n+p: VecN{al, ar} = a VecN{bl, br} = b l r = VecT.zip(~f, p, al, bl) VecT.zip(~f, p, ar, br) VecN{l, r}def VecT.add(+n: Nat, a: VecT(n), b: VecT(n)) -> VecT(n): VecT.zip(~F32.add, n, a, b)def VecT.sub(+n: Nat, a: VecT(n), b: VecT(n)) -> VecT(n): VecT.zip(~F32.sub, n, a, b)# elementwise productdef VecT.mul(+n: Nat, a: VecT(n), b: VecT(n)) -> VecT(n): VecT.zip(~F32.mul, n, a, b)def VecT.exp(+n: Nat, v: VecT(n)) -> VecT(n): VecT.map(~F32.exp, n, v)def VecT.scale(+n: Nat, +k: F32, v: VecT(n)) -> VecT(n): match n: case 0n: VecL{x} = v VecL{Vec4.scale(k, x)} case 1n+p: VecN{l, r} = v a b = VecT.scale(p, k, l) VecT.scale(p, k, r) VecN{a, b}# k added to every elementdef VecT.shift(+n: Nat, +k: F32, v: VecT(n)) -> VecT(n): match n: case 0n: VecL{x} = v Vec4{x0, x1, x2, x3} = x VecL{Vec4{(x0 + k : F32), (x1 + k : F32), (x2 + k : F32), (x3 + k : F32)}} case 1n+p: VecN{l, r} = v a b = VecT.shift(p, k, l) VecT.shift(p, k, r) VecN{a, b}# the sum, left half plus right halfdef VecT.sum(+n: Nat, v: VecT(n)) -> F32: match n: case 0n: VecL{x} = v Vec4.sum(x) case 1n+p: VecN{l, r} = v a b = VecT.sum(p, l) VecT.sum(p, r) (a + b : F32)def VecT.dot(+n: Nat, a: VecT(n), b: VecT(n)) -> F32: match n: case 0n: VecL{x} = a VecL{y} = b Vec4.dot(x, y) case 1n+p: VecN{al, ar} = a VecN{bl, br} = b l r = VecT.dot(p, al, bl) VecT.dot(p, ar, br) (l + r : F32)def VecT.max(+n: Nat, v: VecT(n)) -> F32: match n: case 0n: VecL{x} = v Vec4.max(x) case 1n+p: VecN{l, r} = v a b = VecT.max(p, l) VecT.max(p, r) F32.max(a, b)def VecT.min(+n: Nat, v: VecT(n)) -> F32: match n: case 0n: VecL{x} = v Vec4.min(x) case 1n+p: VecN{l, r} = v a b = VecT.min(p, l) VecT.min(p, r) F32.min(a, b)# the two halves' winners, the right one's index offset by the half's sizedef VecT.argmax.two(+off: U32, ix: U32 & F32, jy: U32 & F32) -> U32 & F32: (i, x) = ix (j, y) = jy Vec4.argmax.two((i, x), ((j + off : U32), y))# the index and value of the largest element; the first wins tiesdef VecT.argmax(+n: Nat, v: VecT(n)) -> U32 & F32: match n: case 0n: VecL{x} = v Vec4.argmax(x) case 1n+p: VecN{l, r} = v a b = VecT.argmax(p, l) VecT.argmax(p, r) VecT.argmax.two(U32.from_nat(VecT.size(p)), a, b)# the elements in order, before acc# the tree is read borrowed (+v, matched in place), never consumed: a# reader that consumes a shared tree makes the compiler copy it, and# bend 2.0.16 miscompiles that copy when the depth is a literal (the# bias of a model: memory faults in the loop; issues/shared_vec_loop_*)def VecT.to_list.go(+n: Nat, +v: VecT(n), acc: List<&2, F32>) -> List<&2, F32>: match n v: case 0n VecL{x}: Vec4.to_list(x, acc) case 1n+p VecN{l, r}: VecT.to_list.go(p, l, VecT.to_list.go(p, r, acc))def VecT.to_list(+n: Nat, +v: VecT(n)) -> List<&2, F32>: VecT.to_list.go(n, v, [])# the first 4 * 2^n elements of xs, padded with 0.0def VecT.from_list(+n: Nat, +xs: List<&2, F32>) -> VecT(n): match n: case 0n: VecL{Vec4.from_list(xs)} case 1n+p: +h = VecT.size(p) l r = VecT.from_list(p, List.take(&2, F32, xs, h)) VecT.from_list(p, List.drop(&2, F32, xs, h)) VecN{l, r}def VecT.softmax.fin(+n: Nat, +e: VecT(n)) -> VecT(n): VecT.scale(n, F32.div(1.0, VecT.sum(n, e)), e)# exp(v - max v) / sumdef VecT.softmax(+n: Nat, +v: VecT(n)) -> VecT(n): VecT.softmax.fin(n, VecT.exp(n, VecT.shift(n, (0.0 - VecT.max(n, v) : F32), v)))def VecT.show(+n: Nat, v: VecT(n)) -> String: List.show(~&2, ~F32, ~F32.show.of, VecT.to_list(n, v))# Matrices# ========## MatT(r, c) has 4 * 2^r rows and 4 * 2^c columns of F32, as a tree of Blk# leaves. While both depths are positive a node is a quadrant split# (MatQ); when only rows remain it splits top and bottom (MatR); when# only columns remain, left and right (MatC). The block product recurses# on that tree, so it is cache-oblivious: every level halves the working# set, down to a 4x4 kernel in registers. Operands are + and shared, never# copied. The shape is the type: MatT.matmul only accepts MatT(r, k) and# MatT(k, c), and answers MatT(r, c).law MatT: for r: Nat for c: Nat Datatype MatT.Leaf is Data: MatB{b: Blk}type MatT.Rows<-p: Nat, -c: Nat> is Data: MatR{top: MatT(p, c), bot: MatT(p, c)}type MatT.Cols<-r: Nat, -q: Nat> is Data: MatC{left: MatT(r, q), right: MatT(r, q)}type MatT.Quad<-p: Nat, -q: Nat> is Data: MatQ{tl: MatT(p, q), tr: MatT(p, q), bl: MatT(p, q), br: MatT(p, q)}def MatT(r, c): match r c: case 0n 0n: MatT.Leaf case 1n+p 0n: MatT.Rows<p, 0n> case 0n 1n+q: MatT.Cols<0n, q> case 1n+p 1n+q: MatT.Quad<p, q>def MatT.full(+r: Nat, +c: Nat, +x: F32) -> MatT(r, c): match r c: case 0n 0n: MatB{Blk.full(x)} case 1n+p 0n: t b = MatT.full(p, 0n, x) MatT.full(p, 0n, x) MatR{t, b} case 0n 1n+q: l r2 = MatT.full(0n, q, x) MatT.full(0n, q, x) MatC{l, r2} case 1n+p 1n+q: a b c2 d = MatT.full(p, q, x) MatT.full(p, q, x) MatT.full(p, q, x) MatT.full(p, q, x) MatQ{a, b, c2, d}def MatT.zeros(+r: Nat, +c: Nat) -> MatT(r, c): MatT.full(r, c, 0.0)def MatT.ones(+r: Nat, +c: Nat) -> MatT(r, c): MatT.full(r, c, 1.0)# the identity of MatT(n, n)def MatT.eye(+n: Nat) -> MatT(n, n): match n: case 0n: MatB{Blk.eye()} case 1n+p: a b c d = MatT.eye(p) MatT.zeros(p, p) MatT.zeros(p, p) MatT.eye(p) MatQ{a, b, c, d}def MatT.map(~f: F32 -> F32, +r: Nat, +c: Nat, m: MatT(r, c)) -> MatT(r, c): match r c: case 0n 0n: MatB{b} = m MatB{Blk.map(~f, b)} case 1n+p 0n: MatR{t, b} = m x y = MatT.map(~f, p, 0n, t) MatT.map(~f, p, 0n, b) MatR{x, y} case 0n 1n+q: MatC{l, r2} = m x y = MatT.map(~f, 0n, q, l) MatT.map(~f, 0n, q, r2) MatC{x, y} case 1n+p 1n+q: MatQ{a, b, c2, d} = m w x y z = MatT.map(~f, p, q, a) MatT.map(~f, p, q, b) MatT.map(~f, p, q, c2) MatT.map(~f, p, q, d) MatQ{w, x, y, z}def MatT.zip(~f: F32 -> F32 -> F32, +r: Nat, +c: Nat, m: MatT(r, c), n: MatT(r, c)) -> MatT(r, c): match r c: case 0n 0n: MatB{x} = m MatB{y} = n MatB{Blk.zip(~f, x, y)} case 1n+p 0n: MatR{mt, mb} = m MatR{nt, nb} = n x y = MatT.zip(~f, p, 0n, mt, nt) MatT.zip(~f, p, 0n, mb, nb) MatR{x, y} case 0n 1n+q: MatC{ml, mr} = m MatC{nl, nr} = n x y = MatT.zip(~f, 0n, q, ml, nl) MatT.zip(~f, 0n, q, mr, nr) MatC{x, y} case 1n+p 1n+q: MatQ{ma, mb, mc, md} = m MatQ{na, nb, nc, nd} = n w x y z = MatT.zip(~f, p, q, ma, na) MatT.zip(~f, p, q, mb, nb) MatT.zip(~f, p, q, mc, nc) MatT.zip(~f, p, q, md, nd) MatQ{w, x, y, z}def MatT.add(+r: Nat, +c: Nat, m: MatT(r, c), n: MatT(r, c)) -> MatT(r, c): MatT.zip(~F32.add, r, c, m, n)def MatT.sub(+r: Nat, +c: Nat, m: MatT(r, c), n: MatT(r, c)) -> MatT(r, c): MatT.zip(~F32.sub, r, c, m, n)# elementwise (Hadamard) product; the matrix product is MatT.matmuldef MatT.mul(+r: Nat, +c: Nat, m: MatT(r, c), n: MatT(r, c)) -> MatT(r, c): MatT.zip(~F32.mul, r, c, m, n)def MatT.exp(+r: Nat, +c: Nat, m: MatT(r, c)) -> MatT(r, c): MatT.map(~F32.exp, r, c, m)def MatT.scale(+r: Nat, +c: Nat, +k: F32, m: MatT(r, c)) -> MatT(r, c): match r c: case 0n 0n: MatB{b} = m MatB{Blk.scale(k, b)} case 1n+p 0n: MatR{t, b} = m x y = MatT.scale(p, 0n, k, t) MatT.scale(p, 0n, k, b) MatR{x, y} case 0n 1n+q: MatC{l, r2} = m x y = MatT.scale(0n, q, k, l) MatT.scale(0n, q, k, r2) MatC{x, y} case 1n+p 1n+q: MatQ{a, b, c2, d} = m w x y z = MatT.scale(p, q, k, a) MatT.scale(p, q, k, b) MatT.scale(p, q, k, c2) MatT.scale(p, q, k, d) MatQ{w, x, y, z}def MatT.transpose(+r: Nat, +c: Nat, m: MatT(r, c)) -> MatT(c, r): match r c: case 0n 0n: MatB{b} = m MatB{Blk.transpose(b)} case 1n+p 0n: MatR{t, b} = m x y = MatT.transpose(p, 0n, t) MatT.transpose(p, 0n, b) MatC{x, y} case 0n 1n+q: MatC{l, r2} = m x y = MatT.transpose(0n, q, l) MatT.transpose(0n, q, r2) MatR{x, y} case 1n+p 1n+q: MatQ{a, b, c2, d} = m w x y z = MatT.transpose(p, q, a) MatT.transpose(p, q, c2) MatT.transpose(p, q, b) MatT.transpose(p, q, d) MatQ{w, x, y, z}# the four quadrant sums of a product node, in one balanced forkdef MatT.add4(+p: Nat, +q: Nat, a0: MatT(p, q), b0: MatT(p, q), a1: MatT(p, q), b1: MatT(p, q), a2: MatT(p, q), b2: MatT(p, q), a3: MatT(p, q), b3: MatT(p, q)) -> MatT(1n+p, 1n+q): w x y z = MatT.add(p, q, a0, b0) MatT.add(p, q, a1, b1) MatT.add(p, q, a2, b2) MatT.add(p, q, a3, b3) MatQ{w, x, y, z}def MatT.add2r(+p: Nat, a0: MatT(p, 0n), b0: MatT(p, 0n), a1: MatT(p, 0n), b1: MatT(p, 0n)) -> MatT(1n+p, 0n): x y = MatT.add(p, 0n, a0, b0) MatT.add(p, 0n, a1, b1) MatR{x, y}def MatT.add2c(+q: Nat, a0: MatT(0n, q), b0: MatT(0n, q), a1: MatT(0n, q), b1: MatT(0n, q)) -> MatT(0n, 1n+q): x y = MatT.add(0n, q, a0, b0) MatT.add(0n, q, a1, b1) MatC{x, y}# a * b for square operands of one depth: the same recursion with one# depth riding instead of three, which measures 3x faster per nodedef MatT.matmul.sq(+d: Nat, +a: MatT(d, d), +b: MatT(d, d)) -> MatT(d, d): match d: case 0n: MatB{x} = a MatB{y} = b MatB{Blk.matmul(x, y)} case 1n+p: MatQ{a0, a1, a2, a3} = a MatQ{b0, b1, b2, b3} = b p0 q0 p1 q1 p2 q2 p3 q3 = MatT.matmul.sq(p, a0, b0) MatT.matmul.sq(p, a1, b2) MatT.matmul.sq(p, a0, b1) MatT.matmul.sq(p, a1, b3) MatT.matmul.sq(p, a2, b0) MatT.matmul.sq(p, a3, b2) MatT.matmul.sq(p, a2, b1) MatT.matmul.sq(p, a3, b3) MatT.add4(p, p, p0, q0, p1, q1, p2, q2, p3, q3)# a * b for any shapes: MatT(r, k) times MatT(k, c) is MatT(r, c). A node# forks its whole burst of products at once (eight at a quadrant), then# joins them with a fork of adds. Both operands are shared through +,# never copied. Three depths ride every call, which is what the square# path below avoids.def MatT.matmul.any(+r: Nat, +k: Nat, +c: Nat, +a: MatT(r, k), +b: MatT(k, c)) -> MatT(r, c): match r k c: case 0n 0n 0n: # the leaf goes through the square path on purpose: compiled # without MatT.matmul.sq reachable, this recursion fails at runtime # from depth 2 on (a bend 2.0.16 bug, see README); with it, it runs MatT.matmul.sq(0n, a, b) case 1n+p 0n 0n: MatR{at, ab} = a x y = MatT.matmul.any(p, 0n, 0n, at, b) MatT.matmul.any(p, 0n, 0n, ab, b) MatR{x, y} case 0n 0n 1n+s: MatC{bl, br} = b x y = MatT.matmul.any(0n, 0n, s, a, bl) MatT.matmul.any(0n, 0n, s, a, br) MatC{x, y} case 0n 1n+q 0n: MatC{al, ar} = a MatR{bt, bb} = b x y = MatT.matmul.any(0n, q, 0n, al, bt) MatT.matmul.any(0n, q, 0n, ar, bb) MatT.add(0n, 0n, x, y) case 1n+p 1n+q 0n: MatQ{a0, a1, a2, a3} = a MatR{bt, bb} = b p0 q0 p1 q1 = MatT.matmul.any(p, q, 0n, a0, bt) MatT.matmul.any(p, q, 0n, a1, bb) MatT.matmul.any(p, q, 0n, a2, bt) MatT.matmul.any(p, q, 0n, a3, bb) MatT.add2r(p, p0, q0, p1, q1) case 0n 1n+q 1n+s: MatC{al, ar} = a MatQ{b0, b1, b2, b3} = b p0 q0 p1 q1 = MatT.matmul.any(0n, q, s, al, b0) MatT.matmul.any(0n, q, s, ar, b2) MatT.matmul.any(0n, q, s, al, b1) MatT.matmul.any(0n, q, s, ar, b3) MatT.add2c(s, p0, q0, p1, q1) case 1n+p 0n 1n+s: MatR{at, ab} = a MatC{bl, br} = b w x y z = MatT.matmul.any(p, 0n, s, at, bl) MatT.matmul.any(p, 0n, s, at, br) MatT.matmul.any(p, 0n, s, ab, bl) MatT.matmul.any(p, 0n, s, ab, br) MatQ{w, x, y, z} case 1n+p 1n+q 1n+s: MatQ{a0, a1, a2, a3} = a MatQ{b0, b1, b2, b3} = b p0 q0 p1 q1 p2 q2 p3 q3 = MatT.matmul.any(p, q, s, a0, b0) MatT.matmul.any(p, q, s, a1, b2) MatT.matmul.any(p, q, s, a0, b1) MatT.matmul.any(p, q, s, a1, b3) MatT.matmul.any(p, q, s, a2, b0) MatT.matmul.any(p, q, s, a3, b2) MatT.matmul.any(p, q, s, a2, b1) MatT.matmul.any(p, q, s, a3, b3) MatT.add4(p, s, p0, q0, p1, q1, p2, q2, p3, q3)# a * b: MatT(r, k) times MatT(k, c) is MatT(r, c). For two square# operands of one depth, MatT.matmul.sq is the same recursion with one# depth riding instead of three, and 3x faster per node; a runtime# dispatch to it is not possible without a cast, and any def that returns# or rebuilds a MatT parameter makes the compiler reference-count every# MatT node in the program (see AGENTS.md), so the choice is the caller's.def MatT.matmul(+r: Nat, +k: Nat, +c: Nat, +a: MatT(r, k), +b: MatT(k, c)) -> MatT(r, c): MatT.matmul.any(r, k, c, a, b)def VecT.add2(+p: Nat, a0: VecT(p), b0: VecT(p), a1: VecT(p), b1: VecT(p)) -> VecT(1n+p): x y = VecT.add(p, a0, b0) VecT.add(p, a1, b1) VecN{x, y}# m * v: MatT(r, c) times VecT(c) is VecT(r)def MatT.matvec(+r: Nat, +c: Nat, +m: MatT(r, c), +v: VecT(c)) -> VecT(r): match r c: case 0n 0n: MatB{b} = m VecL{x} = v VecL{Blk.matvec(b, x)} case 1n+p 0n: MatR{t, b} = m x y = MatT.matvec(p, 0n, t, v) MatT.matvec(p, 0n, b, v) VecN{x, y} case 0n 1n+q: MatC{l, r2} = m VecN{vl, vr} = v x y = MatT.matvec(0n, q, l, vl) MatT.matvec(0n, q, r2, vr) VecT.add(0n, x, y) case 1n+p 1n+q: MatQ{a, b, c2, d} = m VecN{vl, vr} = v p0 q0 p1 q1 = MatT.matvec(p, q, a, vl) MatT.matvec(p, q, b, vr) MatT.matvec(p, q, c2, vl) MatT.matvec(p, q, d, vr) VecT.add2(p, p0, q0, p1, q1)# the sum of every celldef MatT.sum(+r: Nat, +c: Nat, m: MatT(r, c)) -> F32: match r c: case 0n 0n: MatB{b} = m Blk.sum(b) case 1n+p 0n: MatR{t, b} = m x y = MatT.sum(p, 0n, t) MatT.sum(p, 0n, b) (x + y : F32) case 0n 1n+q: MatC{l, r2} = m x y = MatT.sum(0n, q, l) MatT.sum(0n, q, r2) (x + y : F32) case 1n+p 1n+q: MatQ{a, b, c2, d} = m w x y z = MatT.sum(p, q, a) MatT.sum(p, q, b) MatT.sum(p, q, c2) MatT.sum(p, q, d) ((w + x) + (y + z) : F32)# each row's sum: VecT(r)def MatT.sum_rows(+r: Nat, +c: Nat, m: MatT(r, c)) -> VecT(r): match r c: case 0n 0n: MatB{b} = m VecL{Blk.sum_rows(b)} case 1n+p 0n: MatR{t, b} = m x y = MatT.sum_rows(p, 0n, t) MatT.sum_rows(p, 0n, b) VecN{x, y} case 0n 1n+q: MatC{l, r2} = m x y = MatT.sum_rows(0n, q, l) MatT.sum_rows(0n, q, r2) VecT.add(0n, x, y) case 1n+p 1n+q: MatQ{a, b, c2, d} = m p0 q0 p1 q1 = MatT.sum_rows(p, q, a) MatT.sum_rows(p, q, b) MatT.sum_rows(p, q, c2) MatT.sum_rows(p, q, d) VecT.add2(p, p0, q0, p1, q1)# each column's sum: VecT(c)def MatT.sum_cols(+r: Nat, +c: Nat, m: MatT(r, c)) -> VecT(c): match r c: case 0n 0n: MatB{b} = m VecL{Blk.sum_cols(b)} case 1n+p 0n: MatR{t, b} = m x y = MatT.sum_cols(p, 0n, t) MatT.sum_cols(p, 0n, b) VecT.add(0n, x, y) case 0n 1n+q: MatC{l, r2} = m x y = MatT.sum_cols(0n, q, l) MatT.sum_cols(0n, q, r2) VecN{x, y} case 1n+p 1n+q: MatQ{a, b, c2, d} = m p0 q0 p1 q1 = MatT.sum_cols(p, q, a) MatT.sum_cols(p, q, c2) MatT.sum_cols(p, q, b) MatT.sum_cols(p, q, d) VecT.add2(q, p0, q0, p1, q1)def VecT.max2(+p: Nat, a0: VecT(p), b0: VecT(p), a1: VecT(p), b1: VecT(p)) -> VecT(1n+p): x y = VecT.zip(~F32.max.of, p, a0, b0) VecT.zip(~F32.max.of, p, a1, b1) VecN{x, y}# each row's largest element: VecT(r)def MatT.max_rows(+r: Nat, +c: Nat, m: MatT(r, c)) -> VecT(r): match r c: case 0n 0n: MatB{b} = m VecL{Blk.max_rows(b)} case 1n+p 0n: MatR{t, b} = m x y = MatT.max_rows(p, 0n, t) MatT.max_rows(p, 0n, b) VecN{x, y} case 0n 1n+q: MatC{l, r2} = m x y = MatT.max_rows(0n, q, l) MatT.max_rows(0n, q, r2) VecT.zip(~F32.max.of, 0n, x, y) case 1n+p 1n+q: MatQ{a, b, c2, d} = m p0 q0 p1 q1 = MatT.max_rows(p, q, a) MatT.max_rows(p, q, b) MatT.max_rows(p, q, c2) MatT.max_rows(p, q, d) VecT.max2(p, p0, q0, p1, q1)# v added to every row: v is a VecT(c)def MatT.add_row(+r: Nat, +c: Nat, m: MatT(r, c), +v: VecT(c)) -> MatT(r, c): match r c: case 0n 0n: MatB{b} = m VecL{x} = v MatB{Blk.add_row(b, x)} case 1n+p 0n: MatR{t, b} = m x y = MatT.add_row(p, 0n, t, v) MatT.add_row(p, 0n, b, v) MatR{x, y} case 0n 1n+q: MatC{l, r2} = m VecN{vl, vr} = v x y = MatT.add_row(0n, q, l, vl) MatT.add_row(0n, q, r2, vr) MatC{x, y} case 1n+p 1n+q: MatQ{a, b, c2, d} = m VecN{vl, vr} = v w x y z = MatT.add_row(p, q, a, vl) MatT.add_row(p, q, b, vr) MatT.add_row(p, q, c2, vl) MatT.add_row(p, q, d, vr) MatQ{w, x, y, z}# row i minus v_i: v is a VecT(r), one scalar per rowdef MatT.sub_col(+r: Nat, +c: Nat, m: MatT(r, c), +v: VecT(r)) -> MatT(r, c): match r c: case 0n 0n: MatB{b} = m VecL{x} = v MatB{Blk.sub_col(b, x)} case 1n+p 0n: MatR{t, b} = m VecN{vt, vb} = v x y = MatT.sub_col(p, 0n, t, vt) MatT.sub_col(p, 0n, b, vb) MatR{x, y} case 0n 1n+q: MatC{l, r2} = m x y = MatT.sub_col(0n, q, l, v) MatT.sub_col(0n, q, r2, v) MatC{x, y} case 1n+p 1n+q: MatQ{a, b, c2, d} = m VecN{vt, vb} = v w x y z = MatT.sub_col(p, q, a, vt) MatT.sub_col(p, q, b, vt) MatT.sub_col(p, q, c2, vb) MatT.sub_col(p, q, d, vb) MatQ{w, x, y, z}# row i times v_i: v is a VecT(r), one scalar per rowdef MatT.scale_rows(+r: Nat, +c: Nat, m: MatT(r, c), +v: VecT(r)) -> MatT(r, c): match r c: case 0n 0n: MatB{b} = m VecL{x} = v MatB{Blk.scale_rows(b, x)} case 1n+p 0n: MatR{t, b} = m VecN{vt, vb} = v x y = MatT.scale_rows(p, 0n, t, vt) MatT.scale_rows(p, 0n, b, vb) MatR{x, y} case 0n 1n+q: MatC{l, r2} = m x y = MatT.scale_rows(0n, q, l, v) MatT.scale_rows(0n, q, r2, v) MatC{x, y} case 1n+p 1n+q: MatQ{a, b, c2, d} = m VecN{vt, vb} = v w x y z = MatT.scale_rows(p, q, a, vt) MatT.scale_rows(p, q, b, vt) MatT.scale_rows(p, q, c2, vb) MatT.scale_rows(p, q, d, vb) MatQ{w, x, y, z}def MatT.softmax_rows.fin(+r: Nat, +c: Nat, +e: MatT(r, c)) -> MatT(r, c): MatT.scale_rows(r, c, e, VecT.map(~(s => F32.div(1.0, s)), r, MatT.sum_rows(r, c, e)))# softmax along each row: exp(row - max row) / sumdef MatT.softmax_rows(+r: Nat, +c: Nat, +m: MatT(r, c)) -> MatT(r, c): MatT.softmax_rows.fin(r, c, MatT.exp(r, c, MatT.sub_col(r, c, m, MatT.max_rows(r, c, m))))# Conversions# -----------# the first w of every rowdef MatT.rows_take(+w: Nat, rows: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>: match rows: case []: [] case row <> rest: List.take(&2, F32, row, w) <> MatT.rows_take(w, rest)# every row past its first wdef MatT.rows_drop(+w: Nat, rows: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>: match rows: case []: [] case row <> rest: List.drop(&2, F32, row, w) <> MatT.rows_drop(w, rest)# the rows of xs side by side with the rows of ysdef MatT.rows_cat(xs: List<&2, List<&2, F32>>, ys: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>: match xs ys: case x <> xt y <> yt: List.append(&2, F32, x, y) <> MatT.rows_cat(xt, yt) case _ _: []# a leaf from up to four rows; missing rows and cells pad with 0.0def MatT.from_rows.leaf(rows: List<&2, List<&2, F32>>) -> Blk: match rows: case r0 <> r1 <> r2 <> r3 <> rest: Blk.from_rows(r0, r1, r2, r3) case r0 <> r1 <> r2 <> []: Blk.from_rows(r0, r1, r2, []) case r0 <> r1 <> []: Blk.from_rows(r0, r1, [], []) case r0 <> []: Blk.from_rows(r0, [], [], []) case []: Blk.zeros()# a MatT(r, c) from its rows, row-major; short or missing rows pad with 0.0def MatT.from_rows(+r: Nat, +c: Nat, +rows: List<&2, List<&2, F32>>) -> MatT(r, c): match r c: case 0n 0n: MatB{MatT.from_rows.leaf(rows)} case 1n+p 0n: +h = VecT.size(p) t b = MatT.from_rows(p, 0n, List.take(&2, List<&2, F32>, rows, h)) MatT.from_rows(p, 0n, List.drop(&2, List<&2, F32>, rows, h)) MatR{t, b} case 0n 1n+q: +w = VecT.size(q) l r2 = MatT.from_rows(0n, q, MatT.rows_take(w, rows)) MatT.from_rows(0n, q, MatT.rows_drop(w, rows)) MatC{l, r2} case 1n+p 1n+q: +h = VecT.size(p) +w = VecT.size(q) +top = List.take(&2, List<&2, F32>, rows, h) +bot = List.drop(&2, List<&2, F32>, rows, h) a b c2 d = MatT.from_rows(p, q, MatT.rows_take(w, top)) MatT.from_rows(p, q, MatT.rows_drop(w, top)) MatT.from_rows(p, q, MatT.rows_take(w, bot)) MatT.from_rows(p, q, MatT.rows_drop(w, bot)) MatQ{a, b, c2, d}# the rows of a MatT(r, c), row-majordef MatT.to_rows(+r: Nat, +c: Nat, m: MatT(r, c)) -> List<&2, List<&2, F32>>: match r c: case 0n 0n: MatB{b} = m Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = b [[a00, a01, a02, a03], [a10, a11, a12, a13], [a20, a21, a22, a23], [a30, a31, a32, a33]] case 1n+p 0n: MatR{t, b} = m List.append(&2, List<&2, F32>, MatT.to_rows(p, 0n, t), MatT.to_rows(p, 0n, b)) case 0n 1n+q: MatC{l, r2} = m MatT.rows_cat(MatT.to_rows(0n, q, l), MatT.to_rows(0n, q, r2)) case 1n+p 1n+q: MatQ{a, b, c2, d} = m List.append(&2, List<&2, F32>, MatT.rows_cat(MatT.to_rows(p, q, a), MatT.to_rows(p, q, b)), MatT.rows_cat(MatT.to_rows(p, q, c2), MatT.to_rows(p, q, d)))# xs cut into n rows of wdef MatT.chunks(+n: Nat, +w: Nat, +xs: List<&2, F32>) -> List<&2, List<&2, F32>>: match n: case 0n: [] case 1n+p: List.take(&2, F32, xs, w) <> MatT.chunks(p, w, List.drop(&2, F32, xs, w))# a MatT(r, c) from 4 * 2^r times 4 * 2^c floats in row-major orderdef MatT.from_list(+r: Nat, +c: Nat, xs: List<&2, F32>) -> MatT(r, c): +r2 = r MatT.from_rows(r2, c, MatT.chunks(VecT.size(r2), VecT.size(c), xs))# the cells in row-major orderdef MatT.to_list(+r: Nat, +c: Nat, m: MatT(r, c)) -> List<&2, F32>: List.concat(&2, F32, MatT.to_rows(r, c, m))def MatT.show.row(row: List<&2, F32>) -> String: List.show(~&2, ~F32, ~F32.show.of, row)def MatT.show(+r: Nat, +c: Nat, m: MatT(r, c)) -> String: List.show(~&2, ~List<&2, F32>, ~MatT.show.row, MatT.to_rows(r, c, m))# Sizes# =====## The public types. A Vec<m> holds exactly m floats and a Mat<rows, cols># exactly rows x cols, in a tree of the smallest depth that fits, padded# with 0.0. This is what every dense library does underneath (BLAS packs# edge tiles into padded buffers; the tensor-core paths want multiples of# 8 or 16; the LLM stacks pad sequence lengths), here made explicit and# kept by every op: an op that could give a padding cell a value other# than 0.0 (exp, shift, map, eye) masks its result, and a reduction that# padding would fool (max, min, argmax, softmax) stops at the size.# Sums, dots and products need nothing: a padding cell adds 0.# the smallest n with 4 * 2^n >= m; the fuel is m itself, log2(m) stepsdef VecT.depth.go(fuel: Nat, +n: Nat, +cap: Nat, +m: Nat, lt: Bool) -> Nat: match fuel lt: case _ False{}: n case 0n _: n case 1n+f True{}: +cap2 = Nat.add(cap, cap) VecT.depth.go(f, 1n+n, cap2, m, Nat.is_lt(cap2, m))def VecT.depth(+m: Nat) -> Nat: VecT.depth.go(m, 0n, 4n, m, Nat.is_lt(4n, m))# Masked tree ops# ---------------# i0 (and j0) is the index of the subtree's first element; m (rows,# cols) the size. A subtree that starts past the size is computed and# discarded, which costs at most the padding's share of the work.# elements at index >= m set to 0.0def VecT.mask(+n: Nat, v: VecT(n), +i0: Nat, +m: Nat) -> VecT(n): match n: case 0n: VecL{x} = v VecL{Vec4.mask(Nat.sub(m, i0), x)} case 1n+p: VecN{l, r} = v a b = VecT.mask(p, l, i0, m) VecT.mask(p, r, Nat.add(i0, VecT.size(p)), m) VecN{a, b}# the largest element among the first mdef VecT.max_upto(+n: Nat, v: VecT(n), +i0: Nat, +m: Nat) -> F32: match n: case 0n: VecL{x} = v Vec4.max_upto(Nat.sub(m, i0), x) case 1n+p: VecN{l, r} = v +h = Nat.add(i0, VecT.size(p)) a b = VecT.max_upto(p, l, i0, m) VecT.max_upto(p, r, h, m) +a2 = a Bool.pick(F32, Nat.is_lt(h, m), F32.max(a2, b), a2)def VecT.min_upto(+n: Nat, v: VecT(n), +i0: Nat, +m: Nat) -> F32: match n: case 0n: VecL{x} = v Vec4.min_upto(Nat.sub(m, i0), x) case 1n+p: VecN{l, r} = v +h = Nat.add(i0, VecT.size(p)) a b = VecT.min_upto(p, l, i0, m) VecT.min_upto(p, r, h, m) +a2 = a Bool.pick(F32, Nat.is_lt(h, m), F32.min(a2, b), a2)# the right half's winner counts only when the right half starts before mdef VecT.argmax_upto.join(lt: Bool, +off: U32, a: U32 & F32, b: U32 & F32) -> U32 & F32: match lt: case True{}: VecT.argmax.two(off, a, b) case False{}: a# the index and value of the largest among the first m; the first wins tiesdef VecT.argmax_upto(+n: Nat, v: VecT(n), +i0: Nat, +m: Nat) -> U32 & F32: match n: case 0n: VecL{x} = v Vec4.argmax_upto(Nat.sub(m, i0), x) case 1n+p: VecN{l, r} = v +h = Nat.add(i0, VecT.size(p)) a b = VecT.argmax_upto(p, l, i0, m) VecT.argmax_upto(p, r, h, m) VecT.argmax_upto.join(Nat.is_lt(h, m), U32.from_nat(VecT.size(p)), a, b)# rows at index >= rows and columns at index >= cols set to 0.0def MatT.mask(+r: Nat, +c: Nat, m: MatT(r, c), +i0: Nat, +j0: Nat, +rows: Nat, +cols: Nat) -> MatT(r, c): match r c: case 0n 0n: MatB{b} = m MatB{Blk.mask(Nat.sub(rows, i0), Nat.sub(cols, j0), b)} case 1n+p 0n: MatR{t, b} = m x y = MatT.mask(p, 0n, t, i0, j0, rows, cols) MatT.mask(p, 0n, b, Nat.add(i0, VecT.size(p)), j0, rows, cols) MatR{x, y} case 0n 1n+q: MatC{l, rt} = m x y = MatT.mask(0n, q, l, i0, j0, rows, cols) MatT.mask(0n, q, rt, i0, Nat.add(j0, VecT.size(q)), rows, cols) MatC{x, y} case 1n+p 1n+q: MatQ{a, b, c2, d} = m +i1 = Nat.add(i0, VecT.size(p)) +j1 = Nat.add(j0, VecT.size(q)) w x y z = MatT.mask(p, q, a, i0, j0, rows, cols) MatT.mask(p, q, b, i0, j1, rows, cols) MatT.mask(p, q, c2, i1, j0, rows, cols) MatT.mask(p, q, d, i1, j1, rows, cols) MatQ{w, x, y, z}# each row's largest among its first cols columns; a padding row answers 0.0def MatT.max_rows_upto(+r: Nat, +c: Nat, m: MatT(r, c), +j0: Nat, +cols: Nat) -> VecT(r): match r c: case 0n 0n: MatB{b} = m VecL{Blk.max_rows_upto(Nat.sub(cols, j0), b)} case 1n+p 0n: MatR{t, b} = m x y = MatT.max_rows_upto(p, 0n, t, j0, cols) MatT.max_rows_upto(p, 0n, b, j0, cols) VecN{x, y} case 0n 1n+q: MatC{l, rt} = m +h = Nat.add(j0, VecT.size(q)) x y = MatT.max_rows_upto(0n, q, l, j0, cols) MatT.max_rows_upto(0n, q, rt, h, cols) +x2 = x Bool.pick(VecT(0n), Nat.is_lt(h, cols), VecT.zip(~F32.max.of, 0n, x2, y), x2) case 1n+p 1n+q: MatQ{a, b, c2, d} = m +h = Nat.add(j0, VecT.size(q)) p0 q0 p1 q1 = MatT.max_rows_upto(p, q, a, j0, cols) MatT.max_rows_upto(p, q, b, h, cols) MatT.max_rows_upto(p, q, c2, j0, cols) MatT.max_rows_upto(p, q, d, h, cols) +p2 = p0 +p3 = p1 +lt = Nat.is_lt(h, cols) VecN{Bool.pick(VecT(p), lt, VecT.zip(~F32.max.of, p, p2, q0), p2), Bool.pick(VecT(p), lt, VecT.zip(~F32.max.of, p, p3, q1), p3)}# Vec<m># ------type Vec<-m: Nat> is Data: VecS{tree: VecT(VecT.depth(m))}def Vec.zeros(+m: Nat) -> Vec<m>: VecS{VecT.zeros(VecT.depth(m))}def Vec.full(+m: Nat, +x: F32) -> Vec<m>: VecS{VecT.mask(VecT.depth(m), VecT.full(VecT.depth(m), x), 0n, m)}def Vec.ones(+m: Nat) -> Vec<m>: Vec.full(m, 1.0)# 0, 1, .., m - 1def Vec.arange(+m: Nat) -> Vec<m>: VecS{VecT.mask(VecT.depth(m), VecT.arange(VecT.depth(m)), 0n, m)}# the first m of xs, padded with 0.0def Vec.from_list(+m: Nat, xs: List<&2, F32>) -> Vec<m>: VecS{VecT.mask(VecT.depth(m), VecT.from_list(VecT.depth(m), List.take(&2, F32, xs, m)), 0n, m)}# read borrowed: a bias is shared by every pass of a modeldef Vec.to_list(+m: Nat, +v: Vec<m>) -> List<&2, F32>: match v: case VecS{t}: List.take(&2, F32, VecT.to_list(VecT.depth(m), t), m)def Vec.map(~f: F32 -> F32, +m: Nat, v: Vec<m>) -> Vec<m>: VecS{t} = v VecS{VecT.mask(VecT.depth(m), VecT.map(~f, VecT.depth(m), t), 0n, m)}def Vec.zip(~f: F32 -> F32 -> F32, +m: Nat, a: Vec<m>, b: Vec<m>) -> Vec<m>: VecS{x} = a VecS{y} = b VecS{VecT.mask(VecT.depth(m), VecT.zip(~f, VecT.depth(m), x, y), 0n, m)}def Vec.add(+m: Nat, a: Vec<m>, b: Vec<m>) -> Vec<m>: VecS{x} = a VecS{y} = b VecS{VecT.add(VecT.depth(m), x, y)}def Vec.sub(+m: Nat, a: Vec<m>, b: Vec<m>) -> Vec<m>: VecS{x} = a VecS{y} = b VecS{VecT.sub(VecT.depth(m), x, y)}def Vec.mul(+m: Nat, a: Vec<m>, b: Vec<m>) -> Vec<m>: VecS{x} = a VecS{y} = b VecS{VecT.mul(VecT.depth(m), x, y)}def Vec.exp(+m: Nat, v: Vec<m>) -> Vec<m>: Vec.map(~F32.exp, m, v)def Vec.scale(+m: Nat, +k: F32, v: Vec<m>) -> Vec<m>: VecS{t} = v VecS{VecT.scale(VecT.depth(m), k, t)}def Vec.shift(+m: Nat, +k: F32, v: Vec<m>) -> Vec<m>: VecS{t} = v VecS{VecT.mask(VecT.depth(m), VecT.shift(VecT.depth(m), k, t), 0n, m)}def Vec.sum(+m: Nat, v: Vec<m>) -> F32: VecS{t} = v VecT.sum(VecT.depth(m), t)def Vec.dot(+m: Nat, a: Vec<m>, b: Vec<m>) -> F32: VecS{x} = a VecS{y} = b VecT.dot(VecT.depth(m), x, y)def Vec.max(+m: Nat, v: Vec<m>) -> F32: VecS{t} = v VecT.max_upto(VecT.depth(m), t, 0n, m)def Vec.min(+m: Nat, v: Vec<m>) -> F32: VecS{t} = v VecT.min_upto(VecT.depth(m), t, 0n, m)def Vec.argmax(+m: Nat, v: Vec<m>) -> U32 & F32: VecS{t} = v VecT.argmax_upto(VecT.depth(m), t, 0n, m)def Vec.mask(+m: Nat, v: Vec<m>) -> Vec<m>: VecS{t} = v VecS{VecT.mask(VecT.depth(m), t, 0n, m)}def Vec.softmax.fin(+m: Nat, +e: Vec<m>) -> Vec<m>: Vec.mask(m, Vec.scale(m, F32.div(1.0, Vec.sum(m, e)), e))# exp(v - max v) / sum, over the m elementsdef Vec.softmax(+m: Nat, +v: Vec<m>) -> Vec<m>: Vec.softmax.fin(m, Vec.exp(m, Vec.shift(m, (0.0 - Vec.max(m, v) : F32), v)))def Vec.show(+m: Nat, v: Vec<m>) -> String: List.show(~&2, ~F32, ~F32.show.of, Vec.to_list(m, v))# Mat<rows, cols># ---------------type Mat<-rows: Nat, -cols: Nat> is Data: MatS{tree: MatT(VecT.depth(rows), VecT.depth(cols))}def Mat.zeros(+rows: Nat, +cols: Nat) -> Mat<rows, cols>: MatS{MatT.zeros(VecT.depth(rows), VecT.depth(cols))}def Mat.full(+rows: Nat, +cols: Nat, +x: F32) -> Mat<rows, cols>: MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), MatT.full(VecT.depth(rows), VecT.depth(cols), x), 0n, 0n, rows, cols)}def Mat.ones(+rows: Nat, +cols: Nat) -> Mat<rows, cols>: Mat.full(rows, cols, 1.0)def Mat.eye(+n: Nat) -> Mat<n, n>: MatS{MatT.mask(VecT.depth(n), VecT.depth(n), MatT.eye(VecT.depth(n)), 0n, 0n, n, n)}# a matrix from its rows, row-major; short or missing rows pad with 0.0def Mat.from_rows(+rows: Nat, +cols: Nat, rs: List<&2, List<&2, F32>>) -> Mat<rows, cols>: MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), MatT.from_rows(VecT.depth(rows), VecT.depth(cols), MatT.rows_take(cols, List.take(&2, List<&2, F32>, rs, rows))), 0n, 0n, rows, cols)}# a matrix from rows * cols floats in row-major orderdef Mat.from_list(+rows: Nat, +cols: Nat, xs: List<&2, F32>) -> Mat<rows, cols>: MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), MatT.from_rows(VecT.depth(rows), VecT.depth(cols), MatT.chunks(rows, cols, xs)), 0n, 0n, rows, cols)}def Mat.to_rows(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> List<&2, List<&2, F32>>: MatS{t} = m MatT.rows_take(cols, List.take(&2, List<&2, F32>, MatT.to_rows(VecT.depth(rows), VecT.depth(cols), t), rows))def Mat.to_list(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> List<&2, F32>: List.concat(&2, F32, Mat.to_rows(rows, cols, m))def Mat.map(~f: F32 -> F32, +rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Mat<rows, cols>: MatS{t} = m MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), MatT.map(~f, VecT.depth(rows), VecT.depth(cols), t), 0n, 0n, rows, cols)}def Mat.zip(~f: F32 -> F32 -> F32, +rows: Nat, +cols: Nat, m: Mat<rows, cols>, n: Mat<rows, cols>) -> Mat<rows, cols>: MatS{x} = m MatS{y} = n MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), MatT.zip(~f, VecT.depth(rows), VecT.depth(cols), x, y), 0n, 0n, rows, cols)}def Mat.add(+rows: Nat, +cols: Nat, m: Mat<rows, cols>, n: Mat<rows, cols>) -> Mat<rows, cols>: MatS{x} = m MatS{y} = n MatS{MatT.add(VecT.depth(rows), VecT.depth(cols), x, y)}def Mat.sub(+rows: Nat, +cols: Nat, m: Mat<rows, cols>, n: Mat<rows, cols>) -> Mat<rows, cols>: MatS{x} = m MatS{y} = n MatS{MatT.sub(VecT.depth(rows), VecT.depth(cols), x, y)}# elementwise (Hadamard) product; the matrix product is Mat.matmuldef Mat.mul(+rows: Nat, +cols: Nat, m: Mat<rows, cols>, n: Mat<rows, cols>) -> Mat<rows, cols>: MatS{x} = m MatS{y} = n MatS{MatT.mul(VecT.depth(rows), VecT.depth(cols), x, y)}def Mat.exp(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Mat<rows, cols>: Mat.map(~F32.exp, rows, cols, m)def Mat.scale(+rows: Nat, +cols: Nat, +k: F32, m: Mat<rows, cols>) -> Mat<rows, cols>: MatS{t} = m MatS{MatT.scale(VecT.depth(rows), VecT.depth(cols), k, t)}def Mat.transpose(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Mat<cols, rows>: MatS{t} = m MatS{MatT.mask(VecT.depth(cols), VecT.depth(rows), MatT.transpose(VecT.depth(rows), VecT.depth(cols), t), 0n, 0n, cols, rows)}# a * b: Mat<rows, k> times Mat<k, cols> is Mat<rows, cols>def Mat.matmul(+rows: Nat, +k: Nat, +cols: Nat, a: Mat<rows, k>, b: Mat<k, cols>) -> Mat<rows, cols>: MatS{x} = a MatS{y} = b MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), MatT.matmul(VecT.depth(rows), VecT.depth(k), VecT.depth(cols), x, y), 0n, 0n, rows, cols)}# a * b for two Mat<n, n>: the one-depth recursion, 3x faster per nodedef Mat.matmul.sq(+n: Nat, a: Mat<n, n>, b: Mat<n, n>) -> Mat<n, n>: MatS{x} = a MatS{y} = b MatS{MatT.mask(VecT.depth(n), VecT.depth(n), MatT.matmul.sq(VecT.depth(n), x, y), 0n, 0n, n, n)}# m * v: Mat<rows, cols> times Vec<cols> is Vec<rows>def Mat.matvec(+rows: Nat, +cols: Nat, m: Mat<rows, cols>, v: Vec<cols>) -> Vec<rows>: MatS{t} = m VecS{x} = v VecS{VecT.mask(VecT.depth(rows), MatT.matvec(VecT.depth(rows), VecT.depth(cols), t, x), 0n, rows)}def Mat.sum(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> F32: MatS{t} = m MatT.sum(VecT.depth(rows), VecT.depth(cols), t)def Mat.sum_rows(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Vec<rows>: MatS{t} = m VecS{VecT.mask(VecT.depth(rows), MatT.sum_rows(VecT.depth(rows), VecT.depth(cols), t), 0n, rows)}def Mat.sum_cols(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Vec<cols>: MatS{t} = m VecS{VecT.mask(VecT.depth(cols), MatT.sum_cols(VecT.depth(rows), VecT.depth(cols), t), 0n, cols)}def Mat.max_rows(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Vec<rows>: MatS{t} = m VecS{VecT.mask(VecT.depth(rows), MatT.max_rows_upto(VecT.depth(rows), VecT.depth(cols), t, 0n, cols), 0n, rows)}# v added to every row. Masked: a padding row plus v is v, not 0.0.def Mat.add_row(+rows: Nat, +cols: Nat, m: Mat<rows, cols>, v: Vec<cols>) -> Mat<rows, cols>: MatS{t} = m VecS{x} = v MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), MatT.add_row(VecT.depth(rows), VecT.depth(cols), t, x), 0n, 0n, rows, cols)}# row i minus v_i. Masked: a padding column minus v_i is -v_i, not 0.0.def Mat.sub_col(+rows: Nat, +cols: Nat, m: Mat<rows, cols>, v: Vec<rows>) -> Mat<rows, cols>: MatS{t} = m VecS{x} = v MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), MatT.sub_col(VecT.depth(rows), VecT.depth(cols), t, x), 0n, 0n, rows, cols)}# row i times v_i. Masked, so its law needs no fact about x * 0.0.def Mat.scale_rows(+rows: Nat, +cols: Nat, m: Mat<rows, cols>, v: Vec<rows>) -> Mat<rows, cols>: MatS{t} = m VecS{x} = v MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), MatT.scale_rows(VecT.depth(rows), VecT.depth(cols), t, x), 0n, 0n, rows, cols)}def Mat.mask(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Mat<rows, cols>: MatS{t} = m MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), t, 0n, 0n, rows, cols)}def Mat.softmax_rows.fin(+rows: Nat, +cols: Nat, +e: Mat<rows, cols>) -> Mat<rows, cols>: Mat.mask(rows, cols, Mat.scale_rows(rows, cols, e, Vec.map(~(s => F32.div(1.0, s)), rows, Mat.sum_rows(rows, cols, e))))# softmax along each row over its cols columnsdef Mat.softmax_rows(+rows: Nat, +cols: Nat, +m: Mat<rows, cols>) -> Mat<rows, cols>: Mat.softmax_rows.fin(rows, cols, Mat.exp(rows, cols, Mat.sub_col(rows, cols, m, Mat.max_rows(rows, cols, m))))def Mat.show(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> String: List.show(~&2, ~List<&2, F32>, ~MatT.show.row, Mat.to_rows(rows, cols, m))# a Mat<rows, cols> around a tree, masked: what the IO products below# answer, and what their padding law is aboutdef Mat.from_tree(+rows: Nat, +cols: Nat, t: MatT(VecT.depth(rows), VecT.depth(cols))) -> Mat<rows, cols>: MatS{MatT.mask(VecT.depth(rows), VecT.depth(cols), t, 0n, 0n, rows, cols)}# Padding# =======## The predicates the padding laws are stated with. Vec4.Padded(k, v) says# the elements of v at index >= k are 0.0, as a nested pair of equalities# (Unit when k covers all four); VecT.Padded and MatT.Padded lift that# over the trees with the same index bookkeeping as the masks, and# Vec.Padded / Mat.Padded state it for the sized types. A proof of one is# a tree of {==}s; the laws in LAWS.bend say every op answers one.def Vec4.Padded(k: Nat, v: Vec4) -> Type: match k v: case 0n Vec4{x0, x1, x2, x3}: {x0 == 0.0 : F32} & {x1 == 0.0 : F32} & {x2 == 0.0 : F32} & {x3 == 0.0 : F32} case 1n Vec4{x0, x1, x2, x3}: {x1 == 0.0 : F32} & {x2 == 0.0 : F32} & {x3 == 0.0 : F32} case 2n Vec4{x0, x1, x2, x3}: {x2 == 0.0 : F32} & {x3 == 0.0 : F32} case 3n Vec4{x0, x1, x2, x3}: {x3 == 0.0 : F32} case _ _: Unitdef VecT.Padded(+n: Nat, v: VecT(n), +i0: Nat, +m: Nat) -> Type: match n: case 0n: VecL{x} = v Vec4.Padded(Nat.sub(m, i0), x) case 1n+p: VecN{l, r} = v VecT.Padded(p, l, i0, m) & VecT.Padded(p, r, Nat.add(i0, VecT.size(p)), m)# row i of a block is padded past kj when i < ki, and entirely when notdef Blk.kcol(ki: Nat, i: Nat, +kj: Nat) -> Nat: Bool.pick(Nat, Nat.is_lt(i, ki), kj, 0n)def Blk.Padded(+ki: Nat, +kj: Nat, b: Blk) -> Type: Blk{a00, a01, a02, a03, a10, a11, a12, a13, a20, a21, a22, a23, a30, a31, a32, a33} = b Vec4.Padded(Blk.kcol(ki, 0n, kj), Vec4{a00, a01, a02, a03}) & Vec4.Padded(Blk.kcol(ki, 1n, kj), Vec4{a10, a11, a12, a13}) & Vec4.Padded(Blk.kcol(ki, 2n, kj), Vec4{a20, a21, a22, a23}) & Vec4.Padded(Blk.kcol(ki, 3n, kj), Vec4{a30, a31, a32, a33})def MatT.Padded(+r: Nat, +c: Nat, m: MatT(r, c), +i0: Nat, +j0: Nat, +rows: Nat, +cols: Nat) -> Type: match r c: case 0n 0n: MatB{b} = m Blk.Padded(Nat.sub(rows, i0), Nat.sub(cols, j0), b) case 1n+p 0n: MatR{t, b} = m MatT.Padded(p, 0n, t, i0, j0, rows, cols) & MatT.Padded(p, 0n, b, Nat.add(i0, VecT.size(p)), j0, rows, cols) case 0n 1n+q: MatC{l, rt} = m MatT.Padded(0n, q, l, i0, j0, rows, cols) & MatT.Padded(0n, q, rt, i0, Nat.add(j0, VecT.size(q)), rows, cols) case 1n+p 1n+q: MatQ{a, b, c2, d} = m (MatT.Padded(p, q, a, i0, j0, rows, cols) & MatT.Padded(p, q, b, i0, Nat.add(j0, VecT.size(q)), rows, cols)) & (MatT.Padded(p, q, c2, Nat.add(i0, VecT.size(p)), j0, rows, cols) & MatT.Padded(p, q, d, Nat.add(i0, VecT.size(p)), Nat.add(j0, VecT.size(q)), rows, cols))def Vec.Padded(+m: Nat, v: Vec<m>) -> Type: VecS{t} = v VecT.Padded(VecT.depth(m), t, 0n, m)def Mat.Padded(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Type: MatS{t} = m MatT.Padded(VecT.depth(rows), VecT.depth(cols), t, 0n, 0n, rows, cols)# Flat# ====## Flat<rows, cols> is rows x cols floats in one row-major Array<F32>,# the layout BLAS reads, so a product hands its operands to the platform# sgemm with no conversion and a model's forward pass never leaves flat# memory. The array has the smallest power-of-two slot count that holds# rows * cols; nothing reads past rows * cols. An Array has one owner, so# a Flat is affine (a Type): the elementwise ops consume their operand# and answer a new one, the reductions consume it, a product consumes# its array operands, and Flat.clone copies. Weights live outside the# heap as a Weights (Data, shared, never consumed) and biases as a Vec# (Data too), so a forward pass is one do block. Every loop here is one# flat tail loop, the reads as stage chains. There are no laws about# Flat beyond its shape: BLAS is a black box and proofs about array# cells are out of reach; the tests check it against numpy.# Effects# -------# C = op(A) op(B) on row-major arrays. m, n, k are the result rows, the# result columns and the inner size; ta and tb (0 or 1) transpose an# operand as stored. Answers a fresh array of the smallest power-of-two# size holding m * n. Blas.sgemm is the platform CPU BLAS (Accelerate,# OpenBLAS), opened at runtime. The _w forms take the right operand as# a registry weight (see Weights) instead of an array. The GPU build# (bend_tensors_gpu.bend) adds Blas.sgemm_gpu: Metal Performance Shaders# on macOS, cuBLAS on Linux.def Blas.sgemm(m: U32, n: U32, k: U32, ta: U32, tb: U32, a: Array<F32>, b: Array<F32>) -> IO(Result<&1, &1, U32 & String, Array<F32>>): import "./effs/sgemm.c" import "./effs/sgemm.js"def Blas.sgemm_w(m: U32, n: U32, k: U32, ta: U32, tb: U32, a: Array<F32>, id: U32) -> IO(Result<&1, &1, U32 & String, Array<F32>>): import "./effs/sgemm_w.c" import "./effs/sgemm_w.js"# n float32 (little-endian, as numpy writes them) from a file, into the# weight registry (answers the id) or into a fresh arraydef Weights.load.eff(path: String, n: U32) -> IO(Result<&1, &1, U32 & String, U32>): import "./effs/weights_load.c" import "./effs/weights_load.js"def Weights.of_array.eff(a: Array<F32>, n: U32) -> IO(Result<&1, &1, U32 & String, U32>): import "./effs/weights_of_array.c" import "./effs/weights_of_array.js"def Flat.load.eff(path: String, n: U32) -> IO(Result<&1, &1, U32 & String, Array<F32>>): import "./effs/flat_load.c" import "./effs/flat_load.js"# Weights# -------# A Weights<rows, cols> is rows x cols floats that live outside the# Bend heap, in the BLAS effects' memory, named by an id. It is Data:# share it freely, use it in every pass, it is never copied or consumed.# The GPU effects keep a device copy beside it after the first use.type Weights<-rows: Nat, -cols: Nat> is Data: W{id: U32}def Weights.wrap(-rows: Nat, -cols: Nat, id: U32) -> IO(Weights<rows, cols>): IO.pure(Weights<rows, cols>, W{id})# rows * cols float32 from a file, as numpy's tofile writes a C-ordered arraydef Weights.load(+rows: Nat, +cols: Nat, path: String) -> IO(Weights<rows, cols>): IO.bind(U32, Weights<rows, cols>, IO.try(U32, Weights.load.eff(path, (U32.from_nat(rows) * U32.from_nat(cols) : U32))), id => Weights.wrap(rows, cols, id))# Flat# ----type Flat<-rows: Nat, -cols: Nat> is Type: FlatS{a: Array<F32>}# the smallest d with 2^d >= n; the fuel bounds the doubling stepsdef Flat.depth.go(fuel: Nat, +d: Nat, +cap: Nat, +n: Nat, lt: Bool) -> Nat: match fuel lt: case _ False{}: d case 0n _: d case 1n+f True{}: +cap2 = Nat.add(cap, cap) Flat.depth.go(f, 1n+d, cap2, n, Nat.is_lt(cap2, n))def Flat.depth(+n: Nat) -> Nat: Flat.depth.go(48n, 0n, 1n, n, Nat.is_lt(1n, n))def Flat.count(+rows: Nat, +cols: Nat) -> Nat: Nat.mul(rows, cols)# Construction# ------------def Flat.full(+rows: Nat, +cols: Nat, +x: F32) -> Flat<rows, cols>: FlatS{Array.new(F32, Flat.depth(Flat.count(rows, cols)), x)}def Flat.zeros(+rows: Nat, +cols: Nat) -> Flat<rows, cols>: Flat.full(rows, cols, 0.0)def Flat.ones(+rows: Nat, +cols: Nat) -> Flat<rows, cols>: Flat.full(rows, cols, 1.0)def Flat.fill.go(xs: List<&2, F32>, +i: U32, a: Array<F32>) -> Array<F32>: match xs: case []: a case x <> rest: Flat.fill.go(rest, (i + 1 : U32), Array.set(F32, a, i, x))# the first rows * cols of xs, row-major, the rest 0.0def Flat.from_list(+rows: Nat, +cols: Nat, xs: List<&2, F32>) -> Flat<rows, cols>: FlatS{Flat.fill.go(List.take(&2, F32, xs, Flat.count(rows, cols)), 0, Array.new(F32, Flat.depth(Flat.count(rows, cols)), 0.0))}def Flat.wrap(-rows: Nat, -cols: Nat, a: Array<F32>) -> IO(Flat<rows, cols>): IO.pure(Flat<rows, cols>, FlatS{a})# rows * cols float32 from a file, row-majordef Flat.from_file(+rows: Nat, +cols: Nat, path: String) -> IO(Flat<rows, cols>): IO.bind(Array<F32>, Flat<rows, cols>, IO.try(Array<F32>, Flat.load.eff(path, (U32.from_nat(rows) * U32.from_nat(cols) : U32))), a => Flat.wrap(rows, cols, a))# a Flat's floats copied into the weight registry; the Flat is consumeddef Weights.of_flat(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> IO(Weights<rows, cols>): FlatS{a} = m IO.bind(U32, Weights<rows, cols>, IO.try(U32, Weights.of_array.eff(a, (U32.from_nat(rows) * U32.from_nat(cols) : U32))), id => Weights.wrap(rows, cols, id))# the cells in row-major order; i counts down from the lastdef Flat.to_list.go(n: Nat, +i: U32, r: Array<F32> & F32, acc: List<&2, F32>) -> List<&2, F32>: match n: case 0n: (a, x) = r x <> acc case 1n+p: (a, x) = r Flat.to_list.go(p, (i - 1 : U32), Array.get(F32, a, (i - 1 : U32)), x <> acc)def Flat.to_list.of(n: Nat, +last: U32, a: Array<F32>) -> List<&2, F32>: match n: case 0n: [] case 1n+p: Flat.to_list.go(p, last, Array.get(F32, a, last), [])def Flat.to_list(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> List<&2, F32>: FlatS{a} = m Flat.to_list.of(Flat.count(rows, cols), (U32.from_nat(Flat.count(rows, cols)) - 1 : U32), a)def Vec.wrap(+m: Nat, a: Array<F32>) -> IO(Vec<m>): IO.pure(Vec<m>, Vec.from_list(m, Flat.to_list.of(m, (U32.from_nat(m) - 1 : U32), a)))# m float32 from a file, as a Vec (Data: share it, it is never consumed)def Vec.from_file(+m: Nat, path: String) -> IO(Vec<m>): IO.bind(Array<F32>, Vec<m>, IO.try(Array<F32>, Flat.load.eff(path, U32.from_nat(m))), a => Vec.wrap(m, a))def Flat.clone.fin(-rows: Nat, -cols: Nat, r: Array<F32> & Array<F32>) -> Flat<rows, cols> & Flat<rows, cols>: (a, b) = r (FlatS{a}, FlatS{b})# two copies; open the pair in a def of your own (a do block cannot)def Flat.clone(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Flat<rows, cols> & Flat<rows, cols>: FlatS{a} = m Flat.clone.fin(rows, cols, Array.clone(F32, a))def Flat.clone.fst(+rows: Nat, +cols: Nat, r: Flat<rows, cols> & Flat<rows, cols>) -> Flat<rows, cols>: (a, b) = r adef Flat.clone.snd(+rows: Nat, +cols: Nat, r: Flat<rows, cols> & Flat<rows, cols>) -> Flat<rows, cols>: (a, b) = r bdef Flat.to_rows(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> List<&2, List<&2, F32>>: MatT.chunks(rows, cols, Flat.to_list(rows, cols, m))def Flat.show(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> String: List.show(~&2, ~List<&2, F32>, ~MatT.show.row, Flat.to_rows(rows, cols, m))# Elementwise# -----------def Flat.map.step(~f: F32 -> F32, +i: U32, r: Array<F32> & F32) -> Array<F32>: (a, x) = r Array.set(F32, a, i, f(x))def Flat.map.go(~f: F32 -> F32, n: Nat, +i: U32, a: Array<F32>) -> Array<F32>: match n: case 0n: a case 1n+p: Flat.map.go(~f, p, (i + 1 : U32), Flat.map.step(~f, i, Array.get(F32, a, i)))def Flat.map(~f: F32 -> F32, +rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Flat<rows, cols>: FlatS{a} = m FlatS{Flat.map.go(~f, Flat.count(rows, cols), 0, a)}def Flat.exp(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Flat<rows, cols>: Flat.map(~F32.exp, rows, cols, m)def F32.relu(x: F32) -> F32: F32.max(0.0, x)def Flat.relu(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Flat<rows, cols>: Flat.map(~F32.relu, rows, cols, m)# gelu, the tanh form: 0.5 x (1 + tanh(sqrt(2 / pi) (x + 0.044715 x^3)))def F32.gelu(x0: F32) -> F32: +x = x0 (0.5 * x * (1.0 + F32.tanh((0.7978845608 * (x + 0.044715 * x * x * x : F32) : F32))) : F32)def Flat.gelu(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Flat<rows, cols>: Flat.map(~F32.gelu, rows, cols, m)def Flat.scale.step(+k: F32, +i: U32, r: Array<F32> & F32) -> Array<F32>: (a, x) = r Array.set(F32, a, i, (k * x : F32))def Flat.scale.go(n: Nat, +k: F32, +i: U32, a: Array<F32>) -> Array<F32>: match n: case 0n: a case 1n+p: Flat.scale.go(p, k, (i + 1 : U32), Flat.scale.step(k, i, Array.get(F32, a, i)))def Flat.scale(+rows: Nat, +cols: Nat, +k: F32, m: Flat<rows, cols>) -> Flat<rows, cols>: FlatS{a} = m FlatS{Flat.scale.go(Flat.count(rows, cols), k, 0, a)}def Flat.zip.step2(~f: F32 -> F32 -> F32, +i: U32, x: F32, a: Array<F32>, r: Array<F32> & F32) -> Array<F32> & Array<F32>: (b, y) = r (Array.set(F32, a, i, f(x, y)), b)def Flat.zip.step(~f: F32 -> F32 -> F32, +i: U32, b: Array<F32>, r: Array<F32> & F32) -> Array<F32> & Array<F32>: (a, x) = r Flat.zip.step2(~f, i, x, a, Array.get(F32, b, i))def Flat.zip.go(~f: F32 -> F32 -> F32, n: Nat, +i: U32, ab: Array<F32> & Array<F32>) -> Array<F32> & Array<F32>: match n: case 0n: ab case 1n+p: (a, b) = ab Flat.zip.go(~f, p, (i + 1 : U32), Flat.zip.step(~f, i, b, Array.get(F32, a, i)))def Flat.zip.fin(-rows: Nat, -cols: Nat, ab: Array<F32> & Array<F32>) -> Flat<rows, cols>: (a, b) = ab FlatS{a}# f over the pairs of cells; both operands are consumeddef Flat.zip(~f: F32 -> F32 -> F32, +rows: Nat, +cols: Nat, m: Flat<rows, cols>, n: Flat<rows, cols>) -> Flat<rows, cols>: FlatS{a} = m FlatS{b} = n Flat.zip.fin(rows, cols, Flat.zip.go(~f, Flat.count(rows, cols), 0, (a, b)))def Flat.add(+rows: Nat, +cols: Nat, m: Flat<rows, cols>, n: Flat<rows, cols>) -> Flat<rows, cols>: Flat.zip(~F32.add, rows, cols, m, n)def Flat.sub(+rows: Nat, +cols: Nat, m: Flat<rows, cols>, n: Flat<rows, cols>) -> Flat<rows, cols>: Flat.zip(~F32.sub, rows, cols, m, n)# elementwise product; the matrix product is Flat.matmuldef Flat.mul(+rows: Nat, +cols: Nat, m: Flat<rows, cols>, n: Flat<rows, cols>) -> Flat<rows, cols>: Flat.zip(~F32.mul, rows, cols, m, n)# cells above the diagonal (j > i) set to -3.4e38, the most negative# finite float, which a softmax turns into 0 with no inf - infdef Flat.mask_causal.step(+idx: U32, +i: U32, +j: U32, r: Array<F32> & F32) -> Array<F32>: (a, x) = r Array.set(F32, a, idx, Bool.pick(F32, U32.is_gt(j, i), (0.0 - 3.4028235e38 : F32), x))def Flat.mask_causal.go(n: Nat, +idx: U32, +colsu: U32, a: Array<F32>) -> Array<F32>: match n: case 0n: a case 1n+p: Flat.mask_causal.go(p, (idx + 1 : U32), colsu, Flat.mask_causal.step(idx, U32.div(idx, colsu), U32.mod(idx, colsu), Array.get(F32, a, idx)))def Flat.mask_causal(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Flat<rows, cols>: FlatS{a} = m FlatS{Flat.mask_causal.go(Flat.count(rows, cols), 0, U32.from_nat(cols), a)}# Rows# ----# The row ops run an inner loop over one row at a time; the row's start# is base, its length cols.def Flat.add_row.step(+idx: U32, +y: F32, r: Array<F32> & F32) -> Array<F32>: (a, x) = r Array.set(F32, a, idx, (x + y : F32))# the bias as a list, walked from its head for every rowdef Flat.add_row.row(n: Nat, +idx: U32, xs: List<&2, F32>, a: Array<F32>) -> Array<F32>: match n xs: case 0n _: a case 1n+p Nil{}: a case 1n+p Con{y, rest}: Flat.add_row.row(p, (idx + 1 : U32), rest, Flat.add_row.step(idx, y, Array.get(F32, a, idx)))def Flat.add_row.go(n: Nat, +cols: Nat, +colsu: U32, +base: U32, +xs: List<&2, F32>, a: Array<F32>) -> Array<F32>: match n: case 0n: a case 1n+p: Flat.add_row.go(p, cols, colsu, (base + colsu : U32), xs, Flat.add_row.row(cols, base, xs, a))# v added to every row; v is a Vec, Data, so it is shared, not consumeddef Flat.add_row(+rows: Nat, +cols: Nat, m: Flat<rows, cols>, +v: Vec<cols>) -> Flat<rows, cols>: FlatS{a} = m FlatS{Flat.add_row.go(rows, cols, U32.from_nat(cols), 0, Vec.to_list(cols, v), a)}# the largest cell of a row segmentdef Flat.rowmax.go(n: Nat, +i: U32, r: Array<F32> & F32, +best: F32) -> Array<F32> & F32: match n: case 0n: (a, x) = r (a, F32.max(best, x)) case 1n+p: (a, x) = r Flat.rowmax.go(p, (i + 1 : U32), Array.get(F32, a, (i + 1 : U32)), F32.max(best, x))# exp(x - mx) written back, the row's sum answereddef Flat.rowexp.step(+i: U32, +mx: F32, r: Array<F32> & F32, +acc: F32) -> Array<F32> & F32: (a, x) = r +y = F32.exp((x - mx : F32)) (Array.set(F32, a, i, y), (acc + y : F32))def Flat.rowexp.go(n: Nat, +i: U32, +mx: F32, r: Array<F32> & F32) -> Array<F32> & F32: match n: case 0n: r case 1n+p: (a, acc) = r Flat.rowexp.go(p, (i + 1 : U32), mx, Flat.rowexp.step(i, mx, Array.get(F32, a, i), acc))def Flat.rowscale.go(n: Nat, +i: U32, +k: F32, a: Array<F32>) -> Array<F32>: match n: case 0n: a case 1n+p: Flat.rowscale.go(p, (i + 1 : U32), k, Flat.scale.step(k, i, Array.get(F32, a, i)))# one row of softmax: max, exp and sum, scaledef Flat.softmax.row3(+n: Nat, +base: U32, r: Array<F32> & F32) -> Array<F32>: (a, s) = r Flat.rowscale.go(n, base, F32.div(1.0, s), a)def Flat.softmax.row2(+n: Nat, +base: U32, r: Array<F32> & F32) -> Array<F32>: (a, mx) = r Flat.softmax.row3(n, base, Flat.rowexp.go(n, base, mx, (a, 0.0)))def Flat.softmax.row(+n: Nat, +base: U32, a: Array<F32>) -> Array<F32>: match n: case 0n: a case 1n+p: Flat.softmax.row2(1n+p, base, Flat.rowmax.go(p, base, Array.get(F32, a, base), (0.0 - 3.4028235e38 : F32)))def Flat.softmax.go(n: Nat, +cols: Nat, +colsu: U32, +base: U32, a: Array<F32>) -> Array<F32>: match n: case 0n: a case 1n+p: Flat.softmax.go(p, cols, colsu, (base + colsu : U32), Flat.softmax.row(cols, base, a))# softmax along each rowdef Flat.softmax_rows(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Flat<rows, cols>: FlatS{a} = m FlatS{Flat.softmax.go(rows, cols, U32.from_nat(cols), 0, a)}# layer norm over each row: (x - mean) / sqrt(var + eps), var the mean# of the squared deviations; a pass for the sum, one for the squares,# one to writedef F32.sqdev(+x: F32, +mean: F32) -> F32: ((x - mean) * (x - mean) : F32)def Flat.rowsq.go(n: Nat, +i: U32, +mean: F32, r: Array<F32> & F32, +acc: F32) -> Array<F32> & F32: match n: case 0n: (a, x) = r (a, (acc + F32.sqdev(x, mean) : F32)) case 1n+p: (a, x) = r Flat.rowsq.go(p, (i + 1 : U32), mean, Array.get(F32, a, (i + 1 : U32)), (acc + F32.sqdev(x, mean) : F32))def Flat.rowsum.go(n: Nat, +i: U32, r: Array<F32> & F32, +acc: F32) -> Array<F32> & F32: match n: case 0n: (a, x) = r (a, (acc + x : F32)) case 1n+p: (a, x) = r Flat.rowsum.go(p, (i + 1 : U32), Array.get(F32, a, (i + 1 : U32)), (acc + x : F32))def Flat.rownorm.step(+i: U32, +mean: F32, +inv: F32, r: Array<F32> & F32) -> Array<F32>: (a, x) = r Array.set(F32, a, i, ((x - mean) * inv : F32))def Flat.rownorm.go(n: Nat, +i: U32, +mean: F32, +inv: F32, a: Array<F32>) -> Array<F32>: match n: case 0n: a case 1n+p: Flat.rownorm.go(p, (i + 1 : U32), mean, inv, Flat.rownorm.step(i, mean, inv, Array.get(F32, a, i)))def Flat.layer_norm.row3(+n: Nat, +base: U32, +mean: F32, +eps: F32, +cnt: F32, r: Array<F32> & F32) -> Array<F32>: (a, sq) = r Flat.rownorm.go(n, base, mean, F32.div(1.0, F32.sqrt((F32.div(sq, cnt) + eps : F32))), a)def Flat.layer_norm.row2(+n: Nat, +p: Nat, +base: U32, +eps: F32, +cnt: F32, r: Array<F32> & F32) -> Array<F32>: (a, s) = r +mean = F32.div(s, cnt) Flat.layer_norm.row3(n, base, mean, eps, cnt, Flat.rowsq.go(p, base, mean, Array.get(F32, a, base), 0.0))def Flat.layer_norm.row(+n: Nat, +base: U32, +eps: F32, a: Array<F32>) -> Array<F32>: match n: case 0n: a case 1n+p: Flat.layer_norm.row2(1n+p, p, base, eps, F32.from_nat(1n+p), Flat.rowsum.go(p, base, Array.get(F32, a, base), 0.0))def Flat.layer_norm.go(n: Nat, +cols: Nat, +colsu: U32, +base: U32, +eps: F32, a: Array<F32>) -> Array<F32>: match n: case 0n: a case 1n+p: Flat.layer_norm.go(p, cols, colsu, (base + colsu : U32), eps, Flat.layer_norm.row(cols, base, eps, a))def Flat.layer_norm_rows(+rows: Nat, +cols: Nat, +eps: F32, m: Flat<rows, cols>) -> Flat<rows, cols>: FlatS{a} = m FlatS{Flat.layer_norm.go(rows, cols, U32.from_nat(cols), 0, eps, a)}# the index of the largest cell of a row segment; the first wins tiesdef Flat.rowarg.pick(c: Bool, +i: U32, +x: F32, +j: U32, +y: F32) -> U32 & F32: match c: case True{}: (i, x) case False{}: (j, y)# the better of the best so far and cell jdef Flat.rowarg.best(+bj: U32, +bx: F32, +j: U32, +x: F32) -> U32 & F32: Flat.rowarg.pick((bx >= x : F32), bj, bx, j, x)def Flat.rowarg.go(n: Nat, +i: U32, +j: U32, r: Array<F32> & F32, best: U32 & F32) -> Array<F32> & (U32 & F32): match n: case 0n: (a, x) = r (bj, bx) = best (a, Flat.rowarg.best(bj, bx, j, x)) case 1n+p: (a, x) = r (bj, bx) = best Flat.rowarg.go(p, (i + 1 : U32), (j + 1 : U32), Array.get(F32, a, (i + 1 : U32)), Flat.rowarg.best(bj, bx, j, x))def Flat.argmax.row(+n: Nat, +base: U32, a: Array<F32>) -> Array<F32> & (U32 & F32): match n: case 0n: (a, (0, 0.0)) case 1n+p: Flat.rowarg.go(p, base, 0, Array.get(F32, a, base), (0, (0.0 - 3.4028235e38 : F32)))# rows from the last down to the first, so the list comes out in row# order; r is the row just donedef Flat.argmax.go(n: Nat, +cols: Nat, +colsu: U32, +base: U32, r: Array<F32> & (U32 & F32), acc: List<&2, U32>) -> List<&2, U32>: match n: case 0n: (a, (j, x)) = r j <> acc case 1n+p: (a, (j, x)) = r Flat.argmax.go(p, cols, colsu, (base - colsu : U32), Flat.argmax.row(cols, (base - colsu : U32), a), j <> acc)def Flat.argmax.of(+rows: Nat, +cols: Nat, a: Array<F32>) -> List<&2, U32>: match rows: case 0n: [] case 1n+p: +colsu = U32.from_nat(cols) +last = (U32.from_nat(p) * colsu : U32) Flat.argmax.go(p, cols, colsu, last, Flat.argmax.row(cols, last, a), [])# each row's argmax, in row order; the matrix is consumeddef Flat.argmax_rows(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> List<&2, U32>: FlatS{a} = m Flat.argmax.of(rows, cols, a)# Reductions# ----------# Each consumes the matrix; clone it to keep it.def Flat.sum.go(n: Nat, +i: U32, r: Array<F32> & F32, +acc: F32) -> F32: match n: case 0n: (a, x) = r (acc + x : F32) case 1n+p: (a, x) = r Flat.sum.go(p, (i + 1 : U32), Array.get(F32, a, (i + 1 : U32)), (acc + x : F32))def Flat.sum.of(n: Nat, a: Array<F32>) -> F32: match n: case 0n: 0.0 case 1n+p: Flat.sum.go(p, 0, Array.get(F32, a, 0), 0.0)def Flat.sum(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> F32: FlatS{a} = m Flat.sum.of(Flat.count(rows, cols), a)def Flat.mean(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> F32: F32.div(Flat.sum(rows, cols, m), F32.from_nat(Flat.count(rows, cols)))def Flat.max.go(n: Nat, +i: U32, r: Array<F32> & F32, +best: F32) -> F32: match n: case 0n: (a, x) = r F32.max(best, x) case 1n+p: (a, x) = r Flat.max.go(p, (i + 1 : U32), Array.get(F32, a, (i + 1 : U32)), F32.max(best, x))def Flat.max.of(n: Nat, a: Array<F32>) -> F32: match n: case 0n: 0.0 case 1n+p: Flat.max.go(p, 0, Array.get(F32, a, 0), (0.0 - 3.4028235e38 : F32))def Flat.max(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> F32: FlatS{a} = m Flat.max.of(Flat.count(rows, cols), a)# Transpose# ---------def Flat.transpose.step2(+dst: U32, x: F32, a: Array<F32>, b: Array<F32>) -> Array<F32> & Array<F32>: (a, Array.set(F32, b, dst, x))def Flat.transpose.step(+dst: U32, b: Array<F32>, r: Array<F32> & F32) -> Array<F32> & Array<F32>: (a, x) = r Flat.transpose.step2(dst, x, a, b)# cell (i, j) of a at idx = i * cols + j goes to j * rows + i in bdef Flat.transpose.go(n: Nat, +idx: U32, +rowsu: U32, +colsu: U32, ab: Array<F32> & Array<F32>) -> Array<F32> & Array<F32>: match n: case 0n: ab case 1n+p: (a, b) = ab +i = U32.div(idx, colsu) +j = U32.mod(idx, colsu) Flat.transpose.go(p, (idx + 1 : U32), rowsu, colsu, Flat.transpose.step((j * rowsu + i : U32), b, Array.get(F32, a, idx)))def Flat.transpose.fin(-rows: Nat, -cols: Nat, ab: Array<F32> & Array<F32>) -> Flat<cols, rows>: (a, b) = ab FlatS{b}def Flat.transpose(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Flat<cols, rows>: FlatS{a} = m Flat.transpose.fin(rows, cols, Flat.transpose.go(Flat.count(rows, cols), 0, U32.from_nat(rows), U32.from_nat(cols), (a, Array.new(F32, Flat.depth(Flat.count(rows, cols)), 0.0))))# Products# --------# Flat<rows, k> times Flat<k, cols> is Flat<rows, cols>. The _t forms# take the right operand transposed as stored, a Flat<cols, k> (q times# k^T, with no transpose copy); the _w forms take it as a Weights, never# consumed. Array operands are consumed: clone what you keep, or load# it as a Weights. The GPU build adds the _gpu forms of all four.def Flat.matmul.fin(-rows: Nat, -cols: Nat, c: Array<F32>) -> IO(Flat<rows, cols>): IO.pure(Flat<rows, cols>, FlatS{c})def Flat.matmul.run(+rows: Nat, +k: Nat, +cols: Nat, ta: U32, tb: U32, x: Array<F32>, y: Array<F32>) -> IO(Flat<rows, cols>): IO.bind(Array<F32>, Flat<rows, cols>, IO.try(Array<F32>, Blas.sgemm(U32.from_nat(rows), U32.from_nat(cols), U32.from_nat(k), ta, tb, x, y)), c => Flat.matmul.fin(rows, cols, c))def Flat.matmul(+rows: Nat, +k: Nat, +cols: Nat, a: Flat<rows, k>, b: Flat<k, cols>) -> IO(Flat<rows, cols>): FlatS{x} = a FlatS{y} = b Flat.matmul.run(rows, k, cols, 0, 0, x, y)def Flat.matmul_t(+rows: Nat, +k: Nat, +cols: Nat, a: Flat<rows, k>, b: Flat<cols, k>) -> IO(Flat<rows, cols>): FlatS{x} = a FlatS{y} = b Flat.matmul.run(rows, k, cols, 0, 1, x, y)def Flat.matmul_w.run(+rows: Nat, +k: Nat, +cols: Nat, ta: U32, tb: U32, x: Array<F32>, id: U32) -> IO(Flat<rows, cols>): IO.bind(Array<F32>, Flat<rows, cols>, IO.try(Array<F32>, Blas.sgemm_w(U32.from_nat(rows), U32.from_nat(cols), U32.from_nat(k), ta, tb, x, id)), c => Flat.matmul.fin(rows, cols, c))def Flat.matmul_w(+rows: Nat, +k: Nat, +cols: Nat, a: Flat<rows, k>, w: Weights<k, cols>) -> IO(Flat<rows, cols>): FlatS{x} = a W{id} = w Flat.matmul_w.run(rows, k, cols, 0, 0, x, id)def Flat.matmul_wt(+rows: Nat, +k: Nat, +cols: Nat, a: Flat<rows, k>, w: Weights<cols, k>) -> IO(Flat<rows, cols>): FlatS{x} = a W{id} = w Flat.matmul_w.run(rows, k, cols, 0, 1, x, id)# Bridges# -------def Flat.to_mat(+rows: Nat, +cols: Nat, m: Flat<rows, cols>) -> Mat<rows, cols>: Mat.from_list(rows, cols, Flat.to_list(rows, cols, m))def Mat.to_flat(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Flat<rows, cols>: Flat.from_list(rows, cols, Mat.to_list(rows, cols, m))