kernel_pipeline.bend
stelliferous@0.0.3.0 · 145 lines · 0 definitions
1# One convolution leaf owns its input, packs it block by block, and computes
2# before returning. The caller validates dimensions and chooses the parts.
3import Base
4import stelliferous@0.0.3.0/kernel_gemm_buffer.bend as GemmBuffer
5import stelliferous@0.0.3.0/kernel_tile.bend as Tile
6import stelliferous@0.0.3.0/kernel_spatial_packing.bend as Packing
7import stelliferous@0.0.3.0/kernel_packing_buffer.bend as PackingBuffer
8import stelliferous@0.0.3.0/storage_buffer.bend as Storage
9import stelliferous@0.0.3.0/tensor.bend as Tensors
10import stelliferous@0.0.3.0/kernel_window.bend as KernelWindow
11import stelliferous@0.0.3.0/traversal_array.bend as Traversal
12import stelliferous@0.0.3.0/math_activation.bend as Activations
13import stelliferous@0.0.3.0/proofs/traversal_certificate.bend as TraversalCertificate
14import stelliferous@0.0.3.0/tensor_operations.bend as Operations
15import stelliferous@0.0.3.0/tensor_gelu.bend as Gelu
16
17type Work is Data:
18 Work{rows: U32,height: U32,width: U32,kernel: U32,stride: U32,padding: U32,output_width: U32,positions: U32,inner: U32}
19
20# The matrix and its source stencil are independent. Forward convolution is
21# the square specialization; an adjoint phase uses rectangular taps and origins.
22type WindowWork is Data:
23 WindowWork{rows: U32,positions: U32,window: Packing.Stencil}
24
25def square_work(work: Work) -> WindowWork:
26 Work{rows,h,w,k,s,p,ow,positions,inner} = work
27 WindowWork{rows,positions,Packing.square(h,w,k,s,p,ow,inner)}
28
29def activated_value(~activation: F32 -> F32,context: Unit,value: F32) -> F32:
30 activation(value)
31
32def mapped(~Context: Data,~operation: Context -> F32 -> F32,+count: U32,+context: Context,buffer: Storage.Buffer<F32>) -> Storage.Buffer<F32>:
33 Storage.Buffer{storage,length,certificate} = buffer
34 Storage.Buffer{Traversal.map_in_place(~F32,~Context,~operation,U32.to_nat(count),0,context,storage),length,
35 TraversalCertificate.map_in_place(~F32,~Context,~operation,U32.to_nat(count),0,context,storage,certificate)}
36
37def leaky_relu_value(slope: F32,value: F32) -> F32:
38 Activations.leaky_relu_value(value,slope)
39
40def unit_mapped(~operation: F32 -> F32,+count: U32,buffer: Storage.Buffer<F32>) -> Storage.Buffer<F32>:
41 mapped(~Unit,~activated_value(~operation),count,Unit{},buffer)
42
43def unit_lanes(~operation: F32 -> F32,+count: U32,buffer: Storage.Buffer<F32>) -> Storage.Buffer<F32>:
44 Operations.lanes(~Unit,~activated_value(~operation),U32.to_nat(count),0,Unit{},buffer)
45
46# One tight pass over each finished leaf while it is still in cache; applying
47# the activation inside the store measured about twice as slow. The built-in
48# is chosen once per leaf, so each pass is specialized; the identity skips it.
49# Sigmoid, SiLU, tanh and GELU pass their values through tiles.
50def activated(~custom: F32 -> F32,activation: Activations.Activation,+count: U32,buffer: Storage.Buffer<F32>) -> Storage.Buffer<F32>:
51 match activation:
52 case Activations.Identity{}: buffer
53 case Activations.ReLU{}: unit_mapped(~Activations.relu_value,count,buffer)
54 case Activations.ReLU6{}: unit_mapped(~Activations.relu6_value,count,buffer)
55 case Activations.LeakyReLU{slope}: mapped(~F32,~leaky_relu_value,count,slope,buffer)
56 case Activations.Sigmoid{}: unit_lanes(~Activations.sigmoid_value,count,buffer)
57 case Activations.SiLU{}: unit_lanes(~Activations.silu_value,count,buffer)
58 case Activations.Tanh{}: unit_lanes(~Activations.tanh_value,count,buffer)
59 case Activations.HardSigmoid{}: unit_mapped(~Activations.hard_sigmoid_value,count,buffer)
60 case Activations.HardSwish{}: unit_mapped(~Activations.hard_swish_value,count,buffer)
61 case Activations.GELUTanh{}: Gelu.mapped(U32.to_nat(count),0,buffer)
62 case Activations.Custom{operations}: unit_mapped(~custom,count,buffer)
63
64# One tile is 32 bytes; 8192 tiles are 256 KB.
65def block_tiles() -> U32:
66 8192
67
68def block_positions(inner: U32) -> U32:
69 (8 * U32.max(1,(block_tiles() / inner : U32)) : U32)
70
71def workspace_prefix(+target: U32,+inner: U32,positions: U32) -> U32:
72 U32.min((8 * U32.max(1,(target / inner : U32)) : U32),KernelWindow.padded(positions))
73
74def workspace_tiles(+target: U32,+inner: U32,+positions: U32) -> U32:
75 (workspace_prefix(target,inner,positions) / 8 * inner : U32)
76
77type BlockState is Type:
78 BlockState{input: Storage.Buffer<F32>,tiles: Storage.Buffer<Tile.Tile8>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,output: Storage.Buffer<F32>}
79
80type Kept is Type:
81 Kept{input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,output: Storage.Buffer<F32>}
82
83def kept_output(state: Kept) -> Storage.Buffer<F32>:
84 Kept{input,weights,bias,output} = state
85 output
86
87def kept_from_blocks(state: BlockState) -> Kept:
88 BlockState{input,tiles,weights,bias,output} = state
89 Kept{input,weights,bias,output}
90
91def kept_finished(~custom: F32 -> F32,activation: Activations.Activation,count: U32,state: Kept) -> Kept:
92 Kept{input,weights,bias,output} = state
93 Kept{input,weights,bias,activated(~custom,activation,count,output)}
94
95def block_kept(input: Storage.Buffer<F32>,kept: GemmBuffer.Kept) -> BlockState:
96 GemmBuffer.Kept{tiles,weights,bias,output} = kept
97 BlockState{input,tiles,weights,bias,output}
98
99def block_packed(+rows: U32,+positions: U32,inner: U32,+start: U32,+span: U32,
100 weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,output: Storage.Buffer<F32>,packed: Storage.Buffer<F32> & Storage.Buffer<Tile.Tile8>) -> BlockState:
101 (input,tiles) = packed
102 block_kept(input,
103 GemmBuffer.run_keeping(rows,tiles,weights,bias,output,span,inner,start,positions))
104
105# Packing addresses the input window; output positions are local to the leaf.
106# Each block overwrites tiles from zero and writes directly into the leaf output.
107def block_step_with(~Context: Data,
108 ~pack: @+context: Context -> @+count: Nat -> @+start: U32 -> Storage.Buffer<F32> -> Storage.Buffer<Tile.Tile8> -> Storage.Buffer<F32> & Storage.Buffer<Tile.Tile8>,
109 rows: U32,positions: U32,+inner: U32,context: Context,pack_start: U32,output_start: U32,+span: U32,state: BlockState) -> BlockState:
110 BlockState{input,tiles,weights,bias,output} = state
111 block_packed(rows,positions,inner,output_start,span,weights,bias,output,
112 pack(context,U32.to_nat((span / 8 * inner : U32)),pack_start,input,tiles))
113
114def block_step_window(work: WindowWork,pack_start: U32,output_start: U32,span: U32,state: BlockState) -> BlockState:
115 WindowWork{rows,positions,+window} = work
116 Packing.Stencil{base,h,w,kr,kc,s,ro,co,ow,inner} = window
117 block_step_with(~Packing.Stencil,~PackingBuffer.run_window,rows,positions,inner,window,pack_start,output_start,span,state)
118
119def blocks_window(remaining: Nat,+work: WindowWork,+pack_start: U32,+output_start: U32,+padded: U32,+span: U32,state: BlockState) -> BlockState:
120 match remaining:
121 case 0n: state
122 case 1n+rest:
123 blocks_window(rest,work,(pack_start + span : U32),(output_start + span : U32),padded,span,
124 block_step_window(work,pack_start,output_start,U32.min(span,(padded - output_start : U32)),state))
125
126# Exactly one reusable tile buffer and one output buffer per leaf.
127# An inner larger than the target still needs one complete eight-position group.
128def blocked_window_with(+work: WindowWork,pack_start: U32,state: BlockState) -> BlockState:
129 WindowWork{rows,+positions,Packing.Stencil{base,h,w,kr,kc,s,ro,co,ow,+inner}} = work
130 +padded = KernelWindow.padded(positions)
131 +span = block_positions(inner)
132 blocks_window(U32.to_nat(((padded + span - 1) / span : U32)),work,pack_start,0,padded,span,state)
133
134def blocked_window_keeping(~custom: F32 -> F32,+activation: Activations.Activation,+work: WindowWork,pack_start: U32,input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>) -> Kept:
135 WindowWork{+rows,+positions,Packing.Stencil{base,h,w,kr,kc,s,ro,co,ow,+inner}} = work
136 kept_finished(~custom,activation,(rows * positions : U32),kept_from_blocks(
137 blocked_window_with(work,pack_start,
138 BlockState{input,Tensors.buffer(~Tile.Tile8,Tile.filled(0.0),workspace_tiles(block_tiles(),inner,positions)),weights,bias,Tensors.buffer(~F32,0.0,(rows * positions : U32))})))
139
140def 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>:
141 kept_output(blocked_window_keeping(~custom,activation,square_work(work),pack_start,input,weights,bias))
142
143# The deepest split that leaves every leaf at least one eight-position block.
144def depth_limit(+padded: U32) -> U32:
145 (U32.from_nat(Tensors.significant_bits(32n,(padded / 8 : U32),0n,U32.is_zero((padded / 8 : U32)))) - 1 : U32)