1 import Base 2 import ../shape/core.bend as Shape 3 import ../shape/LAWS.bend as ShapeLaws 4 import ../shape/PROOF.bend as ShapeProof 5 import ../storage/LAWS.bend as StorageLaws 6 import ../storage/arith.bend as Arith 7 import ../tensor/core.bend as Core 8 import ../tensor/lifecycle.bend as Lifecycle 9 import ../tensor/word_conversion.bend as Conversion 10 import ../tensor/projection_LAWS.bend as Projection 11 import ../tensor/contracts/replayed.bend as Replayed 12 import ../tensor/contracts/pointwise.bend as Pointwise 13 import ../tensor/contracts/pointwise_proof.bend as PointwiseProof 14 import ../tensor/contracts/broadcast.bend as Broadcast 15 import ../tensor/contracts/broadcast_proof.bend as BroadcastProof 16 import ../tensor/contracts/fold.bend as Fold 17 import ../tensor/contracts/fold_records.bend as Records 18 import ../tensor/contracts/fold_proof.bend as FoldProof 19 import ../tensor/contracts/sum.bend as Sum 20 import ../tensor/contracts/LAWS.bend as Contracts 21 import ../tensor/contracts/PROOF.bend as ContractsProof 22 import ./reduce_LAWS.bend as ReduceLaws 23 import ./reduce_fold.bend as ReduceFold 24 import ./reduce_fold_LAWS.bend as Laws 25 26 def offsets_walk(+plan: Shape.BroadcastPlan, 27 +map: Shape.Index, 28 +position: Nat, +bound: {True{} == Nat.is_lt(position, Shape.elements(Shape.output_shape(plan))) : Bool}) -> 29 {Projection.offset_spec(plan, position) == 30 Some{(U32.from_nat(Broadcast.left_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), 31 Shape.extents(Shape.output_shape(plan)), map, position)), 32 U32.from_nat(Broadcast.right_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), 33 Shape.extents(Shape.output_shape(plan)), map, position)))} : Maybe<&1, U32 & U32>}: 34 match map: 35 case Shape.Indexed{+rows, +exact}: 36 +left = Shape.extents(Shape.left_shape(plan)) 37 +right = Shape.extents(Shape.right_shape(plan)) 38 +output = Shape.extents(Shape.output_shape(plan)) 39 +index = {Shape.FIndex{position, bound} : Shape.FlatIndex} 40 +slot = ShapeLaws.fin_of(Shape.numel(output), position) 41 +x = U32.from_nat(Shape.walk_left(rows, position)) 42 +y = U32.from_nat(Shape.walk_right(rows, position)) 43 %ShapeLaws.project_left_index_refines(plan, Shape.Indexed{rows, exact}, index) : 44 {Projection.pair_offsets(_, ShapeLaws.flatten_spec(right, ShapeLaws.project_right_coordinate(left, right, output, Shape.alignment(plan), slot))) == 45 Some{(x, y)} : Maybe<&1, U32 & U32>} 46 %ShapeLaws.project_right_index_refines(plan, Shape.Indexed{rows, exact}, index) : 47 {Projection.pair_offsets(Some{x}, _) == Some{(x, y)} : Maybe<&1, U32 & U32>} 48 {==} 49 case Shape.Unindexed{zero}: 50 +output = Shape.output_shape(plan) 51 Empty.absurd({Projection.offset_spec(plan, position) == Some{(U32.from_nat(0n), U32.from_nat(0n))} : Maybe<&1, U32 & U32>}, 52 Arith.lt_zero(position, 53 %Equal.trans(Nat, Shape.elements(output), Shape.numel(Shape.extents(output)), 0n, Shape.elements_exact(output), 54 Equal.sym(Nat, 0n, Shape.numel(Shape.extents(output)), zero)) : {Nat.is_lt(position, _) == True{} : Bool} 55 Equal.sym(Bool, True{}, Nat.is_lt(position, Shape.elements(output)), bound))) 56 57 def length_snoc(+values: List<&2, F32>, +word: F32) -> 58 {List.length(&2, F32, List.append(&2, F32, values, [word])) == 1n+List.length(&2, F32, values) : Nat}: 59 match values: 60 case Nil{}: 61 {==} 62 case +value <> +rest: 63 Equal.cong(Nat, Nat, n => 1n+n, List.length(&2, F32, List.append(&2, F32, rest, [word])), 1n+List.length(&2, F32, rest), 64 length_snoc(rest, word)) 65 66 def words_length(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat) -> 67 {List.length(&2, F32, Contracts.words(shape, fill, ops, count)) == count : Nat}: 68 match count: 69 case 0n: 70 {==} 71 case 1n+previous: 72 +before = Contracts.words(shape, fill, ops, previous) 73 Equal.trans(Nat, List.length(&2, F32, List.append(&2, F32, before, [Replayed.value(shape, fill, ops, U32.from_nat(previous))])), 74 1n+List.length(&2, F32, before), 1n+previous, 75 length_snoc(before, Replayed.value(shape, fill, ops, U32.from_nat(previous))), 76 Equal.cong(Nat, Nat, n => 1n+n, List.length(&2, F32, before), previous, words_length(shape, fill, ops, previous))) 77 78 def word_at_front(+values: List<&2, F32>, +word: F32, +position: Nat, 79 below: {Nat.is_lt(position, List.length(&2, F32, values)) == True{} : Bool}) -> 80 {ReduceLaws.word_at(List.append(&2, F32, values, [word]), position) == ReduceLaws.word_at(values, position) : Maybe<&2, F32>}: 81 match values position: 82 case Nil{} _: 83 Empty.absurd({ReduceLaws.word_at([word], position) == None{} : Maybe<&2, F32>}, Arith.lt_zero(position, below)) 84 case +value <> +rest 0n: 85 {==} 86 case +value <> +rest 1n+next: 87 word_at_front(rest, word, next, below) 88 89 def word_at_last(+values: List<&2, F32>, +word: F32) -> 90 {ReduceLaws.word_at(List.append(&2, F32, values, [word]), List.length(&2, F32, values)) == Some{word} : Maybe<&2, F32>}: 91 match values: 92 case Nil{}: 93 {==} 94 case +value <> +rest: 95 word_at_last(rest, word) 96 97 def words_step(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +previous: Nat, +position: Nat, 98 below: {Nat.is_lt(position, 1n+previous) == True{} : Bool}, decision: Shape.BoolCase, 99 earlier: @inside: {Nat.is_lt(position, previous) == True{} : Bool} -> 100 {ReduceLaws.word_at(Contracts.words(shape, fill, ops, previous), position) == 101 Some{Replayed.value(shape, fill, ops, U32.from_nat(position))} : Maybe<&2, F32>}) -> 102 {ReduceLaws.word_at(Contracts.words(shape, fill, ops, 1n+previous), position) == 103 Some{Replayed.value(shape, fill, ops, U32.from_nat(position))} : Maybe<&2, F32>}: 104 match decision: 105 case Shape.BFalse{equal}: 106 +before = Contracts.words(shape, fill, ops, previous) 107 +last = Replayed.value(shape, fill, ops, U32.from_nat(previous)) 108 +inside = PointwiseProof.lt_succ_ne(position, previous, below, equal) 109 Equal.trans(Maybe<&2, F32>, ReduceLaws.word_at(List.append(&2, F32, before, [last]), position), ReduceLaws.word_at(before, position), 110 Some{Replayed.value(shape, fill, ops, U32.from_nat(position))}, 111 word_at_front(before, last, position, 112 %Equal.sym(Nat, List.length(&2, F32, before), previous, words_length(shape, fill, ops, previous)) : 113 {Nat.is_lt(position, _) == True{} : Bool} 114 inside), 115 earlier(inside)) 116 case Shape.BTrue{equal}: 117 +before = Contracts.words(shape, fill, ops, previous) 118 +last = Replayed.value(shape, fill, ops, U32.from_nat(previous)) 119 %Equal.sym(Nat, position, previous, Arith.eq_true(position, previous, Equal.sym(Bool, True{}, Nat.is_eq(position, previous), equal))) : 120 {ReduceLaws.word_at(List.append(&2, F32, before, [last]), _) == Some{Replayed.value(shape, fill, ops, U32.from_nat(_))} : Maybe<&2, F32>} 121 %words_length(shape, fill, ops, previous) : 122 {ReduceLaws.word_at(List.append(&2, F32, before, [last]), _) == Some{last} : Maybe<&2, F32>} 123 word_at_last(before, last) 124 125 def words_at(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat, +position: Nat, 126 below: {Nat.is_lt(position, count) == True{} : Bool}) -> 127 {ReduceLaws.word_at(Contracts.words(shape, fill, ops, count), position) == 128 Some{Replayed.value(shape, fill, ops, U32.from_nat(position))} : Maybe<&2, F32>}: 129 match count: 130 case 0n: 131 Empty.absurd({ReduceLaws.word_at(Nil{}, position) == Some{Replayed.value(shape, fill, ops, U32.from_nat(position))} : Maybe<&2, F32>}, 132 Arith.lt_zero(position, below)) 133 case 1n+previous: 134 words_step(shape, fill, ops, previous, position, below, Shape.bool_case(Nat.is_eq(position, previous)), 135 inside => words_at(shape, fill, ops, previous, position, inside)) 136 137 def matched_chosen(+same: Bool, +total: F32, +word: F32) -> 138 {ReduceLaws.matched(Some{total}, Some{word}, same) == Some{Fold.chosen(~Unit, ~Sum.add, Unit{}, same, total, word)} : Maybe<&2, F32>}: 139 match same: 140 case False{}: 141 {==} 142 case True{}: 143 {==} 144 145 def fiber_total(+plan: Shape.BroadcastPlan, +index: U32, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat, 146 +room: {Nat.is_le(count, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) -> 147 {ReduceLaws.fiber(plan, 148 Contracts.words(Shape.right_shape(plan), fill, ops, Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))), index, count) == 149 Some{Fold.total(~Unit, ~Sum.add, ~ReduceFold.Cell, ~ReduceFold.source, ~ReduceFold.lands, Unit{}, ReduceFold.Cell{plan, index}, 150 Shape.right_shape(plan), fill, ops, count)} : Maybe<&2, F32>}: 151 match count: 152 case 0n: 153 {==} 154 case 1n+k: 155 +right = Shape.right_shape(plan) 156 +output = Shape.output_shape(plan) 157 +limit = Shape.logical_numel(Shape.dimensions(output)) 158 +values = Contracts.words(right, fill, ops, Shape.logical_numel(Shape.dimensions(right))) 159 +cell : ReduceFold.Cell = ReduceFold.Cell{plan, index} 160 +lm = ReduceFold.target_offset(plan, k) 161 +rm = ReduceFold.contribution_offset(plan, k) 162 +before = Fold.total(~Unit, ~Sum.add, ~ReduceFold.Cell, ~ReduceFold.source, ~ReduceFold.lands, Unit{}, cell, right, fill, ops, k) 163 +word = Replayed.value(right, fill, ops, U32.from_nat(rm)) 164 +lands = Nat.is_eq(lm, U32.to_nat(index)) 165 +after : Maybe<&2, F32> = Some{Fold.chosen(~Unit, ~Sum.add, Unit{}, lands, before, word)} 166 +below = Conversion.le_successor_lt(k, limit, room) 167 +rm_below = BroadcastProof.right_walk_bound(plan, Shape.index_map(plan), k, below) 168 +lm_below = BroadcastProof.left_walk_bound(plan, Shape.index_map(plan), k, below) 169 %Equal.sym(Maybe<&1, U32 & U32>, Projection.offset_spec(plan, k), Some{(U32.from_nat(lm), U32.from_nat(rm))}, 170 offsets_walk(plan, Shape.index_map(plan), k, BroadcastProof.output_inside(plan, k, below))) : 171 {ReduceLaws.contribute(values, index, ReduceLaws.fiber(plan, values, index, k), _) == after : Maybe<&2, F32>} 172 %Equal.sym(Maybe<&2, F32>, ReduceLaws.fiber(plan, values, index, k), Some{before}, 173 fiber_total(plan, index, fill, ops, k, Conversion.le_predecessor(k, limit, room))) : 174 {ReduceLaws.matched(_, ReduceLaws.word_at(values, U32.to_nat(U32.from_nat(rm))), Nat.is_eq(U32.to_nat(U32.from_nat(lm)), U32.to_nat(index))) == 175 after : Maybe<&2, F32>} 176 %Equal.sym(Nat, U32.to_nat(U32.from_nat(rm)), rm, PointwiseProof.word_exact(rm, right, rm_below)) : 177 {ReduceLaws.matched(Some{before}, ReduceLaws.word_at(values, _), Nat.is_eq(U32.to_nat(U32.from_nat(lm)), U32.to_nat(index))) == 178 after : Maybe<&2, F32>} 179 %Equal.sym(Maybe<&2, F32>, ReduceLaws.word_at(values, rm), Some{word}, 180 words_at(right, fill, ops, Shape.logical_numel(Shape.dimensions(right)), rm, rm_below)) : 181 {ReduceLaws.matched(Some{before}, _, Nat.is_eq(U32.to_nat(U32.from_nat(lm)), U32.to_nat(index))) == after : Maybe<&2, F32>} 182 %Equal.sym(Nat, U32.to_nat(U32.from_nat(lm)), lm, PointwiseProof.word_exact(lm, Shape.left_shape(plan), lm_below)) : 183 {ReduceLaws.matched(Some{before}, Some{word}, Nat.is_eq(_, U32.to_nat(index))) == after : Maybe<&2, F32>} 184 matched_chosen(lands, before, word) 185 186 def source_inside(+plan: Shape.BroadcastPlan, +index: U32, +count: Nat, 187 +room: {Nat.is_le(count, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) -> 188 {Pointwise.inside(~ReduceFold.Cell, ~ReduceFold.source, ReduceFold.Cell{plan, index}, count, 189 Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))) == True{} : Bool}: 190 match count: 191 case 0n: 192 {==} 193 case 1n+previous: 194 +cell : ReduceFold.Cell = ReduceFold.Cell{plan, index} 195 +limit = Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) 196 +bound = Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) 197 %Equal.sym(Bool, Pointwise.inside(~ReduceFold.Cell, ~ReduceFold.source, cell, previous, bound), True{}, 198 source_inside(plan, index, previous, Conversion.le_predecessor(previous, limit, room))) : 199 {Bool.and(_, Nat.is_lt(ReduceFold.source(cell, previous), bound)) == True{} : Bool} 200 %Equal.sym(Bool, Nat.is_lt(ReduceFold.source(cell, previous), bound), True{}, 201 BroadcastProof.right_walk_bound(plan, Shape.index_map(plan), previous, Conversion.le_successor_lt(previous, limit, room))) : 202 {Bool.and(True{}, _) == True{} : Bool} 203 {==} 204 205 def contribution_room(+plan: Shape.BroadcastPlan, 206 same: {Shape.extents(Shape.output_shape(plan)) == Shape.extents(Shape.right_shape(plan)) : List<&2, Nat>}) -> 207 {Nat.is_le(Shape.elements(Shape.right_shape(plan)), Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}: 208 +right = Shape.right_shape(plan) 209 +limit = Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) 210 %Equal.sym(Nat, Shape.elements(right), Shape.numel(Shape.extents(right)), Shape.elements_exact(right)) : 211 {Nat.is_le(_, limit) == True{} : Bool} 212 %Equal.cong(List<&2, Nat>, Nat, extents => Shape.numel(extents), Shape.extents(Shape.output_shape(plan)), Shape.extents(right), same) : 213 {Nat.is_le(_, limit) == True{} : Bool} 214 Conversion.le_reflexive(limit) 215 216 def fiber_contract(+plan: Shape.BroadcastPlan, +index: U32, +fill: F32, +ops: List<&2, StorageLaws.Op>, 217 -fits: {Lifecycle.logical_length(Replayed.storage(Shape.right_shape(plan), fill, ops)) == 218 Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) : Nat}, 219 +expected: Maybe<&2, F32>, +total: F32, +count: Nat, 220 evidence: Records.FoldContract Replayed.value(Shape.right_shape(plan), fill, ops, other), 222 Fold.run(~Unit, ~Sum.add, ~ReduceFold.Cell, ~ReduceFold.source, ~ReduceFold.lands, Unit{}, ReduceFold.Cell{plan, index}, count, 223 Core.Packed{Shape.right_shape(plan), Core.TTensor{Replayed.storage(Shape.right_shape(plan), fill, ops), fits}})>, 224 equation: {expected == Some{total} : Maybe<&2, F32>}) -> 225 Laws.FiberContract Replayed.value(Shape.right_shape(plan), fill, ops, other), 227 Fold.run(~Unit, ~Sum.add, ~ReduceFold.Cell, ~ReduceFold.source, ~ReduceFold.lands, Unit{}, ReduceFold.Cell{plan, index}, count, 228 Core.Packed{Shape.right_shape(plan), Core.TTensor{Replayed.storage(Shape.right_shape(plan), fill, ops), fits}})>: 229 Records.FoldContract{value, kept} = evidence 230 %Equal.sym(Maybe<&2, F32>, expected, Some{total}, equation) : 231 Laws.FiberContract Replayed.value(Shape.right_shape(plan), fill, ops, other), 233 Fold.run(~Unit, ~Sum.add, ~ReduceFold.Cell, ~ReduceFold.source, ~ReduceFold.lands, Unit{}, ReduceFold.Cell{plan, index}, count, 234 Core.Packed{Shape.right_shape(plan), Core.TTensor{Replayed.storage(Shape.right_shape(plan), fill, ops), fits}})> 235 Laws.FiberContract{value, kept} 236 237 def reduce_fold_fiber(+plan: Shape.BroadcastPlan, 238 same: {Shape.extents(Shape.output_shape(plan)) == Shape.extents(Shape.right_shape(plan)) : List<&2, Nat>}, 239 +index: U32, +fill: F32, +ops: List<&2, StorageLaws.Op>, 240 -fits: {Lifecycle.logical_length(Replayed.storage(Shape.right_shape(plan), fill, ops)) == 241 Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) : Nat}) -> 242 Laws.FiberContract Replayed.value(Shape.right_shape(plan), fill, ops, other), 247 ReduceFold.cell(plan, index, 248 Core.Packed{Shape.right_shape(plan), Core.TTensor{Replayed.storage(Shape.right_shape(plan), fill, ops), fits}})>: 249 +right = Shape.right_shape(plan) 250 +count = Shape.elements(right) 251 +values = Contracts.words(right, fill, ops, Shape.logical_numel(Shape.dimensions(right))) 252 +cell : ReduceFold.Cell = ReduceFold.Cell{plan, index} 253 +total = Fold.total(~Unit, ~Sum.add, ~ReduceFold.Cell, ~ReduceFold.source, ~ReduceFold.lands, Unit{}, cell, right, fill, ops, count) 254 +room = contribution_room(plan, same) 255 fiber_contract(plan, index, fill, ops, fits, 256 ReduceLaws.fiber(plan, values, index, Shape.logical_numel(Shape.dimensions(right))), total, count, 257 FoldProof.contract(~Unit, ~Sum.add, ~ReduceFold.Cell, ~ReduceFold.source, ~ReduceFold.lands, Unit{}, cell, right, fill, ops, count, 258 source_inside(plan, index, count, room), fits), 259 %Shape.elements_exact(right) : {ReduceLaws.fiber(plan, values, index, _) == Some{total} : Maybe<&2, F32>} 260 fiber_total(plan, index, fill, ops, count, room)) 261 262 def Laws.reduce_fold_fiber(plan, same, index, fill, ops, fits): 263 reduce_fold_fiber(plan, same, index, fill, ops, fits)