1	import Base
2	import ../shape/core.bend as Shape
3	import ../storage/model.bend as Model
4	import ../storage/refinement.bend as Refinement
5	import ./core.bend as Core
6	import ./layout.bend as Layout
7	import ./elementwise_totality.bend as Totality
8	import ./slice.bend as Slice
9	import ./signed.bend as Signed
10	import ./gather.bend as Gather
11	
12	# The laws of index_select and gather (gather.bend), written from PyTorch's
13	# definitions: which requests they refuse and in which order, the plan and
14	# the output shape of the ones they take, what every cell of the forward
15	# copy holds, and what every cell of the backward, index_add and
16	# scatter_add into +0.0, holds.
17	
18	type Holds<-P: Type> is Data:
19	  Holds{-proof: P}
20	
21	# The axis a request names: its signed dim wrapped among the input's rank.
22	def axis(+input: Shape.Shape, dim: Signed.Position) -> Nat:
23	  U32.to_nat(Signed.wrap(Shape.rank(input), dim))
24	
25	def rank(+input: Shape.Shape) -> Nat:
26	  Shape.rank_dimensions(Shape.dimensions(input))
27	
28	def extent(+input: Shape.Shape, +at: Nat) -> Nat:
29	  Slice.axis_extent(Shape.dimensions(input), at)
30	
31	# How many entries an index holds.
32	def length(entries: List<&2, U32>) -> Nat:
33	  match entries:
34	    case Nil{}:
35	      0n
36	    case head <> rest:
37	      1n+length(rest)
38	
39	# Whether every entry lies below the extent.
40	def inside(entries: List<&2, U32>, +bound: Nat) -> Bool:
41	  match entries:
42	    case Nil{}:
43	      True{}
44	    case head <> rest:
45	      Bool.and(Nat.is_lt(U32.to_nat(head), bound), inside(rest, bound))
46	
47	# Whether two lists of extents are equal, extent by extent.
48	def equal(left: List<&2, U32>, right: List<&2, U32>) -> Bool:
49	  match left right:
50	    case Nil{} Nil{}:
51	      True{}
52	    case x <> xs y <> ys:
53	      Bool.and(Nat.is_eq(U32.to_nat(x), U32.to_nat(y)), equal(xs, ys))
54	    case _ _:
55	      False{}
56	
57	# Whether a gather index has the input's extent at every axis but `at`.
58	def others(dimensions: List<&2, U32>, index: List<&2, U32>, +at: Nat) -> Bool:
59	  match dimensions index at:
60	    case x <> xs y <> ys 0n:
61	      equal(xs, ys)
62	    case x <> xs y <> ys 1n+next:
63	      Bool.and(Nat.is_eq(U32.to_nat(x), U32.to_nat(y)), others(xs, ys, next))
64	    case _ _ _:
65	      False{}
66	
67	# The plan of a request the checks take: the row-major view of the input
68	# around the axis, slice.bend's, and the length of the output's axis ...
69	def viewed(+kind: Gather.Lookup, +given: Slice.View, +count: Nat) -> Gather.Plan:
70	  Slice.View{+outer, +rows, +inner} = given
71	  Gather.Plan{kind, outer, rows, inner, count}
72	
73	# ... and its output, the input's dimensions with that length at the axis,
74	# or the refusal of the shape's own checks (rank and capacity).
75	def output(-dimensions: List<&2, U32>, +plan: Gather.Plan, built: Shape.ConstructResult<dimensions>) ->
76	  Result<&2, &2, Gather.Error, Gather.Planned>:
77	  match built:
78	    case Shape.SError{error}:
79	      Fail{Gather.TensorError{Core.ShapeError{error}}}
80	    case Shape.SBuilt{+shape, preserved}:
81	      Done{Gather.Planned{plan, shape}}
82	
83	def taken(+kind: Gather.Lookup, +dimensions: List<&2, U32>, +at: Nat, +count: Nat) -> Result<&2, &2, Gather.Error, Gather.Planned>:
84	  output(Slice.resized(dimensions, at, U32.from_nat(count)),
85	    viewed(kind, Slice.sized(dimensions, at, Slice.holds_zero(dimensions)), count),
86	    Shape.construct(Slice.resized(dimensions, at, U32.from_nat(count))))
87	
88	# index_select(dim, index): an axis outside -rank..rank-1 is refused first,
89	# then an entry at or past the extent; any other request is planned.
90	law index_select_axis_refused:
91	  for +input: Shape.Shape
92	  for +dim: Signed.Position
93	  for +entries: List<&2, U32>
94	  for outside: {Nat.is_lt(axis(input, dim), rank(input)) == False{} : Bool}
95	  {Gather.select_plan(input, dim, entries) == Fail{Gather.LayoutError{Layout.AxisOutOfBounds{}}} :
96	    Result<&2, &2, Gather.Error, Gather.Planned>}
97	
98	law index_select_out_of_range:
99	  for +input: Shape.Shape
100	  for +dim: Signed.Position
101	  for +entries: List<&2, U32>
102	  for in_rank: {Nat.is_lt(axis(input, dim), rank(input)) == True{} : Bool}
103	  for past: {inside(entries, extent(input, axis(input, dim))) == False{} : Bool}
104	  {Gather.select_plan(input, dim, entries) == Fail{Gather.OutOfRange{}} : Result<&2, &2, Gather.Error, Gather.Planned>}
105	
106	law index_select_planned:
107	  for +input: Shape.Shape
108	  for +dim: Signed.Position
109	  for +entries: List<&2, U32>
110	  for in_rank: {Nat.is_lt(axis(input, dim), rank(input)) == True{} : Bool}
111	  for within: {inside(entries, extent(input, axis(input, dim))) == True{} : Bool}
112	  {Gather.select_plan(input, dim, entries) == taken(Gather.Selected{}, Shape.dimensions(input), axis(input, dim), length(entries)) :
113	    Result<&2, &2, Gather.Error, Gather.Planned>}
114	
115	# gather(dim, index): an entry count other than the element count of the
116	# index's shape is refused first, then the axis, an index of another rank, one
117	# with another extent at another axis, and an entry at or past the extent;
118	# any other request is planned, with the index's extent at the axis.
119	law gather_count_refused:
120	  for +input: Shape.Shape
121	  for +dim: Signed.Position
122	  for +index: Shape.Shape
123	  for +entries: List<&2, U32>
124	  for miscounted: {Nat.is_eq(length(entries), Shape.elements(index)) == False{} : Bool}
125	  {Gather.gather_plan(input, dim, Gather.Index{index, entries}) == Fail{Gather.TensorError{Core.ValueCountMismatch{}}} :
126	    Result<&2, &2, Gather.Error, Gather.Planned>}
127	
128	law gather_axis_refused:
129	  for +input: Shape.Shape
130	  for +dim: Signed.Position
131	  for +index: Shape.Shape
132	  for +entries: List<&2, U32>
133	  for counted: {Nat.is_eq(length(entries), Shape.elements(index)) == True{} : Bool}
134	  for outside: {Nat.is_lt(axis(input, dim), rank(input)) == False{} : Bool}
135	  {Gather.gather_plan(input, dim, Gather.Index{index, entries}) == Fail{Gather.LayoutError{Layout.AxisOutOfBounds{}}} :
136	    Result<&2, &2, Gather.Error, Gather.Planned>}
137	
138	law gather_rank_refused:
139	  for +input: Shape.Shape
140	  for +dim: Signed.Position
141	  for +index: Shape.Shape
142	  for +entries: List<&2, U32>
143	  for counted: {Nat.is_eq(length(entries), Shape.elements(index)) == True{} : Bool}
144	  for in_rank: {Nat.is_lt(axis(input, dim), rank(input)) == True{} : Bool}
145	  for ranked: {Nat.is_eq(Shape.rank_dimensions(Shape.dimensions(index)), rank(input)) == False{} : Bool}
146	  {Gather.gather_plan(input, dim, Gather.Index{index, entries}) == Fail{Gather.IndexRankMismatch{}} :
147	    Result<&2, &2, Gather.Error, Gather.Planned>}
148	
149	law gather_shape_refused:
150	  for +input: Shape.Shape
151	  for +dim: Signed.Position
152	  for +index: Shape.Shape
153	  for +entries: List<&2, U32>
154	  for counted: {Nat.is_eq(length(entries), Shape.elements(index)) == True{} : Bool}
155	  for in_rank: {Nat.is_lt(axis(input, dim), rank(input)) == True{} : Bool}
156	  for ranked: {Nat.is_eq(Shape.rank_dimensions(Shape.dimensions(index)), rank(input)) == True{} : Bool}
157	  for unmatched: {others(Shape.dimensions(input), Shape.dimensions(index), axis(input, dim)) == False{} : Bool}
158	  {Gather.gather_plan(input, dim, Gather.Index{index, entries}) == Fail{Gather.IndexShapeMismatch{}} :
159	    Result<&2, &2, Gather.Error, Gather.Planned>}
160	
161	law gather_out_of_range:
162	  for +input: Shape.Shape
163	  for +dim: Signed.Position
164	  for +index: Shape.Shape
165	  for +entries: List<&2, U32>
166	  for counted: {Nat.is_eq(length(entries), Shape.elements(index)) == True{} : Bool}
167	  for in_rank: {Nat.is_lt(axis(input, dim), rank(input)) == True{} : Bool}
168	  for ranked: {Nat.is_eq(Shape.rank_dimensions(Shape.dimensions(index)), rank(input)) == True{} : Bool}
169	  for matched: {others(Shape.dimensions(input), Shape.dimensions(index), axis(input, dim)) == True{} : Bool}
170	  for past: {inside(entries, extent(input, axis(input, dim))) == False{} : Bool}
171	  {Gather.gather_plan(input, dim, Gather.Index{index, entries}) == Fail{Gather.OutOfRange{}} :
172	    Result<&2, &2, Gather.Error, Gather.Planned>}
173	
174	law gather_planned:
175	  for +input: Shape.Shape
176	  for +dim: Signed.Position
177	  for +index: Shape.Shape
178	  for +entries: List<&2, U32>
179	  for counted: {Nat.is_eq(length(entries), Shape.elements(index)) == True{} : Bool}
180	  for in_rank: {Nat.is_lt(axis(input, dim), rank(input)) == True{} : Bool}
181	  for ranked: {Nat.is_eq(Shape.rank_dimensions(Shape.dimensions(index)), rank(input)) == True{} : Bool}
182	  for matched: {others(Shape.dimensions(input), Shape.dimensions(index), axis(input, dim)) == True{} : Bool}
183	  for within: {inside(entries, extent(input, axis(input, dim))) == True{} : Bool}
184	  {Gather.gather_plan(input, dim, Gather.Index{index, entries}) ==
185	    taken(Gather.Gathered{}, Shape.dimensions(input), axis(input, dim), Slice.axis_extent(Shape.dimensions(index), axis(input, dim))) :
186	    Result<&2, &2, Gather.Error, Gather.Planned>}
187	
188	# A gather index is built as a tensor of its dimensions is: refused past
189	# rank 4 or 2^20 elements with the shape's own error, else it holds the
190	# shape of those dimensions.
191	def index_of(+entries: List<&2, U32>, -dimensions: List<&2, U32>, built: Shape.ConstructResult<dimensions>) ->
192	  Result<&2, &2, Gather.Error, Gather.Index>:
193	  match built:
194	    case Shape.SError{error}:
195	      Fail{Gather.TensorError{Core.ShapeError{error}}}
196	    case Shape.SBuilt{+shape, preserved}:
197	      Done{Gather.Index{shape, entries}}
198	
199	law gather_index_shaped:
200	  for +dimensions: List<&2, U32>
201	  for +entries: List<&2, U32>
202	  {Gather.index(dimensions, entries) == index_of(entries, dimensions, Shape.construct(dimensions)) : Result<&2, &2, Gather.Error, Gather.Index>}
203	
204	# A refused plan fails the copy with its error and hands the input back.
205	law index_select_refused:
206	  for input: Core.Dynamic
207	  for +dim: Signed.Position
208	  for +entries: List<&2, U32>
209	  for error: Gather.Error
210	  for refused: {Gather.select_plan(Totality.dynamic_shape(input), dim, entries) == Fail{error} : Result<&2, &2, Gather.Error, Gather.Planned>}
211	  {Gather.index_select(input, dim, entries) == (input, Fail{error}) : Core.Dynamic & Result<&2, &1, Gather.Error, Core.Dynamic>}
212	
213	law gather_refused:
214	  for input: Core.Dynamic
215	  for +dim: Signed.Position
216	  for +index: Shape.Shape
217	  for +entries: List<&2, U32>
218	  for error: Gather.Error
219	  for refused: {Gather.gather_plan(Totality.dynamic_shape(input), dim, Gather.Index{index, entries}) == Fail{error} :
220	    Result<&2, &2, Gather.Error, Gather.Planned>}
221	  {Gather.gather(input, dim, Gather.Index{index, entries}) == (input, Fail{error}) : Core.Dynamic & Result<&2, &1, Gather.Error, Core.Dynamic>}
222	
223	# The cells of the builds, in the coordinates of the views: cell (o, k, r)
224	# of [outer, rows, inner] sits at (o * rows + k) * inner + r.
225	def position(+rows: Nat, +inner: Nat, +o: Nat, +k: Nat, +r: Nat) -> Nat:
226	  Nat.add(Nat.mul(Nat.add(Nat.mul(o, rows), k), inner), r)
227	
228	# The entry output cell (o, k, r) of [outer, length, inner] reads: entry k
229	# of an index_select's index, the entry at the cell itself of a gather's.
230	def entry_at(+kind: Gather.Lookup, +count: Nat, +inner: Nat, +o: Nat, +k: Nat, +r: Nat) -> Nat:
231	  match kind:
232	    case Gather.Selected{}:
233	      k
234	    case Gather.Gathered{}:
235	      position(count, inner, o, k, r)
236	
237	# For a balanced input array A, a balanced index array I and every split,
238	# the slot of the forward build that holds output cell (o, k, r), below
239	# the count, holds input cell (o, j, r), where j is the entry the cell
240	# reads: the word PyTorch's index_select and gather put there.
241	law gather_build_cells:
242	  for +depth: Nat
243	  for +splits: Nat
244	  for +kind: Gather.Lookup
245	  for +outer: Nat
246	  for +rows: Nat
247	  for +inner: Nat
248	  for +count: Nat
249	  for +cells: Nat
250	  for +base: Nat
251	  for +da: Nat
252	  for -A: Array<F32>
253	  for a: Array<F32>
254	  for same: {a == A : Array<F32>}
255	  for valid: Model.Balanced(F32, da, A)
256	  for +di: Nat
257	  for -I: Array<U32>
258	  for index: Array<U32>
259	  for same_index: {index == I : Array<U32>}
260	  for index_valid: Model.Balanced(U32, di, I)
261	  for +path: Model.Path(depth)
262	  for +o: Nat
263	  for +k: Nat
264	  for +r: Nat
265	  for k_below: {Nat.is_lt(k, count) == True{} : Bool}
266	  for r_below: {Nat.is_lt(r, inner) == True{} : Bool}
267	  for at: {Nat.add(Model.offset(depth, path), base) == position(count, inner, o, k, r) : Nat}
268	  for in_count: {Nat.is_lt(Nat.add(Model.offset(depth, path), base), cells) == True{} : Bool}
269	  for +j: U32
270	  for read: {Model.observe(U32, di, I, Refinement.route(di, U32.from_nat(entry_at(kind, count, inner, o, k, r)))) == Some{j} : Maybe<&2, U32>}
271	  Holds<{Model.observe(F32, depth, Gather.build(depth, splits, Gather.Plan{kind, outer, rows, inner, count}, cells, base, a, index), path) ==
272	    Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(position(rows, inner, o, U32.to_nat(j), r)))) : Maybe<&2, F32>}>
273	
274	# The backward of input cell (o, i, r): from gradient cell (o, k, r) on, for
275	# `left` more k, a cell whose entry is i adds its word to the sum, in
276	# increasing k; a cell whose entry is another skips. `value` is the sum at
277	# the end.
278	def Added(+dg: Nat, -G: Array<F32>, +di: Nat, -I: Array<U32>, +kind: Gather.Lookup, +count: Nat, +inner: Nat, +o: Nat, +i: Nat, +r: Nat,
279	  left: Nat, +k: Nat, +acc: F32, +value: F32) -> Type:
280	  match left:
281	    case 0n:
282	      {value == acc : F32}
283	    case 1n+rest:
284	      AddStep<dg, G, di, I, kind, count, inner, o, i, r, rest, k, acc, value>
285	
286	type AddStep<-dg: Nat, -G: Array<F32>, -di: Nat, -I: Array<U32>, -kind: Gather.Lookup, -count: Nat, -inner: Nat, -o: Nat, -i: Nat, -r: Nat,
287	  -rest: Nat, -k: Nat, -acc: F32, -value: F32> is Data:
288	  Skips{-j: U32,
289	    -read: {Model.observe(U32, di, I, Refinement.route(di, U32.from_nat(entry_at(kind, count, inner, o, k, r)))) == Some{j} : Maybe<&2, U32>},
290	    -other: {Nat.is_eq(U32.to_nat(j), i) == False{} : Bool},
291	    -next: Added(dg, G, di, I, kind, count, inner, o, i, r, rest, 1n+k, acc, value)}
292	  Adds{-j: U32,
293	    -read: {Model.observe(U32, di, I, Refinement.route(di, U32.from_nat(entry_at(kind, count, inner, o, k, r)))) == Some{j} : Maybe<&2, U32>},
294	    -same: {Nat.is_eq(U32.to_nat(j), i) == True{} : Bool},
295	    -x: F32,
296	    -word: {Model.observe(F32, dg, G, Refinement.route(dg, U32.from_nat(position(count, inner, o, k, r)))) == Some{x} : Maybe<&2, F32>},
297	    -next: Added(dg, G, di, I, kind, count, inner, o, i, r, rest, 1n+k, F32.add(acc, x), value)}
298	
299	# A slot that holds `value`, the sum Added names from +0.0 over every k.
300	type Summed<-depth: Nat, -out: Array<F32>, -path: Model.Path(depth), -dg: Nat, -G: Array<F32>, -di: Nat, -I: Array<U32>,
301	  -kind: Gather.Lookup, -count: Nat, -inner: Nat, -o: Nat, -i: Nat, -r: Nat> is Data:
302	  Summed{-value: F32, -observed: {Model.observe(F32, depth, out, path) == Some{value} : Maybe<&2, F32>},
303	    -added: Added(dg, G, di, I, kind, count, inner, o, i, r, count, 0n, 0.0, value)}
304	
305	# For a balanced gradient array G, a balanced index array I and every split,
306	# the slot of the backward build that holds input cell (o, i, r), below the
307	# count, holds the gradient words of the cells (o, k, r) whose entry is i,
308	# added to +0.0 in increasing k, and +0.0 where no entry is i: PyTorch's
309	# one-thread index_add and scatter_add.
310	law gather_scatter_cells:
311	  for +depth: Nat
312	  for +splits: Nat
313	  for +kind: Gather.Lookup
314	  for +outer: Nat
315	  for +rows: Nat
316	  for +inner: Nat
317	  for +count: Nat
318	  for +cells: Nat
319	  for +base: Nat
320	  for +dg: Nat
321	  for -G: Array<F32>
322	  for g: Array<F32>
323	  for same: {g == G : Array<F32>}
324	  for valid: Model.Balanced(F32, dg, G)
325	  for +di: Nat
326	  for -I: Array<U32>
327	  for index: Array<U32>
328	  for same_index: {index == I : Array<U32>}
329	  for index_valid: Model.Balanced(U32, di, I)
330	  for +path: Model.Path(depth)
331	  for +o: Nat
332	  for +i: Nat
333	  for +r: Nat
334	  for i_below: {Nat.is_lt(i, rows) == True{} : Bool}
335	  for r_below: {Nat.is_lt(r, inner) == True{} : Bool}
336	  for at: {Nat.add(Model.offset(depth, path), base) == position(rows, inner, o, i, r) : Nat}
337	  for in_count: {Nat.is_lt(Nat.add(Model.offset(depth, path), base), cells) == True{} : Bool}
338	  Summed<depth, Gather.scatter_build(depth, splits, Gather.Plan{kind, outer, rows, inner, count}, cells, base, g, index), path,
339	    dg, G, di, I, kind, count, inner, o, i, r>