proofs/kernel_read_witness.bend
stelliferous@0.0.3.0 · 55 lines · 0 definitions
1import Base
2import stelliferous@0.0.3.0/kernel_gemm.bend as Gemm
3import stelliferous@0.0.3.0/kernel_rows.bend as Rows
4import stelliferous@0.0.3.0/kernel_tile.bend as Tile
5import stelliferous@0.0.3.0/traversal_loop.bend as Loop
6import stelliferous@0.0.3.0/proofs/storage_witness.bend as Witness
7
8# The k loop returns every owner it received. Index, context and the erased
9# accumulators are whatever the steps produced.
10type Looped<-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-output: Array<F32>,-state: Gemm.Loop> is Data:
11 Looped{index: U32,context: Gemm.Context,-accumulators: Rows.Rows<Tile.Tile8>,
12 equation: {state == Gemm.Loop{input,weights,bias,output,index,context,accumulators} : Gemm.Loop}}
13
14def after_input(-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-output: Array<F32>,+index: U32,+context: Gemm.Context,
15 -accumulators: Rows.Rows<Tile.Tile8>,-samples: Rows.Rows<F32>,+address: U32,read: Witness.Read<Tile.Tile8,input,Array.get(Tile.Tile8,input,address)>) ->
16 Looped<input,weights,bias,output,Gemm.with_input(weights,bias,output,index,context,accumulators,samples,Array.get(Tile.Tile8,input,address))>:
17 Witness.Read{value,equation} = read
18 %Equal.sym(Array<Tile.Tile8> & Tile.Tile8,Array.get(Tile.Tile8,input,address),(input,value),equation) :
19 Looped<input,weights,bias,output,Gemm.with_input(weights,bias,output,index,context,accumulators,samples,_)>
20 Looped{(index + 1 : U32),context,Gemm.update(accumulators,value,samples),{==}}
21
22def after_weights(~upper: Bool,-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-output: Array<F32>,+index: U32,+context: Gemm.Context,
23 +input_base: U32,-accumulators: Rows.Rows<Tile.Tile8>,+address: U32,input_evidence: Witness.ArrayWitness<Tile.Tile8,input>,
24 read: Witness.Read<Tile.Tile8,weights,Array.get(Tile.Tile8,weights,address)>) ->
25 Looped<input,weights,bias,output,Gemm.with_weights(~upper,input,bias,output,index,context,input_base,accumulators,Array.get(Tile.Tile8,weights,address))>:
26 Witness.Read{value,equation} = read
27 %Equal.sym(Array<Tile.Tile8> & Tile.Tile8,Array.get(Tile.Tile8,weights,address),(weights,value),equation) :
28 Looped<input,weights,bias,output,Gemm.with_weights(~upper,input,bias,output,index,context,input_base,accumulators,_)>
29 after_input(input,weights,bias,output,index,context,accumulators,Rows.half(upper,value),(input_base + index : U32),
30 Witness.observed(Tile.Tile8,input,(input_base + index : U32),input_evidence))
31
32def step_at(~upper: Bool,-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-output: Array<F32>,+index: U32,context: Gemm.Context,
33 -accumulators: Rows.Rows<Tile.Tile8>,+input_evidence: Witness.ArrayWitness<Tile.Tile8,input>,+weight_evidence: Witness.ArrayWitness<Tile.Tile8,weights>) ->
34 Looped<input,weights,bias,output,Gemm.step(~upper,Gemm.Loop{input,weights,bias,output,index,context,accumulators})>:
35 Gemm.Context{+row,+position,+rows,+columns,+input_base,+weight_base} = context
36 after_weights(~upper,input,weights,bias,output,index,Gemm.Context{row,position,rows,columns,input_base,weight_base},input_base,accumulators,
37 (weight_base + index : U32),input_evidence,Witness.observed(Tile.Tile8,weights,(weight_base + index : U32),weight_evidence))
38
39def step(~upper: Bool,-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-output: Array<F32>,-state: Gemm.Loop,
40 +input_evidence: Witness.ArrayWitness<Tile.Tile8,input>,+weight_evidence: Witness.ArrayWitness<Tile.Tile8,weights>,
41 evidence: Looped<input,weights,bias,output,state>) -> Looped<input,weights,bias,output,Gemm.step(~upper,state)>:
42 Looped{+index,+context,accumulators,equation} = evidence
43 %Equal.sym(Gemm.Loop,state,Gemm.Loop{input,weights,bias,output,index,context,accumulators},equation) :
44 Looped<input,weights,bias,output,Gemm.step(~upper,_)>
45 step_at(~upper,input,weights,bias,output,index,context,accumulators,input_evidence,weight_evidence)
46
47def run(~upper: Bool,count: Nat,-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-output: Array<F32>,-state: Gemm.Loop,
48 +input_evidence: Witness.ArrayWitness<Tile.Tile8,input>,+weight_evidence: Witness.ArrayWitness<Tile.Tile8,weights>,
49 evidence: Looped<input,weights,bias,output,state>) ->
50 Looped<input,weights,bias,output,Loop.run(~Gemm.Loop,~Gemm.step(~upper),count,state)>:
51 match count:
52 case 0n: evidence
53 case 1n+rest:
54 run(~upper,rest,input,weights,bias,output,Gemm.step(~upper,state),input_evidence,weight_evidence,
55 step(~upper,input,weights,bias,output,state,input_evidence,weight_evidence,evidence))