tensor_view.bend
stelliferous@0.0.3.0 · 299 lines · 0 definitions
1# Rank-independent forward-strided metadata. The Array owner is passed separately.
2# A vector, matrix and tensor use the same descriptor and lookup operation.
3import Base
4import stelliferous@0.0.3.0/tensor_shape.bend as Shapes
5import stelliferous@0.0.3.0/machine_limits.bend as Machine
6
7type Axis is Data:
8 Axis{extent: U32,stride: U32}
9
10type View is Data:
11 View{axes: List<&2,Axis>,offset: U32}
12
13# Successful addresses are below logical_length. The checked product and the
14# remaining-room comparison prevent multiplication and addition from wrapping.
15def address_when(valid: Bool,offset: U32,increment: U32) -> Maybe<&2,U32>:
16 match valid:
17 case False{}: None{}
18 case True{}: Some{(offset + increment : U32)}
19
20def finish_address(valid: Bool,offset: U32) -> Maybe<&2,U32>:
21 match valid:
22 case False{}: None{}
23 case True{}: Some{offset}
24
25def advance_product(product: Maybe<&2,U32>,offset: U32,room: U32) -> Maybe<&2,U32>:
26 match product:
27 case None{}: None{}
28 case Some{+increment}: address_when(U32.is_le(increment,room),offset,increment)
29
30def advance_inside(inside: Bool,+offset: U32,coordinate: U32,+stride: U32,logical_length: U32) -> Maybe<&2,U32>:
31 match inside:
32 case False{}: None{}
33 case True{}: advance_product(Shapes.checked_multiply(coordinate,stride),offset,(logical_length - 1 - offset : U32))
34
35def advance(+offset: U32,coordinate: U32,stride: U32,+logical_length: U32) -> Maybe<&2,U32>:
36 advance_inside(U32.is_lt(offset,logical_length),offset,coordinate,stride,logical_length)
37
38# The recursive step consumes one axis and one coordinate; no flattening/copy.
39def locate(axes: List<&2,Axis>,coordinates: List<&2,U32>,address: Maybe<&2,U32>,+logical_length: U32) -> Maybe<&2,U32>:
40 match axes coordinates address:
41 case Nil{} Nil{} Some{+offset}:
42 finish_address(U32.is_lt(offset,logical_length),offset)
43 case Con{Axis{+extent,stride},tail} Con{+coordinate,rest} Some{+offset}:
44 locate(tail,rest,advance_inside(U32.is_lt(coordinate,extent) && U32.is_lt(offset,logical_length),offset,coordinate,stride,logical_length),logical_length)
45 case axes coordinates address: None{}
46
47def address(view: View,coordinates: List<&2,U32>,logical_length: U32) -> Maybe<&2,U32>:
48 View{axes,offset} = view
49 locate(axes,coordinates,Some{offset},logical_length)
50
51def axis_extents(axes: List<&2,Axis>) -> List<&2,U32>:
52 match axes:
53 case Nil{}: Nil{}
54 case Con{Axis{extent,stride},tail}: extent <> axis_extents(tail)
55
56def shape(view: View) -> Shapes.Shape:
57 View{axes,offset} = view
58 Shapes.Shape{axis_extents(axes)}
59
60def numel(view: View) -> Maybe<&2,U32>:
61 Shapes.numel(shape(view))
62
63def rank(view: View) -> Nat:
64 Shapes.rank(shape(view))
65
66def extent(view: View,axis: Nat) -> Maybe<&2,U32>:
67 Shapes.extent(shape(view),axis)
68
69def build_axes(reversed: List<&2,U32>,stride: Maybe<&2,U32>,axes: List<&2,Axis>,offset: U32) -> Maybe<&2,View>:
70 match reversed stride:
71 case Nil{} Some{value}: Some{View{axes,offset}}
72 case Con{+size,tail} Some{+value}:
73 build_axes(tail,Shapes.checked_multiply(value,size),Axis{size,value} <> axes,offset)
74 case reversed None{}: None{}
75
76def contiguous_count(count: Maybe<&2,U32>,extents: List<&2,U32>,offset: U32) -> Maybe<&2,View>:
77 match count:
78 case None{}: None{}
79 case Some{+value}:
80 build_axes(List.reverse(&2,U32,extents),Some{Bool.pick(U32,U32.is_zero(value),0,1)},Nil{},offset)
81
82def contiguous(+dimensions: Shapes.Shape,offset: U32) -> Maybe<&2,View>:
83 Shapes.Shape{extents} = dimensions
84 contiguous_count(Shapes.numel(dimensions),extents,offset)
85
86def contiguous_product(valid: Bool,extent: U32,inner_count: U32) -> Maybe<&2,U32>:
87 match valid:
88 case False{}: None{}
89 case True{}: Shapes.checked_multiply(extent,inner_count)
90
91def contiguous_axis(+extent: U32,stride: U32,inner: Maybe<&2,U32>) -> Maybe<&2,U32>:
92 match inner:
93 case None{}: None{}
94 case Some{+count}: contiguous_product(U32.is_le(extent,1) || U32.is_eq(stride,count),extent,count)
95
96# Validate strides from the innermost axis outward. This avoids allocating a
97# second canonical View just to compare its metadata. Empty views bypass it.
98def contiguous_size(axes: List<&2,Axis>) -> Maybe<&2,U32>:
99 match axes:
100 case []: Some{1}
101 case Axis{extent,stride} <> tail: contiguous_axis(extent,stride,contiguous_size(tail))
102
103def has_count(count: Maybe<&2,U32>) -> Bool:
104 match count:
105 case None{}: False{}
106 case Some{value}: True{}
107
108def contiguous_nonempty(empty: Bool,+axes: List<&2,Axis>) -> Bool:
109 match empty:
110 case True{}: True{}
111 case False{}: has_count(contiguous_size(axes))
112
113def is_contiguous(view: View) -> Bool:
114 View{+axes,offset} = view
115 contiguous_nonempty(Shapes.has_zero(axis_extents(axes)),axes)
116
117def last_coordinates(axes: List<&2,Axis>) -> List<&2,U32>:
118 match axes:
119 case Nil{}: Nil{}
120 case Con{Axis{size,stride},tail}: (size - 1 : U32) <> last_coordinates(tail)
121
122def found(address: Maybe<&2,U32>) -> Bool:
123 match address:
124 case None{}: False{}
125 case Some{value}: True{}
126
127def fits_nonempty(empty: Bool,+view: View,+logical_length: U32) -> Bool:
128 match empty:
129 case True{}:
130 View{axes,offset} = view
131 U32.is_le(offset,logical_length)
132 case False{}:
133 View{+axes,offset} = view
134 found(address(View{axes,offset},last_coordinates(axes),logical_length))
135
136# Forward strides are nonnegative: the last coordinate is the maximum address.
137# Empty views may point exactly one past the logical end, and perform no reads.
138def axis_extents_of(view: View) -> List<&2,U32>:
139 View{axes,offset} = view
140 axis_extents(axes)
141
142def fits(view: View,logical_length: U32) -> Bool:
143 +view = view
144 fits_nonempty(Shapes.has_zero(axis_extents_of(view)),view,logical_length)
145
146# The view read with the flagged axes in reverse: such an axis starts at its
147# last element and steps by the negated stride, in wrapping U32 arithmetic.
148# Every address a traversal of it reaches is an address of the view. Its
149# strides are not forward strides, so no check of this file accepts it.
150def flag(flags: List<&2,Bool>) -> Bool:
151 match flags:
152 case Nil{}: False{}
153 case Con{head,tail}: head
154
155def later_flags(flags: List<&2,Bool>) -> List<&2,Bool>:
156 match flags:
157 case Nil{}: Nil{}
158 case Con{head,tail}: tail
159
160# The flags after the first count, the first count flags, and whether none of
161# those is set.
162def drop_flags(count: Nat,flags: List<&2,Bool>) -> List<&2,Bool>:
163 match count:
164 case 0n: flags
165 case 1n+rest: drop_flags(rest,later_flags(flags))
166
167def fitted_flags(count: Nat,+flags: List<&2,Bool>) -> List<&2,Bool>:
168 match count:
169 case 0n: []
170 case 1n+rest: flag(flags) <> fitted_flags(rest,later_flags(flags))
171
172def forward(count: Nat,+flags: List<&2,Bool>) -> Bool:
173 match count:
174 case 0n: True{}
175 case 1n+rest: Bool.not(flag(flags)) && forward(rest,later_flags(flags))
176
177def directed(reverse: Bool,stride: U32) -> U32:
178 match reverse:
179 case False{}: stride
180 case True{}: (0 - stride : U32)
181
182def mirrored_start(reverse: Bool,+extent: U32,stride: U32,+offset: U32) -> U32:
183 match reverse:
184 case False{}: offset
185 case True{}: Bool.pick(U32,U32.is_lt(1,extent),(offset + (extent - 1) * stride : U32),offset)
186
187def mirrored_axes(axes: List<&2,Axis>,+flags: List<&2,Bool>) -> List<&2,Axis>:
188 match axes:
189 case Nil{}: Nil{}
190 case Con{Axis{extent,stride},tail}: Axis{extent,directed(flag(flags),stride)} <> mirrored_axes(tail,later_flags(flags))
191
192def mirrored_offset(axes: List<&2,Axis>,+flags: List<&2,Bool>,offset: U32) -> U32:
193 match axes:
194 case Nil{}: offset
195 case Con{Axis{extent,+stride},tail}: mirrored_offset(tail,later_flags(flags),mirrored_start(flag(flags),extent,stride,offset))
196
197def mirrored(view: View,+flags: List<&2,Bool>) -> View:
198 View{+axes,offset} = view
199 View{mirrored_axes(axes,flags),mirrored_offset(axes,flags,offset)}
200
201def contains_index(indices: List<&2,Nat>,+axis: Nat) -> Bool:
202 match indices:
203 case Nil{}: False{}
204 case Con{index,tail}: Nat.is_eq(index,axis) || contains_index(tail,axis)
205
206def unique_axes(order: List<&2,Nat>) -> Bool:
207 match order:
208 case Nil{}: True{}
209 case Con{axis,+tail}: Bool.not(contains_index(tail,axis)) && unique_axes(tail)
210
211def prepend_axis(axis: Maybe<&2,Axis>,view: Maybe<&2,View>) -> Maybe<&2,View>:
212 match axis view:
213 case Some{item} Some{View{axes,offset}}: Some{View{item <> axes,offset}}
214 case axis view: None{}
215
216def reorder(order: List<&2,Nat>,+axes: List<&2,Axis>,offset: U32) -> Maybe<&2,View>:
217 match order:
218 case Nil{}: Some{View{Nil{},offset}}
219 case Con{index,tail}:
220 prepend_axis(List.get(&2,Axis,axes,index),reorder(tail,axes,offset))
221
222def permute_when(valid: Bool,view: View,order: List<&2,Nat>) -> Maybe<&2,View>:
223 match valid:
224 case False{}: None{}
225 case True{}:
226 View{axes,offset} = view
227 reorder(order,axes,offset)
228
229# Order must contain each axis exactly once, including singleton axes.
230def permute(+view: View,+order: List<&2,Nat>) -> Maybe<&2,View>:
231 permute_when(Nat.is_eq(rank(view),List.length(&2,Nat,order)) && unique_axes(order),view,order)
232
233def slice_step(zero: Bool,+start: U32,+count: U32,+step: U32,+size: U32) -> Bool:
234 match zero:
235 case True{}: False{}
236 case False{}:
237 U32.is_le(start,size) && (U32.is_zero(count) ||
238 (U32.is_lt(start,size) && U32.is_le((count - 1 : U32),((size - 1 - start) / step : U32))))
239
240def sliced_axis(valid: Bool,count: U32,stride: Maybe<&2,U32>,offset: Maybe<&2,U32>,tail: List<&2,Axis>) -> Maybe<&2,View>:
241 match valid stride offset:
242 case True{} Some{stride} Some{offset}: Some{View{Axis{count,stride} <> tail,offset}}
243 case valid stride offset: None{}
244
245def add_offset(shift: Maybe<&2,U32>,+offset: U32) -> Maybe<&2,U32>:
246 match shift:
247 case None{}: None{}
248 case Some{+value}: address_when(U32.is_le(value,(Machine.word_max() - offset : U32)),offset,value)
249
250def slice_axes(axes: List<&2,Axis>,axis: Nat,+start: U32,+count: U32,+step: U32,+offset: U32) -> Maybe<&2,View>:
251 match axes axis:
252 case Con{Axis{size,+stride},tail} 0n:
253 sliced_axis(slice_step(U32.is_zero(step),start,count,step,size),count,Shapes.checked_multiply(stride,step),
254 add_offset(Shapes.checked_multiply(start,stride),offset),tail)
255 case Con{head,tail} 1n+rest:
256 prepend_axis(Some{head},slice_axes(tail,rest,start,count,step,offset))
257 case Nil{} axis: None{}
258
259# Explicit start/count/positive step. The resulting view is validated against
260# the actual Buffer at execution time; arithmetic overflow is rejected here.
261def slice(view: View,axis: Nat,start: U32,count: U32,step: U32) -> Maybe<&2,View>:
262 View{axes,offset} = view
263 slice_axes(axes,axis,start,count,step,offset)
264
265# One axis read as extent elements: unchanged when the extents agree, the
266# single element repeated (stride 0) when the axis has extent 1.
267def stretched_axis(same: Bool,single: Bool,size: U32,stride: U32,extent: U32) -> Maybe<&2,Axis>:
268 match same single:
269 case True{} single: Some{Axis{size,stride}}
270 case False{} True{}: Some{Axis{extent,0}}
271 case False{} False{}: None{}
272
273def stretched_prepend(axis: Maybe<&2,Axis>,rest: Maybe<&2,List<&2,Axis>>) -> Maybe<&2,List<&2,Axis>>:
274 match axis rest:
275 case Some{axis} Some{rest}: Some{axis <> rest}
276 case axis rest: None{}
277
278def stretch_axes(axes: List<&2,Axis>,target: List<&2,U32>) -> Maybe<&2,List<&2,Axis>>:
279 match axes target:
280 case Nil{} Nil{}: Some{[]}
281 case Con{Axis{+size,stride},tail} Con{+extent,rest}:
282 stretched_prepend(stretched_axis(U32.is_eq(size,extent),U32.is_eq(size,1),size,stride,extent),stretch_axes(tail,rest))
283 case axes target: None{}
284
285# The axes with leading axes of extent 1 added up to rank; the same elements.
286def with_rank(+axes: List<&2,Axis>,rank: Nat) -> List<&2,Axis>:
287 List.append(&2,Axis,List.replicate(Axis,Nat.sub(rank,List.length(&2,Axis,axes)),Axis{1,0}),axes)
288
289def stretched_view(axes: Maybe<&2,List<&2,Axis>>,offset: U32) -> Maybe<&2,View>:
290 match axes:
291 case Some{axes}: Some{View{axes,offset}}
292 case None{}: None{}
293
294# The view read with the extents target under NumPy broadcasting. No element
295# moves: an axis of extent 1 is repeated with stride 0, and so is an axis the
296# view lacks.
297def stretch(view: View,+target: List<&2,U32>) -> Maybe<&2,View>:
298 View{axes,offset} = view
299 stretched_view(stretch_axes(with_rank(axes,List.length(&2,U32,target)),target),offset)