1 import Base 2 import ../shape/core.bend as Shape 3 import ../storage/LAWS.bend as StorageLaws 4 import ../tensor/core.bend as Core 5 import ../tensor/lifecycle.bend as Lifecycle 6 import ../tensor/contracts/LAWS.bend as Contracts 7 import ../tensor/contracts/replayed.bend as Replayed 8 import ../tensor/contracts/broadcast.bend as PointwiseBroadcast 9 import ./broadcast.bend as Broadcast 10 11 def succeeded(result: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>) -> Bool: 12 (input, outcome) = result 13 Result.is_done(&2, &1, Broadcast.Error, Core.Dynamic, outcome) 14 15 def output_dimensions(result: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>) -> Maybe<&2, List<&2, U32>>: 16 (input, outcome) = result 17 match outcome: 18 case Fail{error}: 19 None{} 20 case Done{Core.Packed{shape, tensor}}: 21 Some{Shape.dimensions(shape)} 22 23 def output_point(result: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>, +index: U32) -> Maybe<&2, F32>: 24 (input, outcome) = result 25 match outcome: 26 case Fail{error}: 27 None{} 28 case Done{output}: 29 Core.dynamic_read_value(Core.read_dynamic(output, index)) 30 31 def input_point(result: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>, +index: U32) -> Maybe<&2, F32>: 32 (input, outcome) = result 33 Core.dynamic_read_value(Core.read_dynamic(input, index)) 34 35 # What broadcast_to promises, stated only through the result: it succeeds 36 # with the target's dimensions, every in-bounds output index i holds cell(i), 37 # and the returned input still holds every word. 38 type BroadcastToContract<-target: Shape.Shape, -input: Shape.Shape, -cell: Nat -> F32, -word: U32 -> F32, 39 -result: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic>> is Type: 40 BroadcastToContract{ 41 successful: {succeeded(result) == True{} : Bool}, 42 dimensions: {output_dimensions(result) == Some{Shape.dimensions(target)} : Maybe<&2, List<&2, U32>>}, 43 value: @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(target))) == True{} : Bool} -> 44 {output_point(result, index) == Some{cell(U32.to_nat(index))} : Maybe<&2, F32>}, 45 kept: @+other: U32 -> @inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool} -> 46 {input_point(result, other) == Some{word(other)} : Maybe<&2, F32>}} 47 48 # The words of broadcast_to's output in index order: word k is the input 49 # word at left_map(plan, k), its right-aligned broadcast projection. 50 def cells(+plan: Shape.BroadcastPlan, +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat) -> List<&2, F32>: 51 match count: 52 case 0n: 53 Nil{} 54 case 1n+previous: 55 List.append(&2, F32, cells(plan, shape, fill, ops, previous), 56 [Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, previous)))]) 57 58 # What the transport form of broadcast_to_replayed asks of P: P holds for the 59 # input, replayed with any history that keeps its words, next to an output of 60 # the plan's shape, replayed over +0.0 with any history whose words are 61 # `cells`. 62 def BroadcastWith(+plan: Shape.BroadcastPlan, +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, 63 -P: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic> -> Type) -> Type: 64 @+kept: List<&2, StorageLaws.Op> -> 65 @kept_words: {Contracts.words(shape, fill, kept, Shape.logical_numel(Shape.dimensions(shape))) == 66 Contracts.words(shape, fill, ops, Shape.logical_numel(Shape.dimensions(shape))) : List<&2, F32>} -> 67 @-input_fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, kept)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} -> 68 @+output_ops: List<&2, StorageLaws.Op> -> 69 @values: {Contracts.words(Shape.output_shape(plan), 0.0, output_ops, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == 70 cells(plan, shape, fill, ops, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) : List<&2, F32>} -> 71 @-output_fits: {Lifecycle.logical_length(Replayed.storage(Shape.output_shape(plan), 0.0, output_ops)) == 72 Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) : Nat} -> 73 P((Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, kept), input_fits}}, 74 Done{Core.Packed{Shape.output_shape(plan), Core.TTensor{Replayed.storage(Shape.output_shape(plan), 0.0, output_ops), output_fits}}})) 75 76 # Pointwise family, broadcast_to instance: when the input shape broadcasts to 77 # exactly the target, broadcasting a replayed input succeeds with the target's 78 # dimensions, output index i holds the input word at left_map(plan, i), its 79 # right-aligned broadcast projection, and the returned input keeps every word. 80 law broadcast_to_pointwise: 81 for +shape: Shape.Shape 82 for +target: Shape.Shape 83 for +plan: Shape.BroadcastPlan 84 for planned: {Shape.broadcast_checked(shape, target) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>} 85 for same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(target) : List<&2, U32>} 86 for +fill: F32 87 for +ops: List<&2, StorageLaws.Op> 88 for -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} 89 BroadcastToContract Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))), 91 other => Replayed.value(shape, fill, ops, other), 92 Broadcast.broadcast_to(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target)> 93 94 # The broadcast_to result is replayed: for every P, if P holds for every 95 # result that hands back the input with its words and gives an output of the 96 # plan's shape whose words are `cells`, then P holds for broadcast_to. The 97 # output word at k is the input word at left_map(plan, k), whatever split of 98 # the output into parallel halves the build takes, and the output is a 99 # replayed tensor, so a later law about a replayed contribution, such as 100 # reduce_to_fiber, applies to it. 101 law broadcast_to_replayed: 102 for +shape: Shape.Shape 103 for +target: Shape.Shape 104 for +plan: Shape.BroadcastPlan 105 for planned: {Shape.broadcast_checked(shape, target) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>} 106 for same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(target) : List<&2, U32>} 107 for +fill: F32 108 for +ops: List<&2, StorageLaws.Op> 109 for -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} 110 for -P: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic> -> Type 111 for h: BroadcastWith(plan, shape, fill, ops, P) 112 P(Broadcast.broadcast_to(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target))