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.
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 ./storage_buffer.bend as Storage
9	import ./tensor.bend as Tensors
10	import ./kernel_window.bend as KernelWindow
11	import ./traversal_array.bend as Traversal
12	import ./math_activation.bend as Activations
13	import ./proofs/traversal_certificate.bend as TraversalCertificate
14	import ./tensor_operations.bend as Operations
15	import ./tensor_gelu.bend as Gelu
16	
17	type 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.
22	type WindowWork is Data:
23	  WindowWork{rows: U32,positions: U32,window: Packing.Stencil}
24	
25	def 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	
29	def activated_value(~activation: F32 -> F32,context: Unit,value: F32) -> F32:
30	  activation(value)
31	
32	def 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	
37	def leaky_relu_value(slope: F32,value: F32) -> F32:
38	  Activations.leaky_relu_value(value,slope)
39	
40	def unit_mapped(~operation: F32 -> F32,+count: U32,buffer: Storage.Buffer<F32>) -> Storage.Buffer<F32>:
41	  mapped(~Unit,~activated_value(~operation),count,Unit{},buffer)
42	
43	def 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.
50	def 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.
65	def block_tiles() -> U32:
66	  8192
67	
68	def block_positions(inner: U32) -> U32:
69	  (8 * U32.max(1,(block_tiles() / inner : U32)) : U32)
70	
71	def workspace_prefix(+target: U32,+inner: U32,positions: U32) -> U32:
72	  U32.min((8 * U32.max(1,(target / inner : U32)) : U32),KernelWindow.padded(positions))
73	
74	def workspace_tiles(+target: U32,+inner: U32,+positions: U32) -> U32:
75	  (workspace_prefix(target,inner,positions) / 8 * inner : U32)
76	
77	type 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	
80	type Kept is Type:
81	  Kept{input: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,output: Storage.Buffer<F32>}
82	
83	def kept_output(state: Kept) -> Storage.Buffer<F32>:
84	  Kept{input,weights,bias,output} = state
85	  output
86	
87	def kept_from_blocks(state: BlockState) -> Kept:
88	  BlockState{input,tiles,weights,bias,output} = state
89	  Kept{input,weights,bias,output}
90	
91	def 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	
95	def 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	
99	def 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.
107	def 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	
114	def 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	
119	def 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.
128	def 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	
134	def 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	
140	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>:
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.
144	def 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)