1	import Base
2	import ../shape/core.bend as Shape
3	import ./core.bend as Core
4	import ./layout.bend as Layout
5	import ./transforms.bend as Transforms
6	import ./transpose_spec.bend as Transpose
7	import ./slice.bend as Slice
8	
9	# PyTorch's signed axes, indices and reshape sizes, parsed into the unsigned
10	# ones that the tensor operations check. A Position counts from either end, as
11	# a PyTorch int does: First{k} is k, and Last{k} is -1 - k, so Last{0} is -1,
12	# the last axis or element. Each integer has exactly one Position.
13	type Position is Data:
14	  First{k: U32}
15	  Last{k: U32}
16	
17	# A reshape size: an exact size, or -1, the one size PyTorch infers.
18	type Extent is Data:
19	  Exact{size: U32}
20	  Inferred{}
21	
22	# A signed reshape's refusals: the unsigned reshape's, more than one Inferred
23	# size, and an Inferred size that any value would fit, on an empty input whose
24	# other sizes multiply to 0.
25	type Error is Data:
26	  TransformError{error: Transforms.Error}
27	  RepeatedInferred{}
28	  AmbiguousInferred{}
29	
30	# PyTorch's wrap of a position among `limit` axes or elements (maybe_wrap_dim):
31	# a position d in -limit..limit-1 becomes d + limit when d is negative and d
32	# otherwise. A position outside that range stays outside it, at or past
33	# `limit`, so the operation refuses it with its own check, code and order:
34	# First{k} stays k, and an out-of-range Last{k} becomes `limit`.
35	def wrap(+limit: U32, position: Position) -> U32:
36	  match position:
37	    case First{k}:
38	      k
39	    case Last{+k}:
40	      Bool.pick(U32, U32.is_lt(k, limit), U32.sub(U32.sub(limit, 1), k), limit)
41	
42	def wrap_all(+limit: U32, positions: List<&2, Position>) -> List<&2, U32>:
43	  match positions:
44	    case Nil{}:
45	      Nil{}
46	    case head <> tail:
47	      wrap(limit, head) <> wrap_all(limit, tail)
48	
49	# Each coordinate wrapped among its own axis's extent. A coordinate past the
50	# rank has no extent and is wrapped among 0: the read refuses the coordinate
51	# count before any range.
52	def wrap_each(dimensions: List<&2, U32>, positions: List<&2, Position>) -> List<&2, U32>:
53	  match dimensions positions:
54	    case Nil{} positions:
55	      wrap_all(0, positions)
56	    case extent <> rest Nil{}:
57	      Nil{}
58	    case extent <> rest head <> tail:
59	      wrap(extent, head) <> wrap_each(rest, tail)
60	
61	# Whether a position lies in PyTorch's range -limit..limit-1. Last{k} is
62	# -1 - k, at least -limit exactly when k < limit.
63	def inside(+limit: U32, position: Position) -> Bool:
64	  match position:
65	    case First{k}:
66	      U32.is_lt(k, limit)
67	    case Last{k}:
68	      U32.is_lt(k, limit)
69	
70	def all_inside(+limit: U32, positions: List<&2, Position>) -> Bool:
71	  match positions:
72	    case Nil{}:
73	      True{}
74	    case head <> tail:
75	      Bool.and(inside(limit, head), all_inside(limit, tail))
76	
77	# The axes of the input for transpose and sum: its rank, and 1 for a rank-0
78	# input, whose only axis is 0 (or -1).
79	def axis_limit(+input: Shape.Shape) -> U32:
80	  U32.from_nat(Transpose.axis_limit(Shape.rank_dimensions(Shape.dimensions(input))))
81	
82	# How many sizes are Inferred.
83	def inferred_count(extents: List<&2, Extent>) -> Nat:
84	  match extents:
85	    case Nil{}:
86	      0n
87	    case Exact{size} <> tail:
88	      inferred_count(tail)
89	    case Inferred{} <> tail:
90	      1n+inferred_count(tail)
91	
92	# The product of the exact sizes, in Nat: what the laws state the inference
93	# against. The inference never forms it, since it can pass the native Nat
94	# limit of 2^48 on an empty input: [-1, 2^24, 2^24] multiplies to 2^48.
95	def exact_product(extents: List<&2, Extent>) -> Nat:
96	  match extents:
97	    case Nil{}:
98	      1n
99	    case Exact{+size} <> tail:
100	      Nat.mul(U32.to_nat(size), exact_product(tail))
101	    case Inferred{} <> tail:
102	      exact_product(tail)
103	
104	# The sizes, each Inferred one replaced by `size`.
105	def filled(extents: List<&2, Extent>, +size: U32) -> List<&2, U32>:
106	  match extents:
107	    case Nil{}:
108	      Nil{}
109	    case Exact{extent} <> tail:
110	      extent <> filled(tail, size)
111	    case Inferred{} <> tail:
112	      size <> filled(tail, size)
113	
114	# The exact product against a bound, the input's element count: ZeroProduct
115	# when it is 0, Within{p} when it is p > 0 at most the bound, and Beyond when
116	# it is past the bound. capped only multiplies a size and a product that are
117	# both at most the bound, and a tensor holds at most 2^20 elements, so no
118	# number it forms passes 2^40.
119	type Product is Data:
120	  ZeroProduct{}
121	  Within{product: Nat}
122	  Beyond{}
123	
124	def bounded(product: Nat, below: Bool) -> Product:
125	  match below:
126	    case True{}:
127	      Within{product}
128	    case False{}:
129	      Beyond{}
130	
131	def cap(+bound: Nat, +product: Nat) -> Product:
132	  bounded(product, Nat.is_le(product, bound))
133	
134	def scaled(+bound: Nat, +size: Nat, +product: Nat, past: Bool) -> Product:
135	  match past:
136	    case True{}:
137	      Beyond{}
138	    case False{}:
139	      cap(bound, Nat.mul(size, product))
140	
141	# A nonzero size times the product of the sizes after it. A size past the
142	# bound makes any nonzero product pass it, so it is never multiplied.
143	def times(+bound: Nat, size: Nat, rest: Product) -> Product:
144	  match size:
145	    case 0n:
146	      ZeroProduct{}
147	    case 1n+(+less):
148	      match rest:
149	        case ZeroProduct{}:
150	          ZeroProduct{}
151	        case Within{+product}:
152	          scaled(bound, 1n+less, product, Nat.is_lt(bound, 1n+less))
153	        case Beyond{}:
154	          Beyond{}
155	
156	def capped(+bound: Nat, extents: List<&2, Extent>) -> Product:
157	  match extents:
158	    case Nil{}:
159	      cap(bound, 1n)
160	    case Exact{+size} <> tail:
161	      times(bound, U32.to_nat(size), capped(bound, tail))
162	    case Inferred{} <> tail:
163	      capped(bound, tail)
164	
165	def count_mismatch() -> Error:
166	  TransformError{Transforms.LayoutError{Layout.ElementCountMismatch{}}}
167	
168	# An Inferred size whose other sizes multiply to 0: ambiguous on an empty
169	# input, a count mismatch on any other.
170	def unfilled(count: Nat) -> Result<&2, &2, Error, List<&2, U32>>:
171	  match count:
172	    case 0n:
173	      Fail{AmbiguousInferred{}}
174	    case 1n+rest:
175	      Fail{count_mismatch()}
176	
177	def divided(+extents: List<&2, Extent>, split: Nat & Nat) -> Result<&2, &2, Error, List<&2, U32>>:
178	  (quotient, remainder) = split
179	  match remainder:
180	    case 0n:
181	      Done{filled(extents, U32.from_nat(quotient))}
182	    case 1n+rest:
183	      Fail{count_mismatch()}
184	
185	# A product past the count: the inferred size is 0 on an empty input, and no
186	# size fits any other.
187	def beyond(count: Nat, +extents: List<&2, Extent>) -> Result<&2, &2, Error, List<&2, U32>>:
188	  match count:
189	    case 0n:
190	      Done{filled(extents, 0)}
191	    case 1n+rest:
192	      Fail{count_mismatch()}
193	
194	def inferred_one(+count: Nat, +extents: List<&2, Extent>, product: Product) -> Result<&2, &2, Error, List<&2, U32>>:
195	  match product:
196	    case ZeroProduct{}:
197	      unfilled(count)
198	    case Within{product}:
199	      divided(extents, Nat.divmod(count, product))
200	    case Beyond{}:
201	      beyond(count, extents)
202	
203	def inferred_from(+count: Nat, +extents: List<&2, Extent>, inferred: Nat) -> Result<&2, &2, Error, List<&2, U32>>:
204	  match inferred:
205	    case 0n:
206	      Done{filled(extents, 0)}
207	    case 1n+more:
208	      match more:
209	        case 0n:
210	          inferred_one(count, extents, capped(count, extents))
211	        case 1n+rest:
212	          Fail{RepeatedInferred{}}
213	
214	# PyTorch's infer_size: the sizes of a reshape of `count` elements. With no
215	# Inferred size, the sizes themselves; with one, it is count / product of
216	# the others when that divides exactly, the product capped at the count;
217	# more than one is refused first.
218	def inferred_dimensions(+count: Nat, +extents: List<&2, Extent>) -> Result<&2, &2, Error, List<&2, U32>>:
219	  inferred_from(count, extents, inferred_count(extents))
220	
221	def transformed(result: Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Result<&2, &1, Error, Core.Dynamic>:
222	  match result:
223	    case Fail{error}:
224	      Fail{TransformError{error}}
225	    case Done{output}:
226	      Done{output}
227	
228	def reshape_inferred(input: Shape.Shape, tensor: Core.Tensor<input>,
229	  inferred: Result<&2, &2, Error, List<&2, U32>>) -> Result<&2, &1, Error, Core.Dynamic>:
230	  match inferred:
231	    case Fail{error}:
232	      Fail{error}
233	    case Done{+dimensions}:
234	      transformed(Transforms.reshape(input, tensor, dimensions))
235	
236	# PyTorch's reshape with -1: the sizes inferred, then the unsigned reshape.
237	def reshape(+input: Shape.Shape, tensor: Core.Tensor<input>, +extents: List<&2, Extent>) -> Result<&2, &1, Error, Core.Dynamic>:
238	  reshape_inferred(input, tensor, inferred_dimensions(Shape.elements(input), extents))
239	
240	def reshape_dynamic(dynamic: Core.Dynamic, extents: List<&2, Extent>) -> Result<&2, &1, Error, Core.Dynamic>:
241	  Core.Packed{+input, tensor} = dynamic
242	  reshape(input, tensor, extents)
243	
244	# PyTorch's permute_copy with signed axes, wrapped among the input's rank.
245	def permute_copy(+input: Shape.Shape, source: Core.Tensor<input>, axes: List<&2, Position>) ->
246	  Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>:
247	  Transforms.permute_copy(input, source, wrap_all(Shape.rank(input), axes))
248	
249	def permute_dynamic(dynamic: Core.Dynamic, axes: List<&2, Position>) -> Core.Dynamic & Result<&2, &1, Transforms.Error, Core.Dynamic>:
250	  Core.Packed{+input, source} = dynamic
251	  Transforms.dynamic_result(input, permute_copy(input, source, axes))
252	
253	# PyTorch's transpose_copy with signed axes, wrapped among axis_limit.
254	def transpose_copy(+input: Shape.Shape, source: Core.Tensor<input>, a: Position, b: Position) ->
255	  Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>:
256	  Transforms.transpose_copy(input, source, wrap(axis_limit(input), a), wrap(axis_limit(input), b))
257	
258	def transpose_dynamic(dynamic: Core.Dynamic, a: Position, b: Position) -> Core.Dynamic & Result<&2, &1, Transforms.Error, Core.Dynamic>:
259	  Core.Packed{+input, source} = dynamic
260	  Transforms.dynamic_result(input, transpose_copy(input, source, a, b))
261	
262	# A[coordinates] with signed coordinates, each wrapped among its axis's
263	# extent; and the write A[coordinates] = value.
264	def read_coordinate(+shape: Shape.Shape, tensor: Core.Tensor<shape>, coordinates: List<&2, Position>) ->
265	  Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>:
266	  Core.read_coordinate(shape, tensor, wrap_each(Shape.dimensions(shape), coordinates))
267	
268	def write_coordinate(+shape: Shape.Shape, tensor: Core.Tensor<shape>, coordinates: List<&2, Position>, value: F32) ->
269	  Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>:
270	  Core.write_coordinate(shape, tensor, wrap_each(Shape.dimensions(shape), coordinates), value)
271	
272	# The element at a signed flat index, PyTorch's A.flatten()[index], wrapped
273	# among the element count; and the write at it.
274	def read_flat(+shape: Shape.Shape, tensor: Core.Tensor<shape>, index: Position) ->
275	  Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>:
276	  Core.read_flat(shape, tensor, wrap(Shape.count(shape), index))
277	
278	def write_flat(+shape: Shape.Shape, tensor: Core.Tensor<shape>, index: Position, +value: F32) ->
279	  Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>:
280	  Core.write_flat(shape, tensor, wrap(Shape.count(shape), index), value)
281	
282	def read_dynamic(dynamic: Core.Dynamic, index: Position) -> Core.Dynamic & Result<&2, &2, Core.Error, F32>:
283	  Core.Packed{+shape, tensor} = dynamic
284	  Core.read_dynamic_result(shape, read_flat(shape, tensor, index))
285	
286	def write_dynamic(dynamic: Core.Dynamic, index: Position, value: F32) -> Core.Dynamic & Result<&2, &2, Core.Error, Unit>:
287	  Core.Packed{+shape, tensor} = dynamic
288	  Core.write_dynamic_result(shape, write_flat(shape, tensor, index, value))
289	
290	# Whether a position lies before -limit: Last{k}, the position -1 - k, with
291	# k at or past limit. No First does.
292	def before(+limit: U32, position: Position) -> Bool:
293	  match position:
294	    case First{k}:
295	      False{}
296	    case Last{k}:
297	      Bool.not(U32.is_lt(k, limit))
298	
299	# PyTorch's slice positions (aten::slice, get_slice_bounds) before its clamp
300	# to the axis: a negative position d becomes d + limit, the wrap, and one
301	# still below 0 becomes 0. First{k} stays k; Slice.bounds then clamps both
302	# ends to the extent.
303	def bound(+limit: U32, +position: Position) -> U32:
304	  Bool.pick(U32, before(limit, position), 0, wrap(limit, position))
305	
306	# The extent of axis `axis`, 0 past the rank.
307	def extent(dimensions: List<&2, U32>, +axis: U32) -> U32:
308	  U32.from_nat(Slice.extent_at(dimensions, axis))
309	
310	# PyTorch's input[..., start:stop:step, ...] with the slice at a signed axis
311	# `dim`, wrapped among the rank, and signed start and stop: the unsigned
312	# slice of the bounds. A rank-0 input has no axis, so every dim is refused.
313	def slice(input: Core.Dynamic, dim: Position, start: Position, stop: Position, +step: U32) ->
314	  Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>:
315	  Core.Packed{+shape, tensor} = input
316	  +axis = wrap(Shape.rank(shape), dim)
317	  +n = extent(Shape.dimensions(shape), axis)
318	  Slice.slice(Core.Packed{shape, tensor}, axis, bound(n, start), bound(n, stop), step)
319	
320	# PyTorch's input.select(dim, index) at a signed axis and a signed index,
321	# wrapped among the extent: an index outside -extent..extent-1 stays at or
322	# past the extent, so the unsigned select refuses it with OutOfRange after
323	# its axis check.
324	def select(input: Core.Dynamic, dim: Position, index: Position) -> Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>:
325	  Core.Packed{+shape, tensor} = input
326	  +axis = wrap(Shape.rank(shape), dim)
327	  Slice.select(Core.Packed{shape, tensor}, axis, wrap(extent(Shape.dimensions(shape), axis), index))
328	
329	# A narrow's start before -extent: the axis check first, as the unsigned
330	# narrow does it, then OutOfRange.
331	def narrow_before(input: Core.Dynamic, refused: Maybe<&2, Slice.Error>) -> Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>:
332	  match refused:
333	    case Some{error}:
334	      (input, Fail{error})
335	    case None{}:
336	      (input, Fail{Slice.OutOfRange{}})
337	
338	def narrow_started(+shape: Shape.Shape, tensor: Core.Tensor<shape>, +axis: U32, +start: U32, +length: U32, early: Bool) ->
339	  Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>:
340	  match early:
341	    case False{}:
342	      Slice.narrow(Core.Packed{shape, tensor}, axis, start, length)
343	    case True{}:
344	      narrow_before(Core.Packed{shape, tensor}, Slice.refusal(Shape.dimensions(shape), axis, 1))
345	
346	# PyTorch's input.narrow(dim, start, length) at a signed axis with a signed
347	# start. PyTorch takes a start in -extent..extent: Last{k} with k below the
348	# extent is extent - 1 - k, and one before -extent is OutOfRange. A First
349	# start past the extent is refused by the unsigned narrow, which checks
350	# start + length against the extent, with the same code. A length is a
351	# U32, so PyTorch's negative lengths cannot be written.
352	def narrow(input: Core.Dynamic, dim: Position, +start: Position, +length: U32) ->
353	  Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>:
354	  Core.Packed{+shape, tensor} = input
355	  +axis = wrap(Shape.rank(shape), dim)
356	  +n = extent(Shape.dimensions(shape), axis)
357	  narrow_started(shape, tensor, axis, wrap(n, start), length, before(n, start))
358	
359	def error_code(error: Error) -> U32:
360	  match error:
361	    case TransformError{error}:
362	      Transforms.error_code(error)
363	    case RepeatedInferred{}:
364	      17
365	    case AmbiguousInferred{}:
366	      18