1 import Base 2 import ../../shape/core.bend as Shape 3 import ../../shape/flat.bend as Flat 4 import ../../storage/LAWS.bend as StorageLaws 5 import ../../storage/arith.bend as Arith 6 import ../core.bend as Core 7 import ../lifecycle.bend as Lifecycle 8 import ../elementwise.bend as Elementwise 9 import ../word_conversion.bend as Conversion 10 import ./replayed.bend as Replayed 11 import ./pointwise.bend as Pointwise 12 import ./pointwise_proof.bend as Proof 13 import ./pointwise_records.bend as Records 14 import ./broadcast.bend as Broadcast 15 16 def output_inside(+plan: Shape.BroadcastPlan, +index: Nat, 17 below: {Nat.is_lt(index, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) -> 18 {True{} == Nat.is_lt(index, Shape.elements(Shape.output_shape(plan))) : Bool}: 19 %Equal.sym(Nat, Shape.elements(Shape.output_shape(plan)), Shape.numel(Shape.extents(Shape.output_shape(plan))), 20 Shape.elements_exact(Shape.output_shape(plan))) : {True{} == Nat.is_lt(index, _) : Bool} 21 Equal.sym(Bool, Nat.is_lt(index, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))), True{}, below) 22 23 def left_walk_bound(+plan: Shape.BroadcastPlan, 24 map: Shape.Index, 25 +index: Nat, below: {Nat.is_lt(index, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) -> 26 {Nat.is_lt(Broadcast.left_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), map, index), 27 Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan)))) == True{} : Bool}: 28 match map: 29 case Shape.Indexed{rows, exact}: 30 Equal.sym(Bool, True{}, Nat.is_lt(Shape.walk_left(rows, index), Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan)))), 31 %Shape.elements_exact(Shape.left_shape(plan)) : {True{} == Nat.is_lt(Shape.walk_left(rows, index), _) : Bool} 32 Flat.left_bound(plan, rows, exact, index, output_inside(plan, index, below))) 33 case Shape.Unindexed{zero}: 34 Empty.absurd({Nat.is_lt(0n, Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan)))) == True{} : Bool}, 35 Arith.lt_zero(index, 36 %Equal.sym(Nat, 0n, Shape.numel(Shape.extents(Shape.output_shape(plan))), zero) : {Nat.is_lt(index, _) == True{} : Bool} 37 below)) 38 39 def right_walk_bound(+plan: Shape.BroadcastPlan, 40 map: Shape.Index, 41 +index: Nat, below: {Nat.is_lt(index, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) -> 42 {Nat.is_lt(Broadcast.right_walk(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan)), map, index), 43 Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))) == True{} : Bool}: 44 match map: 45 case Shape.Indexed{rows, exact}: 46 Equal.sym(Bool, True{}, Nat.is_lt(Shape.walk_right(rows, index), Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))), 47 %Shape.elements_exact(Shape.right_shape(plan)) : {True{} == Nat.is_lt(Shape.walk_right(rows, index), _) : Bool} 48 Flat.right_bound(plan, rows, exact, index, output_inside(plan, index, below))) 49 case Shape.Unindexed{zero}: 50 Empty.absurd({Nat.is_lt(0n, Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))) == True{} : Bool}, 51 Arith.lt_zero(index, 52 %Equal.sym(Nat, 0n, Shape.numel(Shape.extents(Shape.output_shape(plan))), zero) : {Nat.is_lt(index, _) == True{} : Bool} 53 below)) 54 55 def left_inside(+plan: Shape.BroadcastPlan, +count: Nat, 56 +room: {Nat.is_le(count, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) -> 57 {Pointwise.inside(~Shape.BroadcastPlan, ~Broadcast.left_map, plan, count, Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan)))) == True{} : Bool}: 58 match count: 59 case 0n: 60 {==} 61 case 1n+previous: 62 +total = Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) 63 +bound = Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan))) 64 %Equal.sym(Bool, Pointwise.inside(~Shape.BroadcastPlan, ~Broadcast.left_map, plan, previous, bound), True{}, 65 left_inside(plan, previous, Conversion.le_predecessor(previous, total, room))) : 66 {Bool.and(_, Nat.is_lt(Broadcast.left_map(plan, previous), bound)) == True{} : Bool} 67 %Equal.sym(Bool, Nat.is_lt(Broadcast.left_map(plan, previous), bound), True{}, 68 left_walk_bound(plan, Shape.index_map(plan), previous, Conversion.le_successor_lt(previous, total, room))) : 69 {Bool.and(True{}, _) == True{} : Bool} 70 {==} 71 72 def right_inside(+plan: Shape.BroadcastPlan, +count: Nat, 73 +room: {Nat.is_le(count, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}) -> 74 {Pointwise.inside(~Shape.BroadcastPlan, ~Broadcast.right_map, plan, count, Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan)))) == True{} : Bool}: 75 match count: 76 case 0n: 77 {==} 78 case 1n+previous: 79 +total = Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) 80 +bound = Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) 81 %Equal.sym(Bool, Pointwise.inside(~Shape.BroadcastPlan, ~Broadcast.right_map, plan, previous, bound), True{}, 82 right_inside(plan, previous, Conversion.le_predecessor(previous, total, room))) : 83 {Bool.and(_, Nat.is_lt(Broadcast.right_map(plan, previous), bound)) == True{} : Bool} 84 %Equal.sym(Bool, Nat.is_lt(Broadcast.right_map(plan, previous), bound), True{}, 85 right_walk_bound(plan, Shape.index_map(plan), previous, Conversion.le_successor_lt(previous, total, room))) : 86 {Bool.and(True{}, _) == True{} : Bool} 87 {==} 88 89 def broadcast_pointwise(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, 90 +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>, 91 -left_fits: {Lifecycle.logical_length(Replayed.storage(Shape.left_shape(plan), left_fill, left_ops)) == 92 Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan))) : Nat}, 93 +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, 94 -right_fits: {Lifecycle.logical_length(Replayed.storage(Shape.right_shape(plan), right_fill, right_ops)) == 95 Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) : Nat}) -> 96 Records.PointwiseContract Broadcast.cell(operation, plan, left_fill, left_ops, right_fill, right_ops, index), 98 other => Replayed.value(Shape.left_shape(plan), left_fill, left_ops, other), 99 other => Replayed.value(Shape.right_shape(plan), right_fill, right_ops, other), 100 Broadcast.binary(operation, plan, 101 Core.Packed{Shape.left_shape(plan), Core.TTensor{Replayed.storage(Shape.left_shape(plan), left_fill, left_ops), left_fits}}, 102 Core.Packed{Shape.right_shape(plan), Core.TTensor{Replayed.storage(Shape.right_shape(plan), right_fill, right_ops), right_fits}})>: 103 +total = Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) 104 Proof.contract(~Elementwise.Operation, ~Elementwise.scalar, ~Shape.BroadcastPlan, ~Broadcast.left_map, ~Broadcast.right_map, 105 operation, plan, Shape.left_shape(plan), left_fill, left_ops, Shape.right_shape(plan), right_fill, right_ops, Shape.output_shape(plan), 106 left_inside(plan, total, Conversion.le_reflexive(total)), right_inside(plan, total, Conversion.le_reflexive(total)), left_fits, right_fits)