def fiber_fast

bendtorch-core@0.1.0.0 · src/autograd/reduce.bend:138

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape

138# `fiber`, the same words added in the same order, except that a row whose
139# rows below all take one digit (`single`, which is `single_below` of the
140# rows; a later digit of the same row passes False) walks its digits in one
141# strided loop: below digit d it reads the one word at base + right * d +
142# single_offset(rest, others, target % left, 0), with no closure per word.
143# Column sums, row sums and batch sums over one axis take that loop.
144def fiber_fast(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, fuel: Nat, single: Bool, +target: Nat, +base: Nat, +digit: Nat,
145  accumulator: F32, array: Array<F32>) -> F32 & Array<F32>:
146  match rows extents fuel single:
147    case Nil{} _ _ single:
148      leaf_add(accumulator, Array.get(F32, array, U32.from_nat(base)))
149    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others fuel True{}:
150      strided(fuel, right, base, digit, single_offset(rest, others, Nat.mod(target, left), 0n), accumulator, array)
151    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 0n False{}:
152      (accumulator, array)
153    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 1n+remaining False{}:
154      +inner = Nat.mod(target, left)
155      walk_next(value => kept => fiber_fast(rows, extents, remaining, False{}, target, base, 1n+digit, value, kept),
156        fiber_fast(rest, others, first_span(rest, others), single_below(rest, others), inner, step_offset(base, right, digit),
157          first_digit(rest, inner), accumulator, array))
158    case Shape.Row{+stride, +left, +right} <> +rest Nil{} _ single:
159      (accumulator, array)