# 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 + 3
def 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 them
def 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 ties
def 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 rest
def 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 v
def 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 sums
def 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 transpose
def 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 row
def 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 row
def 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_i
def 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 acc
def 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 rows
def 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.0
def 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.0
def 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 elements
def 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 elements
def 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 elements
def 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.0
def 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.0
def 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.0
def 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 columns
def 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
  Data

type 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^n
def 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 product
def 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 element
def 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 half
def 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 size
def 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 ties
def 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.0
def 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) / sum
def 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
  Data

type 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.matmul
def 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 fork
def 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 node
def 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 cell
def 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 row
def 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 row
def 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) / sum
def 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 row
def 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 w
def 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 ys
def 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.0
def 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.0
def 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-major
def 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 w
def 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 order
def 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 order
def 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) steps
def 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.0
def 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 m
def 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 m
def 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 ties
def 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.0
def 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.0
def 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 - 1
def 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.0
def 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 model
def 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 elements
def 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.0
def 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 order
def 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.matmul
def 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 node
def 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 columns
def 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 about
def 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 _ _:
      Unit

def 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 not
def 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 array
def 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 array
def 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 steps
def 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.0
def 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-major
def 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 consumed
def 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 last
def 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
  a

def Flat.clone.snd(+rows: Nat, +cols: Nat, r: Flat<rows, cols> & Flat<rows, cols>) -> Flat<rows, cols>:
  (a, b) = r
  b

def 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 consumed
def 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.matmul
def 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 - inf
def 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 row
def 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 consumed
def 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 segment
def 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 answered
def 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, scale
def 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 row
def 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 write
def 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 ties
def 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 j
def 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 done
def 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 consumed
def 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 b
def 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))