def single_fiber

bendtorch-core@0.1.0.0 · src/autograd/reduce_parallel_proof.bend:390

import Base
import bendtorch-core@0.1.0.0/src/tensor/parallel_proofs.bend as Generic
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce

390# Rows of span 1 read one word: the one at `single_offset`.
391def single_fiber(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, +single: {Reduce.single_rows(rows, extents) == True{} : Bool},
392  +target: Nat, +base: Nat, +accumulator: F32, -array: Array<F32>) ->
393  Generic.Fact<{Reduce.fiber(rows, extents, Reduce.first_span(rows, extents), target, base, Reduce.first_digit(rows, target), accumulator, array) ==
394    Reduce.leaf_add(accumulator, Array.get(F32, array, U32.from_nat(Reduce.single_offset(rows, extents, target, base)))) : F32 & Array<F32>}>:
395  match rows extents:
396    case Nil{} +extents:
397      Generic.Fact{{==}}
398    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others:
399      +one = span_one(left, extent, Arith.and_left(Nat.is_eq(Reduce.span(left, extent), 1n), Reduce.single_rows(rest, others), single))
400      +inner = Nat.mod(target, left)
401      +offset = Nat.add(base, Nat.mul(right, Nat.div(target, left)))
402      Generic.Fact{
403        %Equal.sym(Nat, Reduce.span(left, extent), 1n, one) :
404          {Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, _, target, base, Nat.div(target, left), accumulator, array) ==
405            Reduce.leaf_add(accumulator, Array.get(F32, array, U32.from_nat(Reduce.single_offset(rest, others, inner, offset)))) : F32 & Array<F32>}
406        Equal.trans(F32 & Array<F32>,
407          Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, 0n, target, base, 1n+Nat.div(target, left), v, k),
408            Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array)),
409          Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, 0n, target, base, 1n+Nat.div(target, left), v, k),
410            Reduce.leaf_add(accumulator, Array.get(F32, array, U32.from_nat(Reduce.single_offset(rest, others, inner, offset))))),
411          Reduce.leaf_add(accumulator, Array.get(F32, array, U32.from_nat(Reduce.single_offset(rest, others, inner, offset)))),
412          Equal.cong(F32 & Array<F32>, F32 & Array<F32>,
413            p => Reduce.walk_next(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, 0n, target, base, 1n+Nat.div(target, left), v, k), p),
414            Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array),
415            Reduce.leaf_add(accumulator, Array.get(F32, array, U32.from_nat(Reduce.single_offset(rest, others, inner, offset)))),
416            {\{Generic.Fact: p => p} : Generic.Fact<({Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array) ==
417              Reduce.leaf_add(accumulator, Array.get(F32, array, U32.from_nat(Reduce.single_offset(rest, others, inner, offset)))) : F32 & Array<F32>})> ->
418              {Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, array) ==
419                Reduce.leaf_add(accumulator, Array.get(F32, array, U32.from_nat(Reduce.single_offset(rest, others, inner, offset)))) : F32 & Array<F32>}}(
420              single_fiber(rest, others, Arith.and_right(Nat.is_eq(Reduce.span(left, extent), 1n), Reduce.single_rows(rest, others), single),
421                inner, offset, accumulator, array))),
422          leaf_kept(v => k => Reduce.fiber(Shape.Row{stride, left, right} <> rest, extent <> others, 0n, target, base, 1n+Nat.div(target, left), v, k),
423            value => kept => {==}, accumulator, Array.get(F32, array, U32.from_nat(Reduce.single_offset(rest, others, inner, offset)))))}
424    case Shape.Row{+stride, +left, +right} <> +rest Nil{}:
425      Empty.absurd(Generic.Fact<{Reduce.fiber(Shape.Row{stride, left, right} <> rest, Nil{}, Reduce.first_span(Shape.Row{stride, left, right} <> rest, Nil{}), target, base,
426        Reduce.first_digit(Shape.Row{stride, left, right} <> rest, target), accumulator, array) ==
427        Reduce.leaf_add(accumulator, Array.get(F32, array, U32.from_nat(Reduce.single_offset(Shape.Row{stride, left, right} <> rest, Nil{}, target, base)))) :
428        F32 & Array<F32>}>, Arith.clash(single))