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]. 3 import Base 4 import ./kernel_parameter_gradient.bend as Parameter 5 import ./kernel_gemm.bend as Gemm 6 import ./kernel_tile.bend as Tile 7 import ./storage_buffer.bend as Storage 8 import ./traversal_array.bend as Traversal 9 import ./traversal_loop.bend as Loop 10 import ./traversal_grouped.bend as Grouped 11 import ./traversal_lanes.bend as Lanes 12 import ./tensor_lanes.bend as TensorLanes 13 import ./storage_clone.bend as StorageClone 14 import ./tensor.bend as Tensors 15 import ./tensor_band.bend as Bands 16 import ./tensor_band_schedule.bend as BandSchedule 17 import ./traversal_partition.bend as Partition 18 import ./kernel_cost.bend as Cost 19 import ./default_kernel_cost.bend as DefaultCost 20 21 type 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 24 def 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 28 def array(~Element: Data,+count: U32,value: Element) -> Array: 29 Array.new(Element,Storage.depth_for(count),value) 30 31 type 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 34 def 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 37 def 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 40 def 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 47 type Patch is Data: 48 Patch{a: Lane,b: Lane,c: Lane,d: Lane,e: Lane,f: Lane,g: Lane,h: Lane} 49 50 def 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 54 def 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 66 type Pixel is Data: 67 Pixel{patch: Patch,row: U32,column: U32,height: U32,width: U32,stride: U32,padding: U32} 68 69 def pixel(input: Storage.Buffer,+row: U32,+column: U32,+height: U32,+width: U32,+stride: U32,+padding: U32,lane: Lane) -> Storage.Buffer & 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 75 def read_pixel(input: Storage.Buffer,context: Pixel,index: Nat) -> Storage.Buffer & F32: 76 Pixel{patch,row,column,height,width,stride,padding} = context 77 pixel(input,row,column,height,width,stride,padding,selected(patch,index)) 78 79 type Packing is Type: 80 Packing{input: Storage.Buffer,output: Array,index: U32,row: U32,column: U32,columns: U32,patch: Patch,height: U32,width: U32,stride: U32,padding: U32} 81 82 def advanced(wrap: Bool,input: Storage.Buffer,output: Array,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 87 def packed(output: Array,+index: U32,row: U32,+column: U32,+columns: U32,patch: Patch,height: U32,width: U32,stride: U32,padding: U32,result: Storage.Buffer & 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 91 def 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,~Pixel,~read_pixel,read_pixel(input,context,0n),context)) 96 97 type Gradient is Data: 98 Gradient{outputs: U32,spatial: U32,start: U32,row: U32,destination: U32} 99 100 def gradient_pixel(input: Storage.Buffer,context: Gradient,lane: Nat,position: U32) -> Storage.Buffer & 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 105 type GradientRead is Data: 106 GradientRead{context: Gradient,position: U32} 107 108 def gradient_lane(input: Storage.Buffer,context: GradientRead,lane: Nat) -> Storage.Buffer & F32: 109 GradientRead{context,position} = context 110 gradient_pixel(input,context,lane,position) 111 112 def gradient_tile(input: Storage.Buffer,context: Gradient,position: U32) -> Storage.Buffer & Tile.Tile8: 113 +context = context 114 +position = position 115 Tile.read_tail(~Storage.Buffer,~GradientRead,~gradient_lane,gradient_pixel(input,context,0n,position),GradientRead{context,position}) 116 117 def gradient_destination(context: Gradient,position: U32) -> U32: 118 position 119 120 type WeightWindow is Data: 121 WeightWindow{geometry: Geometry,image: U32,tap: U32,output: U32,row: U32,column: U32,first: U32} 122 123 def window_input(window: WeightWindow,+count: U32,input: Storage.Buffer) -> 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 129 def window_gradient(window: WeightWindow,+count: U32,input: Storage.Buffer) -> Storage.Buffer & Array: 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,~Tile.Tile8,~Gradient, 133 Traversal.run(~Storage.Buffer,~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 137 type ChunkState is Type: 138 ChunkState{input: Storage.Buffer,gradient: Storage.Buffer,values: Array,remaining: U32,window: WeightWindow} 139 140 def window_gemm(continued: Bool,count: U32,input: Array,weights: Array,output: Array) -> Array: 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 145 def chunk_result(~continued: Bool,+count: U32,+remaining: U32,+window: WeightWindow,values: Array,packed: Packing,read: Storage.Buffer & Array) -> 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 152 def 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 157 def chunk_values(state: ChunkState) -> (Storage.Buffer & Storage.Buffer) & Array: 158 ChunkState{input,gradient,values,remaining,window} = state 159 ((input,gradient),values) 160 161 def window_product(+count: U32,+window: WeightWindow,input: Storage.Buffer,gradient: Storage.Buffer) -> 162 (Storage.Buffer & Storage.Buffer) & Array: 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 167 def copied(valid: Bool,+address: U32,output: Array,read: Array & F32) -> Array & Array: 168 (input,value) = read 169 (input,Gemm.store_lane(valid,output,address,value)) 170 171 type BlockCopy is Data: 172 BlockCopy{tap: U32,output: U32,taps: U32,outputs: U32} 173 174 def copy_cell(index: U32,context: BlockCopy,owners: Array & Array) -> Array & Array: 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 182 def block_result(+context: BlockCopy,output: Array,result: (Storage.Buffer & Storage.Buffer) & Array) -> 183 (Storage.Buffer & Storage.Buffer) & Array: 184 (sources,input) = result 185 (sources,Pair.snd(Array,Array,Traversal.indexed_fold(~(Array & Array),~BlockCopy,~copy_cell,64n,0,context,(input,output)))) 186 187 type GroupBlocks is Data: 188 GroupBlocks{window: WeightWindow,count: U32,output_blocks: U32} 189 190 def block(index: U32,context: GroupBlocks,owners: (Storage.Buffer & Storage.Buffer) & Array) -> 191 (Storage.Buffer & Storage.Buffer) & Array: 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 200 def group(+geometry: Geometry,+image: U32,range: Grouped.Range,input: Storage.Buffer,gradient: Storage.Buffer) -> 201 (Storage.Buffer & Storage.Buffer) & Array: 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 & Storage.Buffer) & Array),~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. 211 def sum_result(+position: U32,left: Array & Tile.Tile8,right: Array & Tile.Tile8) -> Array & Array: 212 (left,a) = left 213 (right,b) = right 214 (Lanes.scatter(left,position,TensorLanes.combine_tile(~F32.add,a,b)),right) 215 216 def sum_step(position: U32,context: Unit,owners: Array & Array) -> Array & Array: 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 221 def sum(+count: U32,left: Array,right: Array) -> Array: 222 Pair.fst(Array,Array,Traversal.indexed_fold(~(Array & Array),~Unit,~sum_step,U32.to_nat(((count + 7) / 8 : U32)),0,Unit{},(left,right))) 223 224 type GroupContext is Data: 225 GroupContext{geometry: Geometry,image: U32} 226 227 def group_read(owners: Storage.Buffer & Storage.Buffer,context: GroupContext,range: Grouped.Range) -> 228 (Storage.Buffer & Storage.Buffer) & Array: 229 (input,gradient) = owners 230 GroupContext{geometry,image} = context 231 group(geometry,image,range,input,gradient) 232 233 def groups(ranges: List<&2,Grouped.Range>,geometry: Geometry,image: U32,input: Storage.Buffer,gradient: Storage.Buffer) -> 234 (Storage.Buffer & Storage.Buffer) & List<&1,Array>: 235 Traversal.collect_list_loop(~&1,~(Storage.Buffer & Storage.Buffer),~Array,~Grouped.Range,~GroupContext,~group_read,ranges,(input,gradient),GroupContext{geometry,image}) 236 237 def reduced(fuel: Nat,count: U32,result: (Storage.Buffer & Storage.Buffer) & List<&1,Array>) -> 238 (Storage.Buffer & Storage.Buffer) & Array: 239 (owners,values) = result 240 (owners,Grouped.levels(~&1,~Array,~U32,~(count => array(~F32,(((count + 7) / 8) * 8 : U32),0.0)),~(count => left => right => sum(count,left,right)),fuel,count,values)) 241 242 def image(+geometry: Geometry,index: U32,+ranges: List<&2,Grouped.Range>,input: Storage.Buffer,gradient: Storage.Buffer) -> 243 (Storage.Buffer & Storage.Buffer) & Array: 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. 249 type ImageContext is Data: 250 ImageContext{geometry: Geometry,image: U32,rows: Nat,columns: Nat,group: Nat} 251 252 def image_collected(context: ImageContext,result: (Storage.Buffer & Storage.Buffer) & Array) -> 253 (Storage.Buffer & Storage.Buffer) & Traversal.Collected<&1,ImageContext,Array>: 254 (owners,value) = result 255 (owners,Traversal.Collected{context,value}) 256 257 def image_read(owners: Storage.Buffer & Storage.Buffer,context: ImageContext) -> 258 (Storage.Buffer & Storage.Buffer) & Traversal.Collected<&1,ImageContext,Array>: 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 263 def images(remaining: Nat,index: U32,geometry: Geometry,rows: Nat,columns: Nat,group: Nat,input: Storage.Buffer,gradient: Storage.Buffer) -> 264 (Storage.Buffer & Storage.Buffer) & List<&1,Array>: 265 Traversal.collect_loop(~&1,~(Storage.Buffer & Storage.Buffer),~Array,~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. 269 type RowCopy is Type: 270 RowCopy{values: Array,result: Storage.Buffer,address: U32,stride: U32,destination: U32} 271 272 def tap_copied(+address: U32,+stride: U32,+destination: U32,result: Storage.Buffer,read: Array & F32) -> RowCopy: 273 (values,value) = read 274 RowCopy{values,Storage.write(~F32,destination,value,result),(address + stride : U32),stride,(destination + 1 : U32)} 275 276 def 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 280 type RowWrite is Type: 281 RowWrite{values: Array,result: Storage.Buffer,row: U32,rows: U32,taps: U32,first: U32} 282 283 def 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 287 def 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 291 def rows_result(state: RowWrite) -> Storage.Buffer: 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. 296 def rows_written(+rows: U32,+taps: U32,+first: U32,values: Array,result: Storage.Buffer) -> Storage.Buffer: 297 rows_result(Loop.run(~RowWrite,~row_step,U32.to_nat(rows),RowWrite{values,result,0,rows,taps,first})) 298 299 def values_selected(+n: Nat,rows: Nat,columns: Nat,g: Nat,+geometry: Geometry,input: Storage.Buffer,gradient: Storage.Buffer) -> 300 (Storage.Buffer & Storage.Buffer) & Array: 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]. 305 def weight_values(+work: Parameter.Work,input: Storage.Buffer,gradient: Storage.Buffer) -> 306 (Storage.Buffer & Storage.Buffer) & Array: 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 310 def written(+outputs: U32,+taps: U32,result: (Storage.Buffer & Storage.Buffer) & Array) -> 311 (Storage.Buffer & Storage.Buffer) & Storage.Buffer: 312 (owners,values) = result 313 (owners,rows_written(outputs,taps,0,values,Tensors.buffer(~F32,0.0,(outputs * taps : U32)))) 314 315 def 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 319 def 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]. 324 def weight(+work: Parameter.Work,input: Storage.Buffer,gradient: Storage.Buffer) -> 325 (Storage.Buffer & Storage.Buffer) & Storage.Buffer: 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. 336 type Owners is Type: 337 Owners{input: Storage.Buffer,gradient: Storage.Buffer} 338 339 def 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 343 def owned(+first: U32,+rows: U32,inputs: Storage.Buffer & Storage.Buffer,gradients: Storage.Buffer & Storage.Buffer) -> 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 349 def 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 357 def part_values(result: (Storage.Buffer & Storage.Buffer) & Array) -> Array: 358 (owners,values) = result 359 values 360 361 def part(work: Parameter.Work,first: U32,rows: U32,owners: Owners,resource: Unit) -> Partition.Tree<&1,Array,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. 367 def unmatched(tree: Partition.Tree<&1,Owners,Unit>) -> Partition.Tree<&1,Array,Unit>: 368 Partition.Leaf{Array.new(F32,0n,0.0),0,0,Unit{}} 369 370 def part_rows(taps: U32,first: U32,rows: U32,values: Array,result: Storage.Buffer) -> Storage.Buffer: 371 rows_written(rows,taps,first,values,result) 372 373 def owners_of(sources: Owners) -> Storage.Buffer & Storage.Buffer: 374 Owners{input,gradient} = sources 375 (input,gradient) 376 377 def parts_run(+depth: Nat,+work: Parameter.Work,made: Owners & Partition.Tree<&1,Owners,Unit>) -> 378 (Storage.Buffer & Storage.Buffer) & Storage.Buffer: 379 (sources,tree) = made 380 +outputs = outputs(work) 381 +taps = taps(work) 382 (owners_of(sources),Bands.fold(~Array,~Storage.Buffer,~U32,~part_rows,taps, 383 Bands.each(~Owners,~Array,~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. 388 def parts(+depth: Nat,+work: Parameter.Work,input: Storage.Buffer,gradient: Storage.Buffer) -> 389 (Storage.Buffer & Storage.Buffer) & Storage.Buffer: 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. 397 def 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. 406 def 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 414 def 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 419 def 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 423 def 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 428 def 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. 433 def 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 438 def routed(depth: Nat,+work: Parameter.Work,input: Storage.Buffer,gradient: Storage.Buffer) -> 439 (Storage.Buffer & Storage.Buffer) & Storage.Buffer: 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. 445 def weight_parts(+work: Parameter.Work,+workers: U32,input: Storage.Buffer,gradient: Storage.Buffer) -> 446 (Storage.Buffer & Storage.Buffer) & Storage.Buffer: 447 routed(depth(work,workers),work,input,gradient)