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.
5	import Base
6	import ./storage_buffer.bend as Storage
7	import ./storage_copy.bend as Copy
8	import ./storage_clone.bend as Clone
9	import ./traversal_partition.bend as Partition
10	import ./traversal_array.bend as Traversal
11	import ./proofs/traversal_certificate.bend as TraversalCertificate
12	import ./tensor_operations.bend as Operations
13	import ./storage_allocation.bend as Allocation
14	import ./tensor_view.bend as Views
15	import ./tensor_indexed.bend as IndexedAccess
16	import ./proofs/tensor_indexed_certificate.bend as IndexedCertificate
17	import ./proofs/storage_certificate.bend as Certificate
18	
19	type Geometry is Data:
20	  Geometry{batch: U32,channels: U32,height: U32,width: U32}
21	
22	def Bands(~Element: Data) -> Type:
23	  Partition.Tree<&1,Storage.Buffer<Element>,Unit>
24	
25	def elements(geometry: Geometry) -> U32:
26	  Geometry{batch,channels,height,width} = geometry
27	  (batch * channels * height * width : U32)
28	
29	def 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.
35	def 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	
40	def 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	
51	def destination(~Element: Data,owners: Storage.Buffer<Element> & Storage.Buffer<Element>) -> Storage.Buffer<Element>:
52	  (source,destination) = owners
53	  destination
54	
55	def 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.
62	def 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	
67	def 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	
71	def 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	
76	def 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.
84	type Plan is Data:
85	  Rows{first: U32,count: U32}
86	  Split{left: Plan,right: Plan}
87	
88	def 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	
93	def 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	
99	def 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	
107	def first(layout: Plan) -> U32:
108	  match layout:
109	    case Rows{first,count}: first
110	    case Split{left,right}: first(left)
111	
112	def last(layout: Plan) -> U32:
113	  match layout:
114	    case Rows{first,count}: (first + count : U32)
115	    case Split{left,right}: last(right)
116	
117	def 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.
123	type Paired<-Left: Type,-Right: Type> is Type:
124	  Both{left: Left,right: Right}
125	  Only{left: Left,first: U32,rows: U32}
126	
127	def 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	
132	def 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	
140	def 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	
145	def 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	
150	def 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	
157	def unshared(resource: Unit) -> Unit & Unit:
158	  (Unit{},Unit{})
159	
160	def 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.
166	def 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.
178	def 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	
208	type Indexed is Data:
209	  Indexed{geometry: Geometry,first: U32,rows: U32}
210	
211	def 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.
217	def 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	
221	def 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.
227	def 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	
232	def scalar_value(~Element: Data,~operation: Element -> Element -> Element,value: Element,element: Element) -> Element:
233	  operation(element,value)
234	
235	def 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	
238	def result(~Element: Data,owners: Storage.Buffer<Element> & Bands(~Element)) -> Bands(~Element):
239	  (source,bands) = owners
240	  bands
241	
242	def 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.
248	def 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	
252	def channels_of(context: Indexed) -> U32:
253	  Indexed{Geometry{batch,channels,height,width},first,rows} = context
254	  channels
255	
256	def plane_of(context: Indexed) -> U32:
257	  Indexed{Geometry{batch,channels,height,width},first,rows} = context
258	  (rows * width : U32)
259	
260	def 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	
267	def 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	
273	type Filled<-Element: Data> is Data:
274	  Filled{geometry: Geometry,fill: Element}
275	
276	def 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	
283	def 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	
288	def 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.
293	def 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.
302	def 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	
310	def 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	
319	def 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	
322	def allocate(~Element: Data,fill: Element,+count: U32) -> Storage.Buffer<Element>:
323	  Allocation.allocate(~Element,Storage.depth_for(count),fill,U32.to_nat(count))
324	
325	def 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	
337	def repartitioned(~Element: Data,owners: Bands(~Element) & Bands(~Element)) -> Bands(~Element):
338	  (source,result) = owners
339	  result
340	
341	def 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	
346	def 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	
350	def 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	
353	def 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	
357	def 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	
363	def 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	
372	def 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	
375	def 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	
384	def 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.
393	def 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	
397	def 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	
402	def 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	
405	def 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	
410	def 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.
419	def 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	
422	def 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	
426	def 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	
431	type Inspection is Data:
432	  Inspection{valid: Bool,first: U32,end: U32,leaves: U32}
433	
434	def 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	
439	def 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	
444	def 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	
451	def 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	
456	def 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))