1	import Base
2	import ../kernel_gemm.bend as Gemm
3	import ../kernel_rows.bend as Rows
4	import ../kernel_tile.bend as Tile
5	import ../traversal_loop.bend as Loop
6	import ./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.
10	type 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	
14	def 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	
22	def 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	
32	def 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	
39	def 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	
47	def 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))