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