1 import Base 2 import ../shape/core.bend as Shape 3 import ../storage/main.bend as Storage 4 import ../storage/LAWS.bend as StorageLaws 5 import ../storage/arith.bend as Arith 6 import ../tensor/core.bend as Core 7 import ../tensor/layout.bend as Layout 8 import ../tensor/lifecycle.bend as Lifecycle 9 import ../tensor/elementwise.bend as Elementwise 10 import ../tensor/word_conversion.bend as Conversion 11 import ../tensor/permutation_mapping.bend as Mapping 12 import ../tensor/projection_LAWS.bend as Projection 13 import ../tensor/projection_PROOF.bend as ProjectionProof 14 import ../tensor/contracts/replayed.bend as Replayed 15 import ../tensor/contracts/pointwise.bend as Pointwise 16 import ../tensor/contracts/pointwise_proof.bend as PointwiseProof 17 import ../tensor/contracts/broadcast_proof.bend as BroadcastProof 18 import ../tensor/contracts/fold.bend as Fold 19 import ../tensor/contracts/sum.bend as Sum 20 import ../tensor/contracts/LAWS.bend as Contracts 21 import ./reduce.bend as Reduce 22 import ./reduce_LAWS.bend as ReduceLaws 23 import ./reduce_fold.bend as ReduceFold 24 import ./reduce_fold_PROOF.bend as FoldProof 25 import ./reduce_fold_LAWS.bend as Laws 26 import ../storage/api.bend as Api 27 import ../storage/array.bend as ArrayCore 28 import ../tensor/read_contract.bend as Reads 29 import ../tensor/contracts/binary_storage_proof.bend as Owned 30 import ./reduce_replayed_proof.bend as Replay 31 32 def planned_shapes(left: Shape.Shape, right: Shape.Shape, result: Result<&2, &2, Shape.Error, Shape.BroadcastPlan>) -> Shape.Shape & Shape.Shape: 33 match result: 34 case Fail{error}: 35 (left, right) 36 case Done{+plan}: 37 (Shape.left_shape(plan), Shape.right_shape(plan)) 38 39 def shapes_construct(+left: Shape.Shape, +right: Shape.Shape, -left_padding: Nat, -right_padding: Nat, 40 -minimal: Shape.PadMinimal, -output: List<&2, U32>, 41 -aligned: Shape.Aligned(Shape.pad(left_padding, Shape.extents(left)), Shape.pad(right_padding, Shape.extents(right)), Shape.naturals(output)), 42 +result: Shape.ConstructResult) -> 43 {planned_shapes(left, right, Shape.broadcast_checked.construct(left, right, left_padding, right_padding, minimal, output, aligned, result)) == 44 (left, right) : Shape.Shape & Shape.Shape}: 45 match result: 46 case Shape.SError{error}: 47 {==} 48 case Shape.SBuilt{shape, preserved}: 49 {==} 50 51 def shapes_align(+left: Shape.Shape, +right: Shape.Shape, +left_padding: Nat, +right_padding: Nat, 52 -minimal: Shape.PadMinimal, 53 +alignment: Shape.Alignment) -> 54 {planned_shapes(left, right, Shape.broadcast_checked.align(left, right, left_padding, right_padding, minimal, alignment)) == 55 (left, right) : Shape.Shape & Shape.Shape}: 56 match alignment: 57 case Shape.Misaligned{}: 58 {==} 59 case Shape.Aligns{+output, aligned}: 60 shapes_construct(left, right, left_padding, right_padding, minimal, output, 61 Shape.broadcast_checked.repad(left_padding, right_padding, Shape.dimensions(left), Shape.dimensions(right), output, aligned), 62 Shape.construct(output)) 63 64 def shapes_checked(+left: Shape.Shape, +right: Shape.Shape) -> 65 {planned_shapes(left, right, Shape.broadcast_checked(left, right)) == (left, right) : Shape.Shape & Shape.Shape}: 66 +left_rank = Shape.rank_dimensions(Shape.dimensions(left)) 67 +right_rank = Shape.rank_dimensions(Shape.dimensions(right)) 68 shapes_align(left, right, Nat.sub(right_rank, left_rank), Nat.sub(left_rank, right_rank), Shape.minimal(left_rank, right_rank), 69 Shape.align(Shape.pad_words(Nat.sub(right_rank, left_rank), Shape.dimensions(left)), 70 Shape.pad_words(Nat.sub(left_rank, right_rank), Shape.dimensions(right)))) 71 72 def first(pair: Shape.Shape & Shape.Shape) -> Shape.Shape: 73 (left, right) = pair 74 left 75 76 def second(pair: Shape.Shape & Shape.Shape) -> Shape.Shape: 77 (left, right) = pair 78 right 79 80 def plan_shapes(+target: Shape.Shape, +shape: Shape.Shape, +plan: Shape.BroadcastPlan, 81 planned: {Shape.broadcast_checked(target, shape) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}) -> 82 {(Shape.left_shape(plan), Shape.right_shape(plan)) == (target, shape) : Shape.Shape & Shape.Shape}: 83 %planned : {planned_shapes(target, shape, _) == (target, shape) : Shape.Shape & Shape.Shape} 84 shapes_checked(target, shape) 85 86 def plan_left(+target: Shape.Shape, +shape: Shape.Shape, +plan: Shape.BroadcastPlan, 87 planned: {Shape.broadcast_checked(target, shape) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}) -> 88 {Shape.left_shape(plan) == target : Shape.Shape}: 89 Equal.cong(Shape.Shape & Shape.Shape, Shape.Shape, first, (Shape.left_shape(plan), Shape.right_shape(plan)), (target, shape), 90 plan_shapes(target, shape, plan, planned)) 91 92 def plan_right(+target: Shape.Shape, +shape: Shape.Shape, +plan: Shape.BroadcastPlan, 93 planned: {Shape.broadcast_checked(target, shape) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}) -> 94 {Shape.right_shape(plan) == shape : Shape.Shape}: 95 Equal.cong(Shape.Shape & Shape.Shape, Shape.Shape, second, (Shape.left_shape(plan), Shape.right_shape(plan)), (target, shape), 96 plan_shapes(target, shape, plan, planned)) 97 98 def target_index(+plan: Shape.BroadcastPlan, +position: Nat) -> U32: 99 U32.from_nat(ReduceFold.target_offset(plan, position)) 100 101 def source_index(+plan: Shape.BroadcastPlan, +position: Nat) -> U32: 102 U32.from_nat(ReduceFold.contribution_offset(plan, position)) 103 104 # Contribution position `position` lands on a target cell below the target's 105 # count. 106 def target_below(+plan: Shape.BroadcastPlan, +target: Shape.Shape, +left: {Shape.left_shape(plan) == target : Shape.Shape}, +position: Nat, 107 below: {Nat.is_lt(position, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) -> 108 {Nat.is_lt(ReduceFold.target_offset(plan, position), Shape.logical_numel(Shape.dimensions(target))) == True{} : Bool}: 109 %left : {Nat.is_lt(ReduceFold.target_offset(plan, position), Shape.logical_numel(Shape.dimensions(_))) == True{} : Bool} 110 BroadcastProof.left_walk_bound(plan, Shape.index_map(plan), position, below) 111 112 def source_below(+plan: Shape.BroadcastPlan, +shape: Shape.Shape, +right: {Shape.right_shape(plan) == shape : Shape.Shape}, +position: Nat, 113 below: {Nat.is_lt(position, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) -> 114 {Nat.is_lt(ReduceFold.contribution_offset(plan, position), Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool}: 115 %right : {Nat.is_lt(ReduceFold.contribution_offset(plan, position), Shape.logical_numel(Shape.dimensions(_))) == True{} : Bool} 116 BroadcastProof.right_walk_bound(plan, Shape.index_map(plan), position, below) 117 118 def some_injective(-A: Type, -x: A, -y: A, equal: {Some{x} == Some{y} : Maybe<&1, A>}) -> {x == y : A}: 119 Equal.cong(Maybe<&1, A>, A, found => Maybe.default(&1, A, found, x), Some{x}, Some{y}, equal) 120 121 def indices_found(+plan: Shape.BroadcastPlan, +cursor: Nat, 122 +below: {Nat.is_lt(cursor, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) -> 123 {Projection.indices_result(Elementwise.input_indices(plan, U32.from_nat(cursor))) == 124 Some{(target_index(plan, cursor), source_index(plan, cursor))} : Maybe<&1, U32 & U32>}: 125 Equal.trans(Maybe<&1, U32 & U32>, Projection.indices_result(Elementwise.input_indices(plan, U32.from_nat(cursor))), 126 Projection.offset_spec(plan, U32.to_nat(U32.from_nat(cursor))), Some{(target_index(plan, cursor), source_index(plan, cursor))}, 127 Projection.input_indices_refines(plan, U32.from_nat(cursor)), 128 %Equal.sym(Nat, U32.to_nat(U32.from_nat(cursor)), cursor, PointwiseProof.word_exact(cursor, Shape.output_shape(plan), below)) : 129 {Projection.offset_spec(plan, _) == Some{(target_index(plan, cursor), source_index(plan, cursor))} : Maybe<&1, U32 & U32>} 130 FoldProof.offsets_walk(plan, Shape.index_map(plan), cursor, BroadcastProof.output_inside(plan, cursor, below))) 131 132 def contribution_count(+shape: Shape.Shape, +plan: Shape.BroadcastPlan, 133 same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(shape) : List<&2, U32>}) -> 134 {Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}: 135 Equal.cong(List<&2, U32>, Nat, Shape.logical_numel, Shape.dimensions(Shape.output_shape(plan)), Shape.dimensions(shape), same) 136 137 def contribution_room(+shape: Shape.Shape, +plan: Shape.BroadcastPlan, 138 same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(shape) : List<&2, U32>}) -> 139 {Nat.is_le(Shape.logical_numel(Shape.dimensions(shape)), Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}: 140 %Equal.sym(Nat, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))), Shape.logical_numel(Shape.dimensions(shape)), 141 contribution_count(shape, plan, same)) : {Nat.is_le(Shape.logical_numel(Shape.dimensions(shape)), _) == True{} : Bool} 142 Conversion.le_reflexive(Shape.logical_numel(Shape.dimensions(shape))) 143 144 def fitted_same(+shape: Shape.Shape, +plan: Shape.BroadcastPlan, 145 same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(shape) : List<&2, U32>}) -> 146 {Layout.words_same(Shape.dimensions(Shape.output_shape(plan)), Shape.dimensions(shape)) == True{} : Bool}: 147 %Equal.sym(List<&2, U32>, Shape.dimensions(Shape.output_shape(plan)), Shape.dimensions(shape), same) : 148 {Layout.words_same(_, Shape.dimensions(shape)) == True{} : Bool} 149 Mapping.words_same_refl(Shape.dimensions(shape)) 150 151 # The parallel reduce over a replayed contribution, in the transport form of 152 # `Replay.Final`. 153 def run_lift(+target: Shape.Shape, +shape: Shape.Shape, +plan: Shape.BroadcastPlan, 154 +planned: {Shape.broadcast_checked(target, shape) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}, 155 +same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(shape) : List<&2, U32>}, 156 +fill: F32, +ops: List<&2, StorageLaws.Op>, 157 -P: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic> -> Type, 158 h: Replay.Final(plan, target, shape, fill, ops, Shape.index_map(plan), StorageLaws.created(Shape.count(target), 0.0), P)) -> 159 @-fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} -> 160 P(Reduce.reduce_to(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target)): 161 +output = Shape.output_shape(plan) 162 +left = plan_left(target, shape, plan, planned) 163 +right = plan_right(target, shape, plan, planned) 164 +count_left : {U32.to_nat(Shape.count(target)) == Shape.numel(Shape.extents(Shape.left_shape(plan))) : Nat} = 165 %Equal.sym(Shape.Shape, Shape.left_shape(plan), target, left) : 166 {U32.to_nat(Shape.count(target)) == Shape.logical_numel(Shape.dimensions(_)) : Nat} 167 Reads.count_exact(target) 168 +facts : Replay.Facts = 169 Replay.Facts{count_left, Equal.cong(Shape.Shape, List<&2, U32>, Shape.dimensions, Shape.right_shape(plan), shape, right)} 170 fits => 171 %Equal.sym(Result<&2, &2, Shape.Error, Shape.BroadcastPlan>, Shape.broadcast_checked(target, shape), Done{plan}, planned) : 172 P(Reduce.planned(shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}, target, _)) 173 %Equal.sym(Bool, Layout.words_same(Shape.dimensions(output), Shape.dimensions(shape)), True{}, fitted_same(shape, plan, same)) : 174 P(Reduce.fitted(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target, plan, _)) 175 Replayed.new_lift(target, 0.0, 176 result => P(Reduce.parallel_allocated(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target, plan, result)), 177 fo => Owned.created_lift(Shape.count(target), 0.0, Conversion.shape_count_within(target), 178 st => (Replay.Final(plan, target, shape, fill, ops, Shape.index_map(plan), st, P) -> 179 @-fo: {Lifecycle.logical_length(st) == Shape.logical_numel(Shape.dimensions(target)) : Nat} -> 180 P(Reduce.parallel_allocated(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target, plan, Done{Core.TTensor{st, fo}}))), 181 zero => final => fo => 182 final(ops)(other => {==})({==})(fits)(Nil{})( 183 index => inside => Empty.absurd( 184 {Replayed.value(target, 0.0, Nil{}, index) == Replay.cell_word(plan, Shape.index_map(plan), shape, fill, ops, U32.to_nat(index)) : F32}, 185 Arith.lt_zero(U32.to_nat(index), 186 %Equal.trans(Nat, Shape.logical_numel(Shape.dimensions(target)), U32.to_nat(Shape.count(target)), 0n, 187 Equal.sym(Nat, U32.to_nat(Shape.count(target)), Shape.logical_numel(Shape.dimensions(target)), Reads.count_exact(target)), zero) : 188 {Nat.is_lt(U32.to_nat(index), _) == True{} : Bool} 189 inside)))(fo), 190 d => fits_d => f => final => fo => 191 Replay.indexed_lift(plan, target, shape, fill, ops, fits, facts, d, fits_d, f, fo, Shape.index_map(plan), P, final))(h)(fo)) 192 193 # A target cell of the replayed reduce holds `ReduceLaws.fiber` over the 194 # contribution's words. 195 def cell_fiber(+target: Shape.Shape, +shape: Shape.Shape, +plan: Shape.BroadcastPlan, 196 +planned: {Shape.broadcast_checked(target, shape) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}, 197 +same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(shape) : List<&2, U32>}, 198 +fill: F32, +ops: List<&2, StorageLaws.Op>, +index: U32, 199 +inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(target))) == True{} : Bool}) -> 200 {Some{Replay.cell_word(plan, Shape.index_map(plan), shape, fill, ops, U32.to_nat(index))} == 201 ReduceLaws.fiber(plan, Contracts.words(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape))), index, 202 Shape.logical_numel(Shape.dimensions(shape))) : Maybe<&2, F32>}: 203 +left = plan_left(target, shape, plan, planned) 204 %contribution_count(shape, plan, same) : 205 {Some{Replay.cell_word(plan, Shape.index_map(plan), shape, fill, ops, U32.to_nat(index))} == 206 ReduceLaws.fiber(plan, Contracts.words(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape))), index, _) : Maybe<&2, F32>} 207 Replay.word_fiber(plan, Shape.index_map(plan), shape, fill, ops, index, 208 %Equal.sym(Shape.Shape, Shape.left_shape(plan), target, left) : {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(_))) == True{} : Bool} 209 inside, 210 Equal.cong(Shape.Shape, List<&2, U32>, Shape.dimensions, Shape.right_shape(plan), shape, plan_right(target, shape, plan, planned))) 211 212 def reduce_to_fiber(+target: Shape.Shape, +shape: Shape.Shape, +plan: Shape.BroadcastPlan, 213 +planned: {Shape.broadcast_checked(target, shape) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}, 214 +same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(shape) : List<&2, U32>}, 215 +fill: F32, +ops: List<&2, StorageLaws.Op>, 216 -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}) -> 217 Laws.ReduceToContract ReduceLaws.fiber(plan, Contracts.words(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape))), 219 index, Shape.logical_numel(Shape.dimensions(shape))), 220 other => Replayed.value(shape, fill, ops, other), 221 Reduce.reduce_to(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target)>: 222 +count = Shape.logical_numel(Shape.dimensions(shape)) 223 Laws.ReduceToContract{ 224 run_lift(target, shape, plan, planned, same, fill, ops, result => {ReduceLaws.succeeded(result) == True{} : Bool}, 225 kept => values => words => kept_fits => out => cells => out_fits => {==})(fits), 226 run_lift(target, shape, plan, planned, same, fill, ops, 227 result => {Laws.output_dimensions(result) == Some{Shape.dimensions(target)} : Maybe<&2, List<&2, U32>>}, 228 kept => values => words => kept_fits => out => cells => out_fits => {==})(fits), 229 index => inside => 230 +inside = inside 231 run_lift(target, shape, plan, planned, same, fill, ops, 232 result => {Laws.output_point(result, index) == 233 ReduceLaws.fiber(plan, Contracts.words(shape, fill, ops, count), index, count) : Maybe<&2, F32>}, 234 kept => values => words => kept_fits => out => cells => out_fits => 235 Equal.trans(Maybe<&2, F32>, 236 Core.dynamic_read_value(Core.read_dynamic(Core.Packed{target, Core.TTensor{Replayed.storage(target, 0.0, out), out_fits}}, index)), 237 Some{Replay.cell_word(plan, Shape.index_map(plan), shape, fill, ops, U32.to_nat(index))}, 238 ReduceLaws.fiber(plan, Contracts.words(shape, fill, ops, count), index, count), 239 Equal.trans(Maybe<&2, F32>, 240 Core.dynamic_read_value(Core.read_dynamic(Core.Packed{target, Core.TTensor{Replayed.storage(target, 0.0, out), out_fits}}, index)), 241 Some{Replayed.value(target, 0.0, out, index)}, Some{Replay.cell_word(plan, Shape.index_map(plan), shape, fill, ops, U32.to_nat(index))}, 242 PointwiseProof.replayed_point(target, 0.0, out, out_fits, index, inside), 243 Equal.cong(F32, Maybe<&2, F32>, word => Some{word}, Replayed.value(target, 0.0, out, index), 244 Replay.cell_word(plan, Shape.index_map(plan), shape, fill, ops, U32.to_nat(index)), cells(index)(inside))), 245 cell_fiber(target, shape, plan, planned, same, fill, ops, index, inside)))(fits), 246 other => inside => 247 run_lift(target, shape, plan, planned, same, fill, ops, 248 result => {Laws.input_point(result, other) == Some{Replayed.value(shape, fill, ops, other)} : Maybe<&2, F32>}, 249 kept => values => words => kept_fits => out => cells => out_fits => 250 Equal.trans(Maybe<&2, F32>, 251 Core.dynamic_read_value(Core.read_dynamic(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, kept), kept_fits}}, other)), 252 Some{Replayed.value(shape, fill, kept, other)}, Some{Replayed.value(shape, fill, ops, other)}, 253 PointwiseProof.replayed_point(shape, fill, kept, kept_fits, other, inside), 254 Equal.cong(F32, Maybe<&2, F32>, word => Some{word}, Replayed.value(shape, fill, kept, other), Replayed.value(shape, fill, ops, other), 255 values(other))))(fits)} 256 257 def Laws.reduce_to_fiber(target, shape, plan, planned, same, fill, ops, fits): 258 reduce_to_fiber(target, shape, plan, planned, same, fill, ops, fits)