1	# The input gradient of a convolution by phases: output positions congruent
2	# modulo the stride read the same taps. Every row of every output channel of
3	# the output gradient is cut into eight-lane tiles once for each column shift
4	# the phases read, zero outside; a phase entry is then a tile index, and a job
5	# stores at the final strided addresses.
6	import Base
7	import ./storage_buffer.bend as Storage
8	import ./storage_allocation.bend as Allocation
9	import ./kernel_tile.bend as Tile
10	import ./kernel_rows.bend as Rows
11	import ./kernel_strided_store.bend as Stores
12	import ./proofs/kernel_strided_store_certificate.bend as StoreCertificate
13	import ./kernel_gemm_buffer.bend as Gemm
14	import ./kernel_spatial_packing.bend as Spatial
15	import ./kernel_window.bend as Window
16	import ./kernel_pipeline.bend as Pipeline
17	import ./traversal_loop.bend as Iteration
18	import ./traversal_array.bend as Traversal
19	import ./traversal_partition.bend as Partition
20	import ./tensor_band.bend as Bands
21	import ./tensor.bend as Tensors
22	import ./kernel_cost.bend as Cost
23	import ./default_kernel_cost.bend as DefaultCost
24	import ./kernel_planning.bend as KernelPlanning
25	import ./kernel_band_planning.bend as BandPlanning
26	
27	type Geometry is Data:
28	  Geometry{channels: U32,outputs: U32,height: U32,width: U32,kernel: U32,stride: U32,padding: U32,rows: U32,columns: U32}
29	
30	type Axis is Data:
31	  Axis{first: U32,count: U32,tap: U32,taps: U32,source: U32}
32	
33	def axis(+first: U32,+size: U32,+kernel: U32,+stride: U32,+padding: U32) -> Axis:
34	  +residue = ((first + padding) % stride : U32)
35	  +taps = Bool.pick(U32,U32.is_lt(residue,kernel),((kernel - 1 - residue) / stride + 1 : U32),0)
36	  +tap = (residue + (taps - 1) * stride : U32)
37	  Axis{first,((size - 1 - first) / stride + 1 : U32),tap,taps,((first + padding) / stride - tap / stride : U32)}
38	
39	type Phase is Data:
40	  Phase{geometry: Geometry,y: Axis,x: Axis,inner: U32}
41	
42	def phase_axes(geometry: Geometry,a: Axis,b: Axis) -> Phase:
43	  Geometry{c,+o,h,w,k,s,p,gh,gw} = geometry
44	  Axis{af,ac,at,+ar,asrc} = a
45	  Axis{bf,bc,bt,+br,bsrc} = b
46	  Phase{geometry,a,b,(o * ar * br : U32)}
47	
48	def phase(+geometry: Geometry,+y: U32,+x: U32) -> Phase:
49	  Geometry{c,o,+h,+w,+k,+s,+p,gh,gw} = geometry
50	  phase_axes(geometry,axis(y,h,k,s,p),axis(x,w,k,s,p))
51	
52	# Tap coordinates and source strides for mapped weight reorder.
53	type WeightWindow is Data:
54	  WeightWindow{kernel: U32,stride: U32,y_tap: U32,y_rows: U32,x_tap: U32,x_columns: U32,tap_area: U32,output_step: U32}
55	
56	def weight_window(context: Phase) -> WeightWindow:
57	  Phase{Geometry{c,o,h,w,+k,s,p,gh,gw},Axis{yf,yc,yt,+yr,ys},Axis{xf,xc,xt,+xr,xs},inner} = context
58	  WeightWindow{k,s,yt,yr,xt,xr,(yr * xr : U32),(c * k * k : U32)}
59	
60	def weight_window_offset(window: WeightWindow,index: U32) -> U32:
61	  WeightWindow{+kernel,+stride,y_tap,y_rows,x_tap,+x_columns,tap_area,output_step} = window
62	  +index = index
63	  +row = (index / x_columns : U32)
64	  (index / tap_area * output_step + (y_tap - row % y_rows * stride) * kernel + x_tap - index % x_columns * stride : U32)
65	
66	def weight_offset(context: Phase,index: U32) -> U32:
67	  weight_window_offset(weight_window(context),index)
68	
69	def window_jobs(geometry: Geometry) -> Nat:
70	  Geometry{c,o,+h,+w,k,+s,p,gh,gw} = geometry
71	  Nat.mul(U32.to_nat(U32.min(s,h)),U32.to_nat(U32.min(s,w)))
72	
73	def phase_of(geometry: Geometry,index: U32) -> Phase:
74	  Geometry{c,o,h,+w,k,+s,p,gh,gw} = geometry
75	  +index = index
76	  +columns = U32.min(s,w)
77	  phase(geometry,(index / columns : U32),(index % columns : U32))
78	
79	def phase_inner(context: Phase) -> U32:
80	  Phase{g,y,x,inner} = context
81	  inner
82	
83	# The tiling cuts every row of every output channel of the output gradient
84	# into eight-lane tiles, once for each column shift the phases read, zero
85	# outside. It is built in chunks of at most chunk_rows output rows and
86	# chunk_groups groups, about one pipeline block of tiles; one output row of one
87	# group is the smallest chunk, and fits bounds it by the capacity. A chunk holds
88	# tile rows * chunk_groups blocks of outputs * shifts tiles (output channel
89	# outer, shift inner) and one unused tile, so that the rows a step reads do
90	# not share cache sets. Tile row t is gradient row t - kernel; shift i
91	# starts at gradient column 8 * group + i + column_low - kernel.
92	type Tiling is Data:
93	  Tiling{groups: U32,shifts: U32,outputs: U32,column_low: U32,kernel: U32,height: U32,width: U32,chunk_groups: U32,chunk_rows: U32}
94	
95	# The outputs [first, last) of an axis of the given size. Over the phases that
96	# read a tap for one of them: what 0 the lowest tile (lowest output + source +
97	# kernel), what 1 one past the highest tile (highest output + source + kernel
98	# + taps), what 2 the most outputs of a phase.
99	type Band is Data:
100	  Band{size: U32,kernel: U32,stride: U32,padding: U32,first: U32,last: U32,what: U32}
101	
102	def band_pick(used: Bool,+low: U32,+high: U32,+count: U32,+value: U32,what: U32) -> U32:
103	  match what:
104	    case 0: Bool.pick(U32,used,U32.min(value,low),value)
105	    case what: Bool.pick(U32,used,U32.max(value,Bool.pick(U32,U32.is_eq(what,1),high,count)),value)
106	
107	def band_measure(+context: Axis,+kernel: U32,+stride: U32,+first: U32,+last: U32,+value: U32,what: U32) -> U32:
108	  Axis{+start,+count,tap,+taps,+source} = context
109	  +low = Bool.pick(U32,U32.is_lt(start,first),((first - start + stride - 1) / stride : U32),0)
110	  +high = U32.min(count,Bool.pick(U32,U32.is_lt(start,last),((last - start + stride - 1) / stride : U32),0))
111	  band_pick(U32.is_lt(0,taps) && U32.is_lt(low,high),(low + source + kernel : U32),(high - 1 + source + kernel + taps : U32),count,value,what)
112	
113	def band_phase(index: U32,context: Band,value: U32) -> U32:
114	  Band{+size,+kernel,+stride,+padding,+first,+last,what} = context
115	  band_measure(axis(index,size,kernel,stride,padding),kernel,stride,first,last,value,what)
116	
117	def band_fold(+context: Band,value: U32) -> U32:
118	  Band{+size,kernel,+stride,padding,first,last,what} = context
119	  Traversal.indexed_fold(~U32,~Band,~band_phase,U32.to_nat(U32.min(stride,size)),0,context,value)
120	
121	# Output rows [first, first + rows) read at most rows - 1 + taps tile rows,
122	# taps = ceil(kernel / stride), and shifts is at most taps + 1.
123	def taps(+kernel: U32,+stride: U32) -> U32:
124	  ((kernel + stride - 1) / stride : U32)
125	
126	def fits(+geometry: Geometry,+capacity: U32) -> Bool:
127	  Geometry{c,+o,h,w,+k,+s,p,gh,gw} = geometry
128	  +taps = Nat.div(Nat.add(U32.to_nat(k),U32.to_nat((s - 1 : U32))),U32.to_nat(s))
129	  Nat.is_le(Nat.mul(taps,Nat.add(Nat.mul(U32.to_nat(o),Nat.add(taps,1n)),1n)),U32.to_nat(capacity))
130	
131	def chunked(+groups: U32,+shifts: U32,+outputs: U32,+column_low: U32,+kernel: U32,+stride: U32,+height: U32,+width: U32,+limit: U32) -> Tiling:
132	  +block = (outputs * shifts + 1 : U32)
133	  +chunk_groups = U32.max(1,U32.min(groups,(limit / (taps(kernel,stride) * block) : U32)))
134	  +fitting = (limit / (chunk_groups * block) + 1 : U32)
135	  Tiling{groups,shifts,outputs,column_low,kernel,height,width,chunk_groups,
136	    Bool.pick(U32,U32.is_lt(taps(kernel,stride),fitting),(fitting - taps(kernel,stride) : U32),1)}
137	
138	# Column shifts come from the first output of every column phase; groups from
139	# the most outputs of a column phase.
140	def tiling(geometry: Geometry) -> Tiling:
141	  Geometry{c,+o,h,+w,+k,+s,+p,+gh,+gw} = geometry
142	  +columns = U32.min(s,w)
143	  +column_low = band_fold(Band{w,k,s,p,0,columns,0},4294967295)
144	  chunked(((band_fold(Band{w,k,s,p,0,w,2},0) + 7) / 8 : U32),(band_fold(Band{w,k,s,p,0,columns,1},0) - column_low : U32),
145	    o,column_low,k,s,gh,gw,Pipeline.block_tiles())
146	
147	type Fill is Type:
148	  Fill{input: Array<F32>,tiles: Array<Tile.Tile8>,index: U32,row: U32,column: U32,first: U32,base: U32,plane: U32,height: U32,width: U32,shifts: Nat}
149	
150	def filled(tiles: Array<Tile.Tile8>,+index: U32,row: U32,+column: U32,first: U32,base: U32,plane: U32,height: U32,width: U32,shifts: Nat,read: Array<F32> & Tile.Tile8) -> Fill:
151	  (input,tile) = read
152	  Fill{input,Array.set(Tile.Tile8,tiles,index,tile),(index + 1 : U32),row,(column + 1 : U32),first,base,plane,height,width,shifts}
153	
154	def fill_shift(state: Fill) -> Fill:
155	  Fill{input,tiles,index,+row,+column,first,+base,plane,+height,+width,shifts} = state
156	  filled(tiles,index,row,column,first,base,plane,height,width,shifts,Spatial.block(input,row,column,base,height,width,1,0,0,0,0,0))
157	
158	def next_channel(state: Fill) -> Fill:
159	  Fill{input,tiles,index,row,column,+first,+base,+plane,height,width,shifts} = state
160	  Fill{input,tiles,index,row,first,first,(base + plane : U32),plane,height,width,shifts}
161	
162	def fill_channel(state: Fill) -> Fill:
163	  Fill{input,tiles,index,row,column,first,base,plane,height,width,+shifts} = state
164	  next_channel(Iteration.run(~Fill,~fill_shift,shifts,Fill{input,tiles,index,row,column,first,base,plane,height,width,shifts}))
165	
166	def fill_finished(state: Fill) -> Array<F32> & Array<Tile.Tile8>:
167	  Fill{input,tiles,index,row,column,first,base,plane,height,width,shifts} = state
168	  (input,tiles)
169	
170	# Rows [first, first + rows) of every output channel of one image.
171	type Source is Type:
172	  Source{data: Array<F32>,base: U32,plane: U32,first: U32,rows: U32}
173	
174	# Tile rows from tile row low and groups [group, group + span) of the tiling,
175	# read from a source.
176	type Filling is Data:
177	  Filling{context: Tiling,low: U32,base: U32,plane: U32,first: U32,rows: U32,group: U32,span: U32}
178	
179	# One tile row and group: every output channel, every shift.
180	def fill_group(index: U32,context: Filling,owners: Array<F32> & Array<Tile.Tile8>) -> Array<F32> & Array<Tile.Tile8>:
181	  Filling{Tiling{groups,+shifts,+outputs,+column_low,+kernel,height,+width,+chunk_groups,chunk_rows},+low,+base,+plane,+first,+source_rows,+group,+span} = context
182	  (input,tiles) = owners
183	  +index = index
184	  +column = ((index % span + group) * 8 + column_low - kernel : U32)
185	  fill_finished(Iteration.run(~Fill,~fill_channel,U32.to_nat(outputs),
186	    Fill{input,tiles,((index / span * chunk_groups + index % span) * (outputs * shifts + 1) : U32),(index / span + low - kernel - first : U32),column,column,
187	      base,plane,source_rows,width,U32.to_nat(shifts)}))
188	
189	type Jobs is Type:
190	  Jobs{tiles: Array<Tile.Tile8>,weights: Array<Tile.Tile8>,offsets: Array<U32>,output: Storage.Buffer<F32>}
191	
192	type Context is Data:
193	  Context{source: U32,weight_base: U32}
194	
195	type Loop is Type:
196	  Loop{tiles: Array<Tile.Tile8>,weights: Array<Tile.Tile8>,offsets: Array<U32>,index: U32,context: Context,accumulators: Rows.Rows<Tile.Tile8>}
197	
198	def with_tile(weights: Array<Tile.Tile8>,offsets: Array<U32>,+index: U32,context: Context,
199	  accumulators: Rows.Rows<Tile.Tile8>,samples: Rows.Rows<F32>,read: Array<Tile.Tile8> & Tile.Tile8) -> Loop:
200	  (tiles,tile) = read
201	  Loop{tiles,weights,offsets,(index + 1 : U32),context,Rows.map(~Tile.Tile8,~Tile.Tile8,~F32,~Tile.update,accumulators,tile,samples)}
202	
203	def with_offset(tiles: Array<Tile.Tile8>,weights: Array<Tile.Tile8>,+index: U32,context: Context,+source: U32,
204	  accumulators: Rows.Rows<Tile.Tile8>,samples: Rows.Rows<F32>,read: Array<U32> & U32) -> Loop:
205	  (offsets,offset) = read
206	  with_tile(weights,offsets,index,context,accumulators,samples,Array.get(Tile.Tile8,tiles,(source + offset : U32)))
207	
208	def with_weights(~upper: Bool,tiles: Array<Tile.Tile8>,offsets: Array<U32>,+index: U32,context: Context,+source: U32,
209	  accumulators: Rows.Rows<Tile.Tile8>,read: Array<Tile.Tile8> & Tile.Tile8) -> Loop:
210	  (weights,tile) = read
211	  with_offset(tiles,weights,index,context,source,accumulators,Rows.half(upper,tile),Array.get(U32,offsets,index))
212	
213	# One ordered step: the weight tile, the tile of the entry through the table,
214	# then the products of four rows at eight outputs.
215	def step(~upper: Bool,state: Loop) -> Loop:
216	  Loop{tiles,weights,offsets,+index,context,accumulators} = state
217	  Context{+source,+weight_base} = context
218	  with_weights(~upper,tiles,offsets,index,Context{source,weight_base},source,accumulators,
219	    Array.get(Tile.Tile8,weights,(weight_base + index : U32)))
220	
221	def execute(upper: Bool,count: Nat,state: Loop) -> Loop:
222	  match upper:
223	    case False{}: Iteration.run(~Loop,~step(~False{}),count,state)
224	    case True{}: Iteration.run(~Loop,~step(~True{}),count,state)
225	
226	# Raw stores within a job, one certificate update after them.
227	def stored(state: Loop,+store: Stores.Store,output: Storage.Buffer<F32>) -> Jobs:
228	  Loop{tiles,weights,offsets,index,context,+accumulators} = state
229	  Storage.Buffer{storage,length,certificate} = output
230	  Jobs{tiles,weights,offsets,Storage.Buffer{Stores.block(store,accumulators,storage),length,StoreCertificate.certificate(store,accumulators,storage,certificate)}}
231	
232	# One job is four input channels at eight consecutive outputs of one phase
233	# row. The loop index runs over the phase table from offset_base; weight_base
234	# is the phase weights less offset_base.
235	type Plan is Data:
236	  Plan{units: U32,groups: U32,inner: U32,source_base: U32,row_step: U32,group_step: U32,
237	    target_base: U32,target_row: U32,step: U32,plane: U32,channels: U32,columns: U32,weight_base: U32,offset_base: U32}
238	
239	def job(index: U32,plan: Plan,state: Jobs) -> Jobs:
240	  Plan{+units,+groups,+inner,+source_base,+row_step,+group_step,+target_base,+target_row,+step,+plane,+channels,+columns,+weight_base,+offset_base} = plan
241	  Jobs{tiles,weights,offsets,output} = state
242	  +index = index
243	  +unit = (index % units : U32)
244	  +group = (index / units % groups : U32)
245	  +line = (index / units / groups : U32)
246	  stored(execute(U32.is_eq((unit % 2 : U32),1),U32.to_nat(inner),
247	    Loop{tiles,weights,offsets,offset_base,Context{(source_base + line * row_step + group * group_step : U32),(weight_base + unit / 2 * inner : U32)},
248	      Rows.filled(~Tile.Tile8,Tile.filled(0.0))}),
249	    Stores.Store{(target_base + unit * 4 * plane + line * target_row + group * 8 * step : U32),step,plane,U32.min(4,(channels - unit * 4 : U32)),U32.min(8,(columns - group * 8 : U32))},output)
250	
251	# The tile of entry (output o, tap row a, tap column b) of a phase, relative to
252	# the tile of the job row and group, at position base + entry of the table.
253	type Table is Data:
254	  Table{tap_rows: U32,tap_columns: U32,row_step: U32,shifts: U32,first_shift: U32,base: U32}
255	
256	def offset(index: U32,context: Table,offsets: Array<U32>) -> Array<U32>:
257	  Table{+tap_rows,+tap_columns,+row_step,+shifts,+first_shift,+base} = context
258	  +index = index
259	  Array.set(U32,offsets,(base + index : U32),(index / tap_columns % tap_rows * row_step + index / (tap_rows * tap_columns) * shifts + index % tap_columns + first_shift : U32))
260	
261	type TileCopy is Type:
262	  TileCopy{source: Array<Tile.Tile8>,target: Array<Tile.Tile8>,from: U32,to: U32}
263	
264	def tile_copied(target: Array<Tile.Tile8>,+from: U32,+to: U32,read: Array<Tile.Tile8> & Tile.Tile8) -> TileCopy:
265	  (source,tile) = read
266	  TileCopy{source,Array.set(Tile.Tile8,target,to,tile),(from + 1 : U32),(to + 1 : U32)}
267	
268	def tile_step(state: TileCopy) -> TileCopy:
269	  TileCopy{source,target,+from,+to} = state
270	  tile_copied(target,from,to,Array.get(Tile.Tile8,source,from))
271	
272	def tile_target(state: TileCopy) -> Array<Tile.Tile8>:
273	  TileCopy{source,target,from,to} = state
274	  target
275	
276	# Every phase weights and table, side by side in phase order.
277	type Packed is Type:
278	  Packed{source: Storage.Buffer<F32>,weights: Array<Tile.Tile8>,offsets: Array<U32>,weight_base: U32,offset_base: U32}
279	
280	def packed_copied(+context: Phase,+layout: Tiling,source: Storage.Buffer<F32>,offsets: Array<U32>,+weight_base: U32,+offset_base: U32,+count: U32,weights: Array<Tile.Tile8>) -> Packed:
281	  Phase{Geometry{c,o,h,w,+k,s,p,gh,gw},Axis{yf,yc,yt,+yr,ys},Axis{xf,xc,xt,+xr,+xs},+inner} = context
282	  Tiling{groups,+shifts,+outputs,+column_low,kernel,height,width,+chunk_groups,chunk_rows} = layout
283	  Packed{source,weights,Traversal.indexed_fold(~Array<U32>,~Table,~offset,U32.to_nat(inner),0,
284	    Table{yr,xr,(chunk_groups * (outputs * shifts + 1) : U32),shifts,(xs + k - column_low : U32),offset_base},offsets),(weight_base + count : U32),(offset_base + inner : U32)}
285	
286	def packed_phase(+context: Phase,+layout: Tiling,weights: Array<Tile.Tile8>,offsets: Array<U32>,+weight_base: U32,+offset_base: U32,read: Storage.Buffer<F32> & Storage.Buffer<Tile.Tile8>) -> Packed:
287	  Phase{Geometry{+c,o,h,w,k,s,p,gh,gw},y,x,+inner} = context
288	  (source,phase_weights) = read
289	  Storage.Buffer{phase_tiles,phase_length,phase_certificate} = phase_weights
290	  +count = ((c + 7) / 8 * inner : U32)
291	  packed_copied(context,layout,source,offsets,weight_base,offset_base,count,
292	    tile_target(Iteration.run(~TileCopy,~tile_step,U32.to_nat(count),TileCopy{phase_tiles,weights,0,weight_base})))
293	
294	def pack_selected(empty: Bool,+context: Phase,+layout: Tiling,state: Packed) -> Packed:
295	  match empty:
296	    case True{}: state
297	    case False{}:
298	      Phase{Geometry{+c,o,h,w,+k,s,p,gh,gw},y,x,+inner} = context
299	      Packed{source,weights,offsets,weight_base,offset_base} = state
300	      packed_phase(context,layout,weights,offsets,weight_base,offset_base,
301	        Gemm.reorder_mapped(~WeightWindow,~weight_window_offset,c,inner,(k * k : U32),weight_window(context),source))
302	
303	type PackContext is Data:
304	  PackContext{geometry: Geometry,layout: Tiling}
305	
306	def pack_phase(index: U32,context: PackContext,state: Packed) -> Packed:
307	  PackContext{+geometry,+layout} = context
308	  +selected = phase_of(geometry,index)
309	  pack_selected(U32.is_eq(phase_inner(selected),0),selected,layout,state)
310	
311	def inner_total(index: U32,geometry: Geometry,total: U32) -> U32:
312	  (total + phase_inner(phase_of(geometry,index)) : U32)
313	
314	def weight_total(index: U32,geometry: Geometry,total: U32) -> U32:
315	  Geometry{+c,o,h,w,k,s,p,gh,gw} = geometry
316	  (total + (c + 7) / 8 * phase_inner(phase_of(geometry,index)) : U32)
317	
318	# The tile rows that output rows [first, last) read.
319	def band_low(+geometry: Geometry,+first: U32,+last: U32) -> U32:
320	  Geometry{c,o,+h,w,+k,+s,+p,gh,gw} = geometry
321	  band_fold(Band{h,k,s,p,first,last,0},4294967295)
322	
323	def band_high(+geometry: Geometry,+first: U32,+last: U32) -> U32:
324	  Geometry{c,o,+h,w,+k,+s,+p,gh,gw} = geometry
325	  band_fold(Band{h,k,s,p,first,last,1},0)
326	
327	def band_rows(+geometry: Geometry,+first: U32,+count: U32) -> U32:
328	  +low = band_low(geometry,first,(first + count : U32))
329	  +high = band_high(geometry,first,(first + count : U32))
330	  Bool.pick(U32,U32.is_lt(low,high),(high - low : U32),0)
331	
332	type Task is Data:
333	  Task{geometry: Geometry,layout: Tiling}
334	
335	type Leaf is Type:
336	  Leaf{tiles: Array<Tile.Tile8>,weights: Array<Tile.Tile8>,offsets: Array<U32>,output: Storage.Buffer<F32>,weight_base: U32,offset_base: U32}
337	
338	# Output rows [first, first + count) go to target + channel * plane +
339	# (row - first) * width + column.
340	type LeafContext is Data:
341	  LeafContext{task: Task,first: U32,count: U32,target: U32,plane: U32}
342	
343	# Output rows [row, row + rows) and groups [group, group + chunk_groups) of a
344	# leaf, over a chunk of the tiling that starts at tile row low.
345	type Chunk is Data:
346	  Chunk{leaf: LeafContext,row: U32,rows: U32,low: U32,group: U32}
347	
348	def leaf_jobs(+context: Phase,+chunk: Chunk,+weight_base: U32,+offset_base: U32,state: Jobs) -> Jobs:
349	  Phase{Geometry{+c,o,+h,+w,+k,+s,p,gh,gw},Axis{+yf,+yc,yt,yr,+ys},Axis{+xf,+xc,xt,xr,xs},+inner} = context
350	  Chunk{LeafContext{Task{geometry,+layout},+first,count,+target,+plane},+row,+rows,+low,+group} = chunk
351	  Tiling{groups,+shifts,+outputs,column_low,kernel,height,width,+chunk_groups,chunk_rows} = layout
352	  +start = Bool.pick(U32,U32.is_lt(yf,row),((row - yf + s - 1) / s : U32),0)
353	  +end = U32.min(yc,Bool.pick(U32,U32.is_lt(yf,(row + rows : U32)),((row + rows - yf + s - 1) / s : U32),0))
354	  +units = ((c + 3) / 4 : U32)
355	  +phase_groups = ((xc + 7) / 8 : U32)
356	  +span = Bool.pick(U32,U32.is_lt(group,phase_groups),U32.min(chunk_groups,(phase_groups - group : U32)),0)
357	  +group_step = (outputs * shifts + 1 : U32)
358	  Traversal.indexed_fold(~Jobs,~Plan,~job,U32.to_nat(Bool.pick(U32,U32.is_lt(start,end),((end - start) * span * units : U32),0)),0,
359	    Plan{units,span,inner,((start + ys + k - low) * chunk_groups * group_step : U32),(chunk_groups * group_step : U32),group_step,
360	      (target + (yf + start * s - first) * w + xf + group * 8 * s : U32),(s * w : U32),s,plane,c,(xc - group * 8 : U32),(weight_base - offset_base : U32),offset_base},state)
361	
362	def leaf_finished(+weight_base: U32,+offset_base: U32,state: Jobs) -> Leaf:
363	  Jobs{tiles,weights,offsets,output} = state
364	  Leaf{tiles,weights,offsets,output,weight_base,offset_base}
365	
366	def leaf_selected(empty: Bool,+context: Phase,+chunk: Chunk,state: Leaf) -> Leaf:
367	  match empty:
368	    case True{}: state
369	    case False{}:
370	      Phase{Geometry{+c,o,h,w,k,s,p,gh,gw},y,x,+inner} = context
371	      Leaf{tiles,weights,offsets,output,+weight_base,+offset_base} = state
372	      leaf_finished((weight_base + (c + 7) / 8 * inner : U32),(offset_base + inner : U32),
373	        leaf_jobs(context,chunk,weight_base,offset_base,Jobs{tiles,weights,offsets,output}))
374	
375	def leaf_phase(index: U32,chunk: Chunk,state: Leaf) -> Leaf:
376	  Chunk{LeafContext{Task{+geometry,layout},first,count,target,plane},row,rows,low,group} = chunk
377	  +selected = phase_of(geometry,index)
378	  leaf_selected(U32.is_eq(phase_inner(selected),0),selected,chunk,state)
379	
380	# The gradient a leaf read and the output it wrote.
381	type Part is Type:
382	  Part{gradient: Array<F32>,output: Storage.Buffer<F32>}
383	
384	# A leaf runs its chunks in turn with the source, the weights and the table.
385	type Sweep is Type:
386	  Sweep{source: Source,weights: Array<Tile.Tile8>,offsets: Array<U32>,output: Storage.Buffer<F32>}
387	
388	def chunk_leaf(+base: U32,+plane: U32,+first: U32,+rows: U32,data: Array<F32>,state: Leaf) -> Sweep:
389	  Leaf{tiles,weights,offsets,output,weight_base,offset_base} = state
390	  Sweep{Source{data,base,plane,first,rows},weights,offsets,output}
391	
392	def chunk_finished(+base: U32,+plane: U32,+first: U32,+rows: U32,pair: Array<F32> & Leaf) -> Sweep:
393	  (data,state) = pair
394	  chunk_leaf(base,plane,first,rows,data,state)
395	
396	def chunk_run(+chunk: Chunk,weights: Array<Tile.Tile8>,offsets: Array<U32>,output: Storage.Buffer<F32>,owners: Array<F32> & Array<Tile.Tile8>) -> Array<F32> & Leaf:
397	  Chunk{LeafContext{Task{+geometry,layout},first,count,target,plane},row,rows,low,group} = chunk
398	  (data,tiles) = owners
399	  (data,Traversal.indexed_fold(~Leaf,~Chunk,~leaf_phase,window_jobs(geometry),0,chunk,Leaf{tiles,weights,offsets,output,0,0}))
400	
401	# Chunk index: row chunk outer, group chunk inner.
402	def chunk_step(index: U32,context: LeafContext,state: Sweep) -> Sweep:
403	  LeafContext{Task{+geometry,+layout},+first,+count,target,plane} = context
404	  Tiling{+groups,+shifts,+outputs,column_low,kernel,height,width,+chunk_groups,+chunk_rows} = layout
405	  Sweep{Source{data,+base,+source_plane,+source_first,+source_rows},weights,offsets,output} = state
406	  +index = index
407	  +group_chunks = ((groups + chunk_groups - 1) / chunk_groups : U32)
408	  +row = (first + index / group_chunks * chunk_rows : U32)
409	  +rows = U32.min(chunk_rows,(first + count - row : U32))
410	  +group = (index % group_chunks * chunk_groups : U32)
411	  +span = U32.min(chunk_groups,(groups - group : U32))
412	  +low = band_low(geometry,row,(row + rows : U32))
413	  +tile_rows = band_rows(geometry,row,rows)
414	  chunk_finished(base,source_plane,source_first,source_rows,
415	    chunk_run(Chunk{context,row,rows,low,group},weights,offsets,output,
416	      Traversal.indexed_fold(~(Array<F32> & Array<Tile.Tile8>),~Filling,~fill_group,U32.to_nat((tile_rows * span : U32)),0,
417	        Filling{layout,low,base,source_plane,source_first,source_rows,group,span},
418	        (data,Allocation.array(~Tile.Tile8,Tensors.depth_for((tile_rows * chunk_groups * (outputs * shifts + 1) + 1 : U32)),Tile.filled(0.0))))))
419	
420	def swept(state: Sweep) -> Part:
421	  Sweep{Source{data,base,plane,first,rows},weights,offsets,output} = state
422	  Part{data,output}
423	
424	# Output rows [first, first + count) from a source holding the gradient rows they read.
425	def leaf_band(+context: LeafContext,source: Source,weights: Array<Tile.Tile8>,offsets: Array<U32>,output: Storage.Buffer<F32>) -> Part:
426	  LeafContext{Task{geometry,+layout},first,+count,target,plane} = context
427	  Tiling{+groups,shifts,outputs,column_low,kernel,height,width,+chunk_groups,+chunk_rows} = layout
428	  swept(Traversal.indexed_fold(~Sweep,~LeafContext,~chunk_step,
429	    U32.to_nat(((count + chunk_rows - 1) / chunk_rows * ((groups + chunk_groups - 1) / chunk_groups) : U32)),0,context,Sweep{source,weights,offsets,output}))
430	
431	# What one part owns: the gradient rows it reads, the weights and the table.
432	type Owners is Type:
433	  Owners{source: Source,weights: Array<Tile.Tile8>,offsets: Array<U32>}
434	
435	# A leaf writes output rows [first, first + count), stored [C][count][W].
436	def leaf(task: Task,first: U32,count: U32,owners: Owners) -> Partition.Tree<&1,Part,Unit>:
437	  Task{+geometry,+layout} = task
438	  Geometry{+c,o,h,+w,k,s,p,gh,gw} = geometry
439	  Owners{source,weights,offsets} = owners
440	  +first = first
441	  +count = count
442	  Partition.Leaf{leaf_band(LeafContext{Task{geometry,layout},first,count,0,(count * w : U32)},
443	    source,weights,offsets,Tensors.buffer(~F32,(0.0 + 0.0 : F32),(c * count * w : U32))),first,count,Unit{}}
444	
445	def copied_pair(+start: U32,+count: U32,+width: U32,+base: U32,+plane: U32,+first: U32,+rows: U32,owners: Array<F32> & Array<F32>) -> Source & Source:
446	  (data,copy) = owners
447	  (Source{copy,0,(count * width : U32),start,count},Source{data,base,plane,first,rows})
448	
449	# The gradient rows [start, start + count) that output rows [first, first + rows) read.
450	def copy_start(+geometry: Geometry,+first: U32,+rows: U32) -> U32:
451	  Geometry{c,o,h,w,+k,s,p,gh,gw} = geometry
452	  +low = band_low(geometry,first,(first + rows : U32))
453	  Bool.pick(U32,U32.is_lt(k,low),(low - k : U32),0)
454	
455	def copy_count(+geometry: Geometry,+first: U32,+rows: U32) -> U32:
456	  Geometry{c,o,h,w,+k,s,p,+gh,gw} = geometry
457	  +start = copy_start(geometry,first,rows)
458	  +high = band_high(geometry,first,(first + rows : U32))
459	  +stop = Bool.pick(U32,U32.is_lt(k,high),U32.min((high - k : U32),gh),0)
460	  Bool.pick(U32,U32.is_lt(start,stop),(stop - start : U32),0)
461	
462	# The left child, output rows [row, row + half), copies the gradient rows it
463	# reads; the right child keeps the source of its parent.
464	def split(task: Task,row: U32,half: U32,source: Source) -> Source & Source:
465	  Task{+geometry,layout} = task
466	  Geometry{c,+o,h,w,k,s,p,gh,+gw} = geometry
467	  Source{data,+base,+plane,+first,+rows} = source
468	  +row = row
469	  +half = half
470	  +start = copy_start(geometry,row,half)
471	  +count = copy_count(geometry,row,half)
472	  copied_pair(start,count,gw,base,plane,first,rows,
473	    Window.copy_planes(U32.to_nat(o),0,(count * gw : U32),plane,(base + (start - first) * gw : U32),(count * gw : U32),
474	      (data,Allocation.array(~F32,Tensors.depth_for((o * count * gw + 1 : U32)),0.0))))
475	
476	def divided(sources: Source & Source,weights: Array<Tile.Tile8> & Array<Tile.Tile8>,offsets: Array<U32> & Array<U32>) -> Bands.Division<Owners>:
477	  (left_source,right_source) = sources
478	  (left_weights,right_weights) = weights
479	  (left_offsets,right_offsets) = offsets
480	  Bands.Divided{Owners{left_source,left_weights,left_offsets},Owners{right_source,right_weights,right_offsets}}
481	
482	# The left half reads the rows of its plan; each half owns the weights and the table.
483	def divide(task: Task,left: Bands.Plan,right: Bands.Plan,owners: Owners) -> Bands.Division<Owners>:
484	  Owners{source,weights,offsets} = owners
485	  +row = Bands.first(left)
486	  divided(split(task,row,(Bands.first(right) - row : U32),source),Array.clone(Tile.Tile8,weights),Array.clone(U32,offsets))
487	
488	# A split always divides the owners; kept whole, they would compute the whole
489	# plan as one part.
490	def whole(task: Task,layout: Bands.Plan,owners: Owners) -> Partition.Tree<&1,Part,Unit>:
491	  +layout = layout
492	  leaf(task,Bands.first(layout),(Bands.last(layout) - Bands.first(layout) : U32),owners)
493	
494	# A leaf output [C][count][W] is the band of its rows of the output, counted
495	# over the batch from start. Its gradient replaces the one carried so far, so
496	# the gradient of the last leaf, the original, comes back.
497	def placed(start: U32,first: U32,count: U32,part: Part,gradient: Array<F32>) -> Bands.Bands(~F32) & Array<F32>:
498	  Part{data,local} = part
499	  (Partition.Leaf{local,(start + first : U32),count,Unit{}},data)
500	
501	# The split depth, priced with the shared cost configuration. A leaf builds
502	# the tiling of the gradient rows it reads, runs its jobs and stores its
503	# rows. A split copies the left band and clones the weights and the table.
504	# A split image is one parallel region with two wake rounds, so the batch
505	# does not change the choice.
506	type Pricing is Data:
507	  Pricing{geometry: Geometry,layout: Tiling,lanes: F32,clone: F32,config: Cost.Configuration}
508	
509	type Movement is Data:
510	  Movement{work: F32,span: F32,leaves: U32}
511	
512	type Choice is Data:
513	  Choice{depth: U32,time: F32}
514	
515	# Lanes are computed in eight-wide groups of each phase row.
516	def phase_lanes(context: Phase,total: F32) -> F32:
517	  Phase{g,Axis{yf,+yc,yt,yr,ys},Axis{xf,+xc,xt,xr,xs},+inner} = context
518	  (total + U32.to_f32(yc) * U32.to_f32(((xc + 7) / 8 * 8 : U32)) * U32.to_f32(inner) : F32)
519	
520	def phase_cost(index: U32,geometry: Geometry,total: F32) -> F32:
521	  phase_lanes(phase_of(geometry,index),total)
522	
523	def leaf_ns(+context: Pricing,+first: U32,+count: U32) -> F32:
524	  Pricing{+geometry,+layout,+lanes,clone,+config} = context
525	  Geometry{+c,o,+h,+w,k,s,p,gh,gw} = geometry
526	  Tiling{+groups,+shifts,+outputs,column_low,kernel,height,width,+chunk_groups,+chunk_rows} = layout
527	  Cost.Configuration{+mac,+byte_ns,operation_ns,grain,wake} = config
528	  # Each row chunk fills the tile rows its output rows read, halo included.
529	  +rows = band_rows(geometry,first,U32.min(count,chunk_rows))
530	  +row_chunks = ((count + chunk_rows - 1) / chunk_rows : U32)
531	  +chunk = (rows * chunk_groups * (outputs * shifts + 1) : U32)
532	  +chunks = U32.to_f32((row_chunks * ((groups + chunk_groups - 1) / chunk_groups) : U32))
533	  ((U32.to_f32(row_chunks) * U32.to_f32(rows) * U32.to_f32(groups) * U32.to_f32((outputs * shifts + 1 : U32)) * 32.0 +
534	    chunks * KernelPlanning.payload_float(Nat.mul(32n,KernelPlanning.capacity(chunk)))) * byte_ns +
535	    U32.to_f32(((c + 3) / 4 * 4 : U32)) * lanes * U32.to_f32(count) / U32.to_f32(U32.max(h,1)) * mac +
536	    U32.to_f32((c * count * w : U32)) * 4.0 * byte_ns : F32)
537	
538	def fork_ns(+context: Pricing,+first: U32,+rows: U32) -> F32:
539	  Pricing{+geometry,layout,lanes,+clone,+config} = context
540	  Geometry{c,+o,h,w,k,s,p,gh,+gw} = geometry
541	  Cost.Configuration{mac,+byte_ns,operation_ns,grain,wake} = config
542	  (BandPlanning.transfer((o * copy_count(geometry,first,rows) * gw : U32),byte_ns) + clone : F32)
543	
544	def merged(+fork: F32,left: Movement,right: Movement) -> Movement:
545	  Movement{left_work,+left_span,left_leaves} = left
546	  Movement{right_work,+right_span,right_leaves} = right
547	  Movement{(fork + left_work + right_work : F32),(fork + F32.max(left_span,right_span) : F32),(left_leaves + right_leaves : U32)}
548	
549	def leaf_movement(context: Pricing,first: U32,count: U32) -> Movement:
550	  +time = leaf_ns(context,first,count)
551	  Movement{time,time,1}
552	
553	def fork_movement(context: Pricing,left: Bands.Plan,right: Bands.Plan,first: Movement,second: Movement) -> Movement:
554	  +row = Bands.first(left)
555	  merged(fork_ns(context,row,(Bands.first(right) - row : U32)),first,second)
556	
557	# The plan the partition runs: the rows of one image halved depth times.
558	def image_rows(depth: Nat,+context: Pricing) -> Bands.Plan:
559	  Pricing{+geometry,layout,lanes,clone,config} = context
560	  Geometry{c,o,+h,w,k,s,p,gh,gw} = geometry
561	  Bands.row_plan(depth,U32.is_lt(1,h),0,h)
562	
563	def movement(depth: Nat,+context: Pricing) -> Movement:
564	  Bands.plan_fold(~Movement,~Pricing,~leaf_movement,~fork_movement,image_rows(depth,context),context)
565	
566	def estimate(+depth: U32,+workers: U32,+wake: F32,cost: Movement) -> F32:
567	  Movement{+work,+span,+leaves} = cost
568	  Bool.pick(F32,U32.is_zero(depth),work,(Cost.region_ns(workers,leaves,span) + 2.0 * wake : F32))
569	
570	def movement_work(cost: Movement) -> F32:
571	  Movement{work,span,leaves} = cost
572	  work
573	
574	def chosen(allowed: Bool,+depth: U32,+time: F32,best: Choice) -> Choice:
575	  Choice{old_depth,+old_time} = best
576	  Bool.pick(Choice,allowed && F32.is_lt(time,old_time),Choice{depth,time},Choice{old_depth,old_time})
577	
578	def candidate(+depth: U32,+workers: U32,+context: Pricing,best: Choice) -> Choice:
579	  Pricing{geometry,layout,lanes,clone,+config} = context
580	  Cost.Configuration{mac,byte_ns,operation_ns,+grain,+wake} = config
581	  +cost = movement(U32.to_nat(depth),context)
582	  chosen(U32.is_zero(depth) || F32.is_le(grain,movement_work(cost)),depth,
583	    estimate(depth,workers,wake,cost),best)
584	
585	# As the other planners: up to five leaves for each worker, at most one leaf a
586	# row, one worker never splits.
587	def depths(remaining: Nat,allowed: Bool,+depth: U32,+tasks: U32,+ceiling: U32,+workers: U32,+context: Pricing,best: Choice) -> Choice:
588	  match remaining allowed:
589	    case 0n allowed: best
590	    case 1n+rest False{}: best
591	    case 1n+rest True{}:
592	      Pricing{+geometry,layout,lanes,clone,config} = context
593	      Geometry{c,o,+h,w,k,s,p,gh,gw} = geometry
594	      +next = (tasks * 2 : U32)
595	      depths(rest,U32.is_le(next,ceiling) && U32.is_le(next,h),(depth + 1 : U32),next,ceiling,workers,context,candidate(depth,workers,context,best))
596	
597	# One worker, no split: the work of the whole image.
598	def baseline(+context: Pricing) -> Choice:
599	  Choice{0,movement_work(movement(0n,context))}
600	
601	def chosen_depth(best: Choice) -> Nat:
602	  Choice{depth,time} = best
603	  U32.to_nat(depth)
604	
605	def byte_cost(config: Cost.Configuration) -> F32:
606	  Cost.Configuration{mac,byte_ns,operation_ns,grain,wake} = config
607	  byte_ns
608	
609	def select(+geometry: Geometry,+layout: Tiling,+workers: U32,+weight_tiles: U32,+entries: U32) -> Nat:
610	  Geometry{c,o,+h,w,k,s,p,gh,gw} = geometry
611	  +config = DefaultCost.configuration()
612	  +workers = U32.max(1,workers)
613	  +context = {Pricing{geometry,layout,Traversal.indexed_fold(~F32,~Geometry,~phase_cost,window_jobs(geometry),0,geometry,0.0),
614	    (KernelPlanning.payload_float(Nat.add(Nat.mul(32n,KernelPlanning.capacity(weight_tiles)),Nat.mul(4n,KernelPlanning.capacity(entries)))) * byte_cost(config) : F32),config} : Pricing}
615	  chosen_depth(depths(32n,U32.is_lt(1,workers) && U32.is_le(2,h),1,2,(workers * 5 : U32),workers,context,baseline(context)))
616	
617	type Images is Data:
618	  Images{geometry: Geometry,layout: Tiling,depth: Nat}
619	
620	type ImageState is Type:
621	  ImageState{gradient: Array<F32>,weights: Array<Tile.Tile8>,offsets: Array<U32>}
622	
623	def image_bands(weights: Array<Tile.Tile8>,offsets: Array<U32>,result: Bands.Bands(~F32) & Array<F32>) -> ImageState & Bands.Bands(~F32):
624	  (bands,gradient) = result
625	  (ImageState{gradient,weights,offsets},bands)
626	
627	def image_offsets(+first: U32,+context: Images,gradient: Array<F32>,weights: Array<Tile.Tile8>,kept_weights: Array<Tile.Tile8>,pair: Array<U32> & Array<U32>) -> ImageState & Bands.Bands(~F32):
628	  Images{+geometry,+layout,+depth} = context
629	  Geometry{+c,+o,+h,+w,k,s,p,+gh,+gw} = geometry
630	  (offsets,kept_offsets) = pair
631	  image_bands(kept_weights,kept_offsets,Bands.walk(~Part,~Storage.Buffer<F32>,~Array<F32>,~U32,~placed,first,
632	    Bands.each_part(~Owners,~Part,~Unit,~Task,~Unit,~Bands.always,~divide,~leaf,~whole,Bands.row_plan(depth,U32.is_lt(1,h),0,h),True{},Unit{},Task{geometry,layout},
633	      Owners{Source{gradient,(first / h * o * gh * gw : U32),(gh * gw : U32),0,gh},weights,offsets}),
634	    Array.new(F32,0n,0.0)))
635	
636	def image_weights(+first: U32,+context: Images,gradient: Array<F32>,offsets: Array<U32>,pair: Array<Tile.Tile8> & Array<Tile.Tile8>) -> ImageState & Bands.Bands(~F32):
637	  (weights,kept_weights) = pair
638	  image_offsets(first,context,gradient,weights,kept_weights,Array.clone(U32,offsets))
639	
640	# The image whose rows start at first runs its plan, a single leaf when the
641	# depth is 0. Its leaves are its bands of the output.
642	def image(context: Images,first: U32,rows: U32,state: ImageState) -> ImageState & Bands.Bands(~F32):
643	  Images{+geometry,+layout,+depth} = context
644	  ImageState{gradient,weights,offsets} = state
645	  image_weights(first,Images{geometry,layout,depth},gradient,offsets,Array.clone(Tile.Tile8,weights))
646	
647	type Outcome is Type:
648	  Outcome{weights: Storage.Buffer<F32>,output: Bands.Bands(~F32)}
649	
650	def finished(source: Storage.Buffer<F32>,result: ImageState & Bands.Bands(~F32)) -> Outcome:
651	  (state,bands) = result
652	  Outcome{source,bands}
653	
654	# The images in order, each one band of rows of the batch.
655	def packed(+geometry: Geometry,+layout: Tiling,+batch: Nat,+depth: Nat,gradient: Array<F32>,state: Packed) -> Outcome:
656	  Geometry{c,o,+h,w,k,s,p,gh,gw} = geometry
657	  Packed{source,weights,offsets,weight_base,offset_base} = state
658	  finished(source,Bands.build(~Storage.Buffer<F32>,~ImageState,~Images,~image,Bands.image_plan(batch,0,0n,h),Images{geometry,layout,depth},ImageState{gradient,weights,offsets}))
659	
660	def tiled(+geometry: Geometry,batch: Nat,+workers: U32,gradient: Array<F32>,weights: Storage.Buffer<F32>) -> Outcome:
661	  +layout = tiling(geometry)
662	  +phases = window_jobs(geometry)
663	  +weight_tiles = Traversal.indexed_fold(~U32,~Geometry,~weight_total,phases,0,geometry,0)
664	  +entries = Traversal.indexed_fold(~U32,~Geometry,~inner_total,phases,0,geometry,0)
665	  packed(geometry,layout,batch,select(geometry,layout,workers,weight_tiles,entries),gradient,
666	    Traversal.indexed_fold(~Packed,~PackContext,~pack_phase,phases,0,PackContext{geometry,layout},
667	      Packed{weights,Allocation.array(~Tile.Tile8,Tensors.depth_for((weight_tiles + 1 : U32)),Tile.filled(0.0)),
668	        Allocation.array(~U32,Tensors.depth_for((entries + 1 : U32)),0),0,0}))
669	
670	def run_selected(empty: Bool,+geometry: Geometry,batch: Nat,+workers: U32,gradient: Storage.Buffer<F32>,weights: Storage.Buffer<F32>) -> Outcome:
671	  match empty:
672	    case True{}: Outcome{weights,Partition.Leaf{Tensors.buffer(~F32,0.0,0),0,0,Unit{}}}
673	    case False{}:
674	      Storage.Buffer{data,length,certificate} = gradient
675	      tiled(geometry,batch,workers,data,weights)
676	
677	# d [N][O][Ho][Wo] and weights [O][C][k][k] give the input gradient
678	# [N][C][H][W] as bands of rows of the batch; count is N * C * H * W.
679	def run(+geometry: Geometry,count: Nat,+workers: U32,gradient: Storage.Buffer<F32>,weights: Storage.Buffer<F32>) -> Outcome:
680	  Geometry{+c,o,+h,+w,k,s,p,gh,gw} = geometry
681	  +batch = Nat.div(count,U32.to_nat((c * h * w : U32)))
682	  run_selected(Nat.is_eq(batch,0n),geometry,batch,workers,gradient,weights)