kernel_gemm_buffer.bend
stelliferous@0.0.3.0 · 60 lines · 0 definitions
1import Base
2import stelliferous@0.0.3.0/kernel_tile.bend as Tile
3import stelliferous@0.0.3.0/kernel_gemm.bend as Gemm
4import stelliferous@0.0.3.0/storage_buffer.bend as Storage
5import stelliferous@0.0.3.0/proofs/kernel_gemm_certificate.bend as GemmCertificate
6import stelliferous@0.0.3.0/kernel_weights.bend as Weights
7import stelliferous@0.0.3.0/proofs/kernel_weights_certificate.bend as WeightCertificate
8
9# The live certificate is updated once, outside every output and MAC loop.
10def 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.
21type Kept is Type:
22 Kept{input: Storage.Buffer<Tile.Tile8>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,output: Storage.Buffer<F32>}
23
24def 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
30def 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
39def 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.
45def 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
51def 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
57def 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))