def dstep

bend-ml-tensor-array@0.1.19.0 · main.bend:46

import Base

46# one step of dotr: adds the pending product to a0, rotates the sums, reads the next pair at ia, ib
47def dstep(+ia: U32, +ib: U32, s: D4) -> D4:
48  match s:
49    case D4{a0, a1, a2, a3, ra, rb}:
50      match ra rb:
51        case Tuple{a, +x} Tuple{b, +y}:
52          D4{a1, a2, a3, (a0 + (x * y : F32) : F32), Array.get(F32, a, ia), Array.get(F32, b, ib)}