kernel_pipeline.bend source
kernel_pipeline.bend on the hub · documented module
# Each leaf owns its input, packs it, and computes before returning. The caller# validates dimensions and chooses depth; the runtime schedules parallel leaves.import Baseimport ./kernel_gemm_buffer.bend as GemmBufferimport ./kernel_tile.bend as Tileimport ./kernel_spatial_packing.bend as Packingimport ./kernel_packing_buffer.bend as PackingBufferimport ./kernel_gather_buffer.bend as GatherBufferimport ./storage_buffer.bend as Storageimport ./storage_clone.bend as Cloneimport ./tensor.bend as Tensorsimport ./traversal_partition.bend as Partitionimport ./storage_copy.bend as Copyimport ./kernel_window.bend as KernelWindowimport ./traversal_array.bend as Traversalimport ./math_activation.bend as Activationsimport ./proofs/traversal_certificate.bend as TraversalCertificateimport ./tensor_operations.bend as Operationstype Work is Data: Work{rows: U32,height: U32,width: U32,kernel: U32,stride: U32,padding: U32,output_width: U32,positions: U32,inner: U32}# A partition node's work and the activation its leaves apply.type Task is Data: Task{work: Work,activation: Activations.Activation}def Outputs() -> Type: Partition.Tree<&1,Storage.Buffer<F32>,U32>def activated_value(~activation: F32 -> F32,context: Unit,value: F32) -> F32: activation(value)def mapped(~Context: Data,~operation: Context -> F32 -> F32,+count: U32,+context: Context,buffer: Storage.Buffer<F32>) -> Storage.Buffer<F32>: Storage.Buffer{storage,length,certificate} = buffer Storage.Buffer{Traversal.map_in_place(~F32,~Context,~operation,U32.to_nat(count),0,context,storage),length, TraversalCertificate.map_in_place(~F32,~Context,~operation,U32.to_nat(count),0,context,storage,certificate)}def leaky_relu_value(slope: F32,value: F32) -> F32: Activations.leaky_relu_value(value,slope)def unit_mapped(~operation: F32 -> F32,+count: U32,buffer: Storage.Buffer<F32>) -> Storage.Buffer<F32>: mapped(~Unit,~activated_value(~operation),count,Unit{},buffer)def unit_lanes(~operation: F32 -> F32,+count: U32,buffer: Storage.Buffer<F32>) -> Storage.Buffer<F32>: Operations.lanes(~Unit,~activated_value(~operation),U32.to_nat(count),0,Unit{},buffer)# One tight pass over each finished leaf while it is still in cache; applying# the activation inside the store measured about twice as slow. The built-in# is chosen once per leaf, so each pass is specialized; the identity skips it.# Sigmoid, SiLU, tanh and GELU pass their values through tiles.def activated(~custom: F32 -> F32,activation: Activations.Activation,+count: U32,buffer: Storage.Buffer<F32>) -> Storage.Buffer<F32>: match activation: case Activations.Identity{}: buffer case Activations.ReLU{}: unit_mapped(~Activations.relu_value,count,buffer) case Activations.ReLU6{}: unit_mapped(~Activations.relu6_value,count,buffer) case Activations.LeakyReLU{slope}: mapped(~F32,~leaky_relu_value,count,slope,buffer) case Activations.Sigmoid{}: unit_lanes(~Activations.sigmoid_value,count,buffer) case Activations.SiLU{}: unit_lanes(~Activations.silu_value,count,buffer) case Activations.Tanh{}: unit_lanes(~Activations.tanh_value,count,buffer) case Activations.HardSigmoid{}: unit_mapped(~Activations.hard_sigmoid_value,count,buffer) case Activations.HardSwish{}: unit_mapped(~Activations.hard_swish_value,count,buffer) case Activations.GELUTanh{}: unit_lanes(~Activations.gelu_tanh_value,count,buffer) case Activations.Custom{operations}: unit_mapped(~custom,count,buffer)# Rows [first, first + rows) of every channel of the image, stored contiguously.type InputWindow is Type: InputWindow{buffer: Storage.Buffer<F32>,first: U32,rows: U32}def copy_rows(remaining: Nat,+channel: U32,+count: U32,+source_stride: U32,+source_start: U32,+target_stride: U32, owners: Storage.Buffer<F32> & Storage.Buffer<F32>) -> Storage.Buffer<F32> & Storage.Buffer<F32>: match remaining: case 0n: owners case 1n+rest: (source,target) = owners copy_rows(rest,(channel + 1 : U32),count,source_stride,source_start,target_stride, Copy.range(~F32,count,(channel * source_stride + source_start : U32),(channel * target_stride : U32),source,target))def windows(+first: U32,+rows: U32,+left_first: U32,+left_rows: U32,owners: Storage.Buffer<F32> & Storage.Buffer<F32>) -> InputWindow & InputWindow: (source,target) = owners (InputWindow{target,left_first,left_rows},InputWindow{source,first,rows})def cloned(+first: U32,+rows: U32,owners: Storage.Buffer<F32> & Storage.Buffer<F32>) -> InputWindow & InputWindow: (left,right) = owners (InputWindow{left,first,rows},InputWindow{right,first,rows})def copied(+channels: U32,+width: U32,+first: U32,+rows: U32,band: KernelWindow.Band,buffer: Storage.Buffer<F32>) -> InputWindow & InputWindow: KernelWindow.Band{+left_first,+left_rows} = band windows(first,rows,left_first,left_rows,copy_rows(U32.to_nat(channels),0,(left_rows * width : U32),(rows * width : U32), ((left_first - first) * width : U32),(left_rows * width : U32),(buffer,Tensors.buffer(~F32,0.0,(channels * left_rows * width : U32)))))def split_when(copy: Bool,+channels: U32,+width: U32,+first: U32,+rows: U32,+band: KernelWindow.Band,buffer: Storage.Buffer<F32>) -> InputWindow & InputWindow: match copy: case False{}: cloned(first,rows,Clone.clone(~F32,buffer)) case True{}: copied(channels,width,first,rows,band,buffer)# The right child keeps the parent's rows; the left child copies only its band.# Unaligned rows would break whole eight-position blocks, so they are cloned;# so is a band that fails its coverage check.def split(work: Work,offset: U32,half: U32,window: InputWindow) -> InputWindow & InputWindow: Work{channels_out,+height,+width,+kernel,+stride,+padding,+output_width,positions,+inner} = work InputWindow{buffer,+first,+rows} = window +offset = offset +half = half +band = KernelWindow.band(first,rows,height,kernel,stride,padding,output_width,offset,half) split_when(KernelWindow.copies_band(first,rows,height,kernel,stride,padding,output_width,offset,half,band), (inner / (kernel * kernel) : U32),width,first,rows,band,buffer)def split_task(task: Task,offset: U32,half: U32,window: InputWindow) -> InputWindow & InputWindow: Task{work,activation} = task split(work,offset,half,window)# One tile is 32 bytes; 8192 tiles are 256 KB.def block_tiles() -> U32: 8192def block_positions(inner: U32) -> U32: (8 * U32.max(1,(block_tiles() / inner : U32)) : U32)type BlockState is Type: BlockState{input: Storage.Buffer<F32>,tiles: Storage.Buffer<Tile.Tile8>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,output: Storage.Buffer<F32>}def block_kept(input: Storage.Buffer<F32>,kept: GemmBuffer.Kept) -> BlockState: GemmBuffer.Kept{tiles,weights,bias,output} = kept BlockState{input,tiles,weights,bias,output}def block_packed(+rows: U32,+positions: U32,inner: U32,+start: U32,+span: U32, weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,output: Storage.Buffer<F32>,packed: Storage.Buffer<F32> & Storage.Buffer<Tile.Tile8>) -> BlockState: (input,tiles) = packed block_kept(input, GemmBuffer.run_keeping(rows,tiles,weights,bias,output,span,inner,start,positions))# Packing addresses the input window; output positions are local to the leaf.# Each block overwrites tiles from zero and writes directly into the leaf output.def block_step(work: Work,pack_start: U32,output_start: U32,+span: U32,state: BlockState) -> BlockState: Work{rows,h,w,k,s,p,ow,positions,+inner} = work BlockState{input,tiles,weights,bias,output} = state block_packed(rows,positions,inner,output_start,span,weights,bias,output, PackingBuffer.run(Packing.Call{U32.to_nat((span / 8 * inner : U32)),pack_start,h,w,k,s,p,ow,inner},input,tiles))def blocks(remaining: Nat,+work: Work,+pack_start: U32,+output_start: U32,+padded: U32,+span: U32,state: BlockState) -> BlockState: match remaining: case 0n: state case 1n+rest: blocks(rest,work,(pack_start + span : U32),(output_start + span : U32),padded,span, block_step(work,pack_start,output_start,U32.min(span,(padded - output_start : U32)),state))def blocks_finished(~custom: F32 -> F32,activation: Activations.Activation,count: U32,state: BlockState) -> Storage.Buffer<F32>: BlockState{input,tiles,weights,bias,output} = state activated(~custom,activation,count,output)# Exactly one reusable tile buffer and one output buffer per leaf.# An inner larger than the target still needs one complete eight-position group.def blocked(~custom: F32 -> F32,+activation: Activations.Activation,+work: Work,pack_start: U32,input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>) -> Storage.Buffer<F32>: Work{+rows,h,w,k,s,p,ow,+positions,+inner} = work +padded = KernelWindow.padded(positions) +span = block_positions(inner) blocks_finished(~custom,activation,(rows * positions : U32), blocks(U32.to_nat(((padded + span - 1) / span : U32)),work,pack_start,0,padded,span, BlockState{input,Tensors.buffer(~Tile.Tile8,Tile.filled(0.0),U32.max(block_tiles(),inner)),weights,bias,Tensors.buffer(~F32,0.0,(rows * positions : U32))}))def leaf(~custom: F32 -> F32,task: Task,offset: U32,span: U32,window: InputWindow,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>) -> Outputs(): Task{Work{+rows,h,w,k,+s,p,+ow,+positions,inner},activation} = task InputWindow{input,+first,height} = window +offset = offset +count = U32.min(span,(positions - offset : U32)) Partition.Leaf{blocked(~custom,activation,Work{rows,height,w,k,s,p,ow,count,inner},KernelWindow.start(offset,s,ow,first),input,weights,bias),offset,count,rows}def partition(~custom: F32 -> F32,depth: Nat,task: Task,offset: U32,span: U32,input: InputWindow,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>) -> Outputs(): Partition.run(~InputWindow,~Storage.Buffer<Tile.Tile8>,~Storage.Buffer<F32>,~Storage.Buffer<F32>,~Task, ~split_task,~Clone.clone(~Tile.Tile8),~Clone.clone(~F32),~leaf(~custom),depth,task,offset,span,input,weights,bias)def gather(partitions: Outputs(),output: Storage.Buffer<F32>,+positions: U32) -> Storage.Buffer<F32>: Partition.gather(~Storage.Buffer<F32>,~(rows => offset => positions => span => input => output => GatherBuffer.run(rows,offset,positions,span,input,output)),partitions,output,positions)def finish_leaf(identity: Bool,input: Storage.Buffer<F32>,+rows: U32,offset: U32,span: U32,+positions: U32) -> Storage.Buffer<F32>: match identity: case True{}: input case False{}: GatherBuffer.run(rows,offset,positions,span,input,Tensors.buffer(~F32,0.0,(rows * positions : U32)))def finish(partitions: Outputs(),+rows: U32,+positions: U32) -> Storage.Buffer<F32>: match partitions: case Partition.Leaf{input,+offset,+span,channels}: finish_leaf(U32.is_zero(offset) && U32.is_eq(span,positions),input,channels,offset,span,positions) case Partition.Fork{left,right}: gather(Partition.Fork{left,right},Tensors.buffer(~F32,0.0,(rows * positions : U32)),positions)# The deepest split that leaves every leaf at least one eight-position block.def depth_limit(+padded: U32) -> U32: (U32.from_nat(Tensors.significant_bits(32n,(padded / 8 : U32),0n,U32.is_zero((padded / 8 : U32)))) - 1 : U32)def run_partitioned(~custom: F32 -> F32,activation: Activations.Activation,depth: U32,work: Work,input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>) -> Storage.Buffer<F32>: Work{+rows,+h,w,k,s,p,ow,+positions,+inner} = work +padded = KernelWindow.padded(positions) +limit = depth_limit(padded) finish(partition(~custom,U32.to_nat(U32.min(depth,limit)),Task{Work{rows,h,w,k,s,p,ow,positions,inner},activation},0,padded,InputWindow{input,0,h},weights,bias),rows,positions)def run_selected(~custom: F32 -> F32,blocked_leaf: Bool,activation: Activations.Activation,depth: U32,work: Work,input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>) -> Storage.Buffer<F32>: match blocked_leaf: case True{}: blocked(~custom,activation,work,0,input,weights,bias) case False{}: run_partitioned(~custom,activation,depth,work,input,weights,bias)def run_packed(~custom: F32 -> F32,activation: Activations.Activation,+depth: U32,work: Work,input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>) -> Storage.Buffer<F32>: run_selected(~custom,U32.is_zero(depth),activation,depth,work,input,weights,bias)def run(~custom: F32 -> F32,activation: Activations.Activation,depth: U32,+work: Work,input: Storage.Buffer<F32>,weights: Storage.Buffer<F32>,bias: Storage.Buffer<F32>) -> Storage.Buffer<F32>: Work{rows,h,w,k,s,p,ow,positions,inner} = work run_packed(~custom,activation,depth,work,input,GemmBuffer.reorder(rows,inner,weights),bias)