1 import Base
2 import ../shape/core.bend as Shape
3 import ./axis_spec.bend as Axis
4 import ./layout.bend as Layout
5 import ./observation.bend as Values
6 import ./permute_build.bend as Build
7 import ./permutation_mapping.bend as Mapping
8 import ./word_conversion.bend as Conversion
9 import ../storage/arith.bend as Arith
10 import ../storage/model.bend as Model
11 import ../storage/main.bend as Storage
12 import ../storage/refinement.bend as Refinement
13 import ../storage/clone.bend as Clone
14 import ../storage/route.bend as Route
15 import ./core.bend as Core
16 import ./lifecycle.bend as Lifecycle
17 import ./transforms.bend as Transforms
18 import ./parallel_proofs.bend as Generic
19 import ./parallel.bend as Parallel
20 import ./transform_totality.bend as Totality
21 import ./read_contract.bend as Reads
22 import ./contracts/replayed.bend as Replayed
23 import ./contracts/binary_build_proof.bend as BinaryBuild
24 import ./contracts/binary_storage_proof.bend as Owned
25 import ./contracts/build_facts.bend as Facts
26 import ../shape/index_proofs.bend as Index
27 import ../storage/LAWS.bend as StorageLaws
28 import ../storage/api.bend as Api
29 import ../storage/array.bend as ArrayCore
30
31 # `Build.source` computes the index `Values.permute_source_index` names. The
32 # proof follows `permute_source` step by step: the coordinate `Axis.select`
33 # takes out of `Shape.digits` is the one its pick computes, the list
34 # `Axis.scatter` builds is the picks' coordinates, and `Shape.offset` of
35 # that list is `Build.offset` with the strides worked out once.
36
37 def coordinate_of(pick: Maybe<&2, Build.Pick>, +x: Nat) -> Maybe<&2, U32>:
38 match pick:
39 case None{}:
40 None{}
41 case Some{found}:
42 Some{U32.from_nat(Build.coordinate(found, x))}
43
44 def coordinates(+picks: List<&2, Build.Pick>, +x: Nat) -> List<&2, U32>:
45 match picks:
46 case Nil{}:
47 Nil{}
48 case +pick <> rest:
49 U32.from_nat(Build.coordinate(pick, x)) <> coordinates(rest, x)
50
51 def coordinates_of(picks: Maybe<&2, List<&2, Build.Pick>>, +x: Nat) -> Maybe<&2, List<&2, U32>>:
52 match picks:
53 case None{}:
54 None{}
55 case Some{+found}:
56 Some{coordinates(found, x)}
57
58 # Widening a pick by an outer count reduces the index modulo that count
59 # first.
60 def widen_coordinate(+outer: Nat, +pick: Maybe<&2, Build.Pick>, +x: Nat) ->
61 {coordinate_of(Build.widen(outer, pick), x) == coordinate_of(pick, Nat.mod(x, outer)) : Maybe<&2, U32>}:
62 match pick:
63 case None{}:
64 {==}
65 case Some{+found}:
66 Build.Pick{+mods, +size} = found
67 {==}
68
69 # Coordinate k of `Shape.digits(dims, x)` is what the pick of k computes.
70 def lookup_digits(+dims: List<&2, Nat>, +x: Nat, +k: Nat) ->
71 {Axis.lookup(U32, Shape.words(Shape.digits(dims, x)), k) == coordinate_of(Build.pick_at(dims, k), x) : Maybe<&2, U32>}:
72 match dims:
73 case Nil{}:
74 {==}
75 case +head <> +tail:
76 match k:
77 case 0n:
78 {==}
79 case 1n+j:
80 Equal.trans(Maybe<&2, U32>, Axis.lookup(U32, Shape.words(Shape.digits(tail, Nat.mod(x, Shape.numel(tail)))), j),
81 coordinate_of(Build.pick_at(tail, j), Nat.mod(x, Shape.numel(tail))),
82 coordinate_of(Build.widen(Shape.numel(tail), Build.pick_at(tail, j)), x),
83 lookup_digits(tail, Nat.mod(x, Shape.numel(tail)), j),
84 Equal.sym(Maybe<&2, U32>, coordinate_of(Build.widen(Shape.numel(tail), Build.pick_at(tail, j)), x),
85 coordinate_of(Build.pick_at(tail, j), Nat.mod(x, Shape.numel(tail))),
86 widen_coordinate(Shape.numel(tail), Build.pick_at(tail, j), x)))
87
88 def select_digits(+dims: List<&2, Nat>, +x: Nat, +position: Maybe<&2, Nat>) ->
89 {Axis.select(U32, Shape.words(Shape.digits(dims, x)), position) == coordinate_of(Build.located(dims, position), x) : Maybe<&2, U32>}:
90 match position:
91 case None{}:
92 {==}
93 case Some{+k}:
94 lookup_digits(dims, x, k)
95
96 def prepend_coordinates(+head: Maybe<&2, Build.Pick>, +tail: Maybe<&2, List<&2, Build.Pick>>, +x: Nat) ->
97 {Axis.prepend(U32, coordinate_of(head, x), coordinates_of(tail, x)) == coordinates_of(Axis.prepend(Build.Pick, head, tail), x) :
98 Maybe<&2, List<&2, U32>>}:
99 match head:
100 case None{}:
101 {==}
102 case Some{+pick}:
103 match tail:
104 case None{}:
105 {==}
106 case Some{+rest}:
107 {==}
108
109 # The list `Axis.scatter` builds from the output coordinates of x is the
110 # picks' coordinates.
111 def scatter_digits(+rank: Nat, +axes: List<&2, Nat>, +out: List<&2, Nat>, +x: Nat, +cursor: Nat) ->
112 {Axis.scatter(U32, rank, axes, Shape.words(Shape.digits(out, x)), cursor) == coordinates_of(Build.picks(rank, axes, out, cursor), x) :
113 Maybe<&2, List<&2, U32>>}:
114 match rank:
115 case 0n:
116 {==}
117 case 1n+p:
118 +position = Axis.position(axes, cursor)
119 %Equal.sym(Maybe<&2, U32>, Axis.select(U32, Shape.words(Shape.digits(out, x)), position), coordinate_of(Build.located(out, position), x),
120 select_digits(out, x, position)) :
121 {Axis.prepend(U32, _, Axis.scatter(U32, p, axes, Shape.words(Shape.digits(out, x)), 1n+cursor)) ==
122 coordinates_of(Axis.prepend(Build.Pick, Build.located(out, position), Build.picks(p, axes, out, 1n+cursor)), x) : Maybe<&2, List<&2, U32>>}
123 %Equal.sym(Maybe<&2, List<&2, U32>>, Axis.scatter(U32, p, axes, Shape.words(Shape.digits(out, x)), 1n+cursor),
124 coordinates_of(Build.picks(p, axes, out, 1n+cursor), x), scatter_digits(p, axes, out, x, 1n+cursor)) :
125 {Axis.prepend(U32, coordinate_of(Build.located(out, position), x), _) ==
126 coordinates_of(Axis.prepend(Build.Pick, Build.located(out, position), Build.picks(p, axes, out, 1n+cursor)), x) : Maybe<&2, List<&2, U32>>}
127 prepend_coordinates(Build.located(out, position), Build.picks(p, axes, out, 1n+cursor), x)
128
129 # `Shape.offset` of the coordinates is `Build.offset` over the strides.
130 def offset_coordinates(+dims: List<&2, Nat>, +picks: List<&2, Build.Pick>, +x: Nat) ->
131 {Shape.offset(dims, Shape.naturals(coordinates(picks, x))) == Build.offset(Build.strides(dims), picks, x) : Nat}:
132 match dims:
133 case Nil{}:
134 {==}
135 case +head <> +tail:
136 match picks:
137 case Nil{}:
138 {==}
139 case +pick <> +rest:
140 Equal.cong(Nat, Nat, rest_offset => Nat.add(Nat.mul(U32.to_nat(U32.from_nat(Build.coordinate(pick, x))), Shape.numel(tail)), rest_offset),
141 Shape.offset(tail, Shape.naturals(coordinates(rest, x))), Build.offset(Build.strides(tail), rest, x),
142 offset_coordinates(tail, rest, x))
143
144 def found_coordinates(+dims: List<&2, U32>, picks: Maybe<&2, List<&2, Build.Pick>>, +index: Nat) ->
145 {U32.to_nat(Maybe.default(&2, U32, Layout.reference_source_coordinates(dims, coordinates_of(picks, U32.to_nat(U32.from_nat(index)))), 0)) ==
146 Build.found(picks, Build.strides(Shape.naturals(dims)), index) : Nat}:
147 match picks:
148 case None{}:
149 {==}
150 case Some{+found}:
151 Equal.cong(Nat, Nat, offset => U32.to_nat(U32.from_nat(offset)),
152 Shape.offset(Shape.naturals(dims), Shape.naturals(coordinates(found, U32.to_nat(U32.from_nat(index))))),
153 Build.offset(Build.strides(Shape.naturals(dims)), found, U32.to_nat(U32.from_nat(index))),
154 offset_coordinates(Shape.naturals(dims), found, U32.to_nat(U32.from_nat(index))))
155
156 # The build's index map is `permute_source_index`.
157 def source_index(+input: Shape.Shape, +axes: List<&2, U32>, +index: Nat) ->
158 {Build.source(Build.recipe(Build.Request{input, axes}), index) == Values.permute_source_index(input, axes, index) : Nat}:
159 +dims = Shape.dimensions(input)
160 +out = Shape.naturals(Axis.order(U32, dims, Shape.naturals(axes)))
161 +x = U32.to_nat(U32.from_nat(index))
162 +picks = Build.picks(Shape.rank_dimensions(dims), Shape.naturals(axes), out, 0n)
163 Equal.sym(Nat, Values.permute_source_index(input, axes, index), Build.found(picks, Build.strides(Shape.naturals(dims)), index),
164 %Equal.sym(Maybe<&2, List<&2, U32>>, Axis.scatter(U32, Shape.rank_dimensions(dims), Shape.naturals(axes), Shape.words(Shape.digits(out, x)), 0n),
165 coordinates_of(picks, x), scatter_digits(Shape.rank_dimensions(dims), Shape.naturals(axes), out, x, 0n)) :
166 {U32.to_nat(Maybe.default(&2, U32, Layout.reference_source_coordinates(dims, _), 0)) ==
167 Build.found(picks, Build.strides(Shape.naturals(dims)), index) : Nat}
168 found_coordinates(dims, picks, index))
169
170 # The word index a cell reads is `permute_source_index` as a word.
171 def source_word(+input: Shape.Shape, +axes: List<&2, U32>, +index: Nat) ->
172 {Build.at(Build.recipe(Build.Request{input, axes}), index) == U32.from_nat(Values.permute_source_index(input, axes, index)) : U32}:
173 Equal.cong(Nat, U32, number => U32.from_nat(number), Build.source(Build.recipe(Build.Request{input, axes}), index),
174 Values.permute_source_index(input, axes, index), source_index(input, axes, index))
175
176 # The mapping a permute's plan gives an output index names the input index
177 # `permute_source` names, and it lies inside the input.
178 def reference(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan, +index: U32,
179 +mapping: Layout.Mapping) ->
180 {Some{Layout.source_index(input, axes, plan, index, mapping)} == Values.permute_source(input, axes, index) : Maybe<&2, U32>}:
181 Equal.trans(Maybe<&2, U32>, Some{Layout.source_index(input, axes, plan, index, mapping)}, Layout.reference_source(input, axes, plan, index),
182 Values.permute_source(input, axes, index),
183 Mapping.mapping_reference(input, axes, plan, index, mapping),
184 Equal.cong(List<&2, U32>, Maybe<&2, U32>,
185 dimensions => Layout.reference_source_coordinates(Shape.dimensions(input),
186 Axis.scatter(U32, Shape.rank_dimensions(Shape.dimensions(input)), Shape.naturals(axes),
187 Shape.words(Shape.digits(Shape.naturals(dimensions), U32.to_nat(index))), 0n)),
188 Shape.dimensions(Layout.output(input, axes, plan)), Axis.order(U32, Shape.dimensions(input), Shape.naturals(axes)),
189 Mapping.permutation_dimensions_order(input, axes, plan)))
190
191 def mapped_index(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan, +index: U32,
192 +mapping: Layout.Mapping) ->
193 {Layout.source_index(input, axes, plan, index, mapping) ==
194 U32.from_nat(U32.to_nat(Maybe.default(&2, U32, Values.permute_source(input, axes, index), 0))) : U32}:
195 +found = Maybe.default(&2, U32, Values.permute_source(input, axes, index), 0)
196 Equal.trans(U32, Layout.source_index(input, axes, plan, index, mapping), found, U32.from_nat(U32.to_nat(found)),
197 Equal.cong(Maybe<&2, U32>, U32, result => Maybe.default(&2, U32, result, 0),
198 Some{Layout.source_index(input, axes, plan, index, mapping)}, Values.permute_source(input, axes, index),
199 reference(input, axes, plan, index, mapping)),
200 Equal.sym(U32, U32.from_nat(U32.to_nat(found)), found, Conversion.word_natural_exact(found)))
201
202 def mapped_bound(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan, +index: U32,
203 mapping: Layout.Mapping) ->
204 {Nat.is_lt(U32.to_nat(Layout.source_index(input, axes, plan, index, mapping)), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool}:
205 Layout.Mapped{oi, oc, ic, found, requested, decoded, inverse, flattened, roundtrip} = mapping
206 Mapping.flat_mathematical_bounded(input, found)
207
208 def mapped_inside(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan, +index: U32,
209 +mapping: Layout.Mapping) ->
210 {Nat.is_lt(U32.to_nat(U32.from_nat(U32.to_nat(Maybe.default(&2, U32, Values.permute_source(input, axes, index), 0)))),
211 Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool}:
212 %mapped_index(input, axes, plan, index, mapping) : {Nat.is_lt(U32.to_nat(_), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool}
213 mapped_bound(input, axes, plan, index, mapping)
214
215 def index_bounded(+shape: Shape.Shape, +index: U32,
216 +inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool}) ->
217 {True{} == U32.is_lt(index, Shape.count(shape)) : Bool}:
218 %Arith.u32_lt(index, Shape.count(shape)) : {True{} == _ : Bool}
219 %Equal.sym(Nat, U32.to_nat(Shape.count(shape)), Shape.elements(shape), Conversion.shape_count_exact(shape)) : {True{} == Nat.is_lt(U32.to_nat(index), _) : Bool}
220 %Equal.sym(Nat, Shape.elements(shape), Shape.logical_numel(Shape.dimensions(shape)), Shape.elements_exact(shape)) : {True{} == Nat.is_lt(U32.to_nat(index), _) : Bool}
221 Equal.sym(Bool, Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(shape))), True{}, inside)
222
223 def mapping_done(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan, +index: U32,
224 +inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) == True{} : Bool}) ->
225 {Result.is_done(&2, &2, Layout.Error, Layout.Mapping, Layout.map_index(input, axes, plan, index)) == True{} : Bool}:
226 Mapping.mapping_decided_success(input, axes, plan, index,
227 Shape.bool_case(U32.is_lt(index, Shape.count(Layout.output(input, axes, plan)))),
228 index_bounded(Layout.output(input, axes, plan), index, inside))
229
230 # The build's cells. The state is the input array next to an empty array,
231 # `(a, Array.new(F32, 0n, 0.0))`. Every cell hands it back, and a split hands it to both halves,
232 # so the generic law of `Parallel.build` gives what every slot holds.
233
234 type Fact<-P: Type> is Data:
235 Fact{-proof: P}
236
237 def take_first(-e: Array, pair: Array & F32) ->
238 {Pair.snd(F32, Array & Array, Build.take(e, pair)) == (Refinement.first(F32, pair), e) : Array & Array}:
239 (x, w) = pair
240 {==}
241
242 def take_word(-e: Array, pair: Array & F32) ->
243 {Pair.fst(F32, Array & Array, Build.take(e, pair)) == Refinement.second(F32, pair) : F32}:
244 (x, w) = pair
245 {==}
246
247 def cell_keeps(+d: Nat, -a: Array, +request: Build.Request, +position: Nat, inside: Bool, -valid: Model.Balanced(F32, d, a)) ->
248 Fact<{Pair.snd(F32, Array & Array, Build.cell(Build.recipe(request), position, (a, Array.new(F32, 0n, 0.0)), inside)) == (a, Array.new(F32, 0n, 0.0)) : Array & Array}>:
249 match inside:
250 case False{}:
251 Fact{{==}}
252 case True{}:
253 +index = Build.at(Build.recipe(request), position)
254 Fact{Equal.trans(Array & Array, Pair.snd(F32, Array & Array, Build.take(Array.new(F32, 0n, 0.0), Array.get(F32, a, index))),
255 (Refinement.first(F32, Array.get(F32, a, index)), Array.new(F32, 0n, 0.0)), (a, Array.new(F32, 0n, 0.0)),
256 take_first(Array.new(F32, 0n, 0.0), Array.get(F32, a, index)),
257 Equal.cong(Array, Array & Array, x => (x, Array.new(F32, 0n, 0.0)), Refinement.first(F32, Array.get(F32, a, index)), a,
258 Refinement.law_get_retains(F32, d, a, index, valid)))}
259
260 # An inside cell holds the input word at `permute_source_index`.
261 def cell_word(+input: Shape.Shape, +axes: List<&2, U32>, -a: Array, +position: Nat) ->
262 Fact<{Pair.fst(F32, Array & Array, Build.cell(Build.recipe(Build.Request{input, axes}), position, (a, Array.new(F32, 0n, 0.0)), True{})) ==
263 Refinement.second(F32, Array.get(F32, a, U32.from_nat(Values.permute_source_index(input, axes, position)))) : F32}>:
264 +index = Build.at(Build.recipe(Build.Request{input, axes}), position)
265 Fact{Equal.trans(F32, Pair.fst(F32, Array & Array, Build.take(Array.new(F32, 0n, 0.0), Array.get(F32, a, index))),
266 Refinement.second(F32, Array.get(F32, a, index)),
267 Refinement.second(F32, Array.get(F32, a, U32.from_nat(Values.permute_source_index(input, axes, position)))),
268 take_word(Array.new(F32, 0n, 0.0), Array.get(F32, a, index)),
269 Equal.cong(U32, F32, word => Refinement.second(F32, Array.get(F32, a, word)), index,
270 U32.from_nat(Values.permute_source_index(input, axes, position)), source_word(input, axes, position)))}
271
272 def clone_first(+d: Nat, -a: Array, pair: Array & Array, -exact: {pair == Array.clone(F32, a) : Array & Array}, -valid: Model.Balanced(F32, d, a)) ->
273 Fact<{Pair.fst(Array, Array, pair) == a : Array}>:
274 Fact{Equal.trans(Array, Pair.fst(Array, Array, pair), Pair.fst(Array, Array, Array.clone(F32, a)), a,
275 Equal.cong(Array & Array, Array, p => Pair.fst(Array, Array, p), pair, Array.clone(F32, a), exact),
276 Clone.first_equal(F32, d, a, Clone.certify(F32, d, a, valid)))}
277
278 def clone_second(+d: Nat, -a: Array, pair: Array & Array, -exact: {pair == Array.clone(F32, a) : Array & Array}, -valid: Model.Balanced(F32, d, a)) ->
279 Fact<{Pair.snd(Array, Array, pair) == a : Array}>:
280 Fact{Equal.trans(Array, Pair.snd(Array, Array, pair), Pair.snd(Array, Array, Array.clone(F32, a)), a,
281 Equal.cong(Array & Array, Array, p => Pair.snd(Array, Array, p), pair, Array.clone(F32, a), exact),
282 Clone.second_equal(F32, d, a, Clone.certify(F32, d, a, valid)))}
283
284 def halves_equal(+d: Nat, -a: Array, lp: Array & Array, -exact: {lp == Array.clone(F32, a) : Array & Array}, -valid: Model.Balanced(F32, d, a)) ->
285 Fact<({Pair.fst(Array & Array, Array & Array, Build.halves(lp, Array.clone(F32, Array.new(F32, 0n, 0.0)))) == (a, Array.new(F32, 0n, 0.0)) : Array & Array} &
286 {Pair.snd(Array & Array, Array & Array, Build.halves(lp, Array.clone(F32, Array.new(F32, 0n, 0.0)))) == (a, Array.new(F32, 0n, 0.0)) : Array & Array})>:
287 (l1, l2) = lp
288 Fact{(Equal.cong(Array, Array & Array, x => (x, Array.new(F32, 0n, 0.0)), l1, a,
289 {\{Fact: q => q} : Fact<{Pair.fst(Array, Array, (l1, l2)) == a : Array}> -> {Pair.fst(Array, Array, (l1, l2)) == a : Array}}(
290 clone_first(d, a, (l1, l2), exact, valid))),
291 Equal.cong(Array, Array & Array, x => (x, Array.new(F32, 0n, 0.0)), l2, a,
292 {\{Fact: q => q} : Fact<{Pair.snd(Array, Array, (l1, l2)) == a : Array}> -> {Pair.snd(Array, Array, (l1, l2)) == a : Array}}(
293 clone_second(d, a, (l1, l2), exact, valid))))}
294
295 def SplitFacts(-a: Array) -> Type:
296 {Pair.fst(Array & Array, Array & Array, Build.split((a, Array.new(F32, 0n, 0.0)))) == (a, Array.new(F32, 0n, 0.0)) : Array & Array} & {Pair.snd(Array & Array, Array & Array, Build.split((a, Array.new(F32, 0n, 0.0)))) == (a, Array.new(F32, 0n, 0.0)) : Array & Array}
297
298 def split_equal(+d: Nat, -a: Array, -valid: Model.Balanced(F32, d, a)) -> Fact:
299 Fact{{\{Fact: q => q} : Fact -> SplitFacts(a)}(halves_equal(d, a, Array.clone(F32, a), {==}, valid))}
300
301 # Every slot of the build holds what its cell function computes.
302 def build_cell(+depth: Nat, +input: Shape.Shape, +axes: List<&2, U32>, +count: Nat, +d: Nat, -a: Array, +path: Model.Path(depth),
303 -valid: Model.Balanced(F32, d, a)) ->
304 Fact<{Model.observe(F32, depth, Build.build(depth, Build.Request{input, axes}, count, a), path) ==
305 Some{Pair.fst(F32, Array & Array, Build.cell(Build.recipe(Build.Request{input, axes}), Nat.add(Model.offset(depth, path), 0n), (a, Array.new(F32, 0n, 0.0)),
306 Nat.is_lt(Nat.add(Model.offset(depth, path), 0n), count)))} : Maybe<&2, F32>}>:
307 +request : Build.Request = Build.Request{input, axes}
308 Fact{{\{Generic.Fact: q => q} :
309 Generic.Fact, ~Array, ~Build.cell, Build.recipe(request), (a, Array.new(F32, 0n, 0.0)), depth, count, 0n,
310 Build.build(depth, request, count, a))> ->
311 Generic.Cells(~Build.Recipe, ~Array, ~Array, ~Build.cell, Build.recipe(request), (a, Array.new(F32, 0n, 0.0)), depth, count, 0n,
312 Build.build(depth, request, count, a))}(
313 Generic.build_cells(~Build.Request, ~Build.Recipe, ~Array, ~Array, ~Build.recipe, ~Build.split, ~Build.cell,
314 request, (a, Array.new(F32, 0n, 0.0)), depth, Build.splits(depth), count, 0n, (a, Array.new(F32, 0n, 0.0)), {==},
315 at => inside => {\{Fact: q => q} : Fact<{Pair.snd(F32, Array & Array, Build.cell(Build.recipe(request), at, (a, Array.new(F32, 0n, 0.0)), inside)) == (a, Array.new(F32, 0n, 0.0)) : Array & Array}> ->
316 {Pair.snd(F32, Array & Array, Build.cell(Build.recipe(request), at, (a, Array.new(F32, 0n, 0.0)), inside)) == (a, Array.new(F32, 0n, 0.0)) : Array & Array}}(cell_keeps(d, a, request, at, inside, valid)),
317 Pair.fst({Pair.fst(Array & Array, Array & Array, Build.split((a, Array.new(F32, 0n, 0.0)))) == (a, Array.new(F32, 0n, 0.0)) : Array & Array}, {Pair.snd(Array & Array, Array & Array, Build.split((a, Array.new(F32, 0n, 0.0)))) == (a, Array.new(F32, 0n, 0.0)) : Array & Array},
318 {\{Fact: q => q} : Fact -> SplitFacts(a)}(split_equal(d, a, valid))),
319 Pair.snd({Pair.fst(Array & Array, Array & Array, Build.split((a, Array.new(F32, 0n, 0.0)))) == (a, Array.new(F32, 0n, 0.0)) : Array & Array}, {Pair.snd(Array & Array, Array & Array, Build.split((a, Array.new(F32, 0n, 0.0)))) == (a, Array.new(F32, 0n, 0.0)) : Array & Array},
320 {\{Fact: q => q} : Fact -> SplitFacts(a)}(split_equal(d, a, valid)))))(path)}
321
322 def plan_elements(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan) ->
323 {Shape.elements(Layout.output(input, axes, plan)) == Shape.elements(input) : Nat}:
324 Layout.Permutation{+output, valid, ordered, +count} = plan
325 Equal.sym(Nat, Shape.elements(input), Shape.elements(output), Lifecycle.natural_equal(Shape.elements(input), Shape.elements(output), count))
326
327 def plan_numel(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan) ->
328 {Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) == Shape.logical_numel(Shape.dimensions(input)) : Nat}:
329 +output = Layout.output(input, axes, plan)
330 Equal.trans(Nat, Shape.logical_numel(Shape.dimensions(output)), Shape.elements(output), Shape.logical_numel(Shape.dimensions(input)),
331 Equal.sym(Nat, Shape.elements(output), Shape.logical_numel(Shape.dimensions(output)), Shape.elements_exact(output)),
332 Equal.trans(Nat, Shape.elements(output), Shape.elements(input), Shape.logical_numel(Shape.dimensions(input)),
333 plan_elements(input, axes, plan), Shape.elements_exact(input)))
334
335 # The parallel copy against the copy loop's meaning: the source comes back
336 # unchanged and every output index below the element count holds the word
337 # `expected_output` names.
338
339 def word_below(+shape: Shape.Shape, +other: Nat, +earlier: {Nat.is_lt(other, Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool}) ->
340 {U32.to_nat(U32.from_nat(other)) == other : Nat}:
341 Conversion.natural_word_exact(other, Shape.count(shape),
342 %Equal.sym(Nat, U32.to_nat(Shape.count(shape)), Shape.logical_numel(Shape.dimensions(shape)),
343 Equal.trans(Nat, U32.to_nat(Shape.count(shape)), Shape.elements(shape), Shape.logical_numel(Shape.dimensions(shape)),
344 Conversion.shape_count_exact(shape), Shape.elements_exact(shape))) : {Nat.is_le(other, _) == True{} : Bool}
345 Conversion.lt_le(other, Shape.logical_numel(Shape.dimensions(shape)), earlier))
346
347 def length_below(+shape: Shape.Shape, +n: Nat, +exact: {n == Shape.logical_numel(Shape.dimensions(shape)) : Nat}, +other: Nat,
348 +earlier: {Nat.is_lt(other, Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool}) ->
349 {Nat.is_lt(U32.to_nat(U32.from_nat(other)), n) == True{} : Bool}:
350 %Equal.sym(Nat, U32.to_nat(U32.from_nat(other)), other, word_below(shape, other, earlier)) : {Nat.is_lt(_, n) == True{} : Bool}
351 %Equal.sym(Nat, n, Shape.logical_numel(Shape.dimensions(shape)), exact) : {Nat.is_lt(other, _) == True{} : Bool}
352 earlier
353
354 def owned_observe(+d: Nat, +n: Nat, -array: Array, +index: U32, +valid: {Storage.bounds(d, n) == True{} : Bool},
355 +inside: {Nat.is_lt(U32.to_nat(index), n) == True{} : Bool}) ->
356 {Lifecycle.checked_observe_array(d, n, array, index) == Model.observe(F32, d, array, Refinement.route(d, index)) : Maybe<&2, F32>}:
357 %Equal.sym(Bool, Nat.is_lt(U32.to_nat(index), n), True{}, inside) :
358 {Lifecycle.observe_allowed(d, n, array, index, _) == Model.observe(F32, d, array, Refinement.route(d, index)) : Maybe<&2, F32>}
359 %Equal.sym(Bool, Storage.index_matches(d, n, index), True{}, Route.index_matched(d, n, index, valid, inside)) :
360 {Lifecycle.observe_index(d, n, array, index, _) == Model.observe(F32, d, array, Refinement.route(d, index)) : Maybe<&2, F32>}
361 {==}
362
363 def balanced_along(-d: Nat, -x: Array, -y: Array, -same: {x == y : Array}, -valid: Model.Balanced(F32, d, y)) ->
364 Fact:
365 Fact{
366 %Equal.sym(Array, x, y, same) : Model.Balanced(F32, d, _)
367 valid}
368
369 # What the copy loop expects at output index `other`: the source word at
370 # `permute_source_index`.
371 def expected_word(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan, +other: Nat,
372 +ds: Nat, +ns: Nat, -a: Array, -b: Model.Balanced(F32, ds, a), +ok: {Storage.bounds(ds, ns) == True{} : Bool},
373 +fits: {ns == Shape.logical_numel(Shape.dimensions(input)) : Nat},
374 -source_fits: {Lifecycle.logical_length(Storage.Owned{ds, ns, a, b, ok}) == Shape.logical_numel(Shape.dimensions(input)) : Nat},
375 result: Result<&2, &2, Layout.Error, Layout.Mapping>,
376 done: {Result.is_done(&2, &2, Layout.Error, Layout.Mapping, result) == True{} : Bool}) ->
377 Fact<{Transforms.expected_mapped(input, axes, plan, U32.from_nat(other), Core.TTensor{Storage.Owned{ds, ns, a, b, ok}, source_fits}, result) ==
378 Some{Refinement.second(F32, Array.get(F32, a, U32.from_nat(Values.permute_source_index(input, axes, other))))} : Maybe<&2, F32>}>:
379 match result:
380 case Fail{error}:
381 Fact{Empty.absurd({None{} == Some{Refinement.second(F32, Array.get(F32, a, U32.from_nat(Values.permute_source_index(input, axes, other))))} : Maybe<&2, F32>}, Arith.clash(done))}
382 case Done{+mapping}:
383 +found = Layout.source_index(input, axes, plan, U32.from_nat(other), mapping)
384 +inside : {Nat.is_lt(U32.to_nat(found), ns) == True{} : Bool} = %Equal.sym(Nat, ns, Shape.logical_numel(Shape.dimensions(input)), fits) : {Nat.is_lt(U32.to_nat(found), _) == True{} : Bool}
385 mapped_bound(input, axes, plan, U32.from_nat(other), mapping)
386 Fact{Equal.trans(Maybe<&2, F32>, Lifecycle.checked_observe_array(ds, ns, a, found), Model.observe(F32, ds, a, Refinement.route(ds, found)),
387 Some{Refinement.second(F32, Array.get(F32, a, U32.from_nat(Values.permute_source_index(input, axes, other))))},
388 owned_observe(ds, ns, a, found, ok, inside),
389 Equal.trans(Maybe<&2, F32>, Model.observe(F32, ds, a, Refinement.route(ds, found)), Some{Refinement.second(F32, Array.get(F32, a, found))},
390 Some{Refinement.second(F32, Array.get(F32, a, U32.from_nat(Values.permute_source_index(input, axes, other))))},
391 Equal.sym(Maybe<&2, F32>, Some{Refinement.second(F32, Array.get(F32, a, found))}, Model.observe(F32, ds, a, Refinement.route(ds, found)),
392 Refinement.law_get_refines(F32, ds, a, found, b)),
393 Equal.cong(U32, Maybe<&2, F32>, index => Some{Refinement.second(F32, Array.get(F32, a, index))}, found, U32.from_nat(Values.permute_source_index(input, axes, other)),
394 mapped_index(input, axes, plan, U32.from_nat(other), mapping))))}
395
396 # What the parallel build holds at output index `other`.
397 def output_word(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan, +d: Nat, +n: Nat,
398 +valid: {Storage.bounds(d, n) == True{} : Bool}, +exact: {n == Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat},
399 +ds: Nat, -c: Array, -bc: Model.Balanced(F32, ds, c), +other: Nat, +earlier: {Nat.is_lt(other, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) == True{} : Bool}) ->
400 Fact<{Lifecycle.checked_observe_array(d, n, Build.build(d, Build.Request{input, axes}, n, c), U32.from_nat(other)) ==
401 Some{Refinement.second(F32, Array.get(F32, c, U32.from_nat(Values.permute_source_index(input, axes, other))))} : Maybe<&2, F32>}>:
402 +index = U32.from_nat(other)
403 +inside = length_below(Layout.output(input, axes, plan), n, exact, other, earlier)
404 +at = Equal.trans(Nat, Nat.add(Model.offset(d, Refinement.route(d, index)), 0n), U32.to_nat(index), other,
405 Facts.output_position(d, n, index, valid, inside), word_below(Layout.output(input, axes, plan), other, earlier))
406 +below : {Nat.is_lt(other, n) == True{} : Bool} = %Equal.sym(Nat, n, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), exact) : {Nat.is_lt(other, _) == True{} : Bool}
407 earlier
408 -built = Build.build(d, Build.Request{input, axes}, n, c)
409 -recipe = Build.recipe(Build.Request{input, axes})
410 Fact{Equal.trans(Maybe<&2, F32>, Lifecycle.checked_observe_array(d, n, built, index), Model.observe(F32, d, built, Refinement.route(d, index)),
411 Some{Refinement.second(F32, Array.get(F32, c, U32.from_nat(Values.permute_source_index(input, axes, other))))},
412 owned_observe(d, n, built, index, valid, inside),
413 Equal.trans(Maybe<&2, F32>, Model.observe(F32, d, built, Refinement.route(d, index)),
414 Some{Pair.fst(F32, Array & Array, Build.cell(recipe, Nat.add(Model.offset(d, Refinement.route(d, index)), 0n), (c, Array.new(F32, 0n, 0.0)),
415 Nat.is_lt(Nat.add(Model.offset(d, Refinement.route(d, index)), 0n), n)))},
416 Some{Refinement.second(F32, Array.get(F32, c, U32.from_nat(Values.permute_source_index(input, axes, other))))},
417 {\{Fact: q => q} : Fact<{Model.observe(F32, d, built, Refinement.route(d, index)) ==
418 Some{Pair.fst(F32, Array & Array, Build.cell(recipe, Nat.add(Model.offset(d, Refinement.route(d, index)), 0n), (c, Array.new(F32, 0n, 0.0)),
419 Nat.is_lt(Nat.add(Model.offset(d, Refinement.route(d, index)), 0n), n)))} : Maybe<&2, F32>}> ->
420 {Model.observe(F32, d, built, Refinement.route(d, index)) ==
421 Some{Pair.fst(F32, Array & Array, Build.cell(recipe, Nat.add(Model.offset(d, Refinement.route(d, index)), 0n), (c, Array.new(F32, 0n, 0.0)),
422 Nat.is_lt(Nat.add(Model.offset(d, Refinement.route(d, index)), 0n), n)))} : Maybe<&2, F32>}}(
423 build_cell(d, input, axes, n, ds, c, Refinement.route(d, index), bc)),
424 %Equal.sym(Nat, Nat.add(Model.offset(d, Refinement.route(d, index)), 0n), other, at) : {Some{Pair.fst(F32, Array & Array, Build.cell(recipe, _, (c, Array.new(F32, 0n, 0.0)), Nat.is_lt(_, n)))} == Some{Refinement.second(F32, Array.get(F32, c, U32.from_nat(Values.permute_source_index(input, axes, other))))} : Maybe<&2, F32>}
425 %Equal.sym(Bool, Nat.is_lt(other, n), True{}, below) :
426 {Some{Pair.fst(F32, Array & Array, Build.cell(recipe, other, (c, Array.new(F32, 0n, 0.0)), _))} == Some{Refinement.second(F32, Array.get(F32, c, U32.from_nat(Values.permute_source_index(input, axes, other))))} : Maybe<&2, F32>}
427 Equal.cong(F32, Maybe<&2, F32>, word => Some{word}, Pair.fst(F32, Array & Array, Build.cell(recipe, other, (c, Array.new(F32, 0n, 0.0)), True{})),
428 Refinement.second(F32, Array.get(F32, c, U32.from_nat(Values.permute_source_index(input, axes, other)))),
429 {\{Fact: q => q} : Fact<{Pair.fst(F32, Array & Array, Build.cell(recipe, other, (c, Array.new(F32, 0n, 0.0)), True{})) == Refinement.second(F32, Array.get(F32, c, U32.from_nat(Values.permute_source_index(input, axes, other)))) : F32}> ->
430 {Pair.fst(F32, Array & Array, Build.cell(recipe, other, (c, Array.new(F32, 0n, 0.0)), True{})) == Refinement.second(F32, Array.get(F32, c, U32.from_nat(Values.permute_source_index(input, axes, other)))) : F32}}(
431 cell_word(input, axes, c, other)))))}
432
433 # The copy from an owned source, whatever the clone hands out.
434 def cloned_meaning(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan, +d: Nat, +n: Nat,
435 -valid: {Storage.bounds(d, n) == True{} : Bool}, -exact: {n == Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat},
436 +ds: Nat, +ns: Nat, -a: Array, pair: Array & Array, -same: {pair == Array.clone(F32, a) : Array & Array},
437 -b: Model.Balanced(F32, ds, a), -ok: {Storage.bounds(ds, ns) == True{} : Bool},
438 -fits: {Lifecycle.logical_length(Storage.Owned{ds, ns, a, b, ok}) == Shape.logical_numel(Shape.dimensions(input)) : Nat}) ->
439 Transforms.LoopMeaning(input, axes, plan, Core.TTensor{Storage.Owned{ds, ns, a, b, ok}, fits}, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))),
440 Transforms.built_cloned(input, axes, plan, d, n, valid, exact, ds, ns, a, pair, same, b, ok, fits)):
441 (kept, copy) = pair
442 -kept_a = {\{Fact: q => q} : Fact<{Pair.fst(Array, Array, (kept, copy)) == a : Array}> -> {kept == a : Array}}(
443 clone_first(ds, a, (kept, copy), same, b))
444 -copy_a = {\{Fact: q => q} : Fact<{Pair.snd(Array, Array, (kept, copy)) == a : Array}> -> {copy == a : Array}}(
445 clone_second(ds, a, (kept, copy), same, b))
446 -bc = {\{Fact: q => q} : Fact -> Model.Balanced(F32, ds, copy)}(balanced_along(ds, copy, a, copy_a, b))
447 Transforms.Prefix{
448 other => Equal.cong(Array, Maybe<&2, F32>, x => Lifecycle.observe_allowed(ds, ns, x, other, Nat.is_lt(U32.to_nat(other), ns)), kept, a, kept_a),
449 other => earlier => word_below(Layout.output(input, axes, plan), other, earlier),
450 other => earlier => Route.index_matched(d, n, U32.from_nat(other), valid, length_below(Layout.output(input, axes, plan), n, exact, other, earlier)),
451 other => earlier => Equal.trans(Maybe<&2, F32>,
452 Lifecycle.checked_observe_array(d, n, Build.build(d, Build.Request{input, axes}, n, copy), U32.from_nat(other)),
453 Some{Refinement.second(F32, Array.get(F32, a, U32.from_nat(Values.permute_source_index(input, axes, other))))},
454 Transforms.expected_output(input, axes, plan, Core.TTensor{Storage.Owned{ds, ns, a, b, ok}, fits}, U32.from_nat(other)),
455 Equal.trans(Maybe<&2, F32>,
456 Lifecycle.checked_observe_array(d, n, Build.build(d, Build.Request{input, axes}, n, copy), U32.from_nat(other)),
457 Some{Refinement.second(F32, Array.get(F32, copy, U32.from_nat(Values.permute_source_index(input, axes, other))))}, Some{Refinement.second(F32, Array.get(F32, a, U32.from_nat(Values.permute_source_index(input, axes, other))))},
458 {\{Fact: q => q} : Fact<{Lifecycle.checked_observe_array(d, n, Build.build(d, Build.Request{input, axes}, n, copy), U32.from_nat(other)) ==
459 Some{Refinement.second(F32, Array.get(F32, copy, U32.from_nat(Values.permute_source_index(input, axes, other))))} : Maybe<&2, F32>}> ->
460 {Lifecycle.checked_observe_array(d, n, Build.build(d, Build.Request{input, axes}, n, copy), U32.from_nat(other)) ==
461 Some{Refinement.second(F32, Array.get(F32, copy, U32.from_nat(Values.permute_source_index(input, axes, other))))} : Maybe<&2, F32>}}(
462 output_word(input, axes, plan, d, n, valid, exact, ds, copy, bc, other, earlier)),
463 Equal.cong(Array, Maybe<&2, F32>, x => Some{Refinement.second(F32, Array.get(F32, x, U32.from_nat(Values.permute_source_index(input, axes, other))))}, copy, a, copy_a)),
464 Equal.sym(Maybe<&2, F32>, Transforms.expected_output(input, axes, plan, Core.TTensor{Storage.Owned{ds, ns, a, b, ok}, fits}, U32.from_nat(other)),
465 Some{Refinement.second(F32, Array.get(F32, a, U32.from_nat(Values.permute_source_index(input, axes, other))))},
466 {\{Fact: q => q} : Fact<{Transforms.expected_mapped(input, axes, plan, U32.from_nat(other), Core.TTensor{Storage.Owned{ds, ns, a, b, ok}, fits},
467 Layout.map_index(input, axes, plan, U32.from_nat(other))) == Some{Refinement.second(F32, Array.get(F32, a, U32.from_nat(Values.permute_source_index(input, axes, other))))} : Maybe<&2, F32>}> ->
468 {Transforms.expected_mapped(input, axes, plan, U32.from_nat(other), Core.TTensor{Storage.Owned{ds, ns, a, b, ok}, fits},
469 Layout.map_index(input, axes, plan, U32.from_nat(other))) == Some{Refinement.second(F32, Array.get(F32, a, U32.from_nat(Values.permute_source_index(input, axes, other))))} : Maybe<&2, F32>}}(
470 expected_word(input, axes, plan, other, ds, ns, a, b, ok, fits, fits, Layout.map_index(input, axes, plan, U32.from_nat(other)),
471 mapping_done(input, axes, plan, U32.from_nat(other),
472 %Equal.sym(Nat, U32.to_nat(U32.from_nat(other)), other, word_below(Layout.output(input, axes, plan), other, earlier)) : {Nat.is_lt(_, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) == True{} : Bool}
473 earlier)))))}
474
475 def owned_meaning(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan, +d: Nat, +n: Nat,
476 -valid: {Storage.bounds(d, n) == True{} : Bool}, -exact: {n == Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat}, source: Core.Tensor) ->
477 Transforms.LoopMeaning(input, axes, plan, source, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Transforms.built_owned(input, axes, plan, d, n, valid, exact, source)):
478 Core.TTensor{storage, fits} = source
479 match storage:
480 case Storage.Empty{}:
481 Unit{}
482 case Storage.Owned{+ds, +ns, a, b, ok}:
483 cloned_meaning(input, axes, plan, d, n, valid, exact, ds, ns, a, Array.clone(F32, a), {==}, b, ok, fits)
484
485 def empty_below(-shape: Shape.Shape, +other: Nat, exact: {Lifecycle.logical_length(Storage.Empty{}) == Shape.logical_numel(Shape.dimensions(shape)) : Nat},
486 earlier: {Nat.is_lt(other, Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool}) -> Empty:
487 Arith.lt_zero(other,
488 %Equal.sym(Nat, Lifecycle.logical_length(Storage.Empty{}), Shape.logical_numel(Shape.dimensions(shape)), exact) : {Nat.is_lt(other, _) == True{} : Bool}
489 earlier)
490
491 # The parallel copy meets the copy loop's meaning on every source and target.
492 def built_meaning(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan, source: Core.Tensor,
493 target: Core.Tensor) ->
494 Transforms.LoopMeaning(input, axes, plan, source, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Transforms.built(input, axes, plan, source, target)):
495 Core.TTensor{storage, exact} = target
496 match storage:
497 case Storage.Empty{}:
498 Transforms.Prefix{other => {==},
499 other => earlier => Empty.absurd({U32.to_nat(U32.from_nat(other)) == other : Nat}, empty_below(Layout.output(input, axes, plan), other, exact, earlier)),
500 other => earlier => Empty.absurd(Core.ObservationChecked(Core.Packed{Layout.output(input, axes, plan), Core.TTensor{Storage.Empty{}, exact}}, U32.from_nat(other)),
501 empty_below(Layout.output(input, axes, plan), other, exact, earlier)),
502 other => earlier => Empty.absurd({Core.checked_observe(Layout.output(input, axes, plan), Core.TTensor{Storage.Empty{}, exact}, U32.from_nat(other)) ==
503 Transforms.expected_output(input, axes, plan, source, U32.from_nat(other)) : Maybe<&2, F32>}, empty_below(Layout.output(input, axes, plan), other, exact, earlier))}
504 case Storage.Owned{+d, +n, array, balanced, valid}:
505 owned_meaning(input, axes, plan, d, n, valid, exact, source)
506
507 def cloned_success(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan, +d: Nat, +n: Nat,
508 -valid: {Storage.bounds(d, n) == True{} : Bool}, -exact: {n == Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat},
509 +ds: Nat, +ns: Nat, -a: Array, pair: Array & Array, -same: {pair == Array.clone(F32, a) : Array & Array},
510 -b: Model.Balanced(F32, ds, a), -ok: {Storage.bounds(ds, ns) == True{} : Bool},
511 -fits: {Lifecycle.logical_length(Storage.Owned{ds, ns, a, b, ok}) == Shape.logical_numel(Shape.dimensions(input)) : Nat}) ->
512 {Totality.copy_succeeded(input, axes, plan, Transforms.built_cloned(input, axes, plan, d, n, valid, exact, ds, ns, a, pair, same, b, ok, fits)) == True{} : Bool}:
513 (kept, copy) = pair
514 {==}
515
516 # An owned output has a cell, so a source of the same element count is owned.
517 def owned_success(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan, +d: Nat, +n: Nat,
518 -valid: {Storage.bounds(d, n) == True{} : Bool}, -exact: {n == Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat}, source: Core.Tensor) ->
519 Fact<{Totality.copy_succeeded(input, axes, plan, Transforms.built_owned(input, axes, plan, d, n, valid, exact, source)) == True{} : Bool}>:
520 Core.TTensor{storage, fits} = source
521 match storage:
522 case Storage.Empty{}:
523 -positive = Arith.and_left(Nat.is_lt(0n, n), Nat.is_le(n, Model.capacity(d)),
524 Arith.and_right(Nat.is_le(d, 20n), Bool.and(Nat.is_lt(0n, n), Nat.is_le(n, Model.capacity(d))), valid))
525 Fact{Empty.absurd({False{} == True{} : Bool},
526 Arith.lt_zero(0n,
527 %Equal.sym(Nat, 0n, n, Equal.trans(Nat, 0n, Shape.logical_numel(Shape.dimensions(input)), n, fits,
528 Equal.sym(Nat, n, Shape.logical_numel(Shape.dimensions(input)), Equal.trans(Nat, n, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Shape.logical_numel(Shape.dimensions(input)), exact, plan_numel(input, axes, plan))))) : {Nat.is_lt(0n, _) == True{} : Bool}
529 positive))}
530 case Storage.Owned{+ds, +ns, a, b, ok}:
531 Fact{cloned_success(input, axes, plan, d, n, valid, exact, ds, ns, a, Array.clone(F32, a), {==}, b, ok, fits)}
532
533 def built_success(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan, source: Core.Tensor,
534 target: Core.Tensor) ->
535 Fact<{Totality.copy_succeeded(input, axes, plan, Transforms.built(input, axes, plan, source, target)) == True{} : Bool}>:
536 Core.TTensor{storage, exact} = target
537 match storage:
538 case Storage.Empty{}:
539 Fact{{==}}
540 case Storage.Owned{+d, +n, array, balanced, valid}:
541 owned_success(input, axes, plan, d, n, valid, exact, source)
542
543 def word(+input: Shape.Shape, +axes: List<&2, U32>, +fill: F32, +ops: List<&2, StorageLaws.Op>, +index: Nat) -> F32:
544 Replayed.value(input, fill, ops, U32.from_nat(Values.permute_source_index(input, axes, index)))
545
546 def leaf(+input: Shape.Shape, +axes: List<&2, U32>, +fill: F32, +ops: List<&2, StorageLaws.Op>, +position: Nat, inside: Bool) -> F32:
547 match inside:
548 case False{}:
549 0.0
550 case True{}:
551 word(input, axes, fill, ops, position)
552
553 def spec(+input: Shape.Shape, +axes: List<&2, U32>, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat, +depth: Nat, +base: Nat) -> Array:
554 match depth:
555 case 0n:
556 ALeaf{leaf(input, axes, fill, ops, base, Nat.is_lt(base, count))}
557 case 1n+inner:
558 ANode{spec(input, axes, fill, ops, count, inner, base), spec(input, axes, fill, ops, count, inner, Nat.add(Model.capacity(inner), base))}
559
560 def spec_balanced(+input: Shape.Shape, +axes: List<&2, U32>, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat, +depth: Nat, +base: Nat) -> Model.Balanced(F32, depth, spec(input, axes, fill, ops, count, depth, base)):
561 match depth:
562 case 0n:
563 Unit{}
564 case 1n+inner:
565 (spec_balanced(input, axes, fill, ops, count, inner, base), spec_balanced(input, axes, fill, ops, count, inner, Nat.add(Model.capacity(inner), base)))
566
567 def new_spec(+input: Shape.Shape, +axes: List<&2, U32>, +fill: F32, +ops: List<&2, StorageLaws.Op>, +depth: Nat, +base: Nat) -> {Array.new(F32, depth, 0.0) == spec(input, axes, fill, ops, 0n, depth, base) : Array}:
568 match depth:
569 case 0n:
570 %Equal.sym(Bool, Nat.is_lt(base, 0n), False{}, Facts.lt_zero_false(base)) :
571 {ALeaf{0.0} == ALeaf{leaf(input, axes, fill, ops, base, _)} : Array}
572 {==}
573 case 1n+inner:
574 %Equal.sym(Array, Array.new(F32, inner, 0.0), spec(input, axes, fill, ops, 0n, inner, base), new_spec(input, axes, fill, ops, inner, base)) :
575 {ANode{_, Array.new(F32, inner, 0.0)} == spec(input, axes, fill, ops, 0n, 1n+inner, base) : Array}
576 %Equal.sym(Array, Array.new(F32, inner, 0.0), spec(input, axes, fill, ops, 0n, inner, Nat.add(Model.capacity(inner), base)),
577 new_spec(input, axes, fill, ops, inner, Nat.add(Model.capacity(inner), base))) :
578 {ANode{spec(input, axes, fill, ops, 0n, inner, base), _} == spec(input, axes, fill, ops, 0n, 1n+inner, base) : Array}
579 {==}
580
581 # A subtree whose positions all lie above k is the same for k and 1 + k.
582 def above(+input: Shape.Shape, +axes: List<&2, U32>, +fill: F32, +ops: List<&2, StorageLaws.Op>, +k: Nat, +depth: Nat, +b: Nat, +h: {Nat.is_lt(k, b) == True{} : Bool}) ->
583 {spec(input, axes, fill, ops, k, depth, b) == spec(input, axes, fill, ops, 1n+k, depth, b) : Array}:
584 match depth:
585 case 0n:
586 %Equal.sym(Bool, Nat.is_lt(b, k), False{}, Facts.lt_asym(k, b, h)) :
587 {ALeaf{leaf(input, axes, fill, ops, b, _)} == ALeaf{leaf(input, axes, fill, ops, b, Nat.is_lt(b, 1n+k))} : Array}
588 %Equal.sym(Bool, Nat.is_lt(b, 1n+k), False{}, Facts.above_false(k, b, h)) :
589 {ALeaf{0.0} == ALeaf{leaf(input, axes, fill, ops, b, _)} : Array}
590 {==}
591 case 1n+inner:
592 +high = Nat.add(Model.capacity(inner), b)
593 %Equal.sym(Array, spec(input, axes, fill, ops, k, inner, b), spec(input, axes, fill, ops, 1n+k, inner, b), above(input, axes, fill, ops, k, inner, b, h)) :
594 {ANode{_, spec(input, axes, fill, ops, k, inner, high)} == spec(input, axes, fill, ops, 1n+k, 1n+inner, b) : Array}
595 %Equal.sym(Array, spec(input, axes, fill, ops, k, inner, high), spec(input, axes, fill, ops, 1n+k, inner, high),
596 above(input, axes, fill, ops, k, inner, high, Arith.lt_le_trans(k, b, high, h, Api.le_add_left(Model.capacity(inner), b)))) :
597 {ANode{spec(input, axes, fill, ops, 1n+k, inner, b), _} == spec(input, axes, fill, ops, 1n+k, 1n+inner, b) : Array}
598 {==}
599
600 # A subtree whose positions all lie below k is the same for k and 1 + k.
601 def below(+input: Shape.Shape, +axes: List<&2, U32>, +fill: F32, +ops: List<&2, StorageLaws.Op>, +k: Nat, +depth: Nat, +b: Nat, +h: {Nat.is_le(Nat.add(Model.capacity(depth), b), k) == True{} : Bool}) ->
602 {spec(input, axes, fill, ops, k, depth, b) == spec(input, axes, fill, ops, 1n+k, depth, b) : Array}:
603 match depth:
604 case 0n:
605 +lower = Facts.below_true(b, k, h)
606 %Equal.sym(Bool, Nat.is_lt(b, k), True{}, lower) :
607 {ALeaf{leaf(input, axes, fill, ops, b, _)} == ALeaf{leaf(input, axes, fill, ops, b, Nat.is_lt(b, 1n+k))} : Array}
608 %Equal.sym(Bool, Nat.is_lt(b, 1n+k), True{}, Facts.lt_succ_right(b, k, lower)) :
609 {ALeaf{word(input, axes, fill, ops, b)} == ALeaf{leaf(input, axes, fill, ops, b, _)} : Array}
610 {==}
611 case 1n+inner:
612 +c = Model.capacity(inner)
613 +high = Nat.add(c, b)
614 +nested = Equal.trans(Bool, Nat.is_le(Nat.add(c, high), k), Nat.is_le(Nat.add(Nat.add(c, c), b), k), True{},
615 Equal.cong(Nat, Bool, x => Nat.is_le(x, k), Nat.add(c, high), Nat.add(Nat.add(c, c), b), Index.add_assoc(c, c, b)), h)
616 +low_h = Facts.le_trans(high, Nat.add(c, high), k, Facts.le_add_cong(c, b, high, Api.le_add_left(c, b)), nested)
617 %Equal.sym(Array, spec(input, axes, fill, ops, k, inner, b), spec(input, axes, fill, ops, 1n+k, inner, b), below(input, axes, fill, ops, k, inner, b, low_h)) :
618 {ANode{_, spec(input, axes, fill, ops, k, inner, high)} == spec(input, axes, fill, ops, 1n+k, 1n+inner, b) : Array}
619 %Equal.sym(Array, spec(input, axes, fill, ops, k, inner, high), spec(input, axes, fill, ops, 1n+k, inner, high), below(input, axes, fill, ops, k, inner, high, nested)) :
620 {ANode{spec(input, axes, fill, ops, 1n+k, inner, b), _} == spec(input, axes, fill, ops, 1n+k, 1n+inner, b) : Array}
621 {==}
622
623 # Writing cell k at the path whose position is k turns the spec for k into
624 # the spec for 1 + k.
625 def write_spec(+input: Shape.Shape, +axes: List<&2, U32>, +fill: F32, +ops: List<&2, StorageLaws.Op>, +k: Nat, +depth: Nat, +base: Nat, +path: Model.Path(depth),
626 +h: {Nat.add(Model.offset(depth, path), base) == k : Nat}) ->
627 {ArrayCore.write_path(F32, depth, spec(input, axes, fill, ops, k, depth, base), path, word(input, axes, fill, ops, k)) == spec(input, axes, fill, ops, 1n+k, depth, base) : Array}:
628 match depth:
629 case 0n:
630 %h : {ALeaf{word(input, axes, fill, ops, _)} == ALeaf{leaf(input, axes, fill, ops, base, Nat.is_lt(base, 1n+_))} : Array}
631 %Equal.sym(Bool, Nat.is_lt(base, 1n+base), True{}, Facts.lt_self_succ(base)) :
632 {ALeaf{word(input, axes, fill, ops, base)} == ALeaf{leaf(input, axes, fill, ops, base, _)} : Array}
633 {==}
634 case 1n+inner:
635 Model.Step{+bit, +rest} = path
636 match bit:
637 case False{}:
638 +c = Model.capacity(inner)
639 +high = Nat.add(c, base)
640 +inside = Equal.trans(Bool, Nat.is_lt(k, high), Nat.is_lt(Nat.add(Model.offset(inner, rest), base), high), True{},
641 Equal.cong(Nat, Bool, x => Nat.is_lt(x, high), k, Nat.add(Model.offset(inner, rest), base),
642 Equal.sym(Nat, Nat.add(Model.offset(inner, rest), base), k, h)),
643 Facts.lt_add_right(Model.offset(inner, rest), c, base, Api.offset_lt(inner, rest)))
644 %Equal.sym(Array, ArrayCore.write_path(F32, inner, spec(input, axes, fill, ops, k, inner, base), rest, word(input, axes, fill, ops, k)),
645 spec(input, axes, fill, ops, 1n+k, inner, base), write_spec(input, axes, fill, ops, k, inner, base, rest, h)) :
646 {ANode{_, spec(input, axes, fill, ops, k, inner, high)} == spec(input, axes, fill, ops, 1n+k, 1n+inner, base) : Array}
647 %Equal.sym(Array, spec(input, axes, fill, ops, k, inner, high), spec(input, axes, fill, ops, 1n+k, inner, high), above(input, axes, fill, ops, k, inner, high, inside)) :
648 {ANode{spec(input, axes, fill, ops, 1n+k, inner, base), _} == spec(input, axes, fill, ops, 1n+k, 1n+inner, base) : Array}
649 {==}
650 case True{}:
651 +c = Model.capacity(inner)
652 +high = Nat.add(c, base)
653 +off = Model.offset(inner, rest)
654 +lower = Equal.trans(Bool, Nat.is_le(high, k), Nat.is_le(high, Nat.add(c, Nat.add(off, base))), True{},
655 Equal.cong(Nat, Bool, x => Nat.is_le(high, x), k, Nat.add(c, Nat.add(off, base)),
656 Equal.trans(Nat, k, Nat.add(Nat.add(c, off), base), Nat.add(c, Nat.add(off, base)),
657 Equal.sym(Nat, Nat.add(Nat.add(c, off), base), k, h),
658 Equal.sym(Nat, Nat.add(c, Nat.add(off, base)), Nat.add(Nat.add(c, off), base), Index.add_assoc(c, off, base)))),
659 Facts.le_add_cong(c, base, Nat.add(off, base), Api.le_add_left(off, base)))
660 +shifted = Equal.trans(Nat, Nat.add(off, high), Nat.add(Nat.add(c, off), base), k, Facts.position_step(c, off, base), h)
661 %Equal.sym(Array, spec(input, axes, fill, ops, k, inner, base), spec(input, axes, fill, ops, 1n+k, inner, base), below(input, axes, fill, ops, k, inner, base, lower)) :
662 {ANode{_, ArrayCore.write_path(F32, inner, spec(input, axes, fill, ops, k, inner, high), rest, word(input, axes, fill, ops, k))} == spec(input, axes, fill, ops, 1n+k, 1n+inner, base) : Array}
663 %Equal.sym(Array, ArrayCore.write_path(F32, inner, spec(input, axes, fill, ops, k, inner, high), rest, word(input, axes, fill, ops, k)),
664 spec(input, axes, fill, ops, 1n+k, inner, high), write_spec(input, axes, fill, ops, k, inner, high, rest, shifted)) :
665 {ANode{spec(input, axes, fill, ops, 1n+k, inner, base), _} == spec(input, axes, fill, ops, 1n+k, 1n+inner, base) : Array}
666 {==}
667
668 def TreeEq(+input: Shape.Shape, +axes: List<&2, U32>, +fill: F32, +ops: List<&2, StorageLaws.Op>, +plan: Layout.Plan, +ds: Nat, +fits: {Storage.bounds(ds, U32.to_nat(Shape.count(input))) == True{} : Bool}, +count: Nat, +exact: {count == Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat}, +depth: Nat, +base: Nat, -result: Array & (Array & Array)) -> Type:
669 {Pair.fst(Array, Array & Array, result) == spec(input, axes, fill, ops, count, depth, base) : Array} &
670 {Pair.snd(Array, Array & Array, result) == (Api.arr((Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{})), ds, fill), Array.new(F32, 0n, 0.0)) : Array & Array}
671
672 def read_inside(+input: Shape.Shape, +axes: List<&2, U32>, +fill: F32, +ops: List<&2, StorageLaws.Op>, +plan: Layout.Plan, +ds: Nat, +fits: {Storage.bounds(ds, U32.to_nat(Shape.count(input))) == True{} : Bool}, +count: Nat, +exact: {count == Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat}, +base: Nat, result: Result<&2, &2, Layout.Error, Layout.Mapping>,
673 done: {Result.is_done(&2, &2, Layout.Error, Layout.Mapping, result) == True{} : Bool}) ->
674 {Nat.is_lt(U32.to_nat(U32.from_nat(Values.permute_source_index(input, axes, base))), U32.to_nat(Shape.count(input))) == True{} : Bool}:
675 match result:
676 case Fail{error}:
677 Empty.absurd({Nat.is_lt(U32.to_nat(U32.from_nat(Values.permute_source_index(input, axes, base))), U32.to_nat(Shape.count(input))) == True{} : Bool}, Arith.clash(done))
678 case Done{+mapping}:
679 %Equal.sym(Nat, U32.to_nat(Shape.count(input)), Shape.logical_numel(Shape.dimensions(input)), Reads.count_exact(input)) : {Nat.is_lt(U32.to_nat(U32.from_nat(Values.permute_source_index(input, axes, base))), _) == True{} : Bool}
680 mapped_inside(input, axes, plan, U32.from_nat(base), mapping)
681
682 # An inside cell reads the replayed word `permute_source_index` names.
683 def word_read(+input: Shape.Shape, +axes: List<&2, U32>, +fill: F32, +ops: List<&2, StorageLaws.Op>, +plan: Layout.Plan, +ds: Nat, +fits: {Storage.bounds(ds, U32.to_nat(Shape.count(input))) == True{} : Bool}, +count: Nat, +exact: {count == Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat}, +base: Nat, +decided: {True{} == Nat.is_lt(base, count) : Bool}, -pair: Array & F32,
684 +same: {pair == Array.get(F32, Api.arr((Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{})), ds, fill), Build.at(Build.recipe(Build.Request{input, axes}), base)) : Array & F32}) ->
685 {Refinement.second(F32, pair) == word(input, axes, fill, ops, base) : F32}:
686 +index = U32.from_nat(Values.permute_source_index(input, axes, base))
687 +below : {Nat.is_lt(base, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) == True{} : Bool} =
688 %exact : {Nat.is_lt(base, _) == True{} : Bool}
689 Equal.sym(Bool, True{}, Nat.is_lt(base, count), decided)
690 +inside = read_inside(input, axes, fill, ops, plan, ds, fits, count, exact, base, Layout.map_index(input, axes, plan, U32.from_nat(base)),
691 mapping_done(input, axes, plan, U32.from_nat(base),
692 %Equal.sym(Nat, U32.to_nat(U32.from_nat(base)), base, word_below(Layout.output(input, axes, plan), base, below)) : {Nat.is_lt(_, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) == True{} : Bool}
693 below))
694 +exact_read = Equal.trans(Array & F32, pair, Array.get(F32, Api.arr((Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{})), ds, fill), Build.at(Build.recipe(Build.Request{input, axes}), base)), Array.get(F32, Api.arr((Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{})), ds, fill), index), same,
695 Equal.cong(U32, Array & F32, word => Array.get(F32, Api.arr((Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{})), ds, fill), word), Build.at(Build.recipe(Build.Request{input, axes}), base), index,
696 source_word(input, axes, base)))
697 Equal.trans(F32, Refinement.second(F32, pair), Api.current((Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{})), index, fill), word(input, axes, fill, ops, base),
698 BinaryBuild.got_value(ds, (Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{})), fill, U32.to_nat(Shape.count(input)), index, pair, exact_read, fits, inside, Owned.norm_written(U32.to_nat(Shape.count(input)), ops, Nil{}, Api.hit_of(U32.to_nat(Shape.count(input)), ops), {==}, {==})),
699 Owned.norm_current(U32.to_nat(Shape.count(input)), ops, Nil{}, index, fill, Api.hit_of(U32.to_nat(Shape.count(input)), ops), {==}))
700
701 def kept_read(+input: Shape.Shape, +axes: List<&2, U32>, +fill: F32, +ops: List<&2, StorageLaws.Op>, +plan: Layout.Plan, +ds: Nat, +fits: {Storage.bounds(ds, U32.to_nat(Shape.count(input))) == True{} : Bool}, +count: Nat, +exact: {count == Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat}, +base: Nat, -pair: Array & F32,
702 +same: {pair == Array.get(F32, Api.arr((Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{})), ds, fill), Build.at(Build.recipe(Build.Request{input, axes}), base)) : Array & F32}) ->
703 {Refinement.first(F32, pair) == Api.arr((Api.Right{} <> Owned.norm(U32.to_nat(Shape.count(input)), ops, Nil{})), ds, fill) : Array