kernel_weight_gradient.bend
stelliferous@0.0.3.0 · 447 lines · 0 definitions
1# Gradient tiles are the GEMM input, so multiplication keeps gradient first.
2# GEMM rows are kernel taps. The final read changes [tap,output] to [output,tap].
3import Base
4import stelliferous@0.0.3.0/kernel_parameter_gradient.bend as Parameter
5import stelliferous@0.0.3.0/kernel_gemm.bend as Gemm
6import stelliferous@0.0.3.0/kernel_tile.bend as Tile
7import stelliferous@0.0.3.0/storage_buffer.bend as Storage
8import stelliferous@0.0.3.0/traversal_array.bend as Traversal
9import stelliferous@0.0.3.0/traversal_loop.bend as Loop
10import stelliferous@0.0.3.0/traversal_grouped.bend as Grouped
11import stelliferous@0.0.3.0/traversal_lanes.bend as Lanes
12import stelliferous@0.0.3.0/tensor_lanes.bend as TensorLanes
13import stelliferous@0.0.3.0/storage_clone.bend as StorageClone
14import stelliferous@0.0.3.0/tensor.bend as Tensors
15import stelliferous@0.0.3.0/tensor_band.bend as Bands
16import stelliferous@0.0.3.0/tensor_band_schedule.bend as BandSchedule
17import stelliferous@0.0.3.0/traversal_partition.bend as Partition
18import stelliferous@0.0.3.0/kernel_cost.bend as Cost
19import stelliferous@0.0.3.0/default_kernel_cost.bend as DefaultCost
20
21type Geometry is Data:
22 Geometry{channels: U32,outputs: U32,height: U32,width: U32,kernel: U32,stride: U32,padding: U32,rows: U32,columns: U32,taps: U32}
23
24def geometry(work: Parameter.Work) -> Geometry:
25 Parameter.Work{n,+c,o,h,w,+k,s,p,rows,columns,g} = work
26 Geometry{U32.from_nat(c),U32.from_nat(o),U32.from_nat(h),U32.from_nat(w),U32.from_nat(k),U32.from_nat(s),U32.from_nat(p),U32.from_nat(rows),U32.from_nat(columns),U32.from_nat(Nat.mul(c,Nat.mul(k,k)))}
27
28def array(~Element: Data,+count: U32,value: Element) -> Array<Element>:
29 Array.new(Element,Storage.depth_for(count),value)
30
31type Lane is Data:
32 Lane{base: U32,row: U32,column: U32,valid: Bool,row_first: U32,row_end: U32,column_first: U32,column_end: U32}
33
34def axis_first(+offset: U32,+stride: U32,padding: U32) -> U32:
35 U32.from_nat(Nat.div(Nat.add(Nat.sub(U32.to_nat(padding),U32.to_nat(offset)),Nat.sub(U32.to_nat(stride),1n)),U32.to_nat(stride)))
36
37def axis_end(+offset: U32,+stride: U32,padding: U32,size: U32,positions: U32) -> U32:
38 U32.from_nat(Nat.min(U32.to_nat(positions),Nat.div(Nat.add(Nat.sub(Nat.add(U32.to_nat(padding),U32.to_nat(size)),U32.to_nat(offset)),Nat.sub(U32.to_nat(stride),1n)),U32.to_nat(stride))))
39
40def lane(+tap: U32,geometry: Geometry,image: U32) -> Lane:
41 Geometry{c,o,+h,+w,+k,+s,+p,+rows,+columns,+taps} = geometry
42 +kr = (tap / k % k : U32)
43 +kc = (tap % k : U32)
44 Lane{((image * c + tap / (k * k)) * h * w : U32),kr,kc,U32.is_lt(tap,taps),
45 axis_first(kr,s,p),axis_end(kr,s,p,h,rows),axis_first(kc,s,p),axis_end(kc,s,p,w,columns)}
46
47type Patch is Data:
48 Patch{a: Lane,b: Lane,c: Lane,d: Lane,e: Lane,f: Lane,g: Lane,h: Lane}
49
50def patch(+tap: U32,+geometry: Geometry,+image: U32) -> Patch:
51 Patch{lane(tap,geometry,image),lane((tap + 1 : U32),geometry,image),lane((tap + 2 : U32),geometry,image),lane((tap + 3 : U32),geometry,image),
52 lane((tap + 4 : U32),geometry,image),lane((tap + 5 : U32),geometry,image),lane((tap + 6 : U32),geometry,image),lane((tap + 7 : U32),geometry,image)}
53
54def selected(patch: Patch,index: Nat) -> Lane:
55 Patch{a,b,c,d,e,f,g,h} = patch
56 match index:
57 case 0n: a
58 case 1n: b
59 case 2n: c
60 case 3n: d
61 case 4n: e
62 case 5n: f
63 case 6n: g
64 case index: h
65
66type Pixel is Data:
67 Pixel{patch: Patch,row: U32,column: U32,height: U32,width: U32,stride: U32,padding: U32}
68
69def pixel(input: Storage.Buffer<F32>,+row: U32,+column: U32,+height: U32,+width: U32,+stride: U32,+padding: U32,lane: Lane) -> Storage.Buffer<F32> & F32:
70 Lane{base,kr,kc,valid,rf,re,cf,ce} = lane
71 +y = (row * stride + kr - padding : U32)
72 +x = (column * stride + kc - padding : U32)
73 Storage.read_masked(~F32,0.0,valid && U32.is_le(rf,row) && U32.is_lt(row,re) && U32.is_le(cf,column) && U32.is_lt(column,ce),input,(base + y * width + x : U32))
74
75def read_pixel(input: Storage.Buffer<F32>,context: Pixel,index: Nat) -> Storage.Buffer<F32> & F32:
76 Pixel{patch,row,column,height,width,stride,padding} = context
77 pixel(input,row,column,height,width,stride,padding,selected(patch,index))
78
79type Packing is Type:
80 Packing{input: Storage.Buffer<F32>,output: Array<Tile.Tile8>,index: U32,row: U32,column: U32,columns: U32,patch: Patch,height: U32,width: U32,stride: U32,padding: U32}
81
82def advanced(wrap: Bool,input: Storage.Buffer<F32>,output: Array<Tile.Tile8>,index: U32,+row: U32,column: U32,columns: U32,patch: Patch,height: U32,width: U32,stride: U32,padding: U32) -> Packing:
83 match wrap:
84 case False{}: Packing{input,output,index,row,column,columns,patch,height,width,stride,padding}
85 case True{}: Packing{input,output,index,(row + 1 : U32),0,columns,patch,height,width,stride,padding}
86
87def packed(output: Array<Tile.Tile8>,+index: U32,row: U32,+column: U32,+columns: U32,patch: Patch,height: U32,width: U32,stride: U32,padding: U32,result: Storage.Buffer<F32> & Tile.Tile8) -> Packing:
88 (input,value) = result
89 advanced(U32.is_eq((column + 1 : U32),columns),input,Array.set(Tile.Tile8,output,index,value),(index + 1 : U32),row,(column + 1 : U32),columns,patch,height,width,stride,padding)
90
91def pack_step(state: Packing) -> Packing:
92 Packing{input,output,index,+row,+column,columns,+patch,+height,+width,+stride,+padding} = state
93 +context = {Pixel{patch,row,column,height,width,stride,padding} : Pixel}
94 packed(output,index,row,column,columns,patch,height,width,stride,padding,
95 Tile.read_tail(~Storage.Buffer<F32>,~Pixel,~read_pixel,read_pixel(input,context,0n),context))
96
97type Gradient is Data:
98 Gradient{outputs: U32,spatial: U32,start: U32,row: U32,destination: U32}
99
100def gradient_pixel(input: Storage.Buffer<F32>,context: Gradient,lane: Nat,position: U32) -> Storage.Buffer<F32> & F32:
101 Gradient{outputs,+spatial,+start,row,destination} = context
102 +row = (row + U32.from_nat(lane) : U32)
103 Storage.read_masked(~F32,0.0,U32.is_lt(row,outputs),input,(row * spatial + start + position : U32))
104
105type GradientRead is Data:
106 GradientRead{context: Gradient,position: U32}
107
108def gradient_lane(input: Storage.Buffer<F32>,context: GradientRead,lane: Nat) -> Storage.Buffer<F32> & F32:
109 GradientRead{context,position} = context
110 gradient_pixel(input,context,lane,position)
111
112def gradient_tile(input: Storage.Buffer<F32>,context: Gradient,position: U32) -> Storage.Buffer<F32> & Tile.Tile8:
113 +context = context
114 +position = position
115 Tile.read_tail(~Storage.Buffer<F32>,~GradientRead,~gradient_lane,gradient_pixel(input,context,0n,position),GradientRead{context,position})
116
117def gradient_destination(context: Gradient,position: U32) -> U32:
118 position
119
120type WeightWindow is Data:
121 WeightWindow{geometry: Geometry,image: U32,tap: U32,output: U32,row: U32,column: U32,first: U32}
122
123def window_input(window: WeightWindow,+count: U32,input: Storage.Buffer<F32>) -> Packing:
124 WeightWindow{+geometry,image,tap,output,row,column,first} = window
125 Geometry{c,o,h,w,k,s,p,rows,columns,taps} = geometry
126 Loop.run(~Packing,~pack_step,U32.to_nat(count),
127 Packing{input,array(~Tile.Tile8,count,Tile.filled(0.0)),0,row,column,columns,patch(tap,geometry,image),h,w,s,p})
128
129def window_gradient(window: WeightWindow,+count: U32,input: Storage.Buffer<F32>) -> Storage.Buffer<F32> & Array<Tile.Tile8>:
130 WeightWindow{geometry,image,tap,output,row,column,first} = window
131 Geometry{c,+o,h,w,k,s,p,+rows,+columns,taps} = geometry
132 Traversal.finish(~Storage.Buffer<F32>,~Tile.Tile8,~Gradient,
133 Traversal.run(~Storage.Buffer<F32>,~Tile.Tile8,~Gradient,~gradient_tile,~gradient_destination,U32.to_nat(count),
134 Traversal.State{input,array(~Tile.Tile8,count,Tile.filled(0.0)),0,
135 Gradient{o,(rows * columns : U32),(image * o * rows * columns + first : U32),output,0}}))
136
137type ChunkState is Type:
138 ChunkState{input: Storage.Buffer<F32>,gradient: Storage.Buffer<F32>,values: Array<F32>,remaining: U32,window: WeightWindow}
139
140def window_gemm(continued: Bool,count: U32,input: Array<Tile.Tile8>,weights: Array<Tile.Tile8>,output: Array<F32>) -> Array<F32>:
141 match continued:
142 case False{}: Gemm.result(Gemm.run_keeping(~False{},8,input,weights,array(~F32,1,0.0),output,8,count,0,8))
143 case True{}: Gemm.result(Gemm.continue_keeping(8,input,weights,array(~F32,1,0.0),output,8,count))
144
145def chunk_result(~continued: Bool,+count: U32,+remaining: U32,+window: WeightWindow,values: Array<F32>,packed: Packing,read: Storage.Buffer<F32> & Array<Tile.Tile8>) -> ChunkState:
146 WeightWindow{geometry,image,tap,output,old_row,old_column,first} = window
147 Packing{input,weights,index,row,column,columns,patch,h,w,s,p} = packed
148 (gradient,tiles) = read
149 ChunkState{input,gradient,window_gemm(continued,count,tiles,weights,values),(remaining - count : U32),
150 WeightWindow{geometry,image,tap,output,row,column,(first + count : U32)}}
151
152def chunk(~continued: Bool,state: ChunkState) -> ChunkState:
153 ChunkState{input,gradient,values,+remaining,+window} = state
154 +count = U32.min(256,remaining)
155 chunk_result(~continued,count,remaining,window,values,window_input(window,count,input),window_gradient(window,count,gradient))
156
157def chunk_values(state: ChunkState) -> (Storage.Buffer<F32> & Storage.Buffer<F32>) & Array<F32>:
158 ChunkState{input,gradient,values,remaining,window} = state
159 ((input,gradient),values)
160
161def window_product(+count: U32,+window: WeightWindow,input: Storage.Buffer<F32>,gradient: Storage.Buffer<F32>) ->
162 (Storage.Buffer<F32> & Storage.Buffer<F32>) & Array<F32>:
163 +turns = (((count - 1) / 256) + 1 : U32)
164 chunk_values(Loop.run(~ChunkState,~chunk(~True{}),U32.to_nat((turns - 1 : U32)),
165 chunk(~False{},ChunkState{input,gradient,array(~F32,64,0.0),count,window})))
166
167def copied(valid: Bool,+address: U32,output: Array<F32>,read: Array<F32> & F32) -> Array<F32> & Array<F32>:
168 (input,value) = read
169 (input,Gemm.store_lane(valid,output,address,value))
170
171type BlockCopy is Data:
172 BlockCopy{tap: U32,output: U32,taps: U32,outputs: U32}
173
174def copy_cell(index: U32,context: BlockCopy,owners: Array<F32> & Array<F32>) -> Array<F32> & Array<F32>:
175 BlockCopy{+tap,+output,taps,+outputs} = context
176 (input,destination) = owners
177 +index = index
178 +row = (tap + index / 8 : U32)
179 +column = (output + index % 8 : U32)
180 copied(U32.is_lt(row,taps) && U32.is_lt(column,outputs),(row * outputs + column : U32),destination,Array.get(F32,input,index))
181
182def block_result(+context: BlockCopy,output: Array<F32>,result: (Storage.Buffer<F32> & Storage.Buffer<F32>) & Array<F32>) ->
183 (Storage.Buffer<F32> & Storage.Buffer<F32>) & Array<F32>:
184 (sources,input) = result
185 (sources,Pair.snd(Array<F32>,Array<F32>,Traversal.indexed_fold(~(Array<F32> & Array<F32>),~BlockCopy,~copy_cell,64n,0,context,(input,output))))
186
187type GroupBlocks is Data:
188 GroupBlocks{window: WeightWindow,count: U32,output_blocks: U32}
189
190def block(index: U32,context: GroupBlocks,owners: (Storage.Buffer<F32> & Storage.Buffer<F32>) & Array<F32>) ->
191 (Storage.Buffer<F32> & Storage.Buffer<F32>) & Array<F32>:
192 GroupBlocks{WeightWindow{+geometry,+image,tap,output,+row,+column,+first},count,+output_blocks} = context
193 Geometry{c,+o,h,w,k,s,p,rows,columns,+taps} = geometry
194 ((input,gradient),values) = owners
195 +index = index
196 +tap = (index / output_blocks * 8 : U32)
197 +output = (index % output_blocks * 8 : U32)
198 block_result(BlockCopy{tap,output,taps,o},values,window_product(count,WeightWindow{geometry,image,tap,output,row,column,first},input,gradient))
199
200def group(+geometry: Geometry,+image: U32,range: Grouped.Range,input: Storage.Buffer<F32>,gradient: Storage.Buffer<F32>) ->
201 (Storage.Buffer<F32> & Storage.Buffer<F32>) & Array<F32>:
202 Geometry{c,+o,h,w,k,s,p,rows,+columns,+taps} = geometry
203 Grouped.Range{+first,count} = range
204 +output_blocks = ((o + 7) / 8 : U32)
205 +context = {GroupBlocks{WeightWindow{geometry,image,0,0,U32.from_nat(Nat.div(first,U32.to_nat(columns))),U32.from_nat(Nat.mod(first,U32.to_nat(columns))),U32.from_nat(first)},U32.from_nat(count),output_blocks} : GroupBlocks}
206 Traversal.indexed_fold(~((Storage.Buffer<F32> & Storage.Buffer<F32>) & Array<F32>),~GroupBlocks,~block,
207 U32.to_nat((((taps + 7) / 8) * output_blocks : U32)),0,context,
208 ((input,gradient),array(~F32,(((taps * o + 7) / 8) * 8 : U32),0.0)))
209
210# Array sums use the same eight-lane gather, addition and scatter as tensors.
211def sum_result(+position: U32,left: Array<F32> & Tile.Tile8,right: Array<F32> & Tile.Tile8) -> Array<F32> & Array<F32>:
212 (left,a) = left
213 (right,b) = right
214 (Lanes.scatter(left,position,TensorLanes.combine_tile(~F32.add,a,b)),right)
215
216def sum_step(position: U32,context: Unit,owners: Array<F32> & Array<F32>) -> Array<F32> & Array<F32>:
217 (left,right) = owners
218 +position = (position * 8 : U32)
219 sum_result(position,Lanes.gather(~Unit,~TensorLanes.identity,left,position,Unit{}),Lanes.gather(~Unit,~TensorLanes.identity,right,position,Unit{}))
220
221def sum(+count: U32,left: Array<F32>,right: Array<F32>) -> Array<F32>:
222 Pair.fst(Array<F32>,Array<F32>,Traversal.indexed_fold(~(Array<F32> & Array<F32>),~Unit,~sum_step,U32.to_nat(((count + 7) / 8 : U32)),0,Unit{},(left,right)))
223
224type GroupContext is Data:
225 GroupContext{geometry: Geometry,image: U32}
226
227def group_read(owners: Storage.Buffer<F32> & Storage.Buffer<F32>,context: GroupContext,range: Grouped.Range) ->
228 (Storage.Buffer<F32> & Storage.Buffer<F32>) & Array<F32>:
229 (input,gradient) = owners
230 GroupContext{geometry,image} = context
231 group(geometry,image,range,input,gradient)
232
233def groups(ranges: List<&2,Grouped.Range>,geometry: Geometry,image: U32,input: Storage.Buffer<F32>,gradient: Storage.Buffer<F32>) ->
234 (Storage.Buffer<F32> & Storage.Buffer<F32>) & List<&1,Array<F32>>:
235 Traversal.collect_list_loop(~&1,~(Storage.Buffer<F32> & Storage.Buffer<F32>),~Array<F32>,~Grouped.Range,~GroupContext,~group_read,ranges,(input,gradient),GroupContext{geometry,image})
236
237def reduced(fuel: Nat,count: U32,result: (Storage.Buffer<F32> & Storage.Buffer<F32>) & List<&1,Array<F32>>) ->
238 (Storage.Buffer<F32> & Storage.Buffer<F32>) & Array<F32>:
239 (owners,values) = result
240 (owners,Grouped.levels(~&1,~Array<F32>,~U32,~(count => array(~F32,(((count + 7) / 8) * 8 : U32),0.0)),~(count => left => right => sum(count,left,right)),fuel,count,values))
241
242def image(+geometry: Geometry,index: U32,+ranges: List<&2,Grouped.Range>,input: Storage.Buffer<F32>,gradient: Storage.Buffer<F32>) ->
243 (Storage.Buffer<F32> & Storage.Buffer<F32>) & Array<F32>:
244 Geometry{c,+o,h,w,k,s,p,rows,columns,+taps} = geometry
245 reduced(List.length(&2,Grouped.Range,ranges),(o * taps : U32),groups(ranges,geometry,index,input,gradient))
246
247# Every image builds its own group ranges: a list kept for the next image
248# would be used twice and make the kernel's constructors shared.
249type ImageContext is Data:
250 ImageContext{geometry: Geometry,image: U32,rows: Nat,columns: Nat,group: Nat}
251
252def image_collected(context: ImageContext,result: (Storage.Buffer<F32> & Storage.Buffer<F32>) & Array<F32>) ->
253 (Storage.Buffer<F32> & Storage.Buffer<F32>) & Traversal.Collected<&1,ImageContext,Array<F32>>:
254 (owners,value) = result
255 (owners,Traversal.Collected{context,value})
256
257def image_read(owners: Storage.Buffer<F32> & Storage.Buffer<F32>,context: ImageContext) ->
258 (Storage.Buffer<F32> & Storage.Buffer<F32>) & Traversal.Collected<&1,ImageContext,Array<F32>>:
259 (input,gradient) = owners
260 ImageContext{+geometry,+index,+rows,+columns,+group} = context
261 image_collected(ImageContext{geometry,(index + 1 : U32),rows,columns,group},image(geometry,index,Grouped.ranges(rows,columns,group),input,gradient))
262
263def images(remaining: Nat,index: U32,geometry: Geometry,rows: Nat,columns: Nat,group: Nat,input: Storage.Buffer<F32>,gradient: Storage.Buffer<F32>) ->
264 (Storage.Buffer<F32> & Storage.Buffer<F32>) & List<&1,Array<F32>>:
265 Traversal.collect_loop(~&1,~(Storage.Buffer<F32> & Storage.Buffer<F32>),~Array<F32>,~ImageContext,~image_read,remaining,(input,gradient),ImageContext{geometry,index,rows,columns,group})
266
267# The values are [tap][output]; the result is [output][tap]. Output row of
268# the values is read at stride rows and written in place at its row of the result.
269type RowCopy is Type:
270 RowCopy{values: Array<F32>,result: Storage.Buffer<F32>,address: U32,stride: U32,destination: U32}
271
272def tap_copied(+address: U32,+stride: U32,+destination: U32,result: Storage.Buffer<F32>,read: Array<F32> & F32) -> RowCopy:
273 (values,value) = read
274 RowCopy{values,Storage.write(~F32,destination,value,result),(address + stride : U32),stride,(destination + 1 : U32)}
275
276def tap_copy(state: RowCopy) -> RowCopy:
277 RowCopy{values,result,+address,+stride,+destination} = state
278 tap_copied(address,stride,destination,result,Array.get(F32,values,address))
279
280type RowWrite is Type:
281 RowWrite{values: Array<F32>,result: Storage.Buffer<F32>,row: U32,rows: U32,taps: U32,first: U32}
282
283def row_copied(+row: U32,+rows: U32,+taps: U32,+first: U32,state: RowCopy) -> RowWrite:
284 RowCopy{values,result,address,stride,destination} = state
285 RowWrite{values,result,(row + 1 : U32),rows,taps,first}
286
287def row_step(state: RowWrite) -> RowWrite:
288 RowWrite{values,result,+row,+rows,+taps,+first} = state
289 row_copied(row,rows,taps,first,Loop.run(~RowCopy,~tap_copy,U32.to_nat(taps),RowCopy{values,result,row,rows,((first + row) * taps : U32)}))
290
291def rows_result(state: RowWrite) -> Storage.Buffer<F32>:
292 RowWrite{values,result,row,rows,taps,first} = state
293 result
294
295# Rows [first, first + rows) of the result from values of rows outputs.
296def rows_written(+rows: U32,+taps: U32,+first: U32,values: Array<F32>,result: Storage.Buffer<F32>) -> Storage.Buffer<F32>:
297 rows_result(Loop.run(~RowWrite,~row_step,U32.to_nat(rows),RowWrite{values,result,0,rows,taps,first}))
298
299def values_selected(+n: Nat,rows: Nat,columns: Nat,g: Nat,+geometry: Geometry,input: Storage.Buffer<F32>,gradient: Storage.Buffer<F32>) ->
300 (Storage.Buffer<F32> & Storage.Buffer<F32>) & Array<F32>:
301 Geometry{gc,+go,gh,gw,gk,gs,gp,gr,gcols,+taps} = geometry
302 reduced(n,(go * taps : U32),images(n,0,geometry,rows,columns,g,input,gradient))
303
304# The values of every tap and output, [tap][output].
305def weight_values(+work: Parameter.Work,input: Storage.Buffer<F32>,gradient: Storage.Buffer<F32>) ->
306 (Storage.Buffer<F32> & Storage.Buffer<F32>) & Array<F32>:
307 Parameter.Work{n,c,o,h,w,k,s,p,rows,columns,g} = work
308 values_selected(n,rows,columns,g,geometry(work),input,gradient)
309
310def written(+outputs: U32,+taps: U32,result: (Storage.Buffer<F32> & Storage.Buffer<F32>) & Array<F32>) ->
311 (Storage.Buffer<F32> & Storage.Buffer<F32>) & Storage.Buffer<F32>:
312 (owners,values) = result
313 (owners,rows_written(outputs,taps,0,values,Tensors.buffer(~F32,0.0,(outputs * taps : U32))))
314
315def taps(work: Parameter.Work) -> U32:
316 Parameter.Work{n,+c,o,h,w,+k,s,p,rows,columns,g} = work
317 U32.from_nat(Nat.mul(c,Nat.mul(k,k)))
318
319def outputs(work: Parameter.Work) -> U32:
320 Parameter.Work{n,c,o,h,w,k,s,p,rows,columns,g} = work
321 U32.from_nat(o)
322
323# The weight gradient of every output, [output][tap].
324def weight(+work: Parameter.Work,input: Storage.Buffer<F32>,gradient: Storage.Buffer<F32>) ->
325 (Storage.Buffer<F32> & Storage.Buffer<F32>) & Storage.Buffer<F32>:
326 written(outputs(work),taps(work),weight_values(work,input,gradient))
327
328# ---- Parts of the outputs ----
329
330# The outputs split at whole blocks of eight; each part computes its own
331# outputs with the serial kernel, so every value keeps its order of sums. A
332# part owns a copy of the input and the output gradient of its outputs, laid
333# out as [images, outputs, rows, columns]: the gradient seen as images of
334# output rows, cut at the rows of the part. The copies are made before the
335# parts run.
336type Owners is Type:
337 Owners{input: Storage.Buffer<F32>,gradient: Storage.Buffer<F32>}
338
339def part_work(+work: Parameter.Work,+outputs: U32) -> Parameter.Work:
340 Parameter.Work{n,c,o,h,w,k,s,p,rows,columns,g} = work
341 Parameter.Work{n,c,U32.to_nat(outputs),h,w,k,s,p,rows,columns,g}
342
343def owned(+first: U32,+rows: U32,inputs: Storage.Buffer<F32> & Storage.Buffer<F32>,gradients: Storage.Buffer<F32> & Storage.Buffer<F32>) ->
344 Owners & Partition.Tree<&1,Owners,Unit>:
345 (input,copy) = inputs
346 (gradient,channels) = gradients
347 (Owners{input,gradient},Partition.Leaf{Owners{copy,channels},first,rows,Unit{}})
348
349def own(work: Parameter.Work,first: U32,rows: U32,sources: Owners) -> Owners & Partition.Tree<&1,Owners,Unit>:
350 Parameter.Work{n,c,o,h,w,k,s,p,height,width,g} = work
351 Owners{input,gradient} = sources
352 +first = first
353 +rows = rows
354 owned(first,rows,StorageClone.clone(~F32,input),
355 Bands.whole_rows(~F32,Bands.Geometry{1,U32.from_nat(n),U32.from_nat(o),U32.from_nat(Nat.mul(height,width))},0.0,first,rows,gradient))
356
357def part_values(result: (Storage.Buffer<F32> & Storage.Buffer<F32>) & Array<F32>) -> Array<F32>:
358 (owners,values) = result
359 values
360
361def part(work: Parameter.Work,first: U32,rows: U32,owners: Owners,resource: Unit) -> Partition.Tree<&1,Array<F32>,Unit>:
362 Owners{input,gradient} = owners
363 +rows = rows
364 Partition.Leaf{part_values(weight_values(part_work(work,rows),input,gradient)),first,rows,Unit{}}
365
366# The owners are made along the plan, so a part never meets another shape.
367def unmatched(tree: Partition.Tree<&1,Owners,Unit>) -> Partition.Tree<&1,Array<F32>,Unit>:
368 Partition.Leaf{Array.new(F32,0n,0.0),0,0,Unit{}}
369
370def part_rows(taps: U32,first: U32,rows: U32,values: Array<F32>,result: Storage.Buffer<F32>) -> Storage.Buffer<F32>:
371 rows_written(rows,taps,first,values,result)
372
373def owners_of(sources: Owners) -> Storage.Buffer<F32> & Storage.Buffer<F32>:
374 Owners{input,gradient} = sources
375 (input,gradient)
376
377def parts_run(+depth: Nat,+work: Parameter.Work,made: Owners & Partition.Tree<&1,Owners,Unit>) ->
378 (Storage.Buffer<F32> & Storage.Buffer<F32>) & Storage.Buffer<F32>:
379 (sources,tree) = made
380 +outputs = outputs(work)
381 +taps = taps(work)
382 (owners_of(sources),Bands.fold(~Array<F32>,~Storage.Buffer<F32>,~U32,~part_rows,taps,
383 Bands.each(~Owners,~Array<F32>,~Unit,~Parameter.Work,~Bands.unshared,~part,~unmatched,Bands.unit_plan(depth,8,0,outputs),True{},work,tree,Unit{}),
384 Tensors.buffer(~F32,0.0,(outputs * taps : U32))))
385
386# The outputs [0, outputs) halved depth times at whole blocks of eight. The
387# rows of the parts are written in output order.
388def parts(+depth: Nat,+work: Parameter.Work,input: Storage.Buffer<F32>,gradient: Storage.Buffer<F32>) ->
389 (Storage.Buffer<F32> & Storage.Buffer<F32>) & Storage.Buffer<F32>:
390 parts_run(depth,work,Bands.build(~Owners,~Owners,~Parameter.Work,~own,Bands.unit_plan(depth,8,0,outputs(work)),work,Owners{input,gradient}))
391
392# ---- How many parts ----
393
394# One output sums its positions in every image: the MACs, and the input and
395# gradient lanes the kernel packs, counted from source as 13 and 6
396# operations.
397def output_ns(+work: Parameter.Work,config: Cost.Configuration) -> F32:
398 Parameter.Work{n,c,o,h,w,k,s,p,rows,columns,g} = work
399 Cost.Configuration{mac_ns,byte_ns,operation_ns,grain,wake} = config
400 +positions = U32.to_f32(U32.from_nat(Nat.mul(n,Nat.mul(rows,columns))))
401 +taps = U32.to_f32(U32.from_nat(Nat.mul(c,Nat.mul(k,k))))
402 (positions * taps * mac_ns + positions * (13.0 + 6.0) * operation_ns : F32)
403
404# A split's time: the parts in waves of the workers, one copy of the input
405# for every part and one of the gradient, and the two wake rounds.
406def split_ns(+work: Parameter.Work,+workers: U32,+config: Cost.Configuration,size: BandSchedule.Size) -> F32:
407 Parameter.Work{n,c,o,h,w,k,s,p,rows,columns,g} = work
408 Cost.Configuration{mac_ns,+byte_ns,operation_ns,grain,+wake} = config
409 BandSchedule.Size{+leaves,minimum,+maximum} = size
410 +input = U32.to_f32(U32.from_nat(Nat.mul(n,Nat.mul(c,Nat.mul(h,w)))))
411 +gradient = U32.to_f32(U32.from_nat(Nat.mul(n,Nat.mul(o,Nat.mul(rows,columns)))))
412 (Cost.region_ns(workers,leaves,(U32.to_f32(maximum) * output_ns(work,config) : F32)) + (U32.to_f32(leaves) * input + gradient) * 4.0 * byte_ns + 2.0 * wake : F32)
413
414def depth_ns(depth: Nat,+work: Parameter.Work,+workers: U32,+config: Cost.Configuration) -> F32:
415 match depth:
416 case 0n: (U32.to_f32(outputs(work)) * output_ns(work,config) : F32)
417 case 1n+ +rest: split_ns(work,workers,config,BandSchedule.size(Bands.unit_plan(1n+rest,8,0,outputs(work))))
418
419def faster(+depth: Nat,+time: F32,best: Nat & F32) -> Nat & F32:
420 (old_depth,+old_time) = best
421 Bool.pick((Nat & F32),F32.is_lt(time,old_time),(depth,time),(old_depth,old_time))
422
423def best_depth(remaining: Nat,+depth: Nat,+work: Parameter.Work,+workers: U32,+config: Cost.Configuration,best: Nat & F32) -> Nat & F32:
424 match remaining:
425 case 0n: best
426 case 1n+rest: best_depth(rest,1n+depth,work,workers,config,faster(depth,depth_ns(depth,work,workers,config),best))
427
428def log2(+value: U32) -> Nat:
429 Nat.sub(Tensors.significant_bits(32n,value,0n,U32.is_zero(value)),1n)
430
431# The fastest depth that leaves every part at least eight outputs and no more
432# parts than workers.
433def depth(+work: Parameter.Work,+workers: U32) -> Nat:
434 +config = DefaultCost.configuration()
435 +limit = Nat.min(log2(U32.max(workers,1)),log2(U32.max((outputs(work) / 8 : U32),1)))
436 Pair.fst(Nat,F32,best_depth(limit,1n,work,workers,config,(0n,depth_ns(0n,work,workers,config))))
437
438def routed(depth: Nat,+work: Parameter.Work,input: Storage.Buffer<F32>,gradient: Storage.Buffer<F32>) ->
439 (Storage.Buffer<F32> & Storage.Buffer<F32>) & Storage.Buffer<F32>:
440 match depth:
441 case 0n: weight(work,input,gradient)
442 case 1n+rest: parts(1n+rest,work,input,gradient)
443
444# The weight gradient on the workers of the call.
445def weight_parts(+work: Parameter.Work,+workers: U32,input: Storage.Buffer<F32>,gradient: Storage.Buffer<F32>) ->
446 (Storage.Buffer<F32> & Storage.Buffer<F32>) & Storage.Buffer<F32>:
447 routed(depth(work,workers),work,input,gradient)