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(~True{},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(~True{},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)}
43	
44	# Mapping runs once for each packed tile; both affine owners return.
45	def mapped_buffers(+input_length: Nat,+output_length: Nat,result: Array<F32> & Array<Tile.Tile8>,
46	  certificates: WeightCertificate.MappedCertificates<result>) -> Storage.Buffer<F32> & Storage.Buffer<Tile.Tile8>:
47	  (input,output) = result
48	  WeightCertificate.MappedCertificates{input_certificate,output_certificate} = certificates
49	  (Storage.Buffer{input,input_length,input_certificate},Storage.Buffer{output,output_length,output_certificate})
50	
51	def reorder_mapped(~Context: Data,~offset: Context -> U32 -> U32,+rows: U32,+inner: U32,+row_stride: U32,+context: Context,weights: Storage.Buffer<F32>) -> Storage.Buffer<F32> & Storage.Buffer<Tile.Tile8>:
52	  Storage.Buffer{weights,+length,certificate} = weights
53	  mapped_buffers(length,U32.to_nat((((rows + 7) / 8) * inner : U32)),Weights.run_mapped(~Context,~offset,rows,inner,row_stride,context,weights),
54	    WeightCertificate.certificates_mapped(~Context,~offset,rows,inner,row_stride,context,weights,certificate))
55	
56	
57	def reorder_mapped_with(~Context: Data,~offset: Context -> U32 -> U32,~row_bases: Context -> @+row_stride: U32 -> @+row: U32 -> Weights.Bases,+rows: U32,+inner: U32,+row_stride: U32,+context: Context,weights: Storage.Buffer<F32>) -> Storage.Buffer<F32> & Storage.Buffer<Tile.Tile8>:
58	  Storage.Buffer{weights,+length,certificate} = weights
59	  mapped_buffers(length,U32.to_nat((((rows + 7) / 8) * inner : U32)),Weights.run_mapped_with(~Context,~offset,~row_bases,rows,inner,row_stride,context,weights),
60	    WeightCertificate.certificates_mapped_with(~Context,~offset,~row_bases,rows,inner,row_stride,context,weights,certificate))