proofs/kernel_strided_store_certificate.bend
stelliferous@0.0.3.0 · 57 lines · 0 definitions
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.
4import Base
5import stelliferous@0.0.3.0/kernel_tile.bend as Tile
6import stelliferous@0.0.3.0/kernel_rows.bend as Rows
7import stelliferous@0.0.3.0/kernel_gemm.bend as Gemm
8import stelliferous@0.0.3.0/kernel_strided_store.bend as Stores
9import stelliferous@0.0.3.0/proofs/storage_witness.bend as Witness
10import stelliferous@0.0.3.0/proofs/storage_certificate.bend as Certificate
11import stelliferous@0.0.3.0/proofs/storage_tree_model.bend as Tree
12import stelliferous@0.0.3.0/proofs/kernel_gemm_certificate.bend as GemmCertificate
13
14def 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
35def 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
46def 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
54def 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)}