1 import Base 2 import ../shape/core.bend as Shape 3 import ../storage/LAWS.bend as StorageLaws 4 import ../storage/arith.bend as Arith 5 import ../tensor/core.bend as Core 6 import ../tensor/layout.bend as Layout 7 import ../tensor/lifecycle.bend as Lifecycle 8 import ../tensor/elementwise.bend as Elementwise 9 import ../tensor/word_conversion.bend as Conversion 10 import ../tensor/contracts/LAWS.bend as Contracts 11 import ../tensor/contracts/replayed.bend as Replayed 12 import ../tensor/contracts/pointwise_proof.bend as PointwiseProof 13 import ../tensor/contracts/broadcast.bend as PointwiseBroadcast 14 import ../tensor/contracts/binary_LAWS.bend as BinaryLaws 15 import ../tensor/contracts/binary_proofs.bend as Binary 16 import ../tensor/contracts/binary_PROOF.bend as BinaryProof 17 import ./broadcast.bend as Broadcast 18 import ./reduce_LAWS.bend as ReduceLaws 19 import ./reduce_fold_PROOF.bend as FoldProof 20 import ./reduce_to_fold_PROOF.bend as ReduceTo 21 import ./broadcast_pointwise_LAWS.bend as Laws 22 23 # The pointwise family's KeepLeft words, with the input on the left and the 24 # +0.0 tensor of the target on the right, are broadcast_to's `cells`. 25 def cells_keep(+plan: Shape.BroadcastPlan, +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +target: Shape.Shape, 26 +count: Nat) -> 27 {BinaryLaws.cells(Elementwise.KeepLeft{}, plan, shape, fill, ops, target, 0.0, Nil{}, count) == 28 Laws.cells(plan, shape, fill, ops, count) : List<&2, F32>}: 29 match count: 30 case 0n: 31 {==} 32 case 1n+previous: 33 Equal.cong(List<&2, F32>, List<&2, F32>, 34 before => List.append(&2, F32, before, [Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, previous)))]), 35 BinaryLaws.cells(Elementwise.KeepLeft{}, plan, shape, fill, ops, target, 0.0, Nil{}, previous), Laws.cells(plan, shape, fill, ops, previous), 36 cells_keep(plan, shape, fill, ops, target, previous)) 37 38 def cells_length(+plan: Shape.BroadcastPlan, +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat) -> 39 {List.length(&2, F32, Laws.cells(plan, shape, fill, ops, count)) == count : Nat}: 40 match count: 41 case 0n: 42 {==} 43 case 1n+previous: 44 +before = Laws.cells(plan, shape, fill, ops, previous) 45 +last = Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, previous))) 46 Equal.trans(Nat, List.length(&2, F32, List.append(&2, F32, before, [last])), 1n+List.length(&2, F32, before), 1n+previous, 47 FoldProof.length_snoc(before, last), 48 Equal.cong(Nat, Nat, n => 1n+n, List.length(&2, F32, before), previous, cells_length(plan, shape, fill, ops, previous))) 49 50 # Word `position` of `cells` is the input word at left_map(plan, position). 51 def cells_step(+plan: Shape.BroadcastPlan, +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +previous: Nat, +position: Nat, 52 below: {Nat.is_lt(position, 1n+previous) == True{} : Bool}, decision: Shape.BoolCase, 53 earlier: @inside: {Nat.is_lt(position, previous) == True{} : Bool} -> 54 {ReduceLaws.word_at(Laws.cells(plan, shape, fill, ops, previous), position) == 55 Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, position)))} : Maybe<&2, F32>}) -> 56 {ReduceLaws.word_at(Laws.cells(plan, shape, fill, ops, 1n+previous), position) == 57 Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, position)))} : Maybe<&2, F32>}: 58 match decision: 59 case Shape.BFalse{equal}: 60 +before = Laws.cells(plan, shape, fill, ops, previous) 61 +last = Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, previous))) 62 +inside = PointwiseProof.lt_succ_ne(position, previous, below, equal) 63 Equal.trans(Maybe<&2, F32>, ReduceLaws.word_at(List.append(&2, F32, before, [last]), position), ReduceLaws.word_at(before, position), 64 Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, position)))}, 65 FoldProof.word_at_front(before, last, position, 66 %Equal.sym(Nat, List.length(&2, F32, before), previous, cells_length(plan, shape, fill, ops, previous)) : 67 {Nat.is_lt(position, _) == True{} : Bool} 68 inside), 69 earlier(inside)) 70 case Shape.BTrue{equal}: 71 +before = Laws.cells(plan, shape, fill, ops, previous) 72 +last = Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, previous))) 73 %Equal.sym(Nat, position, previous, Arith.eq_true(position, previous, Equal.sym(Bool, True{}, Nat.is_eq(position, previous), equal))) : 74 {ReduceLaws.word_at(List.append(&2, F32, before, [last]), _) == 75 Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, _)))} : Maybe<&2, F32>} 76 %cells_length(plan, shape, fill, ops, previous) : 77 {ReduceLaws.word_at(List.append(&2, F32, before, [last]), _) == Some{last} : Maybe<&2, F32>} 78 FoldProof.word_at_last(before, last) 79 80 def cells_at(+plan: Shape.BroadcastPlan, +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat, +position: Nat, 81 below: {Nat.is_lt(position, count) == True{} : Bool}) -> 82 {ReduceLaws.word_at(Laws.cells(plan, shape, fill, ops, count), position) == 83 Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, position)))} : Maybe<&2, F32>}: 84 match count: 85 case 0n: 86 Empty.absurd({ReduceLaws.word_at(Nil{}, position) == 87 Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, position)))} : Maybe<&2, F32>}, 88 Arith.lt_zero(position, below)) 89 case 1n+previous: 90 cells_step(plan, shape, fill, ops, previous, position, below, Shape.bool_case(Nat.is_eq(position, previous)), 91 inside => cells_at(plan, shape, fill, ops, previous, position, inside)) 92 93 # A replayed tensor whose first `count` words are `values` holds, below 94 # `count`, the word `values` lists there. 95 def listed_value(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat, +values: List<&2, F32>, +index: U32, 96 words: {Contracts.words(shape, fill, ops, count) == values : List<&2, F32>}, 97 below: {Nat.is_lt(U32.to_nat(index), count) == True{} : Bool}) -> 98 {Some{Replayed.value(shape, fill, ops, index)} == ReduceLaws.word_at(values, U32.to_nat(index)) : Maybe<&2, F32>}: 99 %Conversion.word_natural_exact(index) : 100 {Some{Replayed.value(shape, fill, ops, _)} == ReduceLaws.word_at(values, U32.to_nat(index)) : Maybe<&2, F32>} 101 Equal.trans(Maybe<&2, F32>, Some{Replayed.value(shape, fill, ops, U32.from_nat(U32.to_nat(index)))}, 102 ReduceLaws.word_at(Contracts.words(shape, fill, ops, count), U32.to_nat(index)), ReduceLaws.word_at(values, U32.to_nat(index)), 103 Equal.sym(Maybe<&2, F32>, ReduceLaws.word_at(Contracts.words(shape, fill, ops, count), U32.to_nat(index)), 104 Some{Replayed.value(shape, fill, ops, U32.from_nat(U32.to_nat(index)))}, 105 FoldProof.words_at(shape, fill, ops, count, U32.to_nat(index), below)), 106 Equal.cong(List<&2, F32>, Maybe<&2, F32>, listed => ReduceLaws.word_at(listed, U32.to_nat(index)), 107 Contracts.words(shape, fill, ops, count), values, words)) 108 109 # broadcast_to on a replayed input whose shape broadcasts to exactly the 110 # target: the planner and the target check pass, the zero tensor of the target 111 # is allocated, and the parallel build runs KeepLeft with the input on the 112 # left and the zero tensor on the right, whose words are `cells`. 113 def run_lift(+shape: Shape.Shape, +target: Shape.Shape, +plan: Shape.BroadcastPlan, 114 +planned: {Shape.broadcast_checked(shape, target) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}, 115 +same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(target) : List<&2, U32>}, 116 +fill: F32, +ops: List<&2, StorageLaws.Op>, 117 -P: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic> -> Type, 118 h: Laws.BroadcastWith(plan, shape, fill, ops, P)) -> 119 @-fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} -> 120 P(Broadcast.broadcast_to(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target)): 121 +output = Shape.output_shape(plan) 122 +count = Shape.logical_numel(Shape.dimensions(output)) 123 fits => 124 -input = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}} : Core.Dynamic} 125 %Equal.sym(Result<&2, &2, Shape.Error, Shape.BroadcastPlan>, Shape.broadcast_checked(shape, target), Done{plan}, planned) : 126 P(Broadcast.planned(input, target, _)) 127 %Equal.sym(Bool, Layout.words_same(Shape.dimensions(output), Shape.dimensions(target)), True{}, ReduceTo.fitted_same(target, plan, same)) : 128 P(Broadcast.fitted(input, target, plan, _)) 129 Replayed.new_lift(target, 0.0, result => P(Broadcast.zeroed(input, target, plan, result)), 130 zeros_fits => Binary.allocated_replayed(Elementwise.KeepLeft{}, plan, shape, fill, ops, target, 0.0, Nil{}, planned, 131 pair => P(Broadcast.finish(pair)), 132 kept => kept_words => input_fits => zeros_kept => zeros_words => zeros_kept_fits => output_ops => values => output_fits => 133 h(kept, kept_words, input_fits, output_ops, 134 Equal.trans(List<&2, F32>, Contracts.words(output, 0.0, output_ops, count), 135 BinaryLaws.cells(Elementwise.KeepLeft{}, plan, shape, fill, ops, target, 0.0, Nil{}, count), Laws.cells(plan, shape, fill, ops, count), 136 values, cells_keep(plan, shape, fill, ops, target, count)), 137 output_fits))(fits, zeros_fits)) 138 139 def output_below(+plan: Shape.BroadcastPlan, +target: Shape.Shape, +index: U32, 140 same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(target) : List<&2, U32>}, 141 inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(target))) == True{} : Bool}) -> 142 {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}: 143 %Equal.sym(List<&2, U32>, Shape.dimensions(Shape.output_shape(plan)), Shape.dimensions(target), same) : 144 {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(_)) == True{} : Bool} 145 inside 146 147 def broadcast_to_pointwise(+shape: Shape.Shape, +target: Shape.Shape, +plan: Shape.BroadcastPlan, 148 +planned: {Shape.broadcast_checked(shape, target) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}, 149 +same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(target) : List<&2, U32>}, 150 +fill: F32, +ops: List<&2, StorageLaws.Op>, 151 -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}) -> 152 Laws.BroadcastToContract Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))), 154 other => Replayed.value(shape, fill, ops, other), 155 Broadcast.broadcast_to(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target)>: 156 +output = Shape.output_shape(plan) 157 +count = Shape.logical_numel(Shape.dimensions(output)) 158 +total = Shape.logical_numel(Shape.dimensions(shape)) 159 run_lift(shape, target, plan, planned, same, fill, ops, 160 result => Laws.BroadcastToContract Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))), 162 other => Replayed.value(shape, fill, ops, other), result>, 163 kept => kept_words => input_fits => output_ops => values => output_fits => 164 Laws.BroadcastToContract{ 165 {==}, 166 Equal.cong(List<&2, U32>, Maybe<&2, List<&2, U32>>, dimensions => Some{dimensions}, Shape.dimensions(output), Shape.dimensions(target), same), 167 index => inside => 168 +below = output_below(plan, target, index, same, inside) 169 Equal.trans(Maybe<&2, F32>, 170 Core.dynamic_read_value(Core.read_dynamic(Core.Packed{output, Core.TTensor{Replayed.storage(output, 0.0, output_ops), output_fits}}, index)), 171 Some{Replayed.value(output, 0.0, output_ops, index)}, 172 Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, U32.to_nat(index))))}, 173 PointwiseProof.replayed_point(output, 0.0, output_ops, output_fits, index, below), 174 Equal.trans(Maybe<&2, F32>, Some{Replayed.value(output, 0.0, output_ops, index)}, 175 ReduceLaws.word_at(Laws.cells(plan, shape, fill, ops, count), U32.to_nat(index)), 176 Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, U32.to_nat(index))))}, 177 listed_value(output, 0.0, output_ops, count, Laws.cells(plan, shape, fill, ops, count), index, values, below), 178 cells_at(plan, shape, fill, ops, count, U32.to_nat(index), below))), 179 other => inside => 180 +within = {inside : {Nat.is_lt(U32.to_nat(other), total) == True{} : Bool}} 181 Equal.trans(Maybe<&2, F32>, 182 Core.dynamic_read_value(Core.read_dynamic(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, kept), input_fits}}, other)), 183 Some{Replayed.value(shape, fill, kept, other)}, Some{Replayed.value(shape, fill, ops, other)}, 184 PointwiseProof.replayed_point(shape, fill, kept, input_fits, other, within), 185 Equal.trans(Maybe<&2, F32>, Some{Replayed.value(shape, fill, kept, other)}, 186 ReduceLaws.word_at(Contracts.words(shape, fill, ops, total), U32.to_nat(other)), Some{Replayed.value(shape, fill, ops, other)}, 187 listed_value(shape, fill, kept, total, Contracts.words(shape, fill, ops, total), other, kept_words, within), 188 Equal.sym(Maybe<&2, F32>, Some{Replayed.value(shape, fill, ops, other)}, 189 ReduceLaws.word_at(Contracts.words(shape, fill, ops, total), U32.to_nat(other)), 190 listed_value(shape, fill, ops, total, Contracts.words(shape, fill, ops, total), other, {==}, within))))})(fits) 191 192 def Laws.broadcast_to_pointwise(shape, target, plan, planned, same, fill, ops, fits): 193 broadcast_to_pointwise(shape, target, plan, planned, same, fill, ops, fits) 194 195 def Laws.broadcast_to_replayed(shape, target, plan, planned, same, fill, ops, fits, P, h): 196 run_lift(shape, target, plan, planned, same, fill, ops, P, h)(fits)