def product_splits

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

import Base

345# A product forks as `splits` says, but at most four levels, sixteen
346# subtrees: each fork copies both operands into each half, and with a row
347# of eight cells running as one loop a subtree's work no longer hides its
348# copies. On medusa2 four levels ran fastest at 8 threads for the products
349# of an MNIST step and for 256 x 256 x 256 (the MNIST step 15.2 ms against
350# 17.5 at six levels) and at one thread (20.1 ms against 21.6).
351def product_splits(+depth: Nat, +inner: Nat) -> Nat:
352  at_most(splits(depth, inner), 4n)