def dot4

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

import bend-ml-tensor-array@0.1.19.0/main.bend as BendMlTensorArray
import Base

59# dotr with its loop unrolled by four: q rounds of four dstep, then the last rem (< 4) steps by dotr itself.
60# Every product goes to the same sum in the same order as in dotr (four rotations bring the names back), so the
61# result has dotr's bits (tests/dot4.bend). The compiled loop runs ~16% fewer instructions; one thread on a
62# 50,257 x 768 matrix · vector goes from 26.7 to 20.0 ms, smaller ones and 16 threads stay within noise (NOTES.md,
63# 2026-11-01). Eight steps a round were no faster.
64def dot4(q: Nat, +rem: Nat, +ia: U32, +ib: U32, +sa: U32, +sb: U32, s: D4) -> R:
65  match q:
66    case 0n:
67      dfin(rem, ia, ib, sa, sb, s)
68    case 1n+p:
69      dot4(p, rem, (ia + (4 * sa : U32) : U32), (ib + (4 * sb : U32) : U32), sa, sb, dstep((ia + (4 * sa : U32) : U32), (ib + (4 * sb : U32) : U32), dstep((ia + (3 * sa : U32) : U32), (ib + (3 * sb : U32) : U32), dstep((ia + (2 * sa : U32) : U32), (ib + (2 * sb : U32) : U32), dstep((ia + sa : U32), (ib + sb : U32), s)))))

Uses (8)

Used by (1)