def MatT.matmul.any

bend-tensors@0.0.0.2 · bend_tensors.bend:842

import Base

842# a * b for any shapes: MatT(r, k) times MatT(k, c) is MatT(r, c). A node
843# forks its whole burst of products at once (eight at a quadrant), then
844# joins them with a fork of adds. Both operands are shared through +,
845# never copied. Three depths ride every call, which is what the square
846# path below avoids.
847def MatT.matmul.any(+r: Nat, +k: Nat, +c: Nat, +a: MatT(r, k), +b: MatT(k, c)) -> MatT(r, c):
848  match r k c:
849    case 0n 0n 0n:
850      # the leaf goes through the square path on purpose: compiled
851      # without MatT.matmul.sq reachable, this recursion fails at runtime
852      # from depth 2 on (a bend 2.0.16 bug, see README); with it, it runs
853      MatT.matmul.sq(0n, a, b)
854    case 1n+p 0n 0n:
855      MatR{at, ab} = a
856      x y = MatT.matmul.any(p, 0n, 0n, at, b) MatT.matmul.any(p, 0n, 0n, ab, b)
857      MatR{x, y}
858    case 0n 0n 1n+s:
859      MatC{bl, br} = b
860      x y = MatT.matmul.any(0n, 0n, s, a, bl) MatT.matmul.any(0n, 0n, s, a, br)
861      MatC{x, y}
862    case 0n 1n+q 0n:
863      MatC{al, ar} = a
864      MatR{bt, bb} = b
865      x y = MatT.matmul.any(0n, q, 0n, al, bt) MatT.matmul.any(0n, q, 0n, ar, bb)
866      MatT.add(0n, 0n, x, y)
867    case 1n+p 1n+q 0n:
868      MatQ{a0, a1, a2, a3} = a
869      MatR{bt, bb} = b
870      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)
871      MatT.add2r(p, p0, q0, p1, q1)
872    case 0n 1n+q 1n+s:
873      MatC{al, ar} = a
874      MatQ{b0, b1, b2, b3} = b
875      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)
876      MatT.add2c(s, p0, q0, p1, q1)
877    case 1n+p 0n 1n+s:
878      MatR{at, ab} = a
879      MatC{bl, br} = b
880      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)
881      MatQ{w, x, y, z}
882    case 1n+p 1n+q 1n+s:
883      MatQ{a0, a1, a2, a3} = a
884      MatQ{b0, b1, b2, b3} = b
885      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)
886      MatT.add4(p, s, p0, q0, p1, q1, p2, q2, p3, q3)