1 # Independent row bands keep their output owners. Fork distributes existing 2 # core/halo owners; only read-only weights and bias need affine clones. 3 import Base 4 import ./tensor_band.bend as Bands 5 import ./tensor_band_window.bend as BandWindow 6 import ./kernel_window.bend as KernelWindow 7 import ./kernel_pipeline.bend as Pipeline 8 import ./kernel_tile.bend as Tile 9 import ./storage_buffer.bend as Storage 10 import ./storage_clone.bend as Clone 11 import ./traversal_partition.bend as Partition 12 import ./math_activation.bend as Activations 13 14 type Context is Data: 15 Context{window: BandWindow.Geometry,channels: U32,inner: U32,parallel: Bool,activation: Activations.Activation} 16 17 # A leaf assembles the rows its outputs read into one contiguous window whose 18 # first row is a stride multiple, then runs the leaf on it: packing a 19 # window is packing a shorter image with output rows shifted by first / stride. 20 def packed_window(~custom: F32 -> F32,activation: Activations.Activation,+geometry: BandWindow.Geometry,+channels: U32,+inner: U32,+first: U32,+rows: U32,+output_first: U32,+output_rows: U32, 21 +window: BandWindow.Rows,+above_range: BandWindow.Rows,+below_range: BandWindow.Rows, 22 core: Storage.Buffer,above: Storage.Buffer,below: Storage.Buffer,weights: Storage.Buffer,bias: Storage.Buffer) -> Storage.Buffer: 23 BandWindow.Geometry{Bands.Geometry{batch,+ci,+height,+width},+kernel,+stride,+padding,+output_height,+output_width} = geometry 24 BandWindow.Rows{+window_first,+window_rows} = window 25 +positions = (output_rows * output_width : U32) 26 Pipeline.blocked(~custom,activation, 27 Pipeline.Work{channels,window_rows,width,kernel,stride,padding,output_width,positions,inner}, 28 KernelWindow.start(((output_first - first / height * output_height) * output_width : U32),stride,output_width,(window_first - first / height * height : U32)), 29 BandWindow.assemble(~F32,0.0,BandWindow.input_geometry(geometry),window,BandWindow.Rows{first,rows},above_range,below_range,core,above,below),weights,bias) 30 31 # A window is empty when the leaf has no outputs or its outputs read only 32 # padding above the image. The leaf then packs the one row at the stride 33 # position of its first output, which no output reads. 34 def core_row(+geometry: BandWindow.Geometry,+first: U32,+output_first: U32) -> BandWindow.Rows: 35 BandWindow.Geometry{Bands.Geometry{batch,ci,+height,width},kernel,+stride,padding,+output_height,output_width} = geometry 36 BandWindow.Rows{(first / height * height + (output_first - first / height * output_height) * stride : U32),1} 37 38 def source_rows(empty: Bool,+geometry: BandWindow.Geometry,+first: U32,+output_first: U32,window: BandWindow.Rows) -> BandWindow.Rows: 39 match empty: 40 case False{}: window 41 case True{}: core_row(geometry,first,output_first) 42 43 def windowed(~custom: F32 -> F32,activation: Activations.Activation,+geometry: BandWindow.Geometry,+channels: U32,+inner: U32,+first: U32,+rows: U32,+rows_read: BandWindow.RowWindow, 44 core: Storage.Buffer,above: Storage.Buffer,below: Storage.Buffer,weights: Storage.Buffer,bias: Storage.Buffer) -> Bands.Bands(~F32): 45 BandWindow.RowWindow{BandWindow.Rows{+output_first,+output_rows},+above_range,+below_range,BandWindow.Rows{+window_first,+window_rows}} = rows_read 46 Partition.Leaf{packed_window(~custom,activation,geometry,channels,inner,first,rows,output_first,output_rows, 47 source_rows(U32.is_zero(window_rows),geometry,first,output_first,BandWindow.Rows{window_first,window_rows}), 48 above_range,below_range,core,above,below,weights,bias),output_first,output_rows,Unit{}} 49 50 def leaf(~custom: F32 -> F32,context: Context,+first: U32,+rows: U32,core: Storage.Buffer,halo: BandWindow.Halo, 51 weights: Storage.Buffer,bias: Storage.Buffer) -> Bands.Bands(~F32): 52 Context{+geometry,+channels,+inner,parallel,activation} = context 53 BandWindow.Halo{above,below,+rows_read} = halo 54 windowed(~custom,activation,geometry,channels,inner,first,rows,rows_read,core,above,below,weights,bias) 55 56 def parallel(context: Context) -> Bool: 57 Context{window,channels,inner,parallel,activation} = context 58 parallel 59 60 def Parameters() -> Type: 61 Storage.Buffer & Storage.Buffer 62 63 def shared(weights: Storage.Buffer & Storage.Buffer,bias: Storage.Buffer & Storage.Buffer) -> Parameters() & Parameters(): 64 (left_weights,right_weights) = weights 65 (left_bias,right_bias) = bias 66 ((left_weights,left_bias),(right_weights,right_bias)) 67 68 # Each half of a split owns its copy of the read-only weights and bias. 69 def share(parameters: Parameters()) -> Parameters() & Parameters(): 70 (weights,bias) = parameters 71 shared(Clone.clone(~Tile.Tile8,weights),Clone.clone(~F32,bias)) 72 73 def band(~custom: F32 -> F32,context: Context,first: U32,rows: U32,core: Storage.Buffer,halo: BandWindow.Halo,parameters: Parameters()) -> Bands.Bands(~F32): 74 (weights,bias) = parameters 75 +first = first 76 +rows = rows 77 leaf(~custom,context,first,rows,core,halo,weights,bias) 78 79 # Every core leaf with its halo, through the shared band recursion. Trees of 80 # another shape keep the cores. 81 def run(~custom: F32 -> F32,plan: Bands.Plan,+context: Context,bands: Bands.Bands(~F32),halos: BandWindow.Halos(~F32), 82 weights: Storage.Buffer,bias: Storage.Buffer) -> Bands.Bands(~F32): 83 Bands.each(~Bands.Paired,BandWindow.Halo>,~Storage.Buffer,~Parameters(),~Context,~share, 84 ~Bands.paired(~Storage.Buffer,~BandWindow.Halo,~Parameters(),~Context,~band(~custom)), 85 ~Bands.left_of(~Storage.Buffer,~BandWindow.Halo), 86 plan,parallel(context),context,Bands.zip(~Storage.Buffer,~BandWindow.Halo,bands,halos),(weights,bias)) 87 88 def prepared(~custom: F32 -> F32,+context: Context,+plan: Bands.Plan,weights: Storage.Buffer,bias: Storage.Buffer, 89 owners: Bands.Bands(~F32) & BandWindow.Halos(~F32)) -> Bands.Bands(~F32): 90 (bands,halos) = owners 91 run(~custom,plan,context,bands,halos,weights,bias) 92 93 def planned(~custom: F32 -> F32,+context: Context,weights: Storage.Buffer,bias: Storage.Buffer, 94 owners: Bands.Bands(~F32) & Bands.Plan) -> Bands.Bands(~F32): 95 Context{window,channels,inner,parallel,activation} = context 96 (bands,+plan) = owners 97 prepared(~custom,context,plan,weights,bias,BandWindow.prepare(~F32,window,0.0,plan,bands)) 98 99 # The weights arrive in the order the kernel reads them. 100 def execute(~custom: F32 -> F32,context: Context,bands: Bands.Bands(~F32),weights: Storage.Buffer,bias: Storage.Buffer) -> Bands.Bands(~F32): 101 planned(~custom,context,weights,bias,Bands.plan(~F32,bands))