kernel_input_gradient.bend
stelliferous@0.0.3.0 · 682 lines · 0 definitions
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.
6import Base
7import stelliferous@0.0.3.0/storage_buffer.bend as Storage
8import stelliferous@0.0.3.0/storage_allocation.bend as Allocation
9import stelliferous@0.0.3.0/kernel_tile.bend as Tile
10import stelliferous@0.0.3.0/kernel_rows.bend as Rows
11import stelliferous@0.0.3.0/kernel_strided_store.bend as Stores
12import stelliferous@0.0.3.0/proofs/kernel_strided_store_certificate.bend as StoreCertificate
13import stelliferous@0.0.3.0/kernel_gemm_buffer.bend as Gemm
14import stelliferous@0.0.3.0/kernel_spatial_packing.bend as Spatial
15import stelliferous@0.0.3.0/kernel_window.bend as Window
16import stelliferous@0.0.3.0/kernel_pipeline.bend as Pipeline
17import stelliferous@0.0.3.0/traversal_loop.bend as Iteration
18import stelliferous@0.0.3.0/traversal_array.bend as Traversal
19import stelliferous@0.0.3.0/traversal_partition.bend as Partition
20import stelliferous@0.0.3.0/tensor_band.bend as Bands
21import stelliferous@0.0.3.0/tensor.bend as Tensors
22import stelliferous@0.0.3.0/kernel_cost.bend as Cost
23import stelliferous@0.0.3.0/default_kernel_cost.bend as DefaultCost
24import stelliferous@0.0.3.0/kernel_planning.bend as KernelPlanning
25import stelliferous@0.0.3.0/kernel_band_planning.bend as BandPlanning
26
27type Geometry is Data:
28 Geometry{channels: U32,outputs: U32,height: U32,width: U32,kernel: U32,stride: U32,padding: U32,rows: U32,columns: U32}
29
30type Axis is Data:
31 Axis{first: U32,count: U32,tap: U32,taps: U32,source: U32}
32
33def 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
39type Phase is Data:
40 Phase{geometry: Geometry,y: Axis,x: Axis,inner: U32}
41
42def 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
48def 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.
53type 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
56def 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
60def 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
66def weight_offset(context: Phase,index: U32) -> U32:
67 weight_window_offset(weight_window(context),index)
68
69def 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
73def 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
79def 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.
92type 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.
99type Band is Data:
100 Band{size: U32,kernel: U32,stride: U32,padding: U32,first: U32,last: U32,what: U32}
101
102def 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
107def 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
113def 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
117def 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.
123def taps(+kernel: U32,+stride: U32) -> U32:
124 ((kernel + stride - 1) / stride : U32)
125
126def 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
131def 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.
140def 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
147type 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
150def 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
154def 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
158def 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
162def 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
166def 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.
171type 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.
176type 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.
180def 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
189type Jobs is Type:
190 Jobs{tiles: Array<Tile.Tile8>,weights: Array<Tile.Tile8>,offsets: Array<U32>,output: Storage.Buffer<F32>}
191
192type Context is Data:
193 Context{source: U32,weight_base: U32}
194
195type 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
198def 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
203def 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
208def 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.
215def 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
221def 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.
227def 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.
235type 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
239def 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.
253type Table is Data:
254 Table{tap_rows: U32,tap_columns: U32,row_step: U32,shifts: U32,first_shift: U32,base: U32}
255
256def 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
261type TileCopy is Type:
262 TileCopy{source: Array<Tile.Tile8>,target: Array<Tile.Tile8>,from: U32,to: U32}
263
264def 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
268def 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
272def 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.
277type Packed is Type:
278 Packed{source: Storage.Buffer<F32>,weights: Array<Tile.Tile8>,offsets: Array<U32>,weight_base: U32,offset_base: U32}
279
280def 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
286def 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
294def 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
303type PackContext is Data:
304 PackContext{geometry: Geometry,layout: Tiling}
305
306def 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
311def inner_total(index: U32,geometry: Geometry,total: U32) -> U32:
312 (total + phase_inner(phase_of(geometry,index)) : U32)
313
314def 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.
319def 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
323def 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
327def 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
332type Task is Data:
333 Task{geometry: Geometry,layout: Tiling}
334
335type 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.
340type 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.
345type Chunk is Data:
346 Chunk{leaf: LeafContext,row: U32,rows: U32,low: U32,group: U32}
347
348def 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
362def 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
366def 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
375def 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.
381type 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.
385type Sweep is Type:
386 Sweep{source: Source,weights: Array<Tile.Tile8>,offsets: Array<U32>,output: Storage.Buffer<F32>}
387
388def 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
392def 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
396def 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.
402def 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
420def 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.
425def 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.
432type 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].
436def 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
445def 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.
450def 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
455def 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.
464def 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
476def 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.
483def 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.
490def 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.
497def 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.
506type Pricing is Data:
507 Pricing{geometry: Geometry,layout: Tiling,lanes: F32,clone: F32,config: Cost.Configuration}
508
509type Movement is Data:
510 Movement{work: F32,span: F32,leaves: U32}
511
512type Choice is Data:
513 Choice{depth: U32,time: F32}
514
515# Lanes are computed in eight-wide groups of each phase row.
516def 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
520def phase_cost(index: U32,geometry: Geometry,total: F32) -> F32:
521 phase_lanes(phase_of(geometry,index),total)
522
523def 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
538def 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
544def 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
549def leaf_movement(context: Pricing,first: U32,count: U32) -> Movement:
550 +time = leaf_ns(context,first,count)
551 Movement{time,time,1}
552
553def 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.
558def 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
563def movement(depth: Nat,+context: Pricing) -> Movement:
564 Bands.plan_fold(~Movement,~Pricing,~leaf_movement,~fork_movement,image_rows(depth,context),context)
565
566def 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
570def movement_work(cost: Movement) -> F32:
571 Movement{work,span,leaves} = cost
572 work
573
574def 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
578def 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.
587def 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.
598def baseline(+context: Pricing) -> Choice:
599 Choice{0,movement_work(movement(0n,context))}
600
601def chosen_depth(best: Choice) -> Nat:
602 Choice{depth,time} = best
603 U32.to_nat(depth)
604
605def byte_cost(config: Cost.Configuration) -> F32:
606 Cost.Configuration{mac,byte_ns,operation_ns,grain,wake} = config
607 byte_ns
608
609def 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
617type Images is Data:
618 Images{geometry: Geometry,layout: Tiling,depth: Nat}
619
620type ImageState is Type:
621 ImageState{gradient: Array<F32>,weights: Array<Tile.Tile8>,offsets: Array<U32>}
622
623def 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
627def 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
636def 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.
642def 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
647type Outcome is Type:
648 Outcome{weights: Storage.Buffer<F32>,output: Bands.Bands(~F32)}
649
650def 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.
655def 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
660def 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
670def 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.
679def 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)