kernel_pipeline.bend

stelliferous@0.0.2.0 · 207 lines · 36 definitions

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.
3import Base
4import stelliferous@0.0.2.0/kernel_gemm_buffer.bend as GemmBuffer
5import stelliferous@0.0.2.0/kernel_tile.bend as Tile
6import stelliferous@0.0.2.0/kernel_spatial_packing.bend as Packing
7import stelliferous@0.0.2.0/kernel_packing_buffer.bend as PackingBuffer
8import stelliferous@0.0.2.0/kernel_gather_buffer.bend as GatherBuffer
9import stelliferous@0.0.2.0/storage_buffer.bend as Storage
10import stelliferous@0.0.2.0/storage_clone.bend as Clone
11import stelliferous@0.0.2.0/tensor.bend as Tensors
12import stelliferous@0.0.2.0/traversal_partition.bend as Partition
13import stelliferous@0.0.2.0/storage_copy.bend as Copy
14import stelliferous@0.0.2.0/kernel_window.bend as KernelWindow
15import stelliferous@0.0.2.0/traversal_array.bend as Traversal
16import stelliferous@0.0.2.0/math_activation.bend as Activations
17import stelliferous@0.0.2.0/proofs/traversal_certificate.bend as TraversalCertificate
18import stelliferous@0.0.2.0/tensor_operations.bend as Operations
19
20type 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.
24type Task is Data:
25  Task{work: Work,activation: Activations.Activation}
26
27def Outputs() -> Type:
28  Partition.Tree<&1,Storage.Buffer<F32>,U32>
29
30def activated_value(~activation: F32 -> F32,context: Unit,value: F32) -> F32:
31  activation(value)
32
33def mapped(~Context: Data,~operation: Context -> F32 -> F32,+count: U32,+context: Context,buffer: Storage.Buffer<F32>) -> Storage.Buffer<F32>:
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
38def leaky_relu_value(slope: F32,value: F32) -> F32:
39  Activations.leaky_relu_value(value,slope)
40
41def unit_mapped(~operation: F32 -> F32,+count: U32,buffer: Storage.Buffer<F32>) -> Storage.Buffer<F32>:
42  mapped(~Unit,~activated_value(~operation),count,Unit{},buffer)
43
44def unit_lanes(~operation: F32 -> F32,+count: U32,buffer: Storage.Buffer<F32>) -> Storage.Buffer<F32>:
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.
51def activated(~custom: F32 -> F32,activation: Activations.Activation,+count: U32,buffer: Storage.Buffer<F32>) -> Storage.Buffer<F32>:
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.
66type InputWindow is Type:
67  InputWindow{buffer: Storage.Buffer<F32>,first: U32,rows: U32}
68
69def copy_rows(remaining: Nat,+channel: U32,+count: U32,+source_stride: U32,+source_start: U32,+target_stride: U32,
70  owners: Storage.Buffer<F32> & Storage.Buffer<F32>) -> Storage.Buffer<F32> & Storage.Buffer<F32>:
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
78def windows(+first: U32,+rows: U32,+left_first: U32,+left_rows: U32,owners: Storage.Buffer<F32> & Storage.Buffer<F32>) -> InputWindow & InputWindow:
79  (source,target) = owners
80  (InputWindow{target,left_first,left_rows},InputWindow{source,first,rows})
81
82def cloned(+first: U32,+rows: U32,owners: Storage.Buffer<F32> & Storage.Buffer<F32>) -> InputWindow & InputWindow:
83  (left,right) = owners
84  (InputWindow{left,first,rows},InputWindow{right,first,rows})
85
86def copied(+channels: U32,+width: U32,+first: U32,+rows: U32,band: KernelWindow.Band,buffer: Storage.Buffer<F32>) -> 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
91def split_when(copy: Bool,+channels: U32,+width: U32,+first: U32,+rows: U32,+band: KernelWindow.Band,buffer: Storage.Buffer<F32>) -> 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.
99def 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
108def 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.
113def block_tiles() -> U32:
114  8192
115
116def block_positions(inner: U32) -> U32:
117  (8 * U32.max(1,(block_tiles() / inner : U32)) : U32)
118
119type BlockState is Type:
120  BlockState{input: Storage.Buffer<F32>,tiles: Storage.Buffer<Tile.Tile8>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,output: Storage.Buffer<F32>}
121
122def block_kept(input: Storage.Buffer<F32>,kept: GemmBuffer.Kept) -> BlockState:
123  GemmBuffer.Kept{tiles,weights,bias,output} = kept
124  BlockState{input,tiles,weights,bias,output}
125
126def block_packed(+rows: U32,+positions: U32,inner: U32,+start: U32,+span: U32,
127  weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,output: Storage.Buffer<F32>,packed: Storage.Buffer<F32> & Storage.Buffer<Tile.Tile8>) -> 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.
134def 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
140def 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
147def blocks_finished(~custom: F32 -> F32,activation: Activations.Activation,count: U32,state: BlockState) -> Storage.Buffer<F32>:
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.
153def 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>:
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
161def leaf(~custom: F32 -> F32,task: Task,offset: U32,span: U32,window: InputWindow,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>) -> 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
168def partition(~custom: F32 -> F32,depth: Nat,task: Task,offset: U32,span: U32,input: InputWindow,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>) -> Outputs():
169  Partition.run(~InputWindow,~Storage.Buffer<Tile.Tile8>,~Storage.Buffer<F32>,~Storage.Buffer<F32>,~Task,
170    ~split_task,~Clone.clone(~Tile.Tile8),~Clone.clone(~F32),~leaf(~custom),depth,task,offset,span,input,weights,bias)
171
172def gather(partitions: Outputs(),output: Storage.Buffer<F32>,+positions: U32) -> Storage.Buffer<F32>:
173  Partition.gather(~Storage.Buffer<F32>,~(rows => offset => positions => span => input => output => GatherBuffer.run(rows,offset,positions,span,input,output)),partitions,output,positions)
174
175def finish_leaf(identity: Bool,input: Storage.Buffer<F32>,+rows: U32,offset: U32,span: U32,+positions: U32) -> Storage.Buffer<F32>:
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
180def finish(partitions: Outputs(),+rows: U32,+positions: U32) -> Storage.Buffer<F32>:
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.
188def 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
191def 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>:
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
197def 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>:
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
202def 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>:
203  run_selected(~custom,U32.is_zero(depth),activation,depth,work,input,weights,bias)
204
205def 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>:
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)