1 import Base 2 import ../../shape/core.bend as Shape 3 import ../../storage/model.bend as Model 4 import ../../storage/main.bend as Storage 5 import ../../storage/refinement.bend as Refinement 6 import ../../storage/clone.bend as Clone 7 import ../../storage/api.bend as Api 8 import ../../storage/arith.bend as Arith 9 import ../elementwise.bend as Elementwise 10 import ../parallel.bend as Parallel 11 import ../elementwise_walk_proof.bend as WalkProof 12 import ../word_conversion.bend as Conversion 13 import ../read_contract.bend as Reads 14 import ./replayed.bend as Replayed 15 import ./broadcast.bend as Broadcast 16 import ./broadcast_proof.bend as BroadcastProof 17 18 # The parallel elementwise build, walked with live code. The two input 19 # arrays are `Api.arr` of a history of storage effects, so every array the 20 # build reads or clones is a term of reusable data, and each step of the proof 21 # can name it again. `spec` is the balanced tree the build must return: cell 22 # `base` holds `scalar(operation, left word, right word)` below `count` and 23 # `+0.0` above it. 24 25 def left_index(+plan: Shape.BroadcastPlan, +position: Nat) -> U32: 26 U32.from_nat(Elementwise.left_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), Shape.index_map(plan), position)) 27 28 def right_index(+plan: Shape.BroadcastPlan, +position: Nat) -> U32: 29 U32.from_nat(Elementwise.right_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), Shape.index_map(plan), position)) 30 31 def word(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +count: Nat, +ha: List<&2, Api.Effect>, +dl: Nat, +fa: F32, +nl: Nat, +hb: List<&2, Api.Effect>, +dr: Nat, +fb: F32, +nr: Nat, +position: Nat) -> F32: 32 Elementwise.scalar(operation, Api.current(ha, left_index(plan, position), fa), Api.current(hb, right_index(plan, position), fb)) 33 34 def leaf(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +count: Nat, +ha: List<&2, Api.Effect>, +dl: Nat, +fa: F32, +nl: Nat, +hb: List<&2, Api.Effect>, +dr: Nat, +fb: F32, +nr: Nat, +position: Nat, inside: Bool) -> F32: 35 match inside: 36 case False{}: 37 0.0 38 case True{}: 39 word(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, position) 40 41 def spec(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +count: Nat, +ha: List<&2, Api.Effect>, +dl: Nat, +fa: F32, +nl: Nat, +hb: List<&2, Api.Effect>, +dr: Nat, +fb: F32, +nr: Nat, +depth: Nat, +base: Nat) -> Array: 42 match depth: 43 case 0n: 44 ALeaf{leaf(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, base, Nat.is_lt(base, count))} 45 case 1n+inner: 46 ANode{spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, base), spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, Nat.add(Model.capacity(inner), base))} 47 48 def spec_balanced(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +count: Nat, +ha: List<&2, Api.Effect>, +dl: Nat, +fa: F32, +nl: Nat, +hb: List<&2, Api.Effect>, +dr: Nat, +fb: F32, +nr: Nat, +depth: Nat, +base: Nat) -> Model.Balanced(F32, depth, spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, depth, base)): 49 match depth: 50 case 0n: 51 Unit{} 52 case 1n+inner: 53 (spec_balanced(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, base), spec_balanced(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, Nat.add(Model.capacity(inner), base))) 54 55 def TreeEq(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +count: Nat, +ha: List<&2, Api.Effect>, +dl: Nat, +fa: F32, +nl: Nat, +hb: List<&2, Api.Effect>, +dr: Nat, +fb: F32, +nr: Nat, +depth: Nat, +base: Nat, -result: Array & (Array & Array)) -> Type: 56 {Pair.fst(Array, Array & Array, result) == spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, depth, base) : Array} & 57 ({Pair.fst(Array, Array, Pair.snd(Array, Array & Array, result)) == Api.arr(ha, dl, fa) : Array} & 58 {Pair.snd(Array, Array, Pair.snd(Array, Array & Array, result)) == Api.arr(hb, dr, fb) : Array}) 59 60 # Every write in the history lands below `n`: `Api.Valid` as a Bool, so a 61 # proof of it can be used many times. 62 def written(+n: Nat, +history: List<&2, Api.Effect>) -> Bool: 63 match history: 64 case Nil{}: 65 True{} 66 case +effect <> +past: 67 match effect: 68 case Api.Wrote{+i, w}: 69 Bool.and(Nat.is_lt(U32.to_nat(i), n), written(n, past)) 70 case Api.Kept{i}: 71 written(n, past) 72 case Api.Left{}: 73 written(n, past) 74 case Api.Right{}: 75 written(n, past) 76 77 def valid_of(+n: Nat, +history: List<&2, Api.Effect>, +all: {written(n, history) == True{} : Bool}) -> Api.Valid(n, history): 78 match history: 79 case Nil{}: 80 Unit{} 81 case +effect <> +past: 82 match effect: 83 case Api.Wrote{+i, +w}: 84 (Arith.and_left(Nat.is_lt(U32.to_nat(i), n), written(n, past), all), 85 valid_of(n, past, Arith.and_right(Nat.is_lt(U32.to_nat(i), n), written(n, past), all))) 86 case Api.Kept{+i}: 87 valid_of(n, past, all) 88 case Api.Left{}: 89 valid_of(n, past, all) 90 case Api.Right{}: 91 valid_of(n, past, all) 92 93 # A read of `arr(h)` keeps the array and returns the word `current` names. 94 def got_kept(+d: Nat, +h: List<&2, Api.Effect>, +f: F32, +i: U32, -pair: Array & F32, 95 exact: {pair == Array.get(F32, Api.arr(h, d, f), i) : Array & F32}) -> 96 {Refinement.first(F32, pair) == Api.arr(h, d, f) : Array}: 97 Equal.trans(Array, Refinement.first(F32, pair), Refinement.first(F32, Array.get(F32, Api.arr(h, d, f), i)), Api.arr(h, d, f), 98 Equal.cong(Array & F32, Array, p => Refinement.first(F32, p), pair, Array.get(F32, Api.arr(h, d, f), i), exact), 99 Refinement.law_get_retains(F32, d, Api.arr(h, d, f), i, Api.bal(h, d, f))) 100 101 def some_word(-word: F32, -other: F32, same: {Some{word} == Some{other} : Maybe<&2, F32>}) -> {word == other : F32}: 102 Equal.cong(Maybe<&2, F32>, F32, found => Maybe.default(&2, F32, found, word), Some{word}, Some{other}, same) 103 104 def got_value(+d: Nat, +h: List<&2, Api.Effect>, +f: F32, +n: Nat, +i: U32, -pair: Array & F32, 105 exact: {pair == Array.get(F32, Api.arr(h, d, f), i) : Array & F32}, 106 +fits: {Storage.bounds(d, n) == True{} : Bool}, +inside: {Nat.is_lt(U32.to_nat(i), n) == True{} : Bool}, +valid: {written(n, h) == True{} : Bool}) -> 107 {Refinement.second(F32, pair) == Api.current(h, i, f) : F32}: 108 some_word(Refinement.second(F32, pair), Api.current(h, i, f), 109 Equal.trans(Maybe<&2, F32>, Some{Refinement.second(F32, pair)}, Some{Refinement.second(F32, Array.get(F32, Api.arr(h, d, f), i))}, 110 Some{Api.current(h, i, f)}, 111 Equal.cong(Array & F32, Maybe<&2, F32>, p => Some{Refinement.second(F32, p)}, pair, Array.get(F32, Api.arr(h, d, f), i), exact), 112 Equal.trans(Maybe<&2, F32>, Some{Refinement.second(F32, Array.get(F32, Api.arr(h, d, f), i))}, 113 Model.observe(F32, d, Api.arr(h, d, f), Refinement.route(d, i)), Some{Api.current(h, i, f)}, 114 Refinement.law_get_refines(F32, d, Api.arr(h, d, f), i, Api.bal(h, d, f)), 115 Equal.sym(Maybe<&2, F32>, Some{Api.current(h, i, f)}, Model.observe(F32, d, Api.arr(h, d, f), Refinement.route(d, i)), 116 Api.model(h, d, n, f, i, fits, inside, valid_of(n, h, valid)))))) 117 118 def walk_left_same(-left: List<&2, Nat>, -right: List<&2, Nat>, -output: List<&2, Nat>, 119 +map: Shape.Index, +index: Nat) -> 120 {Elementwise.left_walk(left, right, output, map, index) == Broadcast.left_walk(left, right, output, map, index) : Nat}: 121 match map: 122 case Shape.Indexed{rows, exact}: 123 {==} 124 case Shape.Unindexed{zero}: 125 {==} 126 127 def walk_right_same(-left: List<&2, Nat>, -right: List<&2, Nat>, -output: List<&2, Nat>, 128 +map: Shape.Index, +index: Nat) -> 129 {Elementwise.right_walk(left, right, output, map, index) == Broadcast.right_walk(left, right, output, map, index) : Nat}: 130 match map: 131 case Shape.Indexed{rows, exact}: 132 {==} 133 case Shape.Unindexed{zero}: 134 {==} 135 136 def output_below(+plan: Shape.BroadcastPlan, +count: Nat, +position: Nat, 137 +count_exact: {count == Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) : Nat}, 138 +decided: {True{} == Nat.is_lt(position, count) : Bool}) -> 139 {Nat.is_lt(position, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}: 140 %count_exact : {Nat.is_lt(position, _) == True{} : Bool} 141 Equal.sym(Bool, True{}, Nat.is_lt(position, count), decided) 142 143 # A walked index below the side's element count is a word index below the 144 # side's storage length. 145 def word_below(+side: Shape.Shape, +n: Nat, +walk: Nat, 146 +side_count: {n == Shape.logical_numel(Shape.dimensions(side)) : Nat}, 147 +below: {Nat.is_lt(walk, Shape.logical_numel(Shape.dimensions(side))) == True{} : Bool}) -> 148 {Nat.is_lt(U32.to_nat(U32.from_nat(walk)), n) == True{} : Bool}: 149 +within = Conversion.lt_le(walk, U32.to_nat(Shape.count(side)), 150 %Equal.sym(Nat, U32.to_nat(Shape.count(side)), Shape.logical_numel(Shape.dimensions(side)), Reads.count_exact(side)) : 151 {Nat.is_lt(walk, _) == True{} : Bool} 152 below) 153 %Equal.sym(Nat, U32.to_nat(U32.from_nat(walk)), walk, Conversion.natural_word_exact(walk, Shape.count(side), within)) : 154 {Nat.is_lt(_, n) == True{} : Bool} 155 %Equal.sym(Nat, n, Shape.logical_numel(Shape.dimensions(side)), side_count) : {Nat.is_lt(walk, _) == True{} : Bool} 156 below 157 158 def left_below(+plan: Shape.BroadcastPlan, +count: Nat, +nl: Nat, +position: Nat, 159 +count_exact: {count == Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) : Nat}, 160 +left_count: {nl == Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan))) : Nat}, 161 +decided: {True{} == Nat.is_lt(position, count) : Bool}) -> 162 {Nat.is_lt(U32.to_nat(left_index(plan, position)), nl) == True{} : Bool}: 163 word_below(Shape.left_shape(plan), nl, Elementwise.left_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), Shape.index_map(plan), position), left_count, 164 %Equal.sym(Nat, Elementwise.left_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), Shape.index_map(plan), position), Broadcast.left_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), Shape.index_map(plan), position), walk_left_same(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), Shape.index_map(plan), position)) : {Nat.is_lt(_, Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan)))) == True{} : Bool} 165 BroadcastProof.left_walk_bound(plan, Shape.index_map(plan), position, output_below(plan, count, position, count_exact, decided))) 166 167 def right_below(+plan: Shape.BroadcastPlan, +count: Nat, +nr: Nat, +position: Nat, 168 +count_exact: {count == Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) : Nat}, 169 +right_count: {nr == Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) : Nat}, 170 +decided: {True{} == Nat.is_lt(position, count) : Bool}) -> 171 {Nat.is_lt(U32.to_nat(right_index(plan, position)), nr) == True{} : Bool}: 172 word_below(Shape.right_shape(plan), nr, Elementwise.right_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), Shape.index_map(plan), position), right_count, 173 %Equal.sym(Nat, Elementwise.right_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), Shape.index_map(plan), position), Broadcast.right_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), Shape.index_map(plan), position), walk_right_same(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), Shape.index_map(plan), position)) : {Nat.is_lt(_, Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))) == True{} : Bool} 174 BroadcastProof.right_walk_bound(plan, Shape.index_map(plan), position, output_below(plan, count, position, count_exact, decided))) 175 176 def cell_right_eq(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +count: Nat, +ha: List<&2, Api.Effect>, +dl: Nat, +fa: F32, +nl: Nat, +hb: List<&2, Api.Effect>, +dr: Nat, +fb: F32, +nr: Nat, +lfits: {Storage.bounds(dl, nl) == True{} : Bool}, +rfits: {Storage.bounds(dr, nr) == True{} : Bool}, +lvalid: {written(nl, ha) == True{} : Bool}, +rvalid: {written(nr, hb) == True{} : Bool}, +count_exact: {count == Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) : Nat}, +left_count: {nl == Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan))) : Nat}, +right_count: {nr == Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) : Nat}, +base: Nat, +decided: {True{} == Nat.is_lt(base, count) : Bool}, 177 -x: F32, -kept: Array, pair: Array & F32, 178 +exact: {pair == Array.get(F32, Api.arr(hb, dr, fb), right_index(plan, base)) : Array & F32}, 179 ex: {x == Api.current(ha, left_index(plan, base), fa) : F32}, ek: {kept == Api.arr(ha, dl, fa) : Array}) -> 180 TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 0n, base, Parallel.leaf(~(Array & Array), Elementwise.cell_right(operation, x, kept, pair))): 181 (other, y) = pair 182 +ey = got_value(dr, hb, fb, nr, right_index(plan, base), (other, y), exact, rfits, 183 right_below(plan, count, nr, base, count_exact, right_count, decided), rvalid) 184 (%decided : 185 {ALeaf{Elementwise.scalar(operation, x, y)} == ALeaf{leaf(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, base, _)} : Array} 186 %Equal.sym(F32, x, Api.current(ha, left_index(plan, base), fa), ex) : {ALeaf{Elementwise.scalar(operation, _, y)} == ALeaf{word(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, base)} : Array} 187 %Equal.sym(F32, y, Api.current(hb, right_index(plan, base), fb), ey) : {ALeaf{Elementwise.scalar(operation, Api.current(ha, left_index(plan, base), fa), _)} == ALeaf{word(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, base)} : Array} 188 {==}, 189 (ek, got_kept(dr, hb, fb, right_index(plan, base), (other, y), exact))) 190 191 def cell_left_eq(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +count: Nat, +ha: List<&2, Api.Effect>, +dl: Nat, +fa: F32, +nl: Nat, +hb: List<&2, Api.Effect>, +dr: Nat, +fb: F32, +nr: Nat, +lfits: {Storage.bounds(dl, nl) == True{} : Bool}, +rfits: {Storage.bounds(dr, nr) == True{} : Bool}, +lvalid: {written(nl, ha) == True{} : Bool}, +rvalid: {written(nr, hb) == True{} : Bool}, +count_exact: {count == Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) : Nat}, +left_count: {nl == Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan))) : Nat}, +right_count: {nr == Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) : Nat}, +base: Nat, +decided: {True{} == Nat.is_lt(base, count) : Bool}, 192 pair: Array & F32, +exact: {pair == Array.get(F32, Api.arr(ha, dl, fa), left_index(plan, base)) : Array & F32}) -> 193 TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 0n, base, Parallel.leaf(~(Array & Array), Elementwise.cell_left(operation, Api.arr(hb, dr, fb), right_index(plan, base), pair))): 194 (kept, x) = pair 195 cell_right_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, lfits, rfits, lvalid, rvalid, count_exact, left_count, right_count, base, decided, x, kept, Array.get(F32, Api.arr(hb, dr, fb), right_index(plan, base)), {==}, 196 got_value(dl, ha, fa, nl, left_index(plan, base), (kept, x), exact, lfits, 197 left_below(plan, count, nl, base, count_exact, left_count, decided), lvalid), 198 got_kept(dl, ha, fa, left_index(plan, base), (kept, x), exact)) 199 200 def cell_eq(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +count: Nat, +ha: List<&2, Api.Effect>, +dl: Nat, +fa: F32, +nl: Nat, +hb: List<&2, Api.Effect>, +dr: Nat, +fb: F32, +nr: Nat, +lfits: {Storage.bounds(dl, nl) == True{} : Bool}, +rfits: {Storage.bounds(dr, nr) == True{} : Bool}, +lvalid: {written(nl, ha) == True{} : Bool}, +rvalid: {written(nr, hb) == True{} : Bool}, +count_exact: {count == Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) : Nat}, +left_count: {nl == Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan))) : Nat}, +right_count: {nr == Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) : Nat}, +base: Nat, +inside: Bool, decided: {inside == Nat.is_lt(base, count) : Bool}) -> 201 TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 0n, base, Parallel.leaf(~(Array & Array), Elementwise.cell(operation, Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), Shape.index_map(plan), base, Api.arr(ha, dl, fa), Api.arr(hb, dr, fb), inside))): 202 match inside: 203 case False{}: 204 (%decided : {ALeaf{0.0} == ALeaf{leaf(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, base, _)} : Array} 205 {==}, ({==}, {==})) 206 case True{}: 207 cell_left_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, lfits, rfits, lvalid, rvalid, count_exact, left_count, right_count, base, decided, Array.get(F32, Api.arr(ha, dl, fa), left_index(plan, base)), {==}) 208 209 # A cell of the build, with or without the walk: `Elementwise.cell` of the plan. 210 def cell_at_eq(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +count: Nat, +ha: List<&2, Api.Effect>, +dl: Nat, +fa: F32, +nl: Nat, +hb: List<&2, Api.Effect>, +dr: Nat, +fb: F32, +nr: Nat, +lfits: {Storage.bounds(dl, nl) == True{} : Bool}, +rfits: {Storage.bounds(dr, nr) == True{} : Bool}, +lvalid: {written(nl, ha) == True{} : Bool}, +rvalid: {written(nr, hb) == True{} : Bool}, +count_exact: {count == Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) : Nat}, +left_count: {nl == Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan))) : Nat}, +right_count: {nr == Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) : Nat}, +base: Nat, +inside: Bool, decided: {inside == Nat.is_lt(base, count) : Bool}) -> 211 TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 0n, base, Parallel.leaf(~(Array & Array), Elementwise.cell_at(Elementwise.walk(Elementwise.Request{operation, plan}), base, (Api.arr(ha, dl, fa), Api.arr(hb, dr, fb)), inside))): 212 %Equal.sym(F32 & (Array & Array), Elementwise.cell_at(Elementwise.walk(Elementwise.Request{operation, plan}), base, (Api.arr(ha, dl, fa), Api.arr(hb, dr, fb)), inside), 213 Elementwise.cell(operation, Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), Shape.index_map(plan), base, Api.arr(ha, dl, fa), Api.arr(hb, dr, fb), inside), 214 WalkProof.cell_at_walk(operation, plan, base, Api.arr(ha, dl, fa), Api.arr(hb, dr, fb), inside)) : 215 TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 0n, base, Parallel.leaf(~(Array & Array), _)) 216 cell_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, lfits, rfits, lvalid, rvalid, count_exact, left_count, right_count, base, inside, decided) 217 218 def tree_join_eq(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +count: Nat, +ha: List<&2, Api.Effect>, +dl: Nat, +fa: F32, +nl: Nat, +hb: List<&2, Api.Effect>, +dr: Nat, +fb: F32, +nr: Nat, +inner: Nat, +base: Nat, -a: Array, high: Array & (Array & Array), 219 ea: {a == spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, base) : Array}, 220 facts: TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, Nat.add(Model.capacity(inner), base), high)) -> 221 TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 1n+inner, base, Parallel.join(~(Array & Array), a, high)): 222 (b, inputs) = high 223 (eb, kept) = facts 224 (%Equal.sym(Array, a, spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, base), ea) : {ANode{_, b} == ANode{spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, base), spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, Nat.add(Model.capacity(inner), base))} : Array} 225 %Equal.sym(Array, b, spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, Nat.add(Model.capacity(inner), base)), eb) : {ANode{spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, base), _} == ANode{spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, base), spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, Nat.add(Model.capacity(inner), base))} : Array} 226 {==}, kept) 227 228 def tree_node_eq(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +count: Nat, +ha: List<&2, Api.Effect>, +dl: Nat, +fa: F32, +nl: Nat, +hb: List<&2, Api.Effect>, +dr: Nat, +fb: F32, +nr: Nat, +inner: Nat, +base: Nat, low: Array & (Array & Array), facts: TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, base, low), 229 next: TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, Nat.add(Model.capacity(inner), base), Parallel.tree(~Elementwise.Walk, ~Array, ~Array, ~Elementwise.cell_at, Elementwise.walk(Elementwise.Request{operation, plan}), inner, count, Nat.add(Model.capacity(inner), base), (Api.arr(ha, dl, fa), Api.arr(hb, dr, fb))))) -> 230 TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 1n+inner, base, Parallel.node(~Array, ~Array, l => r => Parallel.tree(~Elementwise.Walk, ~Array, ~Array, ~Elementwise.cell_at, Elementwise.walk(Elementwise.Request{operation, plan}), inner, count, Nat.add(Model.capacity(inner), base), (l, r)), low)): 231 (a, inputs) = low 232 (la, ra) = inputs 233 (ea, kept) = facts 234 (el, er) = kept 235 %Equal.sym(Array, la, Api.arr(ha, dl, fa), el) : 236 TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 1n+inner, base, Parallel.join(~(Array & Array), a, Parallel.tree(~Elementwise.Walk, ~Array, ~Array, ~Elementwise.cell_at, Elementwise.walk(Elementwise.Request{operation, plan}), inner, count, Nat.add(Model.capacity(inner), base), (_, ra)))) 237 %Equal.sym(Array, ra, Api.arr(hb, dr, fb), er) : 238 TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 1n+inner, base, Parallel.join(~(Array & Array), a, Parallel.tree(~Elementwise.Walk, ~Array, ~Array, ~Elementwise.cell_at, Elementwise.walk(Elementwise.Request{operation, plan}), inner, count, Nat.add(Model.capacity(inner), base), (Api.arr(ha, dl, fa), _)))) 239 tree_join_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, base, a, Parallel.tree(~Elementwise.Walk, ~Array, ~Array, ~Elementwise.cell_at, Elementwise.walk(Elementwise.Request{operation, plan}), inner, count, Nat.add(Model.capacity(inner), base), (Api.arr(ha, dl, fa), Api.arr(hb, dr, fb))), ea, next) 240 241 # The depth-1 tree: the leaf of cell `base` and the leaf of the cell after 242 # it, which reads the arrays the first hands back. 243 def pair_join_eq(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +count: Nat, +ha: List<&2, Api.Effect>, +dl: Nat, +fa: F32, +nl: Nat, +hb: List<&2, Api.Effect>, +dr: Nat, +fb: F32, +nr: Nat, +base: Nat, x: F32, high: F32 & (Array & Array), 244 ex: {ALeaf{x} == spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 0n, base) : Array}, 245 leaf_facts: TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 0n, Nat.add(1n, base), Parallel.leaf(~(Array & Array), high))) -> 246 TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 1n, base, Parallel.pair_join(~(Array & Array), x, high)): 247 (y, inputs) = high 248 (ey, kept) = leaf_facts 249 (%Equal.sym(Array, ALeaf{x}, spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 0n, base), ex) : {ANode{_, ALeaf{y}} == ANode{spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 0n, base), spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 0n, Nat.add(1n, base))} : Array} 250 %Equal.sym(Array, ALeaf{y}, spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 0n, Nat.add(1n, base)), ey) : {ANode{spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 0n, base), _} == ANode{spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 0n, base), spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 0n, Nat.add(1n, base))} : Array} 251 {==}, kept) 252 253 def pair_eq(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +count: Nat, +ha: List<&2, Api.Effect>, +dl: Nat, +fa: F32, +nl: Nat, +hb: List<&2, Api.Effect>, +dr: Nat, +fb: F32, +nr: Nat, +lfits: {Storage.bounds(dl, nl) == True{} : Bool}, +rfits: {Storage.bounds(dr, nr) == True{} : Bool}, +lvalid: {written(nl, ha) == True{} : Bool}, +rvalid: {written(nr, hb) == True{} : Bool}, +count_exact: {count == Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) : Nat}, +left_count: {nl == Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan))) : Nat}, +right_count: {nr == Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) : Nat}, +base: Nat, low: F32 & (Array & Array), 254 leaf_facts: TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 0n, base, Parallel.leaf(~(Array & Array), low))) -> 255 TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 1n, base, Parallel.pair(~Elementwise.Walk, ~Array, ~Array, ~Elementwise.cell_at, Elementwise.walk(Elementwise.Request{operation, plan}), count, base, low)): 256 (x, inputs) = low 257 (la, ra) = inputs 258 (ex, kept) = leaf_facts 259 (el, er) = kept 260 %Equal.sym(Array, la, Api.arr(ha, dl, fa), el) : 261 TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 1n, base, Parallel.pair_join(~(Array & Array), x, Elementwise.cell_at(Elementwise.walk(Elementwise.Request{operation, plan}), Nat.add(1n, base), (_, ra), Nat.is_lt(Nat.add(1n, base), count)))) 262 %Equal.sym(Array, ra, Api.arr(hb, dr, fb), er) : 263 TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 1n, base, Parallel.pair_join(~(Array & Array), x, Elementwise.cell_at(Elementwise.walk(Elementwise.Request{operation, plan}), Nat.add(1n, base), (Api.arr(ha, dl, fa), _), Nat.is_lt(Nat.add(1n, base), count)))) 264 pair_join_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, base, x, Elementwise.cell_at(Elementwise.walk(Elementwise.Request{operation, plan}), Nat.add(1n, base), (Api.arr(ha, dl, fa), Api.arr(hb, dr, fb)), Nat.is_lt(Nat.add(1n, base), count)), ex, 265 cell_at_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, lfits, rfits, lvalid, rvalid, count_exact, left_count, right_count, Nat.add(1n, base), Nat.is_lt(Nat.add(1n, base), count), {==})) 266 267 def tree_eq(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +count: Nat, +ha: List<&2, Api.Effect>, +dl: Nat, +fa: F32, +nl: Nat, +hb: List<&2, Api.Effect>, +dr: Nat, +fb: F32, +nr: Nat, +lfits: {Storage.bounds(dl, nl) == True{} : Bool}, +rfits: {Storage.bounds(dr, nr) == True{} : Bool}, +lvalid: {written(nl, ha) == True{} : Bool}, +rvalid: {written(nr, hb) == True{} : Bool}, +count_exact: {count == Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) : Nat}, +left_count: {nl == Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan))) : Nat}, +right_count: {nr == Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) : Nat}, +depth: Nat, +base: Nat) -> 268 TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, depth, base, Parallel.tree(~Elementwise.Walk, ~Array, ~Array, ~Elementwise.cell_at, Elementwise.walk(Elementwise.Request{operation, plan}), depth, count, base, (Api.arr(ha, dl, fa), Api.arr(hb, dr, fb)))): 269 match depth: 270 case 0n: 271 cell_at_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, lfits, rfits, lvalid, rvalid, count_exact, left_count, right_count, base, Nat.is_lt(base, count), {==}) 272 case 1n+inner: 273 match inner: 274 case 0n: 275 pair_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, lfits, rfits, lvalid, rvalid, count_exact, left_count, right_count, base, Elementwise.cell_at(Elementwise.walk(Elementwise.Request{operation, plan}), base, (Api.arr(ha, dl, fa), Api.arr(hb, dr, fb)), Nat.is_lt(base, count)), 276 cell_at_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, lfits, rfits, lvalid, rvalid, count_exact, left_count, right_count, base, Nat.is_lt(base, count), {==})) 277 case 1n+below: 278 tree_node_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, base, Parallel.tree(~Elementwise.Walk, ~Array, ~Array, ~Elementwise.cell_at, Elementwise.walk(Elementwise.Request{operation, plan}), inner, count, base, (Api.arr(ha, dl, fa), Api.arr(hb, dr, fb))), tree_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, lfits, rfits, lvalid, rvalid, count_exact, left_count, right_count, inner, base), 279 tree_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, lfits, rfits, lvalid, rvalid, count_exact, left_count, right_count, inner, Nat.add(Model.capacity(inner), base))) 280 281 def tree_first(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +count: Nat, +ha: List<&2, Api.Effect>, +dl: Nat, +fa: F32, +nl: Nat, +hb: List<&2, Api.Effect>, +dr: Nat, +fb: F32, +nr: Nat, +depth: Nat, +base: Nat, -result: Array & (Array & Array), facts: TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, depth, base, result)) -> 282 {Pair.fst(Array, Array & Array, result) == spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, depth, base) : Array}: 283 (first, kept) = facts 284 first 285 286 def clone_first_eq(+d: Nat, +h: List<&2, Api.Effect>, +f: F32, -pair: Array & Array, 287 exact: {pair == Array.clone(F32, Api.arr(h, d, f)) : Array & Array}) -> 288 {Pair.fst(Array, Array, pair) == Api.arr(h, d, f) : Array}: 289 Equal.trans(Array, Pair.fst(Array, Array, pair), Pair.fst(Array, Array, Array.clone(F32, Api.arr(h, d, f))), Api.arr(h, d, f), 290 Equal.cong(Array & Array, Array, p => Pair.fst(Array, Array, p), pair, Array.clone(F32, Api.arr(h, d, f)), exact), 291 Clone.first_equal(F32, d, Api.arr(h, d, f), Clone.certify(F32, d, Api.arr(h, d, f), Api.bal(h, d, f)))) 292 293 def clone_second_eq(+d: Nat, +h: List<&2, Api.Effect>, +f: F32, -pair: Array & Array, 294 exact: {pair == Array.clone(F32, Api.arr(h, d, f)) : Array & Array}) -> 295 {Pair.snd(Array, Array, pair) == Api.arr(h, d, f) : Array}: 296 Equal.trans(Array, Pair.snd(Array, Array, pair), Pair.snd(Array, Array, Array.clone(F32, Api.arr(h, d, f))), Api.arr(h, d, f), 297 Equal.cong(Array & Array, Array, p => Pair.snd(Array, Array, p), pair, Array.clone(F32, Api.arr(h, d, f)), exact), 298 Clone.second_equal(F32, d, Api.arr(h, d, f), Clone.certify(F32, d, Api.arr(h, d, f), Api.bal(h, d, f)))) 299 300 def fork_eq(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +count: Nat, +ha: List<&2, Api.Effect>, +dl: Nat, +fa: F32, +nl: Nat, +hb: List<&2, Api.Effect>, +dr: Nat, +fb: F32, +nr: Nat, +inner: Nat, +rest: Nat, +base: Nat, 301 low: {Elementwise.build(inner, rest, operation, plan, count, base, Api.arr(ha, dl, fa), Api.arr(hb, dr, fb)) == spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, base) : Array}, 302 high: {Elementwise.build(inner, rest, operation, plan, count, Nat.add(Model.capacity(inner), base), Api.arr(ha, dl, fa), Api.arr(hb, dr, fb)) == spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, Nat.add(Model.capacity(inner), base)) : Array}, 303 lp: Array & Array, +le: {lp == Array.clone(F32, Api.arr(ha, dl, fa)) : Array & Array}, 304 rp: Array & Array, +re: {rp == Array.clone(F32, Api.arr(hb, dr, fb)) : Array & Array}) -> 305 {Parallel.fork(~Array, ~Array, Elementwise.halves(lp, rp), 306 l => r => Parallel.build(~Elementwise.Request, ~Elementwise.Walk, ~Array, ~Array, ~Elementwise.walk, ~Elementwise.split, ~Elementwise.cell_at, 307 Elementwise.Request{operation, plan}, inner, rest, count, base, (l, r)), 308 l => r => Parallel.build(~Elementwise.Request, ~Elementwise.Walk, ~Array, ~Array, ~Elementwise.walk, ~Elementwise.split, ~Elementwise.cell_at, 309 Elementwise.Request{operation, plan}, inner, rest, count, Nat.add(Model.capacity(inner), base), (l, r))) == 310 spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 1n+inner, base) : Array}: 311 (l1, l2) = lp 312 (r1, r2) = rp 313 +upper = Nat.add(Model.capacity(inner), base) 314 %Equal.sym(Array, l1, Api.arr(ha, dl, fa), clone_first_eq(dl, ha, fa, (l1, l2), le)) : 315 {ANode{Elementwise.build(inner, rest, operation, plan, count, base, _, r1), Elementwise.build(inner, rest, operation, plan, count, upper, l2, r2)} == spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 1n+inner, base) : Array} 316 %Equal.sym(Array, r1, Api.arr(hb, dr, fb), clone_first_eq(dr, hb, fb, (r1, r2), re)) : 317 {ANode{Elementwise.build(inner, rest, operation, plan, count, base, Api.arr(ha, dl, fa), _), Elementwise.build(inner, rest, operation, plan, count, upper, l2, r2)} == spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 1n+inner, base) : Array} 318 %Equal.sym(Array, l2, Api.arr(ha, dl, fa), clone_second_eq(dl, ha, fa, (l1, l2), le)) : 319 {ANode{Elementwise.build(inner, rest, operation, plan, count, base, Api.arr(ha, dl, fa), Api.arr(hb, dr, fb)), Elementwise.build(inner, rest, operation, plan, count, upper, _, r2)} == spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 1n+inner, base) : Array} 320 %Equal.sym(Array, r2, Api.arr(hb, dr, fb), clone_second_eq(dr, hb, fb, (r1, r2), re)) : 321 {ANode{Elementwise.build(inner, rest, operation, plan, count, base, Api.arr(ha, dl, fa), Api.arr(hb, dr, fb)), Elementwise.build(inner, rest, operation, plan, count, upper, Api.arr(ha, dl, fa), _)} == spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 1n+inner, base) : Array} 322 %Equal.sym(Array, Elementwise.build(inner, rest, operation, plan, count, base, Api.arr(ha, dl, fa), Api.arr(hb, dr, fb)), spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, base), low) : 323 {ANode{_, Elementwise.build(inner, rest, operation, plan, count, upper, Api.arr(ha, dl, fa), Api.arr(hb, dr, fb))} == spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 1n+inner, base) : Array} 324 %Equal.sym(Array, Elementwise.build(inner, rest, operation, plan, count, upper, Api.arr(ha, dl, fa), Api.arr(hb, dr, fb)), spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, upper), high) : 325 {ANode{spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, base), _} == spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 1n+inner, base) : Array} 326 {==} 327 328 # Whatever the split, the parallel build returns `spec`. 329 def build_eq(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +count: Nat, +ha: List<&2, Api.Effect>, +dl: Nat, +fa: F32, +nl: Nat, +hb: List<&2, Api.Effect>, +dr: Nat, +fb: F32, +nr: Nat, +lfits: {Storage.bounds(dl, nl) == True{} : Bool}, +rfits: {Storage.bounds(dr, nr) == True{} : Bool}, +lvalid: {written(nl, ha) == True{} : Bool}, +rvalid: {written(nr, hb) == True{} : Bool}, +count_exact: {count == Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) : Nat}, +left_count: {nl == Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan))) : Nat}, +right_count: {nr == Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) : Nat}, +depth: Nat, +splits: Nat, +base: Nat) -> 330 {Elementwise.build(depth, splits, operation, plan, count, base, Api.arr(ha, dl, fa), Api.arr(hb, dr, fb)) == spec(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, depth, base) : Array}: 331 match depth splits: 332 case 1n+inner 1n+rest: 333 fork_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, rest, base, build_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, lfits, rfits, lvalid, rvalid, count_exact, left_count, right_count, inner, rest, base), 334 build_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, lfits, rfits, lvalid, rvalid, count_exact, left_count, right_count, inner, rest, Nat.add(Model.capacity(inner), base)), Array.clone(F32, Api.arr(ha, dl, fa)), {==}, Array.clone(F32, Api.arr(hb, dr, fb)), {==}) 335 case 0n _: 336 tree_first(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 0n, base, Parallel.tree(~Elementwise.Walk, ~Array, ~Array, ~Elementwise.cell_at, Elementwise.walk(Elementwise.Request{operation, plan}), 0n, count, base, (Api.arr(ha, dl, fa), Api.arr(hb, dr, fb))), tree_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, lfits, rfits, lvalid, rvalid, count_exact, left_count, right_count, 0n, base)) 337 case 1n+inner 0n: 338 tree_first(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, 1n+inner, base, Parallel.tree(~Elementwise.Walk, ~Array, ~Array, ~Elementwise.cell_at, Elementwise.walk(Elementwise.Request{operation, plan}), 1n+inner, count, base, (Api.arr(ha, dl, fa), Api.arr(hb, dr, fb))), tree_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, lfits, rfits, lvalid, rvalid, count_exact, left_count, right_count, 1n+inner, base))