1	import Base
2	import ./kernel_tile.bend as Tile
3	import ./kernel_gemm.bend as Gemm
4	import ./storage_buffer.bend as Storage
5	import ./proofs/kernel_gemm_certificate.bend as GemmCertificate
6	import ./kernel_weights.bend as Weights
7	import ./proofs/kernel_weights_certificate.bend as WeightCertificate
8	
9	# The live certificate is updated once, outside every output and MAC loop.
10	def run(+rows: U32,input: Storage.Buffer<Tile.Tile8>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,
11	  output: Storage.Buffer<F32>,+columns: U32,+inner: U32) -> Storage.Buffer<F32>:
12	  Storage.Buffer{input,input_length,input_certificate} = input
13	  Storage.Buffer{weights,weight_length,weight_certificate} = weights
14	  Storage.Buffer{bias,bias_length,bias_certificate} = bias
15	  Storage.Buffer{output,length,certificate} = output
16	  Storage.Buffer{Gemm.run(rows,input,weights,bias,output,columns,inner),length,
17	    GemmCertificate.certificate(rows,input,weights,bias,output,columns,inner,0,columns,input_certificate,weight_certificate,bias_certificate,certificate)}
18	
19	# The tiles, the weights and the bias come back with the output, for a caller
20	# that uses them again.
21	type Kept is Type:
22	  Kept{input: Storage.Buffer<Tile.Tile8>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,output: Storage.Buffer<F32>}
23	
24	def kept(state: Gemm.Buffers,input_length: Nat,weight_length: Nat,bias_length: Nat,length: Nat,certificates: GemmCertificate.Kept<state>) -> Kept:
25	  Gemm.Buffers{input,weights,bias,output} = state
26	  GemmCertificate.Kept{input_certificate,weight_certificate,bias_certificate,certificate} = certificates
27	  Kept{Storage.Buffer{input,input_length,input_certificate},Storage.Buffer{weights,weight_length,weight_certificate},
28	    Storage.Buffer{bias,bias_length,bias_certificate},Storage.Buffer{output,length,certificate}}
29	
30	def run_keeping(+rows: U32,input: Storage.Buffer<Tile.Tile8>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,
31	  output: Storage.Buffer<F32>,+columns: U32,+inner: U32,+start: U32,+output_columns: U32) -> Kept:
32	  Storage.Buffer{input,input_length,input_certificate} = input
33	  Storage.Buffer{weights,weight_length,weight_certificate} = weights
34	  Storage.Buffer{bias,bias_length,bias_certificate} = bias
35	  Storage.Buffer{output,length,certificate} = output
36	  kept(Gemm.run_keeping(rows,input,weights,bias,output,columns,inner,start,output_columns),input_length,weight_length,bias_length,length,
37	    GemmCertificate.kept(rows,input,weights,bias,output,columns,inner,start,output_columns,input_certificate,weight_certificate,bias_certificate,certificate))
38	
39	def reorder(+rows: U32,+inner: U32,weights: Storage.Buffer<F32>) -> Storage.Buffer<Tile.Tile8>:
40	  Storage.Buffer{weights,length,certificate} = weights
41	  Storage.Buffer{Weights.run(rows,inner,weights),U32.to_nat((((rows + 7) / 8) * inner : U32)),
42	    WeightCertificate.certificate(rows,inner,weights,certificate)}