tensor_layout.bend

stelliferous@0.0.2.0 · 565 lines · 81 definitions

1# Layout operations consume their arguments. Copy and nearest upsampling
2# preserve row bands. Other layout operations first obtain one dense owner;
3# reshape relabels it, while the remaining operations use checked view copies.
4import Base
5import stelliferous@0.0.2.0/storage_buffer.bend as Storage
6import stelliferous@0.0.2.0/tensor_shape.bend as Shapes
7import stelliferous@0.0.2.0/tensor_view.bend as Views
8import stelliferous@0.0.2.0/tensor_operations.bend as Operations
9import stelliferous@0.0.2.0/tensor.bend as Tensors
10import stelliferous@0.0.2.0/tensor_band.bend as Bands
11import stelliferous@0.0.2.0/traversal_partition.bend as Partition
12import stelliferous@0.0.2.0/tensor_band_concat.bend as BandConcat
13import stelliferous@0.0.2.0/tensor_band_schedule.bend as BandSchedule
14
15def extents(view: Maybe<&2,Views.View>) -> List<&2,U32>:
16  match view:
17    case Some{view}: Views.axis_extents_of(view)
18    case None{}: []
19
20def naturals(values: List<&2,U32>) -> List<&2,Nat>:
21  match values:
22    case []: []
23    case value <> rest: U32.to_nat(value) <> naturals(rest)
24
25def copied(~Element: Data,name: String,+source_shape: List<&2,U32>,+target_shape: List<&2,U32>,+fill: Element,+workers: U32,
26  result: Operations.Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:
27  match result:
28    case Operations.Accepted{(input,output)}:
29      (Tensors.Tensor{output,target_shape,fill,Tensors.Usable{},workers},Tensors.Tensor{input,source_shape,fill,Tensors.Usable{},workers})
30    case Operations.Rejected{(input,output)}:
31      (Tensors.Tensor{output,target_shape,fill,Tensors.Failed{Tensors.Problem{name,Tensors.InvalidArgument{source_shape}}},workers},
32        Tensors.Tensor{input,source_shape,fill,Tensors.Usable{},workers})
33
34def gather_views(~Element: Data,name: String,source_view: Maybe<&2,Views.View>,target_view: Maybe<&2,Views.View>,count: Maybe<&2,U32>,
35  input: Storage.Buffer<Element>,+source_shape: List<&2,U32>,+target_shape: List<&2,U32>,+fill: Element,+workers: U32) ->
36  Tensors.Tensor<Element> & Tensors.Tensor<Element>:
37  match source_view target_view count:
38    case Some{source_view} Some{target_view} Some{count}:
39      copied(~Element,name,source_shape,target_shape,fill,workers,Operations.copy(~Element,input,source_view,Tensors.buffer(~Element,fill,count),target_view))
40    case source_view target_view count:
41      (Tensors.failed(~Element,name,Tensors.InvalidArgument{source_shape},target_shape,fill,workers),Tensors.Tensor{input,source_shape,fill,Tensors.Usable{},workers})
42
43# Copy the elements selected by source_view, in its row-major order, into a
44# new tensor whose storage is laid out as view_shape and labelled target_shape.
45def gather_status(~Element: Data,name: String,status: Tensors.Status,source_view: Maybe<&2,Views.View>,+view_shape: List<&2,U32>,
46  +target_shape: List<&2,U32>,input: Storage.Buffer<Element>,+source_shape: List<&2,U32>,+fill: Element,+workers: U32) ->
47  Tensors.Tensor<Element> & Tensors.Tensor<Element>:
48  match status:
49    case Tensors.Usable{}:
50      gather_views(~Element,name,source_view,Tensors.view_of(view_shape),Tensors.count_of(view_shape),input,source_shape,target_shape,fill,workers)
51    case Tensors.Failed{+problem}:
52      (Tensors.Tensor{Tensors.buffer(~Element,fill,0),target_shape,fill,Tensors.Failed{problem},workers},
53        Tensors.Tensor{input,source_shape,fill,Tensors.Failed{problem},workers})
54
55def result_of(~Element: Data,pair: Tensors.Tensor<Element> & Tensors.Tensor<Element>) -> Tensors.Tensor<Element>:
56  (result,source) = pair
57  result
58
59def swapped(~Element: Data,pair: Tensors.Tensor<Element> & Tensors.Tensor<Element>) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:
60  (result,source) = pair
61  (source,result)
62
63# The only layout operation that keeps its argument: (original, copy).
64def copy_bands(~Element: Data,status: Tensors.Status,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,+shape: List<&2,U32>,+fill: Element,+workers: U32) ->
65  Tensors.Tensor<Element> & Tensors.Tensor<Element>:
66  match status:
67    case Tensors.Usable{}:
68      Storage.bind_pair(~Bands.Bands(~Element),~Bands.Bands(~Element),~(Tensors.Tensor<Element> & Tensors.Tensor<Element>),
69        Bands.copy(~Element,bands),left => right =>
70          (Tensors.Banded{left,geometry,shape,fill,Tensors.Usable{},workers},Tensors.Banded{right,geometry,shape,fill,Tensors.Usable{},workers}))
71    case Tensors.Failed{+problem}:
72      (Tensors.Banded{bands,geometry,shape,fill,Tensors.Failed{problem},workers},
73        Tensors.Tensor{Tensors.buffer(~Element,fill,0),shape,fill,Tensors.Failed{problem},workers})
74
75def copy(~Element: Data,tensor: Tensors.Tensor<Element>) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:
76  match tensor:
77    case Tensors.Tensor{storage,+shape,fill,status,workers}:
78      swapped(~Element,gather_status(~Element,"copy",status,Tensors.view_of(shape),shape,shape,storage,shape,fill,workers))
79    case Tensors.Banded{bands,geometry,shape,fill,status,workers}: copy_bands(~Element,status,bands,geometry,shape,fill,workers)
80
81def permuted(view: Maybe<&2,Views.View>,order: List<&2,Nat>) -> Maybe<&2,Views.View>:
82  match view:
83    case Some{view}: Views.permute(view,order)
84    case None{}: None{}
85
86def transpose_view(~Element: Data,+source: Maybe<&2,Views.View>,status: Tensors.Status,storage: Storage.Buffer<Element>,+shape: List<&2,U32>,fill: Element,
87  +workers: U32) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:
88  gather_status(~Element,"transpose",status,source,extents(source),extents(source),storage,shape,fill,workers)
89
90# The result's axis i is the source axis order[i].
91def transpose(~Element: Data,tensor: Tensors.Tensor<Element>,order: List<&2,U32>) -> Tensors.Tensor<Element>:
92  Tensors.with_dense(~Element,~Tensors.Tensor<Element>,tensor,storage => shape => fill => status => workers =>
93    +shape = shape
94    result_of(~Element,transpose_view(~Element,permuted(Tensors.view_of(shape),naturals(order)),status,storage,shape,fill,workers)))
95
96def sliced(view: Maybe<&2,Views.View>,axis: U32,start: U32,count: U32,step: U32) -> Maybe<&2,Views.View>:
97  match view:
98    case Some{view}: Views.slice(view,U32.to_nat(axis),start,count,step)
99    case None{}: None{}
100
101def slice_view(~Element: Data,+source: Maybe<&2,Views.View>,status: Tensors.Status,storage: Storage.Buffer<Element>,+shape: List<&2,U32>,fill: Element,
102  +workers: U32) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:
103  gather_status(~Element,"slice",status,source,extents(source),extents(source),storage,shape,fill,workers)
104
105# Elements start, start + step, ... (count of them) along one axis.
106def slice(~Element: Data,tensor: Tensors.Tensor<Element>,axis: U32,start: U32,count: U32,step: U32) -> Tensors.Tensor<Element>:
107  Tensors.with_dense(~Element,~Tensors.Tensor<Element>,tensor,storage => shape => fill => status => workers =>
108    +shape = shape
109    result_of(~Element,slice_view(~Element,sliced(Tensors.view_of(shape),axis,start,count,step),status,storage,shape,fill,workers)))
110
111def same_count(left: Maybe<&2,U32>,right: Maybe<&2,U32>) -> Bool:
112  match left right:
113    case Some{left} Some{right}: U32.is_eq(left,right)
114    case left right: False{}
115
116def relabel_ready(~Element: Data,accepted: Bool,storage: Storage.Buffer<Element>,+shape: List<&2,U32>,+target: List<&2,U32>,fill: Element,
117  workers: U32) -> Tensors.Tensor<Element>:
118  match accepted:
119    case True{}: Tensors.Tensor{storage,target,fill,Tensors.Usable{},workers}
120    case False{}: Tensors.Tensor{storage,shape,fill,Tensors.Failed{Tensors.Problem{"reshape",Tensors.ShapeMismatch{shape,target}}},workers}
121
122def relabel(~Element: Data,accepted: Bool,current: Tensors.Status,storage: Storage.Buffer<Element>,+shape: List<&2,U32>,+target: List<&2,U32>,fill: Element,
123  workers: U32) -> Tensors.Tensor<Element>:
124  match current:
125    case Tensors.Usable{}: relabel_ready(~Element,accepted,storage,shape,target,fill,workers)
126    case Tensors.Failed{problem}: Tensors.Tensor{storage,shape,fill,Tensors.Failed{problem},workers}
127
128# Same elements in the same row-major order under new extents.
129def reshape(~Element: Data,tensor: Tensors.Tensor<Element>,+shape: List<&2,U32>) -> Tensors.Tensor<Element>:
130  Tensors.with_dense(~Element,~Tensors.Tensor<Element>,tensor,storage => source_shape => fill => status => workers =>
131    +source_shape = source_shape
132    relabel(~Element,same_count(Tensors.count_of(source_shape),Tensors.count_of(shape)),status,storage,source_shape,shape,fill,workers))
133
134# Nearest-neighbour upsampling repeats each element of the last two axes
135# scale times, by copying from a view whose repeated axes have stride zero.
136def prepend_repeated(axis: Views.Axis,rest: Maybe<&2,List<&2,Views.Axis>>) -> Maybe<&2,List<&2,Views.Axis>>:
137  match rest:
138    case Some{axes}: Some{axis <> axes}
139    case None{}: None{}
140
141def repeated_axes(axes: List<&2,Views.Axis>,+scale: U32) -> Maybe<&2,List<&2,Views.Axis>>:
142  match axes:
143    case [Views.Axis{height,height_stride},Views.Axis{width,width_stride}]:
144      Some{[Views.Axis{height,height_stride},Views.Axis{scale,0},Views.Axis{width,width_stride},Views.Axis{scale,0}]}
145    case [axis] : None{}
146    case []: None{}
147    case axis <> rest: prepend_repeated(axis,repeated_axes(rest,scale))
148
149def with_offset(axes: Maybe<&2,List<&2,Views.Axis>>,offset: U32) -> Maybe<&2,Views.View>:
150  match axes:
151    case Some{axes}: Some{Views.View{axes,offset}}
152    case None{}: None{}
153
154def repeated(view: Maybe<&2,Views.View>,scale: U32) -> Maybe<&2,Views.View>:
155  match view:
156    case Some{Views.View{axes,offset}}: with_offset(repeated_axes(axes,scale),offset)
157    case None{}: None{}
158
159def scaled_extents(extents: List<&2,U32>,+scale: U32) -> List<&2,U32>:
160  match extents:
161    case [height,width]: [(height * scale : U32),(width * scale : U32)]
162    case extent <> rest: extent <> scaled_extents(rest,scale)
163    case []: []
164
165def upsample_view(~Element: Data,+source: Maybe<&2,Views.View>,status: Tensors.Status,storage: Storage.Buffer<Element>,+shape: List<&2,U32>,+scale: U32,fill: Element,
166  +workers: U32) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:
167  gather_status(~Element,"upsample_nearest",status,source,extents(source),scaled_extents(shape,scale),storage,shape,fill,workers)
168
169type Upsample<-Element: Data> is Data:
170  Upsample{geometry: Bands.Geometry,scale: U32,fill: Element}
171
172def upsample_leaf(~Element: Data,context: Upsample<Element>,first: U32,rows: U32,input: Storage.Buffer<Element>) -> Bands.Bands(~Element):
173  Upsample{Bands.Geometry{batch,channels,height,width},+scale,fill} = context
174  +rows = rows
175  +shape = {[channels,rows,width] : List<&2,U32>}
176  +first = (first * scale : U32)
177  +rows = (rows * scale : U32)
178  Tensors.with_dense(~Element,~Bands.Bands(~Element),
179    result_of(~Element,upsample_view(~Element,repeated(Tensors.view_of(shape),scale),Tensors.Usable{},input,shape,scale,fill,1)),
180    output => shape => fill => status => workers => Partition.Leaf{output,first,rows,Unit{}})
181
182def upsample_supported(~Element: Data,valid: Bool,bands: Bands.Bands(~Element),geometry: Bands.Geometry,+shape: List<&2,U32>,+scale: U32,
183  +fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
184  match valid geometry:
185    case False{} geometry:
186      Tensors.with_dense(~Element,~Tensors.Tensor<Element>,Tensors.Banded{bands,geometry,shape,fill,Tensors.Usable{},workers},
187        storage => shape => fill => status => workers =>
188          +shape = shape
189          result_of(~Element,upsample_view(~Element,repeated(Tensors.view_of(shape),scale),status,storage,shape,scale,fill,workers)))
190    case True{} Bands.Geometry{+batch,+channels,+height,+width}:
191      Tensors.Banded{BandSchedule.transform(~Element,~Upsample<Element>,~upsample_leaf(~Element),Upsample{Bands.Geometry{batch,channels,height,width},scale,fill},
192          Bands.Geometry{batch,channels,height,width},workers,(U32.to_f32(scale) * U32.to_f32(scale) * 12.0 : F32),bands),
193        Bands.Geometry{batch,channels,(height * scale : U32),(width * scale : U32)},scaled_extents(shape,scale),fill,Tensors.Usable{},workers}
194
195def upsample_bands(~Element: Data,status: Tensors.Status,bands: Bands.Bands(~Element),geometry: Bands.Geometry,+shape: List<&2,U32>,+scale: U32,
196  +fill: Element,workers: U32) -> Tensors.Tensor<Element>:
197  match status geometry:
198    case Tensors.Failed{+problem} geometry:
199      Tensors.with_dense(~Element,~Tensors.Tensor<Element>,Tensors.Banded{bands,geometry,shape,fill,Tensors.Failed{problem},workers},
200        storage => shape => fill => status => workers =>
201          +shape = shape
202          result_of(~Element,upsample_view(~Element,repeated(Tensors.view_of(shape),scale),status,storage,shape,scale,fill,workers)))
203    case Tensors.Usable{} Bands.Geometry{batch,channels,+height,+width}:
204      upsample_supported(~Element,Bool.not(U32.is_zero(scale)) && Views.has_count(Shapes.checked_multiply(height,scale)) &&
205        Views.has_count(Shapes.checked_multiply(width,scale)) && Views.has_count(Tensors.count_of(scaled_extents(shape,scale))) &&
206        Bool.not(Shapes.has_zero(scaled_extents(shape,scale))),
207        bands,Bands.Geometry{batch,channels,height,width},shape,scale,fill,workers)
208
209def upsample_nearest(~Element: Data,tensor: Tensors.Tensor<Element>,+scale: U32) -> Tensors.Tensor<Element>:
210  match tensor:
211    case Tensors.Tensor{storage,+shape,fill,status,workers}:
212      result_of(~Element,upsample_view(~Element,repeated(Tensors.view_of(shape),scale),status,storage,shape,scale,fill,workers))
213    case Tensors.Banded{bands,geometry,shape,fill,status,workers}: upsample_bands(~Element,status,bands,geometry,shape,scale,fill,workers)
214
215def settle(~Element: Data,valid: Bool,result: Operations.Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>) ->
216  (Storage.Buffer<Element> & Storage.Buffer<Element>) & Bool:
217  match result:
218    case Operations.Accepted{owners}: (owners,valid)
219    case Operations.Rejected{owners}: (owners,False{})
220
221def count_or_zero(count: Maybe<&2,U32>) -> U32:
222  match count:
223    case Some{count}: count
224    case None{}: 0
225
226def extent_at(extents: List<&2,U32>,axis: Nat) -> U32:
227  count_or_zero(List.get(&2,U32,extents,axis))
228
229def settled(~Element: Data,name: String,valid: Bool,part: Storage.Buffer<Element>,target: Storage.Buffer<Element>,+target_shape: List<&2,U32>,+part_shape: List<&2,U32>,
230  target_fill: Element,part_fill: Element,+workers: U32) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:
231  match valid:
232    case True{}: (Tensors.Tensor{target,target_shape,target_fill,Tensors.Usable{},workers},Tensors.Tensor{part,part_shape,part_fill,Tensors.Usable{},workers})
233    case False{}:
234      (Tensors.Tensor{target,target_shape,target_fill,Tensors.Failed{Tensors.Problem{name,Tensors.ShapeMismatch{target_shape,part_shape}}},workers},
235        Tensors.Tensor{part,part_shape,part_fill,Tensors.Usable{},workers})
236
237def placed(~Element: Data,name: String,+target_shape: List<&2,U32>,+part_shape: List<&2,U32>,target_fill: Element,part_fill: Element,+workers: U32,
238  result: (Storage.Buffer<Element> & Storage.Buffer<Element>) & Bool) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:
239  (owners,valid) = result
240  (part,target) = owners
241  settled(~Element,name,valid,part,target,target_shape,part_shape,target_fill,part_fill,workers)
242
243def same_extents(slice: Views.View,part_shape: List<&2,U32>) -> Bool:
244  Shapes.same(Views.shape(slice),Shapes.Shape{part_shape})
245
246# The part written through a view of the target storage: one checked copy,
247# which rejects a view whose extents differ from the part. A part with no
248# element writes nothing; its extents are compared here.
249def place(~Element: Data,name: String,empty: Bool,slice: Maybe<&2,Views.View>,source: Maybe<&2,Views.View>,
250  target: Storage.Buffer<Element>,+target_shape: List<&2,U32>,target_fill: Element,
251  part: Storage.Buffer<Element>,+part_shape: List<&2,U32>,part_fill: Element,+workers: U32) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:
252  match empty slice source:
253    case False{} Some{slice} Some{source}:
254      placed(~Element,name,target_shape,part_shape,target_fill,part_fill,workers,settle(~Element,True{},Operations.scatter(~Element,part,source,target,slice)))
255    case True{} Some{slice} source:
256      settled(~Element,name,same_extents(slice,part_shape),part,target,target_shape,part_shape,target_fill,part_fill,workers)
257    case empty slice source: settled(~Element,name,False{},part,target,target_shape,part_shape,target_fill,part_fill,workers)
258
259def place_status(~Element: Data,name: String,target_status: Tensors.Status,part_status: Tensors.Status,slice: Maybe<&2,Views.View>,
260  target: Storage.Buffer<Element>,+target_shape: List<&2,U32>,+target_fill: Element,
261  part: Storage.Buffer<Element>,+part_shape: List<&2,U32>,+part_fill: Element,+workers: U32) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:
262  match target_status part_status:
263    case Tensors.Usable{} Tensors.Usable{}:
264      place(~Element,name,Shapes.has_zero(part_shape),slice,Tensors.view_of(part_shape),target,target_shape,target_fill,part,part_shape,part_fill,workers)
265    case Tensors.Usable{} Tensors.Failed{+problem}:
266      (Tensors.Tensor{target,target_shape,target_fill,Tensors.Failed{problem},workers},Tensors.Tensor{part,part_shape,part_fill,Tensors.Failed{problem},workers})
267    case Tensors.Failed{problem} part_status:
268      (Tensors.Tensor{target,target_shape,target_fill,Tensors.Failed{problem},workers},Tensors.Tensor{part,part_shape,part_fill,part_status,workers})
269
270# target[..., start : start + part extent, ...] = part along axis: the part is
271# written through the slice of the target that it replaces.
272def set_slice(~Element: Data,target: Tensors.Tensor<Element>,+axis: U32,start: U32,part: Tensors.Tensor<Element>) -> Tensors.Tensor<Element>:
273  Tensors.with_dense(~Element,~Tensors.Tensor<Element>,target,target => target_shape => target_fill => target_status => target_workers =>
274    +target_shape = target_shape
275    Tensors.with_dense(~Element,~Tensors.Tensor<Element>,part,part => part_shape => part_fill => part_status => part_workers =>
276      +part_shape = part_shape
277      result_of(~Element,place_status(~Element,"set_slice",target_status,part_status,
278        sliced(Tensors.view_of(target_shape),axis,start,extent_at(part_shape,U32.to_nat(axis)),1),target,target_shape,target_fill,part,part_shape,part_fill,
279        U32.max(target_workers,part_workers)))))
280
281# Concatenation along an axis.
282type Survey is Data:
283  Survey{valid: Bool,total: U32,status: Tensors.Status,workers: U32}
284
285def surveyed(survey: Survey,+first: List<&2,U32>,+shape: List<&2,U32>,+axis: Nat,status: Tensors.Status,workers: U32) -> Survey:
286  Survey{valid,+total,previous,most} = survey
287  +width = extent_at(shape,axis)
288  Survey{valid && Nat.is_eq(List.length(&2,U32,first),List.length(&2,U32,shape)) &&
289      Shapes.same_extents(List.take(&2,U32,first,axis),List.take(&2,U32,shape,axis)) &&
290      Shapes.same_extents(List.drop(&2,U32,first,1n+axis),List.drop(&2,U32,shape,1n+axis)) &&
291      U32.is_le(width,(4294967295 - total : U32)),
292    (total + width : U32),Tensors.first_failure(previous,status),U32.max(most,workers)}
293
294def prepend_survey(~Element: Data,tensor: Tensors.Tensor<Element>,result: Tensors.Parts<Element> & Survey) -> Tensors.Parts<Element> & Survey:
295  (rest,summary) = result
296  (Tensors.Cons{tensor,rest},summary)
297
298# Summarizes the extents of every part and returns the parts in order.
299def survey(~Element: Data,parts: Tensors.Parts<Element>,+first: List<&2,U32>,+axis: Nat,summary: Survey) -> Tensors.Parts<Element> & Survey:
300  match parts:
301    case Tensors.Single{tensor}:
302      Tensors.with_metadata(~Element,~(Tensors.Parts<Element> & Survey),tensor,tensor => shape => fill => status => workers =>
303        +shape = shape
304        +status = status
305        +workers = workers
306        (Tensors.Single{tensor},surveyed(summary,first,shape,axis,status,workers)))
307    case Tensors.Cons{tensor,rest}:
308      Tensors.with_metadata(~Element,~(Tensors.Parts<Element> & Survey),tensor,tensor => shape => fill => status => workers =>
309        +shape = shape
310        +status = status
311        +workers = workers
312        prepend_survey(~Element,tensor,survey(~Element,rest,first,axis,surveyed(summary,first,shape,axis,status,workers))))
313
314def with_extent(extents: List<&2,U32>,axis: Nat,+extent: U32) -> List<&2,U32>:
315  List.set(&2,U32,extents,axis,extent)
316
317def append_at(~Element: Data,+axis: U32,+start: U32,output: Tensors.Tensor<Element>,part: Tensors.Tensor<Element>) -> Tensors.Tensor<Element> & U32:
318  Tensors.with_dense(~Element,~(Tensors.Tensor<Element> & U32),part,storage => shape => fill => status => workers =>
319    +shape = shape
320    (set_slice(~Element,output,axis,start,Tensors.Tensor{storage,shape,fill,status,workers}),(start + extent_at(shape,U32.to_nat(axis)) : U32)))
321
322def append_part(~Element: Data,+axis: U32,state: Tensors.Tensor<Element> & U32,part: Tensors.Tensor<Element>) -> Tensors.Tensor<Element> & U32:
323  (output,start) = state
324  append_at(~Element,axis,start,output,part)
325
326def appended_all(~Element: Data,state: Tensors.Tensor<Element> & U32) -> Tensors.Tensor<Element>:
327  (output,start) = state
328  output
329
330def append_parts(~Element: Data,parts: Tensors.Parts<Element>,+axis: U32,state: Tensors.Tensor<Element> & U32) -> Tensors.Tensor<Element>:
331  match parts:
332    case Tensors.Single{part}: appended_all(~Element,append_part(~Element,axis,state,part))
333    case Tensors.Cons{part,rest}: append_parts(~Element,rest,axis,append_part(~Element,axis,state,part))
334
335def join_selected(~Element: Data,geometry: Maybe<&2,Bands.Geometry>,count: Maybe<&2,U32>,shape: List<&2,U32>,axis: U32,fill: Element,workers: U32,
336  layout: Maybe<&2,Bands.Plan>,parts: Tensors.Parts<Element>) -> Tensors.Tensor<Element>:
337  match geometry count layout:
338    case Some{geometry} Some{count} Some{layout}: BandConcat.join(~Element,layout,geometry,parts,shape,fill,workers)
339    case geometry count layout: append_parts(~Element,parts,axis,(Tensors.with_workers(~Element,Tensors.full(~Element,shape,fill),workers),0))
340
341def join_discovered(~Element: Data,+shape: List<&2,U32>,+axis: U32,fill: Element,workers: U32,
342  discovered: Tensors.Parts<Element> & Maybe<&2,Bands.Plan>) -> Tensors.Tensor<Element>:
343  (parts,layout) = discovered
344  join_selected(~Element,BandConcat.geometry(shape,axis),Tensors.count_of(shape),shape,axis,fill,workers,layout,parts)
345
346# The final branch is one Bool. Plan comparison stays in the metadata pass,
347# outside every element-copy loop. Dense append gathers each band source once.
348def join_compatible_selected(~Element: Data,compatible: Bool,+shape: List<&2,U32>,+axis: U32,+fill: Element,+workers: U32,
349  discovered: Tensors.Parts<Element> & Maybe<&2,Bands.Plan>) -> Tensors.Tensor<Element>:
350  match compatible:
351    case True{}: join_discovered(~Element,shape,axis,fill,workers,discovered)
352    case False{}:
353      (parts,layout) = discovered
354      append_parts(~Element,parts,axis,(Tensors.with_workers(~Element,Tensors.full(~Element,shape,fill),workers),0))
355
356def join_compatible(~Element: Data,+shape: List<&2,U32>,+axis: U32,+fill: Element,+workers: U32,
357  discovered: (Tensors.Parts<Element> & Maybe<&2,Bands.Plan>) & Bool) -> Tensors.Tensor<Element>:
358  (parts_and_layout,compatible) = discovered
359  join_compatible_selected(~Element,compatible && Bool.not(Shapes.has_zero(shape)),shape,axis,fill,workers,parts_and_layout)
360
361def join_status(~Element: Data,status: Tensors.Status,valid: Bool,+first: List<&2,U32>,+shape: List<&2,U32>,+axis: U32,+fill: Element,+workers: U32,
362  parts: Tensors.Parts<Element>) -> Tensors.Tensor<Element>:
363  match status valid:
364    case Tensors.Usable{} True{}: join_compatible(~Element,shape,axis,fill,workers,BandConcat.discover_compatible(~Element,parts))
365    case Tensors.Usable{} False{}: Tensors.failed(~Element,"concat",Tensors.InvalidArgument{first},first,fill,workers)
366    case Tensors.Failed{problem} other: Tensors.Tensor{Tensors.buffer(~Element,fill,0),first,fill,Tensors.Failed{problem},workers}
367
368def joined(~Element: Data,+first: List<&2,U32>,+axis: U32,+fill: Element,surveyed: Tensors.Parts<Element> & Survey) -> Tensors.Tensor<Element>:
369  (parts,summary) = surveyed
370  Survey{valid,total,status,workers} = summary
371  join_status(~Element,status,valid,first,with_extent(first,U32.to_nat(axis),total),axis,fill,workers,parts)
372
373def lead(~Element: Data,parts: Tensors.Parts<Element>) -> Tensors.Parts<Element> & (List<&2,U32> & Element):
374  match parts:
375    case Tensors.Single{tensor}:
376      Tensors.with_metadata(~Element,~(Tensors.Parts<Element> & (List<&2,U32> & Element)),tensor,tensor => shape => fill => status => workers =>
377        +shape = shape
378        +fill = fill
379        (Tensors.Single{tensor},(shape,fill)))
380    case Tensors.Cons{tensor,rest}:
381      Tensors.with_metadata(~Element,~(Tensors.Parts<Element> & (List<&2,U32> & Element)),tensor,tensor => shape => fill => status => workers =>
382        +shape = shape
383        +fill = fill
384        (Tensors.Cons{tensor,rest},(shape,fill)))
385
386def concat_first(~Element: Data,+axis: U32,parts: Tensors.Parts<Element>,+shape: List<&2,U32>,+fill: Element) -> Tensors.Tensor<Element>:
387  joined(~Element,shape,axis,fill,
388    survey(~Element,parts,shape,U32.to_nat(axis),Survey{Nat.is_lt(U32.to_nat(axis),List.length(&2,U32,shape)),0,Tensors.Usable{},1}))
389
390def concat_led(~Element: Data,+axis: U32,led: Tensors.Parts<Element> & (List<&2,U32> & Element)) -> Tensors.Tensor<Element>:
391  (parts,first) = led
392  (shape,fill) = first
393  concat_first(~Element,axis,parts,shape,fill)
394
395# Parts agree on every extent except axis; the result stacks them in order.
396def concat(~Element: Data,parts: Tensors.Parts<Element>,+axis: U32) -> Tensors.Tensor<Element>:
397  concat_led(~Element,axis,lead(~Element,parts))
398
399# Broadcasting to given extents: one checked copy from a stretched view.
400def broadcast_status(~Element: Data,status: Tensors.Status,view: Maybe<&2,Views.View>,storage: Storage.Buffer<Element>,+shape: List<&2,U32>,
401  +target: List<&2,U32>,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
402  match status view:
403    case Tensors.Usable{} None{}: Tensors.failed(~Element,"broadcast_to",Tensors.ShapeMismatch{shape,target},target,fill,workers)
404    case status view: result_of(~Element,gather_status(~Element,"broadcast_to",status,view,target,target,storage,shape,fill,workers))
405
406# The tensor repeated to the extents target under NumPy broadcasting.
407def broadcast_to(~Element: Data,tensor: Tensors.Tensor<Element>,+target: List<&2,U32>) -> Tensors.Tensor<Element>:
408  Tensors.with_dense(~Element,~Tensors.Tensor<Element>,tensor,storage => shape => fill => status => workers =>
409    +shape = shape
410    broadcast_status(~Element,status,Tensors.stretched(Tensors.view_of(shape),target),storage,shape,target,fill,workers))
411
412# Reduction along one axis. The output starts at the initial value and takes
413# the slices of the axis in increasing index, so each output element is
414# operation(... operation(operation(initial, x[0]), x[1]) ..., x[extent - 1]).
415def fold_slice(~Element: Data,~operation: Element -> Element -> Element,valid: Bool,slice: Maybe<&2,Views.View>,target: Maybe<&2,Views.View>,
416  output: Storage.Buffer<Element>,input: Storage.Buffer<Element>) -> (Storage.Buffer<Element> & Storage.Buffer<Element>) & Bool:
417  match slice target:
418    case Some{slice} Some{target}: settle(~Element,valid,Operations.combine(~Element,~operation,output,target,input,slice))
419    case slice target: ((output,input),False{})
420
421# The views are built from the extents at every step: a loop that carried one
422# View value through its iterations would make the compiler share every View
423# constructor of the program.
424def fold_step(~Element: Data,~operation: Element -> Element -> Element,index: U32,axis: U32,shape: List<&2,U32>,unit: List<&2,U32>,
425  state: (Storage.Buffer<Element> & Storage.Buffer<Element>) & Bool) -> (Storage.Buffer<Element> & Storage.Buffer<Element>) & Bool:
426  (owners,valid) = state
427  (output,input) = owners
428  fold_slice(~Element,~operation,valid,sliced(Tensors.view_of(shape),axis,index,1,1),Tensors.view_of(unit),output,input)
429
430def fold_slices(~Element: Data,~operation: Element -> Element -> Element,remaining: Nat,+index: U32,+axis: U32,+shape: List<&2,U32>,+unit: List<&2,U32>,
431  state: (Storage.Buffer<Element> & Storage.Buffer<Element>) & Bool) -> (Storage.Buffer<Element> & Storage.Buffer<Element>) & Bool:
432  match remaining:
433    case 0n: state
434    case 1n+rest:
435      fold_slices(~Element,~operation,rest,(index + 1 : U32),axis,shape,unit,fold_step(~Element,~operation,index,axis,shape,unit,state))
436
437def folded_valid(~Element: Data,name: String,valid: Bool,output: Storage.Buffer<Element>,shape: List<&2,U32>,result_shape: List<&2,U32>,fill: Element,
438  workers: U32) -> Tensors.Tensor<Element>:
439  match valid:
440    case True{}: Tensors.Tensor{output,result_shape,fill,Tensors.Usable{},workers}
441    case False{}: Tensors.Tensor{output,result_shape,fill,Tensors.Failed{Tensors.Problem{name,Tensors.InvalidArgument{shape}}},workers}
442
443def folded(~Element: Data,name: String,shape: List<&2,U32>,result_shape: List<&2,U32>,fill: Element,workers: U32,
444  result: (Storage.Buffer<Element> & Storage.Buffer<Element>) & Bool) -> Tensors.Tensor<Element>:
445  (owners,valid) = result
446  (output,input) = owners
447  folded_valid(~Element,name,valid,output,shape,result_shape,fill,workers)
448
449# unit is the shape with extent 1 at the axis: the layout of the output and
450# the extents of every slice.
451def reduce_axis_count(~Element: Data,~operation: Element -> Element -> Element,name: String,count: Maybe<&2,U32>,+axis: U32,input: Storage.Buffer<Element>,
452  +shape: List<&2,U32>,unit: List<&2,U32>,+initial: Element,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
453  match count:
454    case Some{count}:
455      folded(~Element,name,shape,Shapes.without_axis(shape,U32.to_nat(axis)),fill,workers,
456        fold_slices(~Element,~operation,U32.to_nat(extent_at(shape,U32.to_nat(axis))),0,axis,shape,unit,((Tensors.buffer(~Element,initial,count),input),True{})))
457    case None{}: Tensors.failed(~Element,name,Tensors.InvalidArgument{shape},Shapes.without_axis(shape,U32.to_nat(axis)),fill,workers)
458
459def reduce_axis_status(~Element: Data,~operation: Element -> Element -> Element,name: String,status: Tensors.Status,valid: Bool,+axis: U32,
460  input: Storage.Buffer<Element>,+shape: List<&2,U32>,+initial: Element,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
461  match status valid:
462    case Tensors.Usable{} True{}:
463      reduce_axis_count(~Element,~operation,name,Tensors.count_of(with_extent(shape,U32.to_nat(axis),1)),axis,input,shape,with_extent(shape,U32.to_nat(axis),1),
464        initial,fill,workers)
465    case Tensors.Usable{} False{}: Tensors.failed(~Element,name,Tensors.InvalidArgument{shape},shape,fill,workers)
466    case Tensors.Failed{problem} valid: Tensors.Tensor{input,shape,fill,Tensors.Failed{problem},workers}
467
468# The result has the extents of the tensor without axis.
469def reduce_axis(~Element: Data,~operation: Element -> Element -> Element,name: String,+initial: Element,tensor: Tensors.Tensor<Element>,+axis: U32) ->
470  Tensors.Tensor<Element>:
471  Tensors.with_dense(~Element,~Tensors.Tensor<Element>,tensor,storage => shape => fill => status => workers =>
472    +shape = shape
473    reduce_axis_status(~Element,~operation,name,status,Nat.is_lt(U32.to_nat(axis),List.length(&2,U32,shape)),axis,storage,shape,initial,fill,workers))
474
475# Padding. Constant fills the new elements with one value and writes the
476# tensor once, whatever the axes; Edge repeats the nearest element of the
477# tensor one axis after another, so a corner holds the nearest corner element.
478type Border<-Element: Data> is Data:
479  Constant{value: Element}
480  Edge{}
481
482def border_name(~Element: Data,border: Border<Element>) -> String:
483  match border:
484    case Constant{value}: "pad"
485    case Edge{}: "pad_edge"
486
487def blank(~Element: Data,count: Maybe<&2,U32>,+shape: List<&2,U32>,+value: Element,+fill: Element,+workers: U32) -> Tensors.Tensor<Element>:
488  match count:
489    case Some{count}: Tensors.Tensor{Tensors.buffer(~Element,value,count),shape,fill,Tensors.Usable{},workers}
490    case None{}: Tensors.failed(~Element,"pad",Tensors.InvalidShape{shape},shape,fill,workers)
491
492# The elements at one end of axis, repeated count times: (strip, tensor).
493def edge_strip(~Element: Data,tensor: Tensors.Tensor<Element>,+axis: U32,from_end: Bool,+count: U32) -> Tensors.Tensor<Element> & Tensors.Tensor<Element>:
494  Tensors.with_dense(~Element,~(Tensors.Tensor<Element> & Tensors.Tensor<Element>),tensor,storage => shape => fill => status => workers =>
495    +shape = shape
496    +strip = with_extent(shape,U32.to_nat(axis),count)
497    gather_status(~Element,"pad_edge",status,
498      Tensors.stretched(sliced(Tensors.view_of(shape),axis,Bool.pick(U32,from_end,(extent_at(shape,U32.to_nat(axis)) - 1 : U32),0),1,1),strip),
499      strip,strip,storage,shape,fill,workers))
500
501def pad_edge_axis(~Element: Data,unchanged: Bool,tensor: Tensors.Tensor<Element>,+axis: U32,+before: U32,+after: U32) -> Tensors.Tensor<Element>:
502  match unchanged:
503    case True{}: tensor
504    case False{}:
505      Tensors.then(~Tensors.Tensor<Element>,~Tensors.Tensor<Element>,~Tensors.Tensor<Element>,edge_strip(~Element,tensor,axis,False{},before),first => tensor =>
506        Tensors.then(~Tensors.Tensor<Element>,~Tensors.Tensor<Element>,~Tensors.Tensor<Element>,edge_strip(~Element,tensor,axis,True{},after),last => tensor =>
507          concat(~Element,Tensors.Cons{first,Tensors.Cons{tensor,Tensors.Single{last}}},axis)))
508
509def pad_edge_axes(~Element: Data,before: List<&2,U32>,after: List<&2,U32>,tensor: Tensors.Tensor<Element>,+axis: U32) -> Tensors.Tensor<Element>:
510  match before after:
511    case Con{+front,before} Con{+back,after}:
512      pad_edge_axes(~Element,before,after,pad_edge_axis(~Element,U32.is_zero(front) && U32.is_zero(back),tensor,axis,front,back),(axis + 1 : U32))
513    case before after: tensor
514
515# The view of the padded extents that holds the tensor: on every axis, the
516# extent of the tensor starting after the elements added in front.
517def interior(before: List<&2,U32>,shape: List<&2,U32>,view: Maybe<&2,Views.View>,+axis: U32) -> Maybe<&2,Views.View>:
518  match before shape:
519    case Con{front,before} Con{extent,shape}: interior(before,shape,sliced(view,axis,front,extent,1),(axis + 1 : U32))
520    case before shape: view
521
522def all_zero(values: List<&2,U32>) -> Bool:
523  match values:
524    case []: True{}
525    case value <> rest: U32.is_zero(value) && all_zero(rest)
526
527# Constant padding: storage of the padded extents filled with the value, then
528# the tensor written through the interior view. One copy whatever the axes.
529def pad_constant(~Element: Data,unchanged: Bool,+value: Element,tensor: Tensors.Tensor<Element>,+before: List<&2,U32>,+padded: List<&2,U32>) -> Tensors.Tensor<Element>:
530  match unchanged:
531    case True{}: tensor
532    case False{}:
533      Tensors.with_dense(~Element,~Tensors.Tensor<Element>,tensor,part => part_shape => part_fill => part_status => part_workers =>
534        +part_shape = part_shape
535        +part_fill = part_fill
536        +part_workers = part_workers
537        Tensors.with_dense(~Element,~Tensors.Tensor<Element>,blank(~Element,Tensors.count_of(padded),padded,value,part_fill,part_workers),
538          target => target_shape => target_fill => target_status => target_workers =>
539            result_of(~Element,place_status(~Element,"pad",target_status,part_status,interior(before,part_shape,Tensors.view_of(padded),0),
540              target,padded,target_fill,part,part_shape,part_fill,part_workers))))
541
542def pad_border(~Element: Data,border: Border<Element>,tensor: Tensors.Tensor<Element>,+before: List<&2,U32>,+after: List<&2,U32>,padded: Maybe<&2,List<&2,U32>>) -> Tensors.Tensor<Element>:
543  match border padded:
544    case Constant{value} Some{padded}: pad_constant(~Element,all_zero(before) && all_zero(after),value,tensor,before,padded)
545    case border padded: pad_edge_axes(~Element,before,after,tensor,0)
546
547def padded_count(shape: Maybe<&2,List<&2,U32>>) -> Maybe<&2,U32>:
548  match shape:
549    case Some{shape}: Tensors.count_of(shape)
550    case None{}: None{}
551
552def pad_valid(~Element: Data,+border: Border<Element>,status: Tensors.Status,valid: Bool,tensor: Tensors.Tensor<Element>,shape: List<&2,U32>,
553  before: List<&2,U32>,after: List<&2,U32>,padded: Maybe<&2,List<&2,U32>>) -> Tensors.Tensor<Element>:
554  match status valid:
555    case Tensors.Usable{} True{}: pad_border(~Element,border,tensor,before,after,padded)
556    case Tensors.Usable{} False{}: Tensors.with_status(~Element,Tensors.Failed{Tensors.Problem{border_name(~Element,border),Tensors.InvalidArgument{shape}}},tensor)
557    case Tensors.Failed{problem} valid: tensor
558
559# before[i] elements are added in front of axis i and after[i] behind it. Both
560# lists have one entry per axis and the padded extents must be allocatable.
561def pad(~Element: Data,+border: Border<Element>,tensor: Tensors.Tensor<Element>,+before: List<&2,U32>,+after: List<&2,U32>) -> Tensors.Tensor<Element>:
562  Tensors.with_metadata(~Element,~Tensors.Tensor<Element>,tensor,tensor => shape => fill => status => workers =>
563    +shape = shape
564    +padded = Shapes.padded(shape,before,after)
565    pad_valid(~Element,border,status,Views.has_count(padded_count(padded)),tensor,shape,before,after,padded))