def product_run

bendtorch-core@0.1.0.0 · src/tensor/matmul.bend:250

import Base

250# The fuel + 1 terms of a cell from the one whose words sit at offsets `lp`
251# and `rp`, added to `accumulator`; `left` and `right` are their reads,
252# already made, and each later term sits `ls` and `rs` further on. Each step
253# reads the next term's words, so the loop is one tail call with no closure
254# and one addition per offset: it compiles to a flat loop.
255def product_run(fuel: Nat, +ls: Nat, +rs: Nat, +lp: Nat, +rp: Nat, accumulator: F32,
256  left: Array<F32> & F32, right: Array<F32> & F32) -> F32 & (Array<F32> & Array<F32>):
257  match fuel left right:
258    case 0n Tuple{la, x} Tuple{ra, y}:
259      (F32.add(accumulator, F32.mul(x, y)), (la, ra))
260    case 1n+rest Tuple{la, x} Tuple{ra, y}:
261      product_run(rest, ls, rs, Nat.add(ls, lp), Nat.add(rs, rp), F32.add(accumulator, F32.mul(x, y)),
262        Array.get(F32, la, U32.from_nat(Nat.add(ls, lp))), Array.get(F32, ra, U32.from_nat(Nat.add(rs, rp))))