1 # Shape evidence follows the real readers and writers. Only the final flat 2 # certificate is live; witness traversals occur beneath erased rewrites. 3 import Base 4 import ../kernel_gemm.bend as Gemm 5 import ../kernel_rows.bend as Rows 6 import ../kernel_tile.bend as Tile 7 import ../traversal_array.bend as Traversal 8 import ../traversal_loop.bend as Loop 9 import ./storage_witness.bend as Witness 10 import ./storage_certificate.bend as Certificate 11 import ./storage_tree_model.bend as Tree 12 import ./loop_projection.bend as Projection 13 import ./kernel_read_witness.bend as ReadWitness 14 15 def store_lane(+depth: Nat,inside: Bool,-output: Array,+offset: U32,-value: F32,evidence: Witness.Witness(F32,depth,output)) -> 16 Witness.Witness(F32,depth,Gemm.store_lane(inside,output,offset,value)): 17 match inside: 18 case True{}: Witness.write(F32,depth,output,offset,value,evidence) 19 case False{}: evidence 20 21 def store_tail(+count: U32,+depth: Nat,-output: Array,+offset: U32,-bias: F32,-accumulator: Tile.Tile8, 22 evidence: Witness.Witness(F32,depth,output)) -> 23 Witness.Witness(F32,depth,Gemm.store_tail(count,output,offset,bias,accumulator)): 24 -a = (Tile.lane0(accumulator) + bias : F32) 25 -b = (Tile.lane1(accumulator) + bias : F32) 26 -c = (Tile.lane2(accumulator) + bias : F32) 27 -d = (Tile.lane3(accumulator) + bias : F32) 28 -e = (Tile.lane4(accumulator) + bias : F32) 29 -f = (Tile.lane5(accumulator) + bias : F32) 30 -g = (Tile.lane6(accumulator) + bias : F32) 31 -h = (Tile.lane7(accumulator) + bias : F32) 32 -o1 = Gemm.store_lane(U32.is_lt(0,count),output,offset,a) 33 -o2 = Gemm.store_lane(U32.is_lt(1,count),o1,(offset + 1 : U32),b) 34 -o3 = Gemm.store_lane(U32.is_lt(2,count),o2,(offset + 2 : U32),c) 35 -o4 = Gemm.store_lane(U32.is_lt(3,count),o3,(offset + 3 : U32),d) 36 -o5 = Gemm.store_lane(U32.is_lt(4,count),o4,(offset + 4 : U32),e) 37 -o6 = Gemm.store_lane(U32.is_lt(5,count),o5,(offset + 5 : U32),f) 38 -o7 = Gemm.store_lane(U32.is_lt(6,count),o6,(offset + 6 : U32),g) 39 evidence = store_lane(depth,U32.is_lt(0,count),output,offset,a,evidence) 40 evidence = store_lane(depth,U32.is_lt(1,count),o1,(offset + 1 : U32),b,evidence) 41 evidence = store_lane(depth,U32.is_lt(2,count),o2,(offset + 2 : U32),c,evidence) 42 evidence = store_lane(depth,U32.is_lt(3,count),o3,(offset + 3 : U32),d,evidence) 43 evidence = store_lane(depth,U32.is_lt(4,count),o4,(offset + 4 : U32),e,evidence) 44 evidence = store_lane(depth,U32.is_lt(5,count),o5,(offset + 5 : U32),f,evidence) 45 evidence = store_lane(depth,U32.is_lt(6,count),o6,(offset + 6 : U32),g,evidence) 46 store_lane(depth,U32.is_lt(7,count),o7,(offset + 7 : U32),h,evidence) 47 48 type Written<-bias: Array,-depth: Nat,-result: Array & Array> is Data: 49 Written{-output: Array,equation: {result == (bias,output) : Array & Array},balanced: Witness.Witness(F32,depth,output)} 50 51 def write_bias(+depth: Nat,-bias: Array,-output: Array,+offset: U32,+row: U32,-accumulator: Tile.Tile8,+count: U32, 52 read: Witness.Read,evidence: Witness.Witness(F32,depth,output)) -> 53 Written: 54 Witness.Read{value,equation} = read 55 Written{Gemm.store_tail(count,output,offset,value,accumulator), 56 %Equal.sym(Array & F32,Array.get(F32,bias,row),(bias,value),equation) : 57 {Gemm.write_bias(output,offset,accumulator,count,_) == (bias,Gemm.store_tail(count,output,offset,value,accumulator)) : Array & Array} 58 {==},store_tail(count,depth,output,offset,value,accumulator,evidence)} 59 60 def write_row(+depth: Nat,valid: Bool,+row: U32,+position: U32,+columns: U32,-accumulator: Tile.Tile8, 61 -bias: Array,-output: Array,bias_evidence: Witness.ArrayWitness,evidence: Witness.Witness(F32,depth,output)) -> 62 Written: 63 match valid: 64 case False{}: Written{output,{==},evidence} 65 case True{}: 66 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) 67 68 def next(+depth: Nat,valid: Bool,row: U32,position: U32,columns: U32,-accumulator: Tile.Tile8, 69 -bias: Array,-owners: Array & Array,bias_evidence: Witness.ArrayWitness,evidence: Written) -> 70 Written: 71 Written{output,equation,balanced} = evidence 72 %Equal.sym(Array & Array,owners,(bias,output),equation) : 73 Written 74 write_row(depth,valid,row,position,columns,accumulator,bias,output,bias_evidence,balanced) 75 76 def write(+depth: Nat,context: Gemm.Context,-accumulators: Rows.Rows, 77 -bias: Array,-output: Array,+bias_evidence: Witness.ArrayWitness,balanced: Witness.Witness(F32,depth,output)) -> 78 Written: 79 Gemm.Context{+row,+position,+rows,+columns,input_base,weight_base} = context 80 -a = Rows.first(~Tile.Tile8,accumulators) 81 -b = Rows.second(~Tile.Tile8,accumulators) 82 -c = Rows.third(~Tile.Tile8,accumulators) 83 -d = Rows.fourth(~Tile.Tile8,accumulators) 84 evidence = write_row(depth,U32.is_lt(row,rows),row,position,columns,a,bias,output,bias_evidence,balanced) 85 -owners = Gemm.write_row(U32.is_lt(row,rows),row,position,columns,a,(bias,output)) 86 evidence = next(depth,U32.is_lt((row + 1 : U32),rows),(row + 1 : U32),position,columns,b,bias,owners,bias_evidence,evidence) 87 -owners = Gemm.write_row(U32.is_lt((row + 1 : U32),rows),(row + 1 : U32),position,columns,b,owners) 88 evidence = next(depth,U32.is_lt((row + 2 : U32),rows),(row + 2 : U32),position,columns,c,bias,owners,bias_evidence,evidence) 89 -owners = Gemm.write_row(U32.is_lt((row + 2 : U32),rows),(row + 2 : U32),position,columns,c,owners) 90 next(depth,U32.is_lt((row + 3 : U32),rows),(row + 3 : U32),position,columns,d,bias,owners,bias_evidence,evidence) 91 92 type Evidence<-input: Array,-weights: Array,-bias: Array,-depth: Nat,-state: Gemm.Buffers> is Data: 93 Evidence{-output: Array,equation: {state == Gemm.Buffers{input,weights,bias,output} : Gemm.Buffers},balanced: Witness.Witness(F32,depth,output)} 94 95 def after_write(-input: Array,-weights: Array,-bias: Array,-depth: Nat,-written: Array & Array, 96 evidence: Written) -> Evidence: 97 Written{output,equation,balanced} = evidence 98 Evidence{output, 99 %Equal.sym(Array & Array,written,(bias,output),equation) : 100 {Gemm.after_write(input,weights,_) == Gemm.Buffers{input,weights,bias,output} : Gemm.Buffers} 101 {==},balanced} 102 103 # The final step stores into the owners the loop returned unchanged. 104 def finished(+depth: Nat,-input: Array,-weights: Array,-bias: Array,-output: Array,-state: Gemm.Loop, 105 +bias_evidence: Witness.ArrayWitness,balanced: Witness.Witness(F32,depth,output), 106 evidence: ReadWitness.Looped) -> Evidence: 107 ReadWitness.Looped{index,context,accumulators,equation} = evidence 108 %Equal.sym(Gemm.Loop,state,Gemm.Loop{input,weights,bias,output,index,context,accumulators},equation) : 109 Evidence 110 after_write(input,weights,bias,depth,Gemm.write(context,accumulators,bias,output), 111 write(depth,context,accumulators,bias,output,bias_evidence,balanced)) 112 113 def reduced(~upper: Bool,+depth: Nat,+count: Nat,-input: Array,-weights: Array,-bias: Array,-output: Array,-state: Gemm.Loop, 114 +input_evidence: Witness.ArrayWitness,+weight_evidence: Witness.ArrayWitness, 115 +bias_evidence: Witness.ArrayWitness,balanced: Witness.Witness(F32,depth,output),evidence: ReadWitness.Looped) -> 116 Evidence: 117 %Equal.sym(Gemm.Buffers,Loop.run_then(~Gemm.Loop,~Gemm.Buffers,~Gemm.step(~upper),~Gemm.finish,count,state), 118 Gemm.finish(Loop.run(~Gemm.Loop,~Gemm.step(~upper),count,state)), 119 Projection.run_then_finishes(~Gemm.Loop,~Gemm.Buffers,~Gemm.step(~upper),~Gemm.finish,count,state)) : 120 Evidence 121 finished(depth,input,weights,bias,output,Loop.run(~Gemm.Loop,~Gemm.step(~upper),count,state),bias_evidence,balanced, 122 ReadWitness.run(~upper,count,input,weights,bias,output,state,input_evidence,weight_evidence,evidence)) 123 124 def reduce(upper: Bool,+depth: Nat,+count: Nat,-input: Array,-weights: Array,-bias: Array,-output: Array,+context: Gemm.Context, 125 +input_evidence: Witness.ArrayWitness,+weight_evidence: Witness.ArrayWitness, 126 +bias_evidence: Witness.ArrayWitness,balanced: Witness.Witness(F32,depth,output)) -> 127 Evidence: 128 match upper: 129 case False{}: 130 reduced(~False{},depth,count,input,weights,bias,output,Gemm.Loop{input,weights,bias,output,0,context,Rows.filled(~Tile.Tile8,Tile.filled(0.0))}, 131 input_evidence,weight_evidence,bias_evidence,balanced,ReadWitness.Looped{0,context,Rows.filled(~Tile.Tile8,Tile.filled(0.0)),{==}}) 132 case True{}: 133 reduced(~True{},depth,count,input,weights,bias,output,Gemm.Loop{input,weights,bias,output,0,context,Rows.filled(~Tile.Tile8,Tile.filled(0.0))}, 134 input_evidence,weight_evidence,bias_evidence,balanced,ReadWitness.Looped{0,context,Rows.filled(~Tile.Tile8,Tile.filled(0.0)),{==}}) 135 136 def execute(+depth: Nat,job: Gemm.Job,-input: Array,-weights: Array,-bias: Array,-state: Gemm.Buffers, 137 +input_evidence: Witness.ArrayWitness,+weight_evidence: Witness.ArrayWitness,+bias_evidence: Witness.ArrayWitness, 138 evidence: Evidence) -> Evidence: 139 Gemm.Job{+context,+upper,+count} = job 140 Evidence{output,equation,balanced} = evidence 141 %Equal.sym(Gemm.Buffers,state,Gemm.Buffers{input,weights,bias,output},equation) : 142 Evidence 143 reduce(upper,depth,count,input,weights,bias,output,context,input_evidence,weight_evidence,bias_evidence,balanced) 144 145 def traverse(+depth: Nat,count: Nat,+index: U32,+layout: Gemm.Layout,-input: Array,-weights: Array,-bias: Array,-state: Gemm.Buffers, 146 +input_evidence: Witness.ArrayWitness,+weight_evidence: Witness.ArrayWitness,+bias_evidence: Witness.ArrayWitness, 147 evidence: Evidence) -> 148 Evidence: 149 match count: 150 case 0n: evidence 151 case 1n+rest: 152 traverse(depth,rest,(index + 1 : U32),layout,input,weights,bias,Gemm.step_job(index,layout,state),input_evidence,weight_evidence,bias_evidence, 153 execute(depth,Gemm.job(layout,index),input,weights,bias,state,input_evidence,weight_evidence,bias_evidence,evidence)) 154 155 def finish(-input: Array,-weights: Array,-bias: Array,-depth: Nat,-state: Gemm.Buffers, 156 evidence: Evidence) -> Witness.Witness(F32,depth,Gemm.result(state)): 157 Evidence{output,equation,balanced} = evidence 158 %Equal.sym(Gemm.Buffers,state,Gemm.Buffers{input,weights,bias,output},equation) : Witness.Witness(F32,depth,Gemm.result(_)) 159 balanced 160 161 def run(+depth: Nat,+rows: U32,-input: Array,-weights: Array,-bias: Array,-output: Array,+columns: U32,+inner: U32,+start: U32,+output_columns: U32, 162 input_evidence: Witness.ArrayWitness,weight_evidence: Witness.ArrayWitness,bias_evidence: Witness.ArrayWitness, 163 balanced: Witness.Witness(F32,depth,output)) -> Witness.Witness(F32,depth,Gemm.run_at(rows,input,weights,bias,output,columns,inner,start,output_columns)): 164 +count = U32.to_nat(Gemm.job_count(rows,columns)) 165 +layout = {Gemm.Layout{rows,columns,inner,start,output_columns} : Gemm.Layout} 166 -state = {Gemm.Buffers{input,weights,bias,output} : Gemm.Buffers} 167 finish(input,weights,bias,depth,Traversal.indexed_fold(~Gemm.Buffers,~Gemm.Layout,~Gemm.step_job,count,0,layout,state), 168 traverse(depth,count,0,layout,input,weights,bias,state,input_evidence,weight_evidence,bias_evidence,Evidence{output,{==},balanced})) 169 170 def equation(+depth: Nat,+rows: U32,-input: Array,-weights: Array,-bias: Array,-output: Array,+columns: U32,+inner: U32,+start: U32,+output_columns: U32, 171 input_certificate: Certificate.Certificate,weight_certificate: Certificate.Certificate,bias_certificate: Certificate.Certificate, 172 balanced: {output == Tree.pack(F32,Certificate.normalized(F32,depth,Tree.reflect(F32,output))) : Array}) -> 173 {Gemm.run_at(rows,input,weights,bias,output,columns,inner,start,output_columns) == 174 Tree.pack(F32,Certificate.normalized(F32,depth,Tree.reflect(F32,Gemm.run_at(rows,input,weights,bias,output,columns,inner,start,output_columns)))) : Array}: 175 %Certificate.restored(F32,depth,Gemm.run_at(rows,input,weights,bias,output,columns,inner,start,output_columns), 176 run(depth,rows,input,weights,bias,output,columns,inner,start,output_columns,Certificate.witness(Tile.Tile8,input,input_certificate), 177 Certificate.witness(Tile.Tile8,weights,weight_certificate),Certificate.witness(F32,bias,bias_certificate),Certificate.witness_at(F32,depth,output,balanced))) : 178 {_ == Tree.pack(F32,Certificate.normalized(F32,depth,Tree.reflect(F32,Gemm.run_at(rows,input,weights,bias,output,columns,inner,start,output_columns)))) : Array} 179 {==} 180 181 def certificate(rows: U32,-input: Array,-weights: Array,-bias: Array,-output: Array,columns: U32,inner: U32,start: U32,output_columns: U32, 182 input_certificate: Certificate.Certificate,weight_certificate: Certificate.Certificate,bias_certificate: Certificate.Certificate, 183 output_certificate: Certificate.Certificate) -> Certificate.Certificate: 184 Certificate.Certificate{+depth,balanced} = output_certificate 185 Certificate.Certificate{depth,equation(depth,rows,input,weights,bias,output,columns,inner,start,output_columns,input_certificate,weight_certificate,bias_certificate,balanced)} 186 187 # The fold gives back the three buffers it only reads, next to the output of run. 188 def kept_state(-input: Array,-weights: Array,-bias: Array,-depth: Nat,-state: Gemm.Buffers, 189 evidence: Evidence) -> {state == Gemm.Buffers{input,weights,bias,Gemm.result(state)} : Gemm.Buffers}: 190 Evidence{output,equation,balanced} = evidence 191 %Equal.sym(Gemm.Buffers,state,Gemm.Buffers{input,weights,bias,output},equation) : 192 {_ == Gemm.Buffers{input,weights,bias,Gemm.result(_)} : Gemm.Buffers} 193 {==} 194 195 def kept_equation(+depth: Nat,+rows: U32,-input: Array,-weights: Array,-bias: Array,-output: Array,+columns: U32,+inner: U32,+start: U32,+output_columns: U32, 196 input_evidence: Witness.ArrayWitness,weight_evidence: Witness.ArrayWitness,bias_evidence: Witness.ArrayWitness, 197 balanced: Witness.Witness(F32,depth,output)) -> 198 {Gemm.run_keeping(rows,input,weights,bias,output,columns,inner,start,output_columns) == 199 Gemm.Buffers{input,weights,bias,Gemm.run_at(rows,input,weights,bias,output,columns,inner,start,output_columns)} : Gemm.Buffers}: 200 +count = U32.to_nat(Gemm.job_count(rows,columns)) 201 +layout = {Gemm.Layout{rows,columns,inner,start,output_columns} : Gemm.Layout} 202 -state = {Gemm.Buffers{input,weights,bias,output} : Gemm.Buffers} 203 kept_state(input,weights,bias,depth,Traversal.indexed_fold(~Gemm.Buffers,~Gemm.Layout,~Gemm.step_job,count,0,layout,state), 204 traverse(depth,count,0,layout,input,weights,bias,state,input_evidence,weight_evidence,bias_evidence,Evidence{output,{==},balanced})) 205 206 type Kept<-state: Gemm.Buffers> is Data: 207 Kept{input: Certificate.Certificate,weights: Certificate.Certificate, 208 bias: Certificate.Certificate,output: Certificate.Certificate} 209 210 # The certificates of the read-only buffers are the ones that came in. 211 def kept(+rows: U32,-input: Array,-weights: Array,-bias: Array,-output: Array,+columns: U32,+inner: U32,+start: U32,+output_columns: U32, 212 input_certificate: Certificate.Certificate,weight_certificate: Certificate.Certificate,bias_certificate: Certificate.Certificate, 213 output_certificate: Certificate.Certificate) -> Kept: 214 Certificate.Certificate{+input_depth,+input_equation} = input_certificate 215 Certificate.Certificate{+weight_depth,+weight_equation} = weight_certificate 216 Certificate.Certificate{+bias_depth,+bias_equation} = bias_certificate 217 Certificate.Certificate{+depth,+balanced} = output_certificate 218 -state = Gemm.run_keeping(rows,input,weights,bias,output,columns,inner,start,output_columns) 219 -result = Gemm.run_at(rows,input,weights,bias,output,columns,inner,start,output_columns) 220 +owners = kept_equation(depth,rows,input,weights,bias,output,columns,inner,start,output_columns, 221 Certificate.witness(Tile.Tile8,input,Certificate.Certificate{input_depth,input_equation}),Certificate.witness(Tile.Tile8,weights,Certificate.Certificate{weight_depth,weight_equation}), 222 Certificate.witness(F32,bias,Certificate.Certificate{bias_depth,bias_equation}),Certificate.witness_at(F32,depth,output,balanced)) 223 Kept{ 224 Certificate.Certificate{input_depth, 225 %Equal.sym(Gemm.Buffers,state,Gemm.Buffers{input,weights,bias,result},owners) : 226 {Gemm.kept_input(_) == Tree.pack(Tile.Tile8,Certificate.normalized(Tile.Tile8,input_depth,Tree.reflect(Tile.Tile8,Gemm.kept_input(_)))) : Array} 227 input_equation}, 228 Certificate.Certificate{weight_depth, 229 %Equal.sym(Gemm.Buffers,state,Gemm.Buffers{input,weights,bias,result},owners) : 230 {Gemm.kept_weights(_) == Tree.pack(Tile.Tile8,Certificate.normalized(Tile.Tile8,weight_depth,Tree.reflect(Tile.Tile8,Gemm.kept_weights(_)))) : Array} 231 weight_equation}, 232 Certificate.Certificate{bias_depth, 233 %Equal.sym(Gemm.Buffers,state,Gemm.Buffers{input,weights,bias,result},owners) : 234 {Gemm.kept_bias(_) == Tree.pack(F32,Certificate.normalized(F32,bias_depth,Tree.reflect(F32,Gemm.kept_bias(_)))) : Array} 235 bias_equation}, 236 Certificate.Certificate{depth,equation(depth,rows,input,weights,bias,output,columns,inner,start,output_columns,Certificate.Certificate{input_depth,input_equation}, 237 Certificate.Certificate{weight_depth,weight_equation},Certificate.Certificate{bias_depth,bias_equation},balanced)}}