proofs/kernel_gemm_certificate.bend source
proofs/kernel_gemm_certificate.bend on the hub · documented module
# Shape evidence follows the real readers and writers. Only the final flat# certificate is live; witness traversals occur beneath erased rewrites.import Baseimport ../kernel_gemm.bend as Gemmimport ../kernel_rows.bend as Rowsimport ../kernel_tile.bend as Tileimport ../traversal_array.bend as Traversalimport ../traversal_loop.bend as Loopimport ./storage_witness.bend as Witnessimport ./storage_certificate.bend as Certificateimport ./storage_tree_model.bend as Treeimport ./loop_projection.bend as Projectionimport ./kernel_read_witness.bend as ReadWitnessdef store_lane(+depth: Nat,inside: Bool,-output: Array<F32>,+offset: U32,-value: F32,evidence: Witness.Witness(F32,depth,output)) -> Witness.Witness(F32,depth,Gemm.store_lane(inside,output,offset,value)): match inside: case True{}: Witness.write(F32,depth,output,offset,value,evidence) case False{}: evidencedef store_tail(+count: U32,+depth: Nat,-output: Array<F32>,+offset: U32,-bias: F32,-accumulator: Tile.Tile8, evidence: Witness.Witness(F32,depth,output)) -> Witness.Witness(F32,depth,Gemm.store_tail(count,output,offset,bias,accumulator)): -a = (Tile.lane0(accumulator) + bias : F32) -b = (Tile.lane1(accumulator) + bias : F32) -c = (Tile.lane2(accumulator) + bias : F32) -d = (Tile.lane3(accumulator) + bias : F32) -e = (Tile.lane4(accumulator) + bias : F32) -f = (Tile.lane5(accumulator) + bias : F32) -g = (Tile.lane6(accumulator) + bias : F32) -h = (Tile.lane7(accumulator) + bias : F32) -o1 = Gemm.store_lane(U32.is_lt(0,count),output,offset,a) -o2 = Gemm.store_lane(U32.is_lt(1,count),o1,(offset + 1 : U32),b) -o3 = Gemm.store_lane(U32.is_lt(2,count),o2,(offset + 2 : U32),c) -o4 = Gemm.store_lane(U32.is_lt(3,count),o3,(offset + 3 : U32),d) -o5 = Gemm.store_lane(U32.is_lt(4,count),o4,(offset + 4 : U32),e) -o6 = Gemm.store_lane(U32.is_lt(5,count),o5,(offset + 5 : U32),f) -o7 = Gemm.store_lane(U32.is_lt(6,count),o6,(offset + 6 : U32),g) evidence = store_lane(depth,U32.is_lt(0,count),output,offset,a,evidence) evidence = store_lane(depth,U32.is_lt(1,count),o1,(offset + 1 : U32),b,evidence) evidence = store_lane(depth,U32.is_lt(2,count),o2,(offset + 2 : U32),c,evidence) evidence = store_lane(depth,U32.is_lt(3,count),o3,(offset + 3 : U32),d,evidence) evidence = store_lane(depth,U32.is_lt(4,count),o4,(offset + 4 : U32),e,evidence) evidence = store_lane(depth,U32.is_lt(5,count),o5,(offset + 5 : U32),f,evidence) evidence = store_lane(depth,U32.is_lt(6,count),o6,(offset + 6 : U32),g,evidence) store_lane(depth,U32.is_lt(7,count),o7,(offset + 7 : U32),h,evidence)type Written<-bias: Array<F32>,-depth: Nat,-result: Array<F32> & Array<F32>> is Data: Written{-output: Array<F32>,equation: {result == (bias,output) : Array<F32> & Array<F32>},balanced: Witness.Witness(F32,depth,output)}def write_bias(+depth: Nat,-bias: Array<F32>,-output: Array<F32>,+offset: U32,+row: U32,-accumulator: Tile.Tile8,+count: U32, read: Witness.Read<F32,bias,Array.get(F32,bias,row)>,evidence: Witness.Witness(F32,depth,output)) -> Written<bias,depth,Gemm.write_bias(output,offset,accumulator,count,Array.get(F32,bias,row))>: Witness.Read{value,equation} = read Written{Gemm.store_tail(count,output,offset,value,accumulator), %Equal.sym(Array<F32> & F32,Array.get(F32,bias,row),(bias,value),equation) : {Gemm.write_bias(output,offset,accumulator,count,_) == (bias,Gemm.store_tail(count,output,offset,value,accumulator)) : Array<F32> & Array<F32>} {==},store_tail(count,depth,output,offset,value,accumulator,evidence)}def write_row(+depth: Nat,valid: Bool,+row: U32,+position: U32,+columns: U32,-accumulator: Tile.Tile8, -bias: Array<F32>,-output: Array<F32>,bias_evidence: Witness.ArrayWitness<F32,bias>,evidence: Witness.Witness(F32,depth,output)) -> Written<bias,depth,Gemm.write_row(valid,row,position,columns,accumulator,(bias,output))>: match valid: case False{}: Written{output,{==},evidence} case True{}: write_bias(depth,bias,output,(row * columns + position : U32),row,accumulator,U32.min(8,(columns - position : U32)),Witness.observed(F32,bias,row,bias_evidence),evidence)def next(+depth: Nat,valid: Bool,row: U32,position: U32,columns: U32,-accumulator: Tile.Tile8, -bias: Array<F32>,-owners: Array<F32> & Array<F32>,bias_evidence: Witness.ArrayWitness<F32,bias>,evidence: Written<bias,depth,owners>) -> Written<bias,depth,Gemm.write_row(valid,row,position,columns,accumulator,owners)>: Written{output,equation,balanced} = evidence %Equal.sym(Array<F32> & Array<F32>,owners,(bias,output),equation) : Written<bias,depth,Gemm.write_row(valid,row,position,columns,accumulator,_)> write_row(depth,valid,row,position,columns,accumulator,bias,output,bias_evidence,balanced)def write(+depth: Nat,context: Gemm.Context,-accumulators: Rows.Rows<Tile.Tile8>, -bias: Array<F32>,-output: Array<F32>,+bias_evidence: Witness.ArrayWitness<F32,bias>,balanced: Witness.Witness(F32,depth,output)) -> Written<bias,depth,Gemm.write(context,accumulators,bias,output)>: Gemm.Context{+row,+position,+rows,+columns,input_base,weight_base} = context -a = Rows.first(~Tile.Tile8,accumulators) -b = Rows.second(~Tile.Tile8,accumulators) -c = Rows.third(~Tile.Tile8,accumulators) -d = Rows.fourth(~Tile.Tile8,accumulators) evidence = write_row(depth,U32.is_lt(row,rows),row,position,columns,a,bias,output,bias_evidence,balanced) -owners = Gemm.write_row(U32.is_lt(row,rows),row,position,columns,a,(bias,output)) evidence = next(depth,U32.is_lt((row + 1 : U32),rows),(row + 1 : U32),position,columns,b,bias,owners,bias_evidence,evidence) -owners = Gemm.write_row(U32.is_lt((row + 1 : U32),rows),(row + 1 : U32),position,columns,b,owners) evidence = next(depth,U32.is_lt((row + 2 : U32),rows),(row + 2 : U32),position,columns,c,bias,owners,bias_evidence,evidence) -owners = Gemm.write_row(U32.is_lt((row + 2 : U32),rows),(row + 2 : U32),position,columns,c,owners) next(depth,U32.is_lt((row + 3 : U32),rows),(row + 3 : U32),position,columns,d,bias,owners,bias_evidence,evidence)type Evidence<-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-depth: Nat,-state: Gemm.Buffers> is Data: Evidence{-output: Array<F32>,equation: {state == Gemm.Buffers{input,weights,bias,output} : Gemm.Buffers},balanced: Witness.Witness(F32,depth,output)}def after_write(-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-depth: Nat,-written: Array<F32> & Array<F32>, evidence: Written<bias,depth,written>) -> Evidence<input,weights,bias,depth,Gemm.after_write(input,weights,written)>: Written{output,equation,balanced} = evidence Evidence{output, %Equal.sym(Array<F32> & Array<F32>,written,(bias,output),equation) : {Gemm.after_write(input,weights,_) == Gemm.Buffers{input,weights,bias,output} : Gemm.Buffers} {==},balanced}# The final step stores into the owners the loop returned unchanged.def finished(+depth: Nat,-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-output: Array<F32>,-state: Gemm.Loop, +bias_evidence: Witness.ArrayWitness<F32,bias>,balanced: Witness.Witness(F32,depth,output), evidence: ReadWitness.Looped<input,weights,bias,output,state>) -> Evidence<input,weights,bias,depth,Gemm.finish(state)>: ReadWitness.Looped{index,context,accumulators,equation} = evidence %Equal.sym(Gemm.Loop,state,Gemm.Loop{input,weights,bias,output,index,context,accumulators},equation) : Evidence<input,weights,bias,depth,Gemm.finish(_)> after_write(input,weights,bias,depth,Gemm.write(context,accumulators,bias,output), write(depth,context,accumulators,bias,output,bias_evidence,balanced))def reduced(~upper: Bool,+depth: Nat,+count: Nat,-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-output: Array<F32>,-state: Gemm.Loop, +input_evidence: Witness.ArrayWitness<Tile.Tile8,input>,+weight_evidence: Witness.ArrayWitness<Tile.Tile8,weights>, +bias_evidence: Witness.ArrayWitness<F32,bias>,balanced: Witness.Witness(F32,depth,output),evidence: ReadWitness.Looped<input,weights,bias,output,state>) -> Evidence<input,weights,bias,depth,Loop.run_then(~Gemm.Loop,~Gemm.Buffers,~Gemm.step(~upper),~Gemm.finish,count,state)>: %Equal.sym(Gemm.Buffers,Loop.run_then(~Gemm.Loop,~Gemm.Buffers,~Gemm.step(~upper),~Gemm.finish,count,state), Gemm.finish(Loop.run(~Gemm.Loop,~Gemm.step(~upper),count,state)), Projection.run_then_finishes(~Gemm.Loop,~Gemm.Buffers,~Gemm.step(~upper),~Gemm.finish,count,state)) : Evidence<input,weights,bias,depth,_> finished(depth,input,weights,bias,output,Loop.run(~Gemm.Loop,~Gemm.step(~upper),count,state),bias_evidence,balanced, ReadWitness.run(~upper,count,input,weights,bias,output,state,input_evidence,weight_evidence,evidence))def reduce(upper: Bool,+depth: Nat,+count: Nat,-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-output: Array<F32>,+context: Gemm.Context, +input_evidence: Witness.ArrayWitness<Tile.Tile8,input>,+weight_evidence: Witness.ArrayWitness<Tile.Tile8,weights>, +bias_evidence: Witness.ArrayWitness<F32,bias>,balanced: Witness.Witness(F32,depth,output)) -> Evidence<input,weights,bias,depth,Gemm.reduce(upper,count,Gemm.Loop{input,weights,bias,output,0,context,Rows.filled(~Tile.Tile8,Tile.filled(0.0))})>: match upper: case False{}: reduced(~False{},depth,count,input,weights,bias,output,Gemm.Loop{input,weights,bias,output,0,context,Rows.filled(~Tile.Tile8,Tile.filled(0.0))}, input_evidence,weight_evidence,bias_evidence,balanced,ReadWitness.Looped{0,context,Rows.filled(~Tile.Tile8,Tile.filled(0.0)),{==}}) case True{}: reduced(~True{},depth,count,input,weights,bias,output,Gemm.Loop{input,weights,bias,output,0,context,Rows.filled(~Tile.Tile8,Tile.filled(0.0))}, input_evidence,weight_evidence,bias_evidence,balanced,ReadWitness.Looped{0,context,Rows.filled(~Tile.Tile8,Tile.filled(0.0)),{==}})def execute(+depth: Nat,job: Gemm.Job,-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-state: Gemm.Buffers, +input_evidence: Witness.ArrayWitness<Tile.Tile8,input>,+weight_evidence: Witness.ArrayWitness<Tile.Tile8,weights>,+bias_evidence: Witness.ArrayWitness<F32,bias>, evidence: Evidence<input,weights,bias,depth,state>) -> Evidence<input,weights,bias,depth,Gemm.execute(job,state)>: Gemm.Job{+context,+upper,+count} = job Evidence{output,equation,balanced} = evidence %Equal.sym(Gemm.Buffers,state,Gemm.Buffers{input,weights,bias,output},equation) : Evidence<input,weights,bias,depth,Gemm.execute(Gemm.Job{context,upper,count},_)> reduce(upper,depth,count,input,weights,bias,output,context,input_evidence,weight_evidence,bias_evidence,balanced)def traverse(+depth: Nat,count: Nat,+index: U32,+layout: Gemm.Layout,-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-state: Gemm.Buffers, +input_evidence: Witness.ArrayWitness<Tile.Tile8,input>,+weight_evidence: Witness.ArrayWitness<Tile.Tile8,weights>,+bias_evidence: Witness.ArrayWitness<F32,bias>, evidence: Evidence<input,weights,bias,depth,state>) -> Evidence<input,weights,bias,depth,Traversal.indexed_fold(~Gemm.Buffers,~Gemm.Layout,~Gemm.step_job,count,index,layout,state)>: match count: case 0n: evidence case 1n+rest: traverse(depth,rest,(index + 1 : U32),layout,input,weights,bias,Gemm.step_job(index,layout,state),input_evidence,weight_evidence,bias_evidence, execute(depth,Gemm.job(layout,index),input,weights,bias,state,input_evidence,weight_evidence,bias_evidence,evidence))def finish(-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-depth: Nat,-state: Gemm.Buffers, evidence: Evidence<input,weights,bias,depth,state>) -> Witness.Witness(F32,depth,Gemm.result(state)): Evidence{output,equation,balanced} = evidence %Equal.sym(Gemm.Buffers,state,Gemm.Buffers{input,weights,bias,output},equation) : Witness.Witness(F32,depth,Gemm.result(_)) balanceddef run(+depth: Nat,+rows: U32,-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-output: Array<F32>,+columns: U32,+inner: U32,+start: U32,+output_columns: U32, input_evidence: Witness.ArrayWitness<Tile.Tile8,input>,weight_evidence: Witness.ArrayWitness<Tile.Tile8,weights>,bias_evidence: Witness.ArrayWitness<F32,bias>, balanced: Witness.Witness(F32,depth,output)) -> Witness.Witness(F32,depth,Gemm.run_at(rows,input,weights,bias,output,columns,inner,start,output_columns)): +count = U32.to_nat(Gemm.job_count(rows,columns)) +layout = {Gemm.Layout{rows,columns,inner,start,output_columns} : Gemm.Layout} -state = {Gemm.Buffers{input,weights,bias,output} : Gemm.Buffers} finish(input,weights,bias,depth,Traversal.indexed_fold(~Gemm.Buffers,~Gemm.Layout,~Gemm.step_job,count,0,layout,state), traverse(depth,count,0,layout,input,weights,bias,state,input_evidence,weight_evidence,bias_evidence,Evidence{output,{==},balanced}))def equation(+depth: Nat,+rows: U32,-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-output: Array<F32>,+columns: U32,+inner: U32,+start: U32,+output_columns: U32, input_certificate: Certificate.Certificate<Tile.Tile8,input>,weight_certificate: Certificate.Certificate<Tile.Tile8,weights>,bias_certificate: Certificate.Certificate<F32,bias>, balanced: {output == Tree.pack(F32,Certificate.normalized(F32,depth,Tree.reflect(F32,output))) : Array<F32>}) -> {Gemm.run_at(rows,input,weights,bias,output,columns,inner,start,output_columns) == Tree.pack(F32,Certificate.normalized(F32,depth,Tree.reflect(F32,Gemm.run_at(rows,input,weights,bias,output,columns,inner,start,output_columns)))) : Array<F32>}: %Certificate.restored(F32,depth,Gemm.run_at(rows,input,weights,bias,output,columns,inner,start,output_columns), run(depth,rows,input,weights,bias,output,columns,inner,start,output_columns,Certificate.witness(Tile.Tile8,input,input_certificate), Certificate.witness(Tile.Tile8,weights,weight_certificate),Certificate.witness(F32,bias,bias_certificate),Certificate.witness_at(F32,depth,output,balanced))) : {_ == Tree.pack(F32,Certificate.normalized(F32,depth,Tree.reflect(F32,Gemm.run_at(rows,input,weights,bias,output,columns,inner,start,output_columns)))) : Array<F32>} {==}def certificate(rows: U32,-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-output: Array<F32>,columns: U32,inner: U32,start: U32,output_columns: U32, input_certificate: Certificate.Certificate<Tile.Tile8,input>,weight_certificate: Certificate.Certificate<Tile.Tile8,weights>,bias_certificate: Certificate.Certificate<F32,bias>, output_certificate: Certificate.Certificate<F32,output>) -> Certificate.Certificate<F32,Gemm.run_at(rows,input,weights,bias,output,columns,inner,start,output_columns)>: Certificate.Certificate{+depth,balanced} = output_certificate Certificate.Certificate{depth,equation(depth,rows,input,weights,bias,output,columns,inner,start,output_columns,input_certificate,weight_certificate,bias_certificate,balanced)}# The fold gives back the three buffers it only reads, next to the output of run.def kept_state(-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-depth: Nat,-state: Gemm.Buffers, evidence: Evidence<input,weights,bias,depth,state>) -> {state == Gemm.Buffers{input,weights,bias,Gemm.result(state)} : Gemm.Buffers}: Evidence{output,equation,balanced} = evidence %Equal.sym(Gemm.Buffers,state,Gemm.Buffers{input,weights,bias,output},equation) : {_ == Gemm.Buffers{input,weights,bias,Gemm.result(_)} : Gemm.Buffers} {==}def kept_equation(+depth: Nat,+rows: U32,-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-output: Array<F32>,+columns: U32,+inner: U32,+start: U32,+output_columns: U32, input_evidence: Witness.ArrayWitness<Tile.Tile8,input>,weight_evidence: Witness.ArrayWitness<Tile.Tile8,weights>,bias_evidence: Witness.ArrayWitness<F32,bias>, balanced: Witness.Witness(F32,depth,output)) -> {Gemm.run_keeping(rows,input,weights,bias,output,columns,inner,start,output_columns) == Gemm.Buffers{input,weights,bias,Gemm.run_at(rows,input,weights,bias,output,columns,inner,start,output_columns)} : Gemm.Buffers}: +count = U32.to_nat(Gemm.job_count(rows,columns)) +layout = {Gemm.Layout{rows,columns,inner,start,output_columns} : Gemm.Layout} -state = {Gemm.Buffers{input,weights,bias,output} : Gemm.Buffers} kept_state(input,weights,bias,depth,Traversal.indexed_fold(~Gemm.Buffers,~Gemm.Layout,~Gemm.step_job,count,0,layout,state), traverse(depth,count,0,layout,input,weights,bias,state,input_evidence,weight_evidence,bias_evidence,Evidence{output,{==},balanced}))type Kept<-state: Gemm.Buffers> is Data: Kept{input: Certificate.Certificate<Tile.Tile8,Gemm.kept_input(state)>,weights: Certificate.Certificate<Tile.Tile8,Gemm.kept_weights(state)>, bias: Certificate.Certificate<F32,Gemm.kept_bias(state)>,output: Certificate.Certificate<F32,Gemm.result(state)>}# The certificates of the read-only buffers are the ones that came in.def kept(+rows: U32,-input: Array<Tile.Tile8>,-weights: Array<Tile.Tile8>,-bias: Array<F32>,-output: Array<F32>,+columns: U32,+inner: U32,+start: U32,+output_columns: U32, input_certificate: Certificate.Certificate<Tile.Tile8,input>,weight_certificate: Certificate.Certificate<Tile.Tile8,weights>,bias_certificate: Certificate.Certificate<F32,bias>, output_certificate: Certificate.Certificate<F32,output>) -> Kept<Gemm.run_keeping(rows,input,weights,bias,output,columns,inner,start,output_columns)>: Certificate.Certificate{+input_depth,+input_equation} = input_certificate Certificate.Certificate{+weight_depth,+weight_equation} = weight_certificate Certificate.Certificate{+bias_depth,+bias_equation} = bias_certificate Certificate.Certificate{+depth,+balanced} = output_certificate -state = Gemm.run_keeping(rows,input,weights,bias,output,columns,inner,start,output_columns) -result = Gemm.run_at(rows,input,weights,bias,output,columns,inner,start,output_columns) +owners = kept_equation(depth,rows,input,weights,bias,output,columns,inner,start,output_columns, Certificate.witness(Tile.Tile8,input,Certificate.Certificate{input_depth,input_equation}),Certificate.witness(Tile.Tile8,weights,Certificate.Certificate{weight_depth,weight_equation}), Certificate.witness(F32,bias,Certificate.Certificate{bias_depth,bias_equation}),Certificate.witness_at(F32,depth,output,balanced)) Kept{ Certificate.Certificate{input_depth, %Equal.sym(Gemm.Buffers,state,Gemm.Buffers{input,weights,bias,result},owners) : {Gemm.kept_input(_) == Tree.pack(Tile.Tile8,Certificate.normalized(Tile.Tile8,input_depth,Tree.reflect(Tile.Tile8,Gemm.kept_input(_)))) : Array<Tile.Tile8>} input_equation}, Certificate.Certificate{weight_depth, %Equal.sym(Gemm.Buffers,state,Gemm.Buffers{input,weights,bias,result},owners) : {Gemm.kept_weights(_) == Tree.pack(Tile.Tile8,Certificate.normalized(Tile.Tile8,weight_depth,Tree.reflect(Tile.Tile8,Gemm.kept_weights(_)))) : Array<Tile.Tile8>} weight_equation}, Certificate.Certificate{bias_depth, %Equal.sym(Gemm.Buffers,state,Gemm.Buffers{input,weights,bias,result},owners) : {Gemm.kept_bias(_) == Tree.pack(F32,Certificate.normalized(F32,bias_depth,Tree.reflect(F32,Gemm.kept_bias(_)))) : Array<F32>} bias_equation}, Certificate.Certificate{depth,equation(depth,rows,input,weights,bias,output,columns,inner,start,output_columns,Certificate.Certificate{input_depth,input_equation}, Certificate.Certificate{weight_depth,weight_equation},Certificate.Certificate{bias_depth,bias_equation},balanced)}}