src/autograd/reduce_to_fold_PROOF.bend
bendtorch-core@0.1.0.0 · 258 lines · 21 definitions
1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/storage/main.bend as Storage
4import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
5import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
6import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
7import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
8import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
9import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
10import bendtorch-core@0.1.0.0/src/tensor/word_conversion.bend as Conversion
11import bendtorch-core@0.1.0.0/src/tensor/permutation_mapping.bend as Mapping
12import bendtorch-core@0.1.0.0/src/tensor/projection_LAWS.bend as Projection
13import bendtorch-core@0.1.0.0/src/tensor/projection_PROOF.bend as ProjectionProof
14import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
15import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise.bend as Pointwise
16import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise_proof.bend as PointwiseProof
17import bendtorch-core@0.1.0.0/src/tensor/contracts/broadcast_proof.bend as BroadcastProof
18import bendtorch-core@0.1.0.0/src/tensor/contracts/fold.bend as Fold
19import bendtorch-core@0.1.0.0/src/tensor/contracts/sum.bend as Sum
20import bendtorch-core@0.1.0.0/src/tensor/contracts/LAWS.bend as Contracts
21import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce
22import bendtorch-core@0.1.0.0/src/autograd/reduce_LAWS.bend as ReduceLaws
23import bendtorch-core@0.1.0.0/src/autograd/reduce_fold.bend as ReduceFold
24import bendtorch-core@0.1.0.0/src/autograd/reduce_fold_PROOF.bend as FoldProof
25import bendtorch-core@0.1.0.0/src/autograd/reduce_fold_LAWS.bend as Laws
26import bendtorch-core@0.1.0.0/src/storage/api.bend as Api
27import bendtorch-core@0.1.0.0/src/storage/array.bend as ArrayCore
28import bendtorch-core@0.1.0.0/src/tensor/read_contract.bend as Reads
29import bendtorch-core@0.1.0.0/src/tensor/contracts/binary_storage_proof.bend as Owned
30import bendtorch-core@0.1.0.0/src/autograd/reduce_replayed_proof.bend as Replay
31
32def 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
39def shapes_construct(+left: Shape.Shape, +right: Shape.Shape, -left_padding: Nat, -right_padding: Nat,
40 -minimal: Shape.PadMinimal<left_padding, right_padding>, -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<output>) ->
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
51def shapes_align(+left: Shape.Shape, +right: Shape.Shape, +left_padding: Nat, +right_padding: Nat,
52 -minimal: Shape.PadMinimal<left_padding, right_padding>,
53 +alignment: Shape.Alignment<Shape.pad_words(left_padding, Shape.dimensions(left)), Shape.pad_words(right_padding, Shape.dimensions(right))>) ->
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
64def 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
72def first(pair: Shape.Shape & Shape.Shape) -> Shape.Shape:
73 (left, right) = pair
74 left
75
76def second(pair: Shape.Shape & Shape.Shape) -> Shape.Shape:
77 (left, right) = pair
78 right
79
80def 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
86def 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
92def 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
98def target_index(+plan: Shape.BroadcastPlan, +position: Nat) -> U32:
99 U32.from_nat(ReduceFold.target_offset(plan, position))
100
101def source_index(+plan: Shape.BroadcastPlan, +position: Nat) -> U32:
102 U32.from_nat(ReduceFold.contribution_offset(plan, position))
103
104
105
106def 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
112def 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
118def 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
121def 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
132def 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
137def 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
144def 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
152
153def 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<plan, shape, U32.to_nat(Shape.count(target))> =
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
194
195def 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
212def 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<target, shape,
218 index => 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
257def Laws.reduce_to_fiber(target, shape, plan, planned, same, fill, ops, fits):
258 reduce_to_fiber(target, shape, plan, planned, same, fill, ops, fits)