tensor_view.bend

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