1	# A block store keeps the storage shape, whatever values it writes. Only the
2	# final flat certificate is live; the witness follows the writes beneath an
3	# erased rewrite.
4	import Base
5	import ../kernel_tile.bend as Tile
6	import ../kernel_rows.bend as Rows
7	import ../kernel_gemm.bend as Gemm
8	import ../kernel_strided_store.bend as Stores
9	import ./storage_witness.bend as Witness
10	import ./storage_certificate.bend as Certificate
11	import ./storage_tree_model.bend as Tree
12	import ./kernel_gemm_certificate.bend as GemmCertificate
13	
14	def row(+depth: Nat,valid: Bool,-output: Array<F32>,+target: U32,+step: U32,+count: U32,-accumulator: Tile.Tile8,
15	  evidence: Witness.Witness(F32,depth,output)) -> Witness.Witness(F32,depth,Stores.row(valid,output,target,step,count,accumulator)):
16	  match valid:
17	    case False{}: evidence
18	    case True{}:
19	      -o1 = Gemm.store_lane(U32.is_lt(0,count),output,target,(Tile.lane0(accumulator) + 0.0 : F32))
20	      -o2 = Gemm.store_lane(U32.is_lt(1,count),o1,(target + step : U32),(Tile.lane1(accumulator) + 0.0 : F32))
21	      -o3 = Gemm.store_lane(U32.is_lt(2,count),o2,(target + 2 * step : U32),(Tile.lane2(accumulator) + 0.0 : F32))
22	      -o4 = Gemm.store_lane(U32.is_lt(3,count),o3,(target + 3 * step : U32),(Tile.lane3(accumulator) + 0.0 : F32))
23	      -o5 = Gemm.store_lane(U32.is_lt(4,count),o4,(target + 4 * step : U32),(Tile.lane4(accumulator) + 0.0 : F32))
24	      -o6 = Gemm.store_lane(U32.is_lt(5,count),o5,(target + 5 * step : U32),(Tile.lane5(accumulator) + 0.0 : F32))
25	      -o7 = Gemm.store_lane(U32.is_lt(6,count),o6,(target + 6 * step : U32),(Tile.lane6(accumulator) + 0.0 : F32))
26	      evidence = GemmCertificate.store_lane(depth,U32.is_lt(0,count),output,target,(Tile.lane0(accumulator) + 0.0 : F32),evidence)
27	      evidence = GemmCertificate.store_lane(depth,U32.is_lt(1,count),o1,(target + step : U32),(Tile.lane1(accumulator) + 0.0 : F32),evidence)
28	      evidence = GemmCertificate.store_lane(depth,U32.is_lt(2,count),o2,(target + 2 * step : U32),(Tile.lane2(accumulator) + 0.0 : F32),evidence)
29	      evidence = GemmCertificate.store_lane(depth,U32.is_lt(3,count),o3,(target + 3 * step : U32),(Tile.lane3(accumulator) + 0.0 : F32),evidence)
30	      evidence = GemmCertificate.store_lane(depth,U32.is_lt(4,count),o4,(target + 4 * step : U32),(Tile.lane4(accumulator) + 0.0 : F32),evidence)
31	      evidence = GemmCertificate.store_lane(depth,U32.is_lt(5,count),o5,(target + 5 * step : U32),(Tile.lane5(accumulator) + 0.0 : F32),evidence)
32	      evidence = GemmCertificate.store_lane(depth,U32.is_lt(6,count),o6,(target + 6 * step : U32),(Tile.lane6(accumulator) + 0.0 : F32),evidence)
33	      GemmCertificate.store_lane(depth,U32.is_lt(7,count),o7,(target + 7 * step : U32),(Tile.lane7(accumulator) + 0.0 : F32),evidence)
34	
35	def block(+depth: Nat,+store: Stores.Store,-accumulators: Rows.Rows<Tile.Tile8>,-output: Array<F32>,
36	  evidence: Witness.Witness(F32,depth,output)) -> Witness.Witness(F32,depth,Stores.block(store,accumulators,output)):
37	  Stores.Store{+target,+step,+plane,+rows,+count} = store
38	  -o1 = Stores.row(U32.is_lt(0,rows),output,target,step,count,Rows.first(~Tile.Tile8,accumulators))
39	  -o2 = Stores.row(U32.is_lt(1,rows),o1,(target + plane : U32),step,count,Rows.second(~Tile.Tile8,accumulators))
40	  -o3 = Stores.row(U32.is_lt(2,rows),o2,(target + 2 * plane : U32),step,count,Rows.third(~Tile.Tile8,accumulators))
41	  evidence = row(depth,U32.is_lt(0,rows),output,target,step,count,Rows.first(~Tile.Tile8,accumulators),evidence)
42	  evidence = row(depth,U32.is_lt(1,rows),o1,(target + plane : U32),step,count,Rows.second(~Tile.Tile8,accumulators),evidence)
43	  evidence = row(depth,U32.is_lt(2,rows),o2,(target + 2 * plane : U32),step,count,Rows.third(~Tile.Tile8,accumulators),evidence)
44	  row(depth,U32.is_lt(3,rows),o3,(target + 3 * plane : U32),step,count,Rows.fourth(~Tile.Tile8,accumulators),evidence)
45	
46	def equation(+depth: Nat,+store: Stores.Store,-accumulators: Rows.Rows<Tile.Tile8>,-output: Array<F32>,
47	  balanced: {output == Tree.pack(F32,Certificate.normalized(F32,depth,Tree.reflect(F32,output))) : Array<F32>}) ->
48	  {Stores.block(store,accumulators,output) == Tree.pack(F32,Certificate.normalized(F32,depth,Tree.reflect(F32,Stores.block(store,accumulators,output)))) : Array<F32>}:
49	  %Certificate.restored(F32,depth,Stores.block(store,accumulators,output),
50	    block(depth,store,accumulators,output,Certificate.witness_at(F32,depth,output,balanced))) :
51	    {_ == Tree.pack(F32,Certificate.normalized(F32,depth,Tree.reflect(F32,Stores.block(store,accumulators,output)))) : Array<F32>}
52	  {==}
53	
54	def certificate(+store: Stores.Store,-accumulators: Rows.Rows<Tile.Tile8>,-output: Array<F32>,certificate: Certificate.Certificate<F32,output>) ->
55	  Certificate.Certificate<F32,Stores.block(store,accumulators,output)>:
56	  Certificate.Certificate{+depth,balanced} = certificate
57	  Certificate.Certificate{depth,equation(depth,store,accumulators,output,balanced)}