tensor_band.bend

stelliferous@0.0.2.0 · 459 lines · 77 definitions

1# Internal H-axis ownership. Offsets and spans are logical image rows, not
2# element positions: offset / height identifies the image. No leaf crosses an
3# image boundary. Its buffer stores [channels][span][width]. The common
4# partition tree owns all buffers; shape and element type remain parameters.
5import Base
6import stelliferous@0.0.2.0/storage_buffer.bend as Storage
7import stelliferous@0.0.2.0/storage_copy.bend as Copy
8import stelliferous@0.0.2.0/storage_clone.bend as Clone
9import stelliferous@0.0.2.0/traversal_partition.bend as Partition
10import stelliferous@0.0.2.0/traversal_array.bend as Traversal
11import stelliferous@0.0.2.0/proofs/traversal_certificate.bend as TraversalCertificate
12import stelliferous@0.0.2.0/tensor_operations.bend as Operations
13import stelliferous@0.0.2.0/storage_allocation.bend as Allocation
14import stelliferous@0.0.2.0/tensor_view.bend as Views
15import stelliferous@0.0.2.0/tensor_indexed.bend as IndexedAccess
16import stelliferous@0.0.2.0/proofs/tensor_indexed_certificate.bend as IndexedCertificate
17import stelliferous@0.0.2.0/proofs/storage_certificate.bend as Certificate
18
19type Geometry is Data:
20  Geometry{batch: U32,channels: U32,height: U32,width: U32}
21
22def Bands(~Element: Data) -> Type:
23  Partition.Tree<&1,Storage.Buffer<Element>,Unit>
24
25def elements(geometry: Geometry) -> U32:
26  Geometry{batch,channels,height,width} = geometry
27  (batch * channels * height * width : U32)
28
29def local_count(geometry: Geometry,rows: U32) -> U32:
30  Geometry{batch,channels,height,width} = geometry
31  (channels * rows * width : U32)
32
33# Channel-major local indices map to the same global NCHW index for every
34# partition. A leaf's first row is global across images; rows are image-local.
35def global_index(geometry: Geometry,+first: U32,+rows: U32,+index: U32) -> U32:
36  Geometry{batch,+channels,+height,+width} = geometry
37  +plane = (rows * width : U32)
38  (((first / height * channels + index / plane) * height + first % height) * width + index % plane : U32)
39
40def planes(~Element: Data,remaining: Nat,+channel: U32,+count: U32,+source_stride: U32,+destination_stride: U32,
41  +source_start: U32,+destination_start: U32,owners: Storage.Buffer<Element> & Storage.Buffer<Element>) ->
42  Storage.Buffer<Element> & Storage.Buffer<Element>:
43  match remaining:
44    case 0n: owners
45    case 1n+rest:
46      (source,destination) = owners
47      planes(~Element,rest,(channel + 1 : U32),count,source_stride,destination_stride,source_start,destination_start,
48        Copy.range(~Element,count,(source_start + channel * source_stride : U32),
49          (destination_start + channel * destination_stride : U32),source,destination))
50
51def destination(~Element: Data,owners: Storage.Buffer<Element> & Storage.Buffer<Element>) -> Storage.Buffer<Element>:
52  (source,destination) = owners
53  destination
54
55def gather_leaf(~Element: Data,geometry: Geometry,+first: U32,+rows: U32,input: Storage.Buffer<Element>,output: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
56  Geometry{batch,+channels,+height,+width} = geometry
57  destination(~Element,planes(~Element,U32.to_nat(channels),0,(rows * width : U32),(rows * width : U32),(height * width : U32),0,
58    ((first / height * channels * height + first % height) * width : U32),(input,output)))
59
60# One allocation is provided by the caller. Each value is copied once, using
61# one contiguous range per channel per band. No intermediate dense images.
62def gather(~Element: Data,+geometry: Geometry,bands: Bands(~Element),output: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
63  match bands:
64    case Partition.Leaf{input,first,rows,details}: gather_leaf(~Element,geometry,first,rows,input,output)
65    case Partition.Fork{left,right}: gather(~Element,geometry,right,gather(~Element,geometry,left,output))
66
67def copied_leaf(~Element: Data,+first: U32,+rows: U32,owners: Storage.Buffer<Element> & Storage.Buffer<Element>) -> Bands(~Element) & Bands(~Element):
68  (left,right) = owners
69  (Partition.Leaf{left,first,rows,Unit{}},Partition.Leaf{right,first,rows,Unit{}})
70
71def copied_fork(~Element: Data,left: Bands(~Element) & Bands(~Element),right: Bands(~Element) & Bands(~Element)) -> Bands(~Element) & Bands(~Element):
72  (left_original,left_copy) = left
73  (right_original,right_copy) = right
74  (Partition.Fork{left_original,right_original},Partition.Fork{left_copy,right_copy})
75
76def copy(~Element: Data,bands: Bands(~Element)) -> Bands(~Element) & Bands(~Element):
77  match bands:
78    case Partition.Leaf{input,+first,+rows,details}: copied_leaf(~Element,first,rows,Clone.clone(~Element,input))
79    case Partition.Fork{left,right}: copied_fork(~Element,copy(~Element,left),copy(~Element,right))
80
81# A pure ownership plan contains no Buffer. It is a concrete Data type:
82# duplicating Tree<&2,...> through a computed type makes all constructors
83# shared in the official compiler, including the affine Buffer tree.
84type Plan is Data:
85  Rows{first: U32,count: U32}
86  Split{left: Plan,right: Plan}
87
88def plan_join(~Element: Data,left: Bands(~Element) & Plan,right: Bands(~Element) & Plan) -> Bands(~Element) & Plan:
89  (left,left_plan) = left
90  (right,right_plan) = right
91  (Partition.Fork{left,right},Split{left_plan,right_plan})
92
93def plan(~Element: Data,bands: Bands(~Element)) -> Bands(~Element) & Plan:
94  match bands:
95    case Partition.Leaf{input,+first,+rows,details}:
96      (Partition.Leaf{input,first,rows,Unit{}},Rows{first,rows})
97    case Partition.Fork{left,right}: plan_join(~Element,plan(~Element,left),plan(~Element,right))
98
99def same_plan(left: Plan,right: Plan) -> Bool:
100  match left right:
101    case Rows{first,rows} Rows{other_first,other_rows}:
102      U32.is_eq(first,other_first) && U32.is_eq(rows,other_rows)
103    case Split{left,right} Split{other_left,other_right}:
104      same_plan(left,other_left) && same_plan(right,other_right)
105    case left right: False{}
106
107def first(layout: Plan) -> U32:
108  match layout:
109    case Rows{first,count}: first
110    case Split{left,right}: first(left)
111
112def last(layout: Plan) -> U32:
113  match layout:
114    case Rows{first,count}: (first + count : U32)
115    case Split{left,right}: last(right)
116
117def same_image(geometry: Geometry,+layout: Plan) -> Bool:
118  Geometry{batch,channels,+height,width} = geometry
119  U32.is_eq((first(layout) / height : U32),((last(layout) - 1) / height : U32))
120
121# Two trees of one shape, joined leaf by leaf. A left leaf without a partner
122# keeps its own row range, so the left tree can be returned unchanged.
123type Paired<-Left: Type,-Right: Type> is Type:
124  Both{left: Left,right: Right}
125  Only{left: Left,first: U32,rows: U32}
126
127def only(~Left: Type,~Right: Type,tree: Partition.Tree<&1,Left,Unit>) -> Partition.Tree<&1,Paired<Left,Right>,Unit>:
128  match tree:
129    case Partition.Leaf{left,+first,+rows,details}: Partition.Leaf{Only{left,first,rows},first,rows,Unit{}}
130    case Partition.Fork{left,right}: Partition.Fork{only(~Left,~Right,left),only(~Left,~Right,right)}
131
132def zip(~Left: Type,~Right: Type,left: Partition.Tree<&1,Left,Unit>,right: Partition.Tree<&1,Right,Unit>) -> Partition.Tree<&1,Paired<Left,Right>,Unit>:
133  match left right:
134    case Partition.Leaf{left,first,rows,details} Partition.Leaf{right,other_first,other_rows,other_details}:
135      Partition.Leaf{Both{left,right},first,rows,Unit{}}
136    case Partition.Fork{left,right} Partition.Fork{other_left,other_right}:
137      Partition.Fork{zip(~Left,~Right,left,other_left),zip(~Left,~Right,right,other_right)}
138    case left right: only(~Left,~Right,left)
139
140def left_leaf(~Left: Type,~Right: Type,pair: Paired<Left,Right>,first: U32,rows: U32) -> Partition.Tree<&1,Left,Unit>:
141  match pair:
142    case Both{left,right}: Partition.Leaf{left,first,rows,Unit{}}
143    case Only{left,own_first,own_rows}: Partition.Leaf{left,own_first,own_rows,Unit{}}
144
145def left_of(~Left: Type,~Right: Type,tree: Partition.Tree<&1,Paired<Left,Right>,Unit>) -> Partition.Tree<&1,Left,Unit>:
146  match tree:
147    case Partition.Leaf{pair,first,rows,details}: left_leaf(~Left,~Right,pair,first,rows)
148    case Partition.Fork{left,right}: Partition.Fork{left_of(~Left,~Right,left),left_of(~Left,~Right,right)}
149
150def paired(~Left: Type,~Right: Type,~Resource: Type,~Context: Data,
151  ~leaf: Context -> U32 -> U32 -> Left -> Right -> Resource -> Partition.Tree<&1,Left,Unit>,
152  context: Context,first: U32,rows: U32,pair: Paired<Left,Right>,resource: Resource) -> Partition.Tree<&1,Left,Unit>:
153  match pair:
154    case Both{left,right}: leaf(context,first,rows,left,right,resource)
155    case Only{left,own_first,own_rows}: Partition.Leaf{left,own_first,own_rows,Unit{}}
156
157def unshared(resource: Unit) -> Unit & Unit:
158  (Unit{},Unit{})
159
160def same(~Block: Type,tree: Partition.Tree<&1,Block,Unit>) -> Partition.Tree<&1,Block,Unit>:
161  tree
162
163# One leaf below a parallel split. It neither forks nor calls a closure, so
164# the runtime keeps it as its own task for the work turn; a leaf inlined into
165# the forking recursion would run in the grow turn on one worker per row.
166def leaf_task(~Block: Type,~Output: Type,~Resource: Type,~Context: Data,
167  ~leaf: Context -> U32 -> U32 -> Block -> Resource -> Partition.Tree<&1,Output,Unit>,
168  ~keep: Partition.Tree<&1,Block,Unit> -> Partition.Tree<&1,Output,Unit>,
169  first: U32,rows: U32,context: Context,tree: Partition.Tree<&1,Block,Unit>,resource: Resource) -> Partition.Tree<&1,Output,Unit>:
170  match tree:
171    case Partition.Leaf{block,offset,span,details}: leaf(context,first,rows,block,resource)
172    case Partition.Fork{left,right}: keep(Partition.Fork{left,right})
173
174# The one band recursion. The plan gives every leaf its rows and decides the
175# structure; a tree of another shape is kept. A split inside one image may
176# run its halves in parallel, and a split across images runs them in order.
177# share gives each half its own owner of a read-only resource.
178def each(~Block: Type,~Output: Type,~Resource: Type,~Context: Data,
179  ~share: Resource -> Resource & Resource,
180  ~leaf: Context -> U32 -> U32 -> Block -> Resource -> Partition.Tree<&1,Output,Unit>,
181  ~keep: Partition.Tree<&1,Block,Unit> -> Partition.Tree<&1,Output,Unit>,
182  layout: Plan,mode: Bool,+enabled: Bool,+geometry: Geometry,+context: Context,
183  tree: Partition.Tree<&1,Block,Unit>,resource: Resource) -> Partition.Tree<&1,Output,Unit>:
184  match layout mode tree:
185    case Rows{first,rows} mode tree: leaf_task(~Block,~Output,~Resource,~Context,~leaf,~keep,first,rows,context,tree,resource)
186    case Split{+left_plan,+right_plan} False{} Partition.Fork{left,right}:
187      Storage.bind_pair(~Resource,~Resource,~Partition.Tree<&1,Output,Unit>,share(resource),left_resource => right_resource =>
188        Partition.Fork{each(~Block,~Output,~Resource,~Context,~share,~leaf,~keep,left_plan,enabled && same_image(geometry,left_plan),enabled,geometry,context,left,left_resource),
189          each(~Block,~Output,~Resource,~Context,~share,~leaf,~keep,right_plan,enabled && same_image(geometry,right_plan),enabled,geometry,context,right,right_resource)})
190    case Split{Rows{left_first,left_rows},Rows{right_first,right_rows}} True{} Partition.Fork{left,right}:
191      Storage.bind_pair(~Resource,~Resource,~Partition.Tree<&1,Output,Unit>,share(resource),left_resource => right_resource =>
192        left right = leaf_task(~Block,~Output,~Resource,~Context,~leaf,~keep,left_first,left_rows,context,left,left_resource) leaf_task(~Block,~Output,~Resource,~Context,~leaf,~keep,right_first,right_rows,context,right,right_resource)
193        Partition.Fork{left,right})
194    case Split{Rows{left_first,left_rows},+right_plan} True{} Partition.Fork{left,right}:
195      Storage.bind_pair(~Resource,~Resource,~Partition.Tree<&1,Output,Unit>,share(resource),left_resource => right_resource =>
196        left right = leaf_task(~Block,~Output,~Resource,~Context,~leaf,~keep,left_first,left_rows,context,left,left_resource) each(~Block,~Output,~Resource,~Context,~share,~leaf,~keep,right_plan,enabled && same_image(geometry,right_plan),enabled,geometry,context,right,right_resource)
197        Partition.Fork{left,right})
198    case Split{+left_plan,Rows{right_first,right_rows}} True{} Partition.Fork{left,right}:
199      Storage.bind_pair(~Resource,~Resource,~Partition.Tree<&1,Output,Unit>,share(resource),left_resource => right_resource =>
200        left right = each(~Block,~Output,~Resource,~Context,~share,~leaf,~keep,left_plan,enabled && same_image(geometry,left_plan),enabled,geometry,context,left,left_resource) leaf_task(~Block,~Output,~Resource,~Context,~leaf,~keep,right_first,right_rows,context,right,right_resource)
201        Partition.Fork{left,right})
202    case Split{+left_plan,+right_plan} True{} Partition.Fork{left,right}:
203      Storage.bind_pair(~Resource,~Resource,~Partition.Tree<&1,Output,Unit>,share(resource),left_resource => right_resource =>
204        left right = each(~Block,~Output,~Resource,~Context,~share,~leaf,~keep,left_plan,enabled && same_image(geometry,left_plan),enabled,geometry,context,left,left_resource) each(~Block,~Output,~Resource,~Context,~share,~leaf,~keep,right_plan,enabled && same_image(geometry,right_plan),enabled,geometry,context,right,right_resource)
205        Partition.Fork{left,right})
206    case layout mode tree: keep(tree)
207
208type Indexed is Data:
209  Indexed{geometry: Geometry,first: U32,rows: U32}
210
211def map_leaf(~Element: Data,~Context: Data,~operation: Context -> Element -> Element,+context: Context,buffer: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
212  Storage.Buffer{storage,+length,certificate} = buffer
213  Storage.Buffer{Traversal.map_in_place(~Element,~Context,~operation,length,0,context,storage),length,
214    TraversalCertificate.map_in_place(~Element,~Context,~operation,length,0,context,storage,certificate)}
215
216# A mapper over the whole of one band.
217def map_whole_context(~Element: Data,~Context: Data,~mapper: @+count: Nat -> @+position: U32 -> @+context: Context -> Storage.Buffer<Element> -> Storage.Buffer<Element>,+context: Context,buffer: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
218  Storage.Buffer{storage,+length,certificate} = buffer
219  mapper(length,0,context,Storage.Buffer{storage,length,certificate})
220
221def map_band(~Element: Data,~Context: Data,~leaf: @+context: Context -> Storage.Buffer<Element> -> Storage.Buffer<Element>,context: Context,first: U32,rows: U32,
222  buffer: Storage.Buffer<Element>,resource: Unit) -> Bands(~Element):
223  Partition.Leaf{leaf(context,buffer),first,rows,Unit{}}
224
225# Parallelism is a planner decision. Sequential execution keeps the bands and
226# their owners: a small operation need not pay either a fork or a gather.
227def map_with(~Element: Data,~Context: Data,~leaf: @+context: Context -> Storage.Buffer<Element> -> Storage.Buffer<Element>,+layout: Plan,+parallel: Bool,+geometry: Geometry,
228  context: Context,bands: Bands(~Element)) -> Bands(~Element):
229  each(~Storage.Buffer<Element>,~Storage.Buffer<Element>,~Unit,~Context,~unshared,~map_band(~Element,~Context,~leaf),~same(~Storage.Buffer<Element>),
230    layout,parallel && same_image(geometry,layout),parallel,geometry,context,bands,Unit{})
231
232def scalar_value(~Element: Data,~operation: Element -> Element -> Element,value: Element,element: Element) -> Element:
233  operation(element,value)
234
235def scalar(~Element: Data,~operation: Element -> Element -> Element,layout: Plan,parallel: Bool,geometry: Geometry,value: Element,bands: Bands(~Element)) -> Bands(~Element):
236  map_with(~Element,~Element,~map_leaf(~Element,~Element,~scalar_value(~Element,~operation)),layout,parallel,geometry,value,bands)
237
238def result(~Element: Data,owners: Storage.Buffer<Element> & Bands(~Element)) -> Bands(~Element):
239  (source,bands) = owners
240  bands
241
242def index_of(context: Indexed,index: U32) -> U32:
243  Indexed{geometry,first,rows} = context
244  global_index(geometry,first,rows,index)
245
246# Channel c of a leaf holds the positions from c * rows * width; their global
247# indices are consecutive from this base.
248def channel_base(context: Indexed,channel: U32) -> U32:
249  Indexed{Geometry{batch,channels,+height,width},+first,rows} = context
250  (((first / height * channels + channel) * height + first % height) * width : U32)
251
252def channels_of(context: Indexed) -> U32:
253  Indexed{Geometry{batch,channels,height,width},first,rows} = context
254  channels
255
256def plane_of(context: Indexed) -> U32:
257  Indexed{Geometry{batch,channels,height,width},first,rows} = context
258  (rows * width : U32)
259
260def indexed_buffers(~Element: Data,~operation: U32 -> Element -> Element,+count: U32,+context: Indexed,input: Storage.Buffer<Element>,output: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
261  Storage.Buffer{input,input_length,input_certificate} = input
262  Storage.Buffer{output,output_length,output_certificate} = output
263  destination(~Element,Operations.pair_buffers(~Element,input_length,output_length,
264    IndexedAccess.run_tiled(~Element,~Indexed,~index_of,~channel_base,~channels_of,~plane_of,~operation,count,context,input,output),
265    IndexedCertificate.tiled_certificates(~Element,~Indexed,~index_of,~channel_base,~channels_of,~plane_of,~operation,count,context,input,output,input_certificate,output_certificate)))
266
267def indexed_leaf(~Element: Data,~operation: U32 -> Element -> Element,
268  +geometry: Geometry,+first: U32,+rows: U32,fill: Element,input: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
269  +count = local_count(geometry,rows)
270  indexed_buffers(~Element,~operation,count,Indexed{geometry,first,rows},input,
271    Allocation.allocate(~Element,Storage.allocation_depth(U32.max(count,1)),fill,U32.to_nat(count)))
272
273type Filled<-Element: Data> is Data:
274  Filled{geometry: Geometry,fill: Element}
275
276def indexed_band(~Element: Data,~operation: U32 -> Element -> Element,context: Filled<Element>,first: U32,rows: U32,
277  input: Storage.Buffer<Element>,resource: Unit) -> Bands(~Element):
278  Filled{geometry,fill} = context
279  +first = first
280  +rows = rows
281  Partition.Leaf{indexed_leaf(~Element,~operation,geometry,first,rows,fill,input),first,rows,Unit{}}
282
283def apply_indexed(~Element: Data,~operation: U32 -> Element -> Element,
284  +layout: Plan,+parallel: Bool,+geometry: Geometry,fill: Element,bands: Bands(~Element)) -> Bands(~Element):
285  each(~Storage.Buffer<Element>,~Storage.Buffer<Element>,~Unit,~Filled<Element>,~unshared,~indexed_band(~Element,~operation),~same(~Storage.Buffer<Element>),
286    layout,parallel && same_image(geometry,layout),parallel,geometry,Filled{geometry,fill},bands,Unit{})
287
288def range_result(~Element: Data,+first: U32,+rows: U32,owners: Storage.Buffer<Element> & Storage.Buffer<Element>) -> Bands(~Element) & Storage.Buffer<Element>:
289  (source,output) = owners
290  (Partition.Leaf{source,first,rows,Unit{}},output)
291
292# An empty intersection copies nothing, so it skips the per-channel loop.
293def range_planes(~Element: Data,empty: Bool,+channels: U32,+count: U32,+source_stride: U32,+destination_stride: U32,
294  +source_start: U32,+destination_start: U32,owners: Storage.Buffer<Element> & Storage.Buffer<Element>) ->
295  Storage.Buffer<Element> & Storage.Buffer<Element>:
296  match empty:
297    case True{}: owners
298    case False{}: planes(~Element,U32.to_nat(channels),0,count,source_stride,destination_stride,source_start,destination_start,owners)
299
300# Repartition and halo extraction share this owned range reader. Empty
301# intersections copy zero elements; each source owner is returned unchanged.
302def range_leaf(~Element: Data,geometry: Geometry,+target_first: U32,+target_rows: U32,target_offset: U32,+first: U32,+rows: U32,
303  source: Storage.Buffer<Element>,output: Storage.Buffer<Element>) -> Bands(~Element) & Storage.Buffer<Element>:
304  Geometry{batch,channels,height,+width} = geometry
305  +begin = U32.max(first,target_first)
306  +end = U32.max(begin,U32.min((first + rows : U32),(target_first + target_rows : U32)))
307  range_result(~Element,first,rows,range_planes(~Element,U32.is_eq(begin,end),channels,((end - begin) * width : U32),
308    (rows * width : U32),(target_rows * width : U32),((begin - first) * width : U32),(target_offset + (begin - target_first) * width : U32),(source,output)))
309
310def range_into(~Element: Data,+geometry: Geometry,+first: U32,+rows: U32,+target_offset: U32,bands: Bands(~Element),output: Storage.Buffer<Element>) -> Bands(~Element) & Storage.Buffer<Element>:
311  match bands:
312    case Partition.Leaf{input,source_first,source_rows,details}: range_leaf(~Element,geometry,first,rows,target_offset,source_first,source_rows,input,output)
313    case Partition.Fork{left,right}:
314      Storage.bind_pair(~Bands(~Element),~Storage.Buffer<Element>,~(Bands(~Element) & Storage.Buffer<Element>),
315        range_into(~Element,geometry,first,rows,target_offset,left,output),left => output =>
316          Storage.bind_pair(~Bands(~Element),~Storage.Buffer<Element>,~(Bands(~Element) & Storage.Buffer<Element>),
317            range_into(~Element,geometry,first,rows,target_offset,right,output),right => output => (Partition.Fork{left,right},output)))
318
319def range(~Element: Data,geometry: Geometry,first: U32,rows: U32,bands: Bands(~Element),output: Storage.Buffer<Element>) -> Bands(~Element) & Storage.Buffer<Element>:
320  range_into(~Element,geometry,first,rows,0,bands,output)
321
322def allocate(~Element: Data,fill: Element,+count: U32) -> Storage.Buffer<Element>:
323  Allocation.allocate(~Element,Storage.depth_for(count),fill,U32.to_nat(count))
324
325def repartition(~Element: Data,+geometry: Geometry,+fill: Element,layout: Plan,bands: Bands(~Element)) -> Bands(~Element) & Bands(~Element):
326  match layout:
327    case Rows{+first,+rows}:
328      Storage.bind_pair(~Bands(~Element),~Storage.Buffer<Element>,~(Bands(~Element) & Bands(~Element)),
329        range(~Element,geometry,first,rows,bands,allocate(~Element,fill,local_count(geometry,rows))),source => output =>
330          (source,Partition.Leaf{output,first,rows,Unit{}}))
331    case Split{left,right}:
332      Storage.bind_pair(~Bands(~Element),~Bands(~Element),~(Bands(~Element) & Bands(~Element)),
333        repartition(~Element,geometry,fill,left,bands),source => left =>
334          Storage.bind_pair(~Bands(~Element),~Bands(~Element),~(Bands(~Element) & Bands(~Element)),
335            repartition(~Element,geometry,fill,right,source),source => right => (source,Partition.Fork{left,right})))
336
337def repartitioned(~Element: Data,owners: Bands(~Element) & Bands(~Element)) -> Bands(~Element):
338  (source,result) = owners
339  result
340
341def align_when(~Element: Data,same: Bool,geometry: Geometry,fill: Element,layout: Plan,bands: Bands(~Element)) -> Bands(~Element):
342  match same:
343    case True{}: bands
344    case False{}: repartitioned(~Element,repartition(~Element,geometry,fill,layout,bands))
345
346def align_planned(~Element: Data,geometry: Geometry,fill: Element,+layout: Plan,planned: Bands(~Element) & Plan) -> Bands(~Element):
347  (bands,other) = planned
348  align_when(~Element,same_plan(layout,other),geometry,fill,layout,bands)
349
350def align(~Element: Data,geometry: Geometry,fill: Element,layout: Plan,bands: Bands(~Element)) -> Bands(~Element):
351  align_planned(~Element,geometry,fill,layout,plan(~Element,bands))
352
353def dense_result(~Element: Data,+first: U32,+rows: U32,owners: Storage.Buffer<Element> & Storage.Buffer<Element>) -> Storage.Buffer<Element> & Bands(~Element):
354  (source,band) = owners
355  (source,Partition.Leaf{band,first,rows,Unit{}})
356
357def dense_leaf(~Element: Data,geometry: Geometry,fill: Element,+first: U32,+rows: U32,input: Storage.Buffer<Element>) -> Storage.Buffer<Element> & Bands(~Element):
358  Geometry{batch,+channels,+height,+width} = geometry
359  dense_result(~Element,first,rows,
360    planes(~Element,U32.to_nat(channels),0,(rows * width : U32),(height * width : U32),(rows * width : U32),
361      ((first / height * channels * height + first % height) * width : U32),0,(input,allocate(~Element,fill,(channels * rows * width : U32)))))
362
363def dense_ranges(~Element: Data,+geometry: Geometry,+fill: Element,layout: Plan,input: Storage.Buffer<Element>) -> Storage.Buffer<Element> & Bands(~Element):
364  match layout:
365    case Rows{first,rows}: dense_leaf(~Element,geometry,fill,first,rows,input)
366    case Split{left,right}:
367      Storage.bind_pair(~Storage.Buffer<Element>,~Bands(~Element),~(Storage.Buffer<Element> & Bands(~Element)),
368        dense_ranges(~Element,geometry,fill,left,input),source => left =>
369          Storage.bind_pair(~Storage.Buffer<Element>,~Bands(~Element),~(Storage.Buffer<Element> & Bands(~Element)),
370            dense_ranges(~Element,geometry,fill,right,source),source => right => (source,Partition.Fork{left,right})))
371
372def from_dense(~Element: Data,geometry: Geometry,fill: Element,layout: Plan,input: Storage.Buffer<Element>) -> Bands(~Element):
373  result(~Element,dense_ranges(~Element,geometry,fill,layout,input))
374
375def row_plan(depth: Nat,split: Bool,+first: U32,+rows: U32) -> Plan:
376  match depth split:
377    case 0n split: Rows{first,rows}
378    case 1n+rest False{}: Rows{first,rows}
379    case 1n+ +rest True{}:
380      +half = (rows / 2 : U32)
381      Split{row_plan(rest,U32.is_lt(1,half),first,half),
382        row_plan(rest,U32.is_lt(1,(rows - half : U32)),(first + half : U32),(rows - half : U32))}
383
384def image_plan(remaining: Nat,+image: U32,+depth: Nat,+height: U32) -> Plan:
385  match remaining:
386    case 0n: Rows{0,0}
387    case 1n: row_plan(depth,U32.is_lt(1,height),(image * height : U32),height)
388    case 2n+rest:
389      Split{row_plan(depth,U32.is_lt(1,height),(image * height : U32),height),image_plan(1n+rest,(image + 1 : U32),depth,height)}
390
391# Plan metadata first, then transfer each dense channel range once. First
392# splits and later changed boundaries share the same copy implementation.
393def split(~Element: Data,depth: Nat,+geometry: Geometry,fill: Element,input: Storage.Buffer<Element>) -> Bands(~Element):
394  Geometry{batch,channels,height,width} = geometry
395  from_dense(~Element,geometry,fill,image_plan(U32.to_nat(batch),0,depth,height),input)
396
397def combined(~Element: Data,result: Operations.Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>) -> Storage.Buffer<Element>:
398  match result:
399    case Operations.Accepted{(left,right)}: left
400    case Operations.Rejected{(left,right)}: left
401
402def combine_leaf_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,+count: U32,left: Storage.Buffer<Element>,right: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
403  combined(~Element,Operations.combine_with(~Element,~kernel,left,Views.View{[Views.Axis{count,1}],0},right,Views.View{[Views.Axis{count,1}],0}))
404
405def combine_band_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,geometry: Geometry,first: U32,rows: U32,
406  left: Storage.Buffer<Element>,right: Storage.Buffer<Element>,resource: Unit) -> Bands(~Element):
407  +rows = rows
408  Partition.Leaf{combine_leaf_with(~Element,~kernel,local_count(geometry,rows),left,right),first,rows,Unit{}}
409
410def combine_planned_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,+geometry: Geometry,+parallel: Bool,planned: Bands(~Element) & Plan,right: Bands(~Element)) -> Bands(~Element):
411  (left,+layout) = planned
412  each(~Paired<Storage.Buffer<Element>,Storage.Buffer<Element>>,~Storage.Buffer<Element>,~Unit,~Geometry,~unshared,
413    ~paired(~Storage.Buffer<Element>,~Storage.Buffer<Element>,~Unit,~Geometry,~combine_band_with(~Element,~kernel)),
414    ~left_of(~Storage.Buffer<Element>,~Storage.Buffer<Element>),
415    layout,parallel && same_image(geometry,layout),parallel,geometry,geometry,zip(~Storage.Buffer<Element>,~Storage.Buffer<Element>,left,right),Unit{})
416
417# The caller supplies the same plan on both sides, after shape and status
418# checks; the plan is read from the left. Trees of another shape keep the left owner.
419def combine_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,geometry: Geometry,parallel: Bool,left: Bands(~Element),right: Bands(~Element)) -> Bands(~Element):
420  combine_planned_with(~Element,~kernel,geometry,parallel,plan(~Element,left),right)
421
422def transform_band(~Element: Data,~Context: Data,~leaf: Context -> U32 -> U32 -> Storage.Buffer<Element> -> Bands(~Element),
423  context: Context,first: U32,rows: U32,input: Storage.Buffer<Element>,resource: Unit) -> Bands(~Element):
424  leaf(context,first,rows,input)
425
426def transform(~Element: Data,~Context: Data,~leaf: Context -> U32 -> U32 -> Storage.Buffer<Element> -> Bands(~Element),
427  +layout: Plan,+parallel: Bool,+geometry: Geometry,context: Context,bands: Bands(~Element)) -> Bands(~Element):
428  each(~Storage.Buffer<Element>,~Storage.Buffer<Element>,~Unit,~Context,~unshared,~transform_band(~Element,~Context,~leaf),~same(~Storage.Buffer<Element>),
429    layout,parallel && same_image(geometry,layout),parallel,geometry,context,bands,Unit{})
430
431type Inspection is Data:
432  Inspection{valid: Bool,first: U32,end: U32,leaves: U32}
433
434def join_inspection(left: Inspection,right: Inspection) -> Inspection:
435  Inspection{lv,lf,+le,ln} = left
436  Inspection{rv,+rf,re,rn} = right
437  Inspection{lv && rv && U32.is_eq(le,rf),lf,re,(ln + rn : U32)}
438
439def inspected_pair(~Element: Data,left: Bands(~Element) & Inspection,right: Bands(~Element) & Inspection) -> Bands(~Element) & Inspection:
440  (left,left_inspection) = left
441  (right,right_inspection) = right
442  (Partition.Fork{left,right},join_inspection(left_inspection,right_inspection))
443
444def leaf_inspection(geometry: Geometry,+length: Nat,depth: Nat,+first: U32,+rows: U32) -> Inspection:
445  Geometry{batch,+channels,+height,+width} = geometry
446  Inspection{U32.is_le(rows,(height - first % height : U32)) &&
447      Nat.is_le(Nat.add(U32.to_nat(first),U32.to_nat(rows)),Nat.mul(U32.to_nat(batch),U32.to_nat(height))) &&
448      Nat.is_eq(length,Nat.mul(U32.to_nat(channels),Nat.mul(U32.to_nat(rows),U32.to_nat(width)))) &&
449      Nat.is_le(length,U32.to_nat(U32.shln(1,depth))),first,(first + rows : U32),1}
450
451def inspect_leaf(~Element: Data,geometry: Geometry,buffer: Storage.Buffer<Element>,+first: U32,+rows: U32) -> Bands(~Element) & Inspection:
452  Storage.Buffer{storage,+length,Certificate.Certificate{+depth,equation}} = buffer
453  (Partition.Leaf{Storage.Buffer{storage,length,Certificate.Certificate{depth,equation}},first,rows,Unit{}},
454    leaf_inspection(geometry,length,depth,first,rows))
455
456def inspect(~Element: Data,+geometry: Geometry,bands: Bands(~Element)) -> Bands(~Element) & Inspection:
457  match bands:
458    case Partition.Leaf{buffer,first,rows,details}: inspect_leaf(~Element,geometry,buffer,first,rows)
459    case Partition.Fork{left,right}: inspected_pair(~Element,inspect(~Element,geometry,left),inspect(~Element,geometry,right))