1 # Each leaf owns its input, packs it, and computes before returning. The caller 2 # validates dimensions and chooses depth; the runtime schedules parallel leaves. 3 import Base 4 import ./kernel_gemm_buffer.bend as GemmBuffer 5 import ./kernel_tile.bend as Tile 6 import ./kernel_spatial_packing.bend as Packing 7 import ./kernel_packing_buffer.bend as PackingBuffer 8 import ./kernel_gather_buffer.bend as GatherBuffer 9 import ./storage_buffer.bend as Storage 10 import ./storage_clone.bend as Clone 11 import ./tensor.bend as Tensors 12 import ./traversal_partition.bend as Partition 13 import ./storage_copy.bend as Copy 14 import ./kernel_window.bend as KernelWindow 15 import ./traversal_array.bend as Traversal 16 import ./math_activation.bend as Activations 17 import ./proofs/traversal_certificate.bend as TraversalCertificate 18 import ./tensor_operations.bend as Operations 19 20 type Work is Data: 21 Work{rows: U32,height: U32,width: U32,kernel: U32,stride: U32,padding: U32,output_width: U32,positions: U32,inner: U32} 22 23 # A partition node's work and the activation its leaves apply. 24 type Task is Data: 25 Task{work: Work,activation: Activations.Activation} 26 27 def Outputs() -> Type: 28 Partition.Tree<&1,Storage.Buffer,U32> 29 30 def activated_value(~activation: F32 -> F32,context: Unit,value: F32) -> F32: 31 activation(value) 32 33 def mapped(~Context: Data,~operation: Context -> F32 -> F32,+count: U32,+context: Context,buffer: Storage.Buffer) -> Storage.Buffer: 34 Storage.Buffer{storage,length,certificate} = buffer 35 Storage.Buffer{Traversal.map_in_place(~F32,~Context,~operation,U32.to_nat(count),0,context,storage),length, 36 TraversalCertificate.map_in_place(~F32,~Context,~operation,U32.to_nat(count),0,context,storage,certificate)} 37 38 def leaky_relu_value(slope: F32,value: F32) -> F32: 39 Activations.leaky_relu_value(value,slope) 40 41 def unit_mapped(~operation: F32 -> F32,+count: U32,buffer: Storage.Buffer) -> Storage.Buffer: 42 mapped(~Unit,~activated_value(~operation),count,Unit{},buffer) 43 44 def unit_lanes(~operation: F32 -> F32,+count: U32,buffer: Storage.Buffer) -> Storage.Buffer: 45 Operations.lanes(~Unit,~activated_value(~operation),U32.to_nat(count),0,Unit{},buffer) 46 47 # One tight pass over each finished leaf while it is still in cache; applying 48 # the activation inside the store measured about twice as slow. The built-in 49 # is chosen once per leaf, so each pass is specialized; the identity skips it. 50 # Sigmoid, SiLU, tanh and GELU pass their values through tiles. 51 def activated(~custom: F32 -> F32,activation: Activations.Activation,+count: U32,buffer: Storage.Buffer) -> Storage.Buffer: 52 match activation: 53 case Activations.Identity{}: buffer 54 case Activations.ReLU{}: unit_mapped(~Activations.relu_value,count,buffer) 55 case Activations.ReLU6{}: unit_mapped(~Activations.relu6_value,count,buffer) 56 case Activations.LeakyReLU{slope}: mapped(~F32,~leaky_relu_value,count,slope,buffer) 57 case Activations.Sigmoid{}: unit_lanes(~Activations.sigmoid_value,count,buffer) 58 case Activations.SiLU{}: unit_lanes(~Activations.silu_value,count,buffer) 59 case Activations.Tanh{}: unit_lanes(~Activations.tanh_value,count,buffer) 60 case Activations.HardSigmoid{}: unit_mapped(~Activations.hard_sigmoid_value,count,buffer) 61 case Activations.HardSwish{}: unit_mapped(~Activations.hard_swish_value,count,buffer) 62 case Activations.GELUTanh{}: unit_lanes(~Activations.gelu_tanh_value,count,buffer) 63 case Activations.Custom{operations}: unit_mapped(~custom,count,buffer) 64 65 # Rows [first, first + rows) of every channel of the image, stored contiguously. 66 type InputWindow is Type: 67 InputWindow{buffer: Storage.Buffer,first: U32,rows: U32} 68 69 def copy_rows(remaining: Nat,+channel: U32,+count: U32,+source_stride: U32,+source_start: U32,+target_stride: U32, 70 owners: Storage.Buffer & Storage.Buffer) -> Storage.Buffer & Storage.Buffer: 71 match remaining: 72 case 0n: owners 73 case 1n+rest: 74 (source,target) = owners 75 copy_rows(rest,(channel + 1 : U32),count,source_stride,source_start,target_stride, 76 Copy.range(~F32,count,(channel * source_stride + source_start : U32),(channel * target_stride : U32),source,target)) 77 78 def windows(+first: U32,+rows: U32,+left_first: U32,+left_rows: U32,owners: Storage.Buffer & Storage.Buffer) -> InputWindow & InputWindow: 79 (source,target) = owners 80 (InputWindow{target,left_first,left_rows},InputWindow{source,first,rows}) 81 82 def cloned(+first: U32,+rows: U32,owners: Storage.Buffer & Storage.Buffer) -> InputWindow & InputWindow: 83 (left,right) = owners 84 (InputWindow{left,first,rows},InputWindow{right,first,rows}) 85 86 def copied(+channels: U32,+width: U32,+first: U32,+rows: U32,band: KernelWindow.Band,buffer: Storage.Buffer) -> InputWindow & InputWindow: 87 KernelWindow.Band{+left_first,+left_rows} = band 88 windows(first,rows,left_first,left_rows,copy_rows(U32.to_nat(channels),0,(left_rows * width : U32),(rows * width : U32), 89 ((left_first - first) * width : U32),(left_rows * width : U32),(buffer,Tensors.buffer(~F32,0.0,(channels * left_rows * width : U32))))) 90 91 def split_when(copy: Bool,+channels: U32,+width: U32,+first: U32,+rows: U32,+band: KernelWindow.Band,buffer: Storage.Buffer) -> InputWindow & InputWindow: 92 match copy: 93 case False{}: cloned(first,rows,Clone.clone(~F32,buffer)) 94 case True{}: copied(channels,width,first,rows,band,buffer) 95 96 # The right child keeps the parent's rows; the left child copies only its band. 97 # Unaligned rows would break whole eight-position blocks, so they are cloned; 98 # so is a band that fails its coverage check. 99 def split(work: Work,offset: U32,half: U32,window: InputWindow) -> InputWindow & InputWindow: 100 Work{channels_out,+height,+width,+kernel,+stride,+padding,+output_width,positions,+inner} = work 101 InputWindow{buffer,+first,+rows} = window 102 +offset = offset 103 +half = half 104 +band = KernelWindow.band(first,rows,height,kernel,stride,padding,output_width,offset,half) 105 split_when(KernelWindow.copies_band(first,rows,height,kernel,stride,padding,output_width,offset,half,band), 106 (inner / (kernel * kernel) : U32),width,first,rows,band,buffer) 107 108 def split_task(task: Task,offset: U32,half: U32,window: InputWindow) -> InputWindow & InputWindow: 109 Task{work,activation} = task 110 split(work,offset,half,window) 111 112 # One tile is 32 bytes; 8192 tiles are 256 KB. 113 def block_tiles() -> U32: 114 8192 115 116 def block_positions(inner: U32) -> U32: 117 (8 * U32.max(1,(block_tiles() / inner : U32)) : U32) 118 119 type BlockState is Type: 120 BlockState{input: Storage.Buffer,tiles: Storage.Buffer,weights: Storage.Buffer,bias: Storage.Buffer,output: Storage.Buffer} 121 122 def block_kept(input: Storage.Buffer,kept: GemmBuffer.Kept) -> BlockState: 123 GemmBuffer.Kept{tiles,weights,bias,output} = kept 124 BlockState{input,tiles,weights,bias,output} 125 126 def block_packed(+rows: U32,+positions: U32,inner: U32,+start: U32,+span: U32, 127 weights: Storage.Buffer,bias: Storage.Buffer,output: Storage.Buffer,packed: Storage.Buffer & Storage.Buffer) -> BlockState: 128 (input,tiles) = packed 129 block_kept(input, 130 GemmBuffer.run_keeping(rows,tiles,weights,bias,output,span,inner,start,positions)) 131 132 # Packing addresses the input window; output positions are local to the leaf. 133 # Each block overwrites tiles from zero and writes directly into the leaf output. 134 def block_step(work: Work,pack_start: U32,output_start: U32,+span: U32,state: BlockState) -> BlockState: 135 Work{rows,h,w,k,s,p,ow,positions,+inner} = work 136 BlockState{input,tiles,weights,bias,output} = state 137 block_packed(rows,positions,inner,output_start,span,weights,bias,output, 138 PackingBuffer.run(Packing.Call{U32.to_nat((span / 8 * inner : U32)),pack_start,h,w,k,s,p,ow,inner},input,tiles)) 139 140 def blocks(remaining: Nat,+work: Work,+pack_start: U32,+output_start: U32,+padded: U32,+span: U32,state: BlockState) -> BlockState: 141 match remaining: 142 case 0n: state 143 case 1n+rest: 144 blocks(rest,work,(pack_start + span : U32),(output_start + span : U32),padded,span, 145 block_step(work,pack_start,output_start,U32.min(span,(padded - output_start : U32)),state)) 146 147 def blocks_finished(~custom: F32 -> F32,activation: Activations.Activation,count: U32,state: BlockState) -> Storage.Buffer: 148 BlockState{input,tiles,weights,bias,output} = state 149 activated(~custom,activation,count,output) 150 151 # Exactly one reusable tile buffer and one output buffer per leaf. 152 # An inner larger than the target still needs one complete eight-position group. 153 def blocked(~custom: F32 -> F32,+activation: Activations.Activation,+work: Work,pack_start: U32,input: Storage.Buffer,weights: Storage.Buffer,bias: Storage.Buffer) -> Storage.Buffer: 154 Work{+rows,h,w,k,s,p,ow,+positions,+inner} = work 155 +padded = KernelWindow.padded(positions) 156 +span = block_positions(inner) 157 blocks_finished(~custom,activation,(rows * positions : U32), 158 blocks(U32.to_nat(((padded + span - 1) / span : U32)),work,pack_start,0,padded,span, 159 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))})) 160 161 def leaf(~custom: F32 -> F32,task: Task,offset: U32,span: U32,window: InputWindow,weights: Storage.Buffer,bias: Storage.Buffer) -> Outputs(): 162 Task{Work{+rows,h,w,k,+s,p,+ow,+positions,inner},activation} = task 163 InputWindow{input,+first,height} = window 164 +offset = offset 165 +count = U32.min(span,(positions - offset : U32)) 166 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} 167 168 def partition(~custom: F32 -> F32,depth: Nat,task: Task,offset: U32,span: U32,input: InputWindow,weights: Storage.Buffer,bias: Storage.Buffer) -> Outputs(): 169 Partition.run(~InputWindow,~Storage.Buffer,~Storage.Buffer,~Storage.Buffer,~Task, 170 ~split_task,~Clone.clone(~Tile.Tile8),~Clone.clone(~F32),~leaf(~custom),depth,task,offset,span,input,weights,bias) 171 172 def gather(partitions: Outputs(),output: Storage.Buffer,+positions: U32) -> Storage.Buffer: 173 Partition.gather(~Storage.Buffer,~(rows => offset => positions => span => input => output => GatherBuffer.run(rows,offset,positions,span,input,output)),partitions,output,positions) 174 175 def finish_leaf(identity: Bool,input: Storage.Buffer,+rows: U32,offset: U32,span: U32,+positions: U32) -> Storage.Buffer: 176 match identity: 177 case True{}: input 178 case False{}: GatherBuffer.run(rows,offset,positions,span,input,Tensors.buffer(~F32,0.0,(rows * positions : U32))) 179 180 def finish(partitions: Outputs(),+rows: U32,+positions: U32) -> Storage.Buffer: 181 match partitions: 182 case Partition.Leaf{input,+offset,+span,channels}: 183 finish_leaf(U32.is_zero(offset) && U32.is_eq(span,positions),input,channels,offset,span,positions) 184 case Partition.Fork{left,right}: 185 gather(Partition.Fork{left,right},Tensors.buffer(~F32,0.0,(rows * positions : U32)),positions) 186 187 # The deepest split that leaves every leaf at least one eight-position block. 188 def depth_limit(+padded: U32) -> U32: 189 (U32.from_nat(Tensors.significant_bits(32n,(padded / 8 : U32),0n,U32.is_zero((padded / 8 : U32)))) - 1 : U32) 190 191 def run_partitioned(~custom: F32 -> F32,activation: Activations.Activation,depth: U32,work: Work,input: Storage.Buffer,weights: Storage.Buffer,bias: Storage.Buffer) -> Storage.Buffer: 192 Work{+rows,+h,w,k,s,p,ow,+positions,+inner} = work 193 +padded = KernelWindow.padded(positions) 194 +limit = depth_limit(padded) 195 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) 196 197 def run_selected(~custom: F32 -> F32,blocked_leaf: Bool,activation: Activations.Activation,depth: U32,work: Work,input: Storage.Buffer,weights: Storage.Buffer,bias: Storage.Buffer) -> Storage.Buffer: 198 match blocked_leaf: 199 case True{}: blocked(~custom,activation,work,0,input,weights,bias) 200 case False{}: run_partitioned(~custom,activation,depth,work,input,weights,bias) 201 202 def run_packed(~custom: F32 -> F32,activation: Activations.Activation,+depth: U32,work: Work,input: Storage.Buffer,weights: Storage.Buffer,bias: Storage.Buffer) -> Storage.Buffer: 203 run_selected(~custom,U32.is_zero(depth),activation,depth,work,input,weights,bias) 204 205 def run(~custom: F32 -> F32,activation: Activations.Activation,depth: U32,+work: Work,input: Storage.Buffer,weights: Storage.Buffer,bias: Storage.Buffer) -> Storage.Buffer: 206 Work{rows,h,w,k,s,p,ow,positions,inner} = work 207 run_packed(~custom,activation,depth,work,input,GemmBuffer.reorder(rows,inner,weights),bias)