def splits

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

import Base

336# Split by work, not by cell count: a cell costs one term per inner index.
337# The build forks until a subtree holds about 2^13 terms, and at most six
338# levels, 64 subtrees, each with its own copy of both inputs: an output of
339# 2^depth cells of `inner` terms forks depth + log2(inner) - 13 levels. On
340# medusa2 at 8 threads a product of 2^15 terms or more ran faster with every
341# extra level up to six, and forking cost nothing measurable on one thread.
342def splits(+depth: Nat, +inner: Nat) -> Nat:
343  at_most(at_most(Nat.sub(Nat.add(depth, halvings(inner, inner)), 13n), 6n), depth)