1 # The packer hands back its input unchanged and keeps the witness of its 2 # output: every loop of kernel_spatial_packing preserves both. 3 import Base 4 import ../traversal_loop.bend as Loop 5 import ../kernel_spatial_packing.bend as Packing 6 import ../kernel_packing_samples.bend as PackingSamples 7 import ../kernel_tile.bend as Tile 8 import ./storage_witness.bend as Witness 9 import ./traversal_witness.bend as TraversalWitness 10 import ./kernel_packing_read_witness.bend as PackingReadWitness 11 import ./kernel_tile_witness.bend as TileWitness 12 13 type RowCursor is Data: 14 RowCursor{index: U32,channel: U32,kernel_row: U32,kernel_column: U32,row: U32,offset: U32,output_width: U32,height: U32,width: U32,stride: U32,padding: U32,taps: Nat} 15 16 def restore_row(input: Array,output: Array,cursor: RowCursor) -> Packing.RowTiles(): 17 RowCursor{i,c,kr,kc,row,offset,ow,h,w,s,p,taps} = cursor 18 Packing.RowTile{input,output,i,c,kr,kc,row,offset,ow,h,w,s,p,taps} 19 20 type RowEvidence<-input: Array,-depth: Nat,-state: Packing.RowTiles()> is Data: 21 RowEvidence{-output: Array,cursor: RowCursor,equation: {state == restore_row(input,output,cursor) : Packing.RowTiles()},balanced: Witness.Witness(Tile.Tile8,depth,output)} 22 23 def stored_row(output: Array,cursor: RowCursor,result: Array & Tile.Tile8) -> Packing.RowTiles(): 24 RowCursor{i,c,kr,kc,row,offset,ow,h,w,s,p,taps} = cursor 25 Packing.store_block(output,i,c,kr,kc,row,offset,ow,h,w,s,p,taps,result) 26 27 def store_row(+depth: Nat,-input: Array,-output: Array,cursor: RowCursor,-result: Array & Tile.Tile8, 28 balanced: Witness.Witness(Tile.Tile8,depth,output),sample: TraversalWitness.Sample,Tile.Tile8,input,result>) -> 29 RowEvidence: 30 RowCursor{+i,c,kr,kc,row,offset,ow,h,w,s,p,taps} = cursor 31 TraversalWitness.Sample{value,equation} = sample 32 %Equal.sym(Array & Tile.Tile8,result,(input,value),equation) : RowEvidence 33 RowEvidence{Array.set(Tile.Tile8,output,i,value),RowCursor{(i + 1 : U32),c,kr,(kc + 1 : U32),row,offset,ow,h,w,s,p,taps},{==},Witness.write(Tile.Tile8,depth,output,i,value,balanced)} 34 35 def column(~block: Packing.Block(), 36 ~sampled: @-input: Array -> @+row: U32 -> @+column: U32 -> @+channel: U32 -> @+height: U32 -> @+width: U32 -> @+stride: U32 -> @+next_row: U32 -> @+offset: U32 -> @+output_width: U32 -> 37 @+kernel_column: U32 -> @+padding: U32 -> 38 @+evidence: PackingReadWitness.family(input) -> TraversalWitness.Sample,Tile.Tile8,input,block(input,row,column,channel,height,width,stride,next_row,offset,output_width,kernel_column,padding)>, 39 +depth: Nat,-input: Array,-state: Packing.RowTiles(),+input_evidence: PackingReadWitness.family(input),evidence: RowEvidence) -> 40 RowEvidence: 41 RowEvidence{output,cursor,equation,balanced} = evidence 42 RowCursor{+i,+c,+kr,+kc,+row,+offset,+ow,+h,+w,+s,+p,+taps} = cursor 43 %Equal.sym(Packing.RowTiles(),state,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,ow,h,w,s,p,taps}),equation) : RowEvidence 44 +iy = (row * s + kr - p : U32) 45 +ix = (offset * s + kc - p : U32) 46 +cb = (c * h * w : U32) 47 store_row(depth,input,output,RowCursor{i,c,kr,kc,row,offset,ow,h,w,s,p,taps},block(input,iy,ix,cb,h,w,s,((row + 1) * s + kr - p : U32),offset,ow,kc,p),balanced, 48 sampled(input,iy,ix,cb,h,w,s,((row + 1) * s + kr - p : U32),offset,ow,kc,p,input_evidence)) 49 50 def column_loop(~block: Packing.Block(), 51 ~sampled: @-input: Array -> @+row: U32 -> @+column: U32 -> @+channel: U32 -> @+height: U32 -> @+width: U32 -> @+stride: U32 -> @+next_row: U32 -> @+offset: U32 -> @+output_width: U32 -> 52 @+kernel_column: U32 -> @+padding: U32 -> 53 @+evidence: PackingReadWitness.family(input) -> TraversalWitness.Sample,Tile.Tile8,input,block(input,row,column,channel,height,width,stride,next_row,offset,output_width,kernel_column,padding)>, 54 +depth: Nat,count: Nat,-input: Array,-state: Packing.RowTiles(),+input_evidence: PackingReadWitness.family(input),evidence: RowEvidence) -> 55 RowEvidence: 56 match count: 57 case 0n: evidence 58 case 1n+rest: column_loop(~block,~sampled,depth,rest,input,Packing.column_step(~block,state),input_evidence,column(~block,~sampled,depth,input,state,input_evidence,evidence)) 59 60 def after_columns(-depth: Nat,-input: Array,-state: Packing.RowTiles(),evidence: RowEvidence) -> 61 RowEvidence: 62 RowEvidence{output,cursor,equation,balanced} = evidence 63 RowCursor{i,c,kr,kc,row,offset,ow,h,w,s,p,taps} = cursor 64 %Equal.sym(Packing.RowTiles(),state,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,ow,h,w,s,p,taps}),equation) : RowEvidence 65 RowEvidence{output,RowCursor{i,c,(kr + 1 : U32),0,row,offset,ow,h,w,s,p,taps},{==},balanced} 66 67 def kernel_row(~block: Packing.Block(), 68 ~sampled: @-input: Array -> @+row: U32 -> @+column: U32 -> @+channel: U32 -> @+height: U32 -> @+width: U32 -> @+stride: U32 -> @+next_row: U32 -> @+offset: U32 -> @+output_width: U32 -> 69 @+kernel_column: U32 -> @+padding: U32 -> 70 @+evidence: PackingReadWitness.family(input) -> TraversalWitness.Sample,Tile.Tile8,input,block(input,row,column,channel,height,width,stride,next_row,offset,output_width,kernel_column,padding)>, 71 +depth: Nat,-input: Array,-state: Packing.RowTiles(),+input_evidence: PackingReadWitness.family(input),evidence: RowEvidence) -> 72 RowEvidence: 73 RowEvidence{output,cursor,equation,balanced} = evidence 74 RowCursor{+i,+c,+kr,+kc,+row,+offset,+ow,+h,+w,+s,+p,+taps} = cursor 75 %Equal.sym(Packing.RowTiles(),state,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,ow,h,w,s,p,taps}),equation) : RowEvidence 76 after_columns(depth,input,Loop.run(~Packing.RowTiles(),~Packing.column_step(~block),taps,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,ow,h,w,s,p,taps})), 77 column_loop(~block,~sampled,depth,taps,input,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,ow,h,w,s,p,taps}),input_evidence, 78 RowEvidence{output,RowCursor{i,c,kr,kc,row,offset,ow,h,w,s,p,taps},{==},balanced})) 79 80 def kernel_row_loop(~block: Packing.Block(), 81 ~sampled: @-input: Array -> @+row: U32 -> @+column: U32 -> @+channel: U32 -> @+height: U32 -> @+width: U32 -> @+stride: U32 -> @+next_row: U32 -> @+offset: U32 -> @+output_width: U32 -> 82 @+kernel_column: U32 -> @+padding: U32 -> 83 @+evidence: PackingReadWitness.family(input) -> TraversalWitness.Sample,Tile.Tile8,input,block(input,row,column,channel,height,width,stride,next_row,offset,output_width,kernel_column,padding)>, 84 +depth: Nat,count: Nat,-input: Array,-state: Packing.RowTiles(),+input_evidence: PackingReadWitness.family(input),evidence: RowEvidence) -> 85 RowEvidence: 86 match count: 87 case 0n: evidence 88 case 1n+rest: kernel_row_loop(~block,~sampled,depth,rest,input,Packing.kernel_row_step(~block,state),input_evidence,kernel_row(~block,~sampled,depth,input,state,input_evidence,evidence)) 89 90 def after_kernel_rows(-depth: Nat,-input: Array,-state: Packing.RowTiles(),evidence: RowEvidence) -> 91 RowEvidence: 92 RowEvidence{output,cursor,equation,balanced} = evidence 93 RowCursor{i,c,kr,kc,row,offset,ow,h,w,s,p,taps} = cursor 94 %Equal.sym(Packing.RowTiles(),state,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,ow,h,w,s,p,taps}),equation) : RowEvidence 95 RowEvidence{output,RowCursor{i,(c + 1 : U32),0,0,row,offset,ow,h,w,s,p,taps},{==},balanced} 96 97 def channel(~block: Packing.Block(), 98 ~sampled: @-input: Array -> @+row: U32 -> @+column: U32 -> @+channel: U32 -> @+height: U32 -> @+width: U32 -> @+stride: U32 -> @+next_row: U32 -> @+offset: U32 -> @+output_width: U32 -> 99 @+kernel_column: U32 -> @+padding: U32 -> 100 @+evidence: PackingReadWitness.family(input) -> TraversalWitness.Sample,Tile.Tile8,input,block(input,row,column,channel,height,width,stride,next_row,offset,output_width,kernel_column,padding)>, 101 +depth: Nat,-input: Array,-state: Packing.RowTiles(),+input_evidence: PackingReadWitness.family(input),evidence: RowEvidence) -> 102 RowEvidence: 103 RowEvidence{output,cursor,equation,balanced} = evidence 104 RowCursor{+i,+c,+kr,+kc,+row,+offset,+ow,+h,+w,+s,+p,+taps} = cursor 105 %Equal.sym(Packing.RowTiles(),state,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,ow,h,w,s,p,taps}),equation) : RowEvidence 106 after_kernel_rows(depth,input,Loop.run(~Packing.RowTiles(),~Packing.kernel_row_step(~block),taps,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,ow,h,w,s,p,taps})), 107 kernel_row_loop(~block,~sampled,depth,taps,input,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,ow,h,w,s,p,taps}),input_evidence, 108 RowEvidence{output,RowCursor{i,c,kr,kc,row,offset,ow,h,w,s,p,taps},{==},balanced})) 109 110 def channel_loop(~block: Packing.Block(), 111 ~sampled: @-input: Array -> @+row: U32 -> @+column: U32 -> @+channel: U32 -> @+height: U32 -> @+width: U32 -> @+stride: U32 -> @+next_row: U32 -> @+offset: U32 -> @+output_width: U32 -> 112 @+kernel_column: U32 -> @+padding: U32 -> 113 @+evidence: PackingReadWitness.family(input) -> TraversalWitness.Sample,Tile.Tile8,input,block(input,row,column,channel,height,width,stride,next_row,offset,output_width,kernel_column,padding)>, 114 +depth: Nat,count: Nat,-input: Array,-state: Packing.RowTiles(),+input_evidence: PackingReadWitness.family(input),evidence: RowEvidence) -> 115 RowEvidence: 116 match count: 117 case 0n: evidence 118 case 1n+rest: channel_loop(~block,~sampled,depth,rest,input,Packing.channel_step(~block,state),input_evidence,channel(~block,~sampled,depth,input,state,input_evidence,evidence)) 119 120 type CrossingCursor is Data: 121 CrossingCursor{index: U32,reduction: U32,position: U32,height: U32,width: U32,kernel: U32,stride: U32,padding: U32,output_width: U32} 122 123 def restore_crossing(input: Array,output: Array,cursor: CrossingCursor) -> Packing.Crossings(): 124 CrossingCursor{i,q,position,h,w,k,s,p,ow} = cursor 125 Packing.Crossing{input,output,i,q,position,h,w,k,s,p,ow} 126 127 type CrossingEvidence<-input: Array,-depth: Nat,-state: Packing.Crossings()> is Data: 128 CrossingEvidence{-output: Array,cursor: CrossingCursor,equation: {state == restore_crossing(input,output,cursor) : Packing.Crossings()}, 129 balanced: Witness.Witness(Tile.Tile8,depth,output)} 130 131 def stored_crossing(output: Array,cursor: CrossingCursor,result: Array & Tile.Tile8) -> Packing.Crossings(): 132 CrossingCursor{i,q,position,h,w,k,s,p,ow} = cursor 133 Packing.store_crossing(output,i,q,position,h,w,k,s,p,ow,result) 134 135 def store_crossing(+depth: Nat,-input: Array,-output: Array,cursor: CrossingCursor,-result: Array & Tile.Tile8, 136 balanced: Witness.Witness(Tile.Tile8,depth,output),sample: TraversalWitness.Sample,Tile.Tile8,input,result>) -> 137 CrossingEvidence: 138 CrossingCursor{+i,q,position,h,w,k,s,p,ow} = cursor 139 TraversalWitness.Sample{value,equation} = sample 140 %Equal.sym(Array & Tile.Tile8,result,(input,value),equation) : CrossingEvidence 141 CrossingEvidence{Array.set(Tile.Tile8,output,i,value),CrossingCursor{(i + 1 : U32),(q + 1 : U32),position,h,w,k,s,p,ow},{==},Witness.write(Tile.Tile8,depth,output,i,value,balanced)} 142 143 def crossing(+depth: Nat,-input: Array,-state: Packing.Crossings(),+input_evidence: PackingReadWitness.family(input),evidence: CrossingEvidence) -> 144 CrossingEvidence: 145 CrossingEvidence{output,cursor,equation,balanced} = evidence 146 CrossingCursor{+i,+q,+position,+h,+w,+k,+s,+p,+ow} = cursor 147 %Equal.sym(Packing.Crossings(),state,restore_crossing(input,output,CrossingCursor{i,q,position,h,w,k,s,p,ow}),equation) : CrossingEvidence 148 +context = PackingSamples.pixel_block(position,q,h,w,k,s,p,ow) 149 -first = PackingSamples.pixel(input,position,q,h,w,k,s,p,ow) 150 store_crossing(depth,input,output,CrossingCursor{i,q,position,h,w,k,s,p,ow},Tile.read_tail(~Array,~PackingSamples.PixelBlock,~PackingSamples.read_block_pixel,first,context),balanced, 151 TileWitness.tail(~Array,~PackingSamples.PixelBlock,~PackingReadWitness.family,~PackingSamples.read_block_pixel,~PackingReadWitness.crossing_read,input,context,first,input_evidence, 152 PackingReadWitness.first(input,position,q,h,w,k,s,p,ow,input_evidence))) 153 154 def crossing_loop(+depth: Nat,count: Nat,-input: Array,-state: Packing.Crossings(),+input_evidence: PackingReadWitness.family(input),evidence: CrossingEvidence) -> 155 CrossingEvidence: 156 match count: 157 case 0n: evidence 158 case 1n+rest: crossing_loop(depth,rest,input,Packing.crossing_step(state),input_evidence,crossing(depth,input,state,input_evidence,evidence)) 159 160 type TileCursor is Data: 161 TileCursor{index: U32,position: U32,height: U32,width: U32,kernel: U32,stride: U32,padding: U32,output_width: U32,channels: Nat,taps: Nat,inner: Nat} 162 163 def restore_tiles(input: Array,output: Array,cursor: TileCursor) -> Packing.Tiled(): 164 TileCursor{i,position,h,w,k,s,p,ow,channels,taps,inner} = cursor 165 Packing.Tiles{input,output,i,position,h,w,k,s,p,ow,channels,taps,inner} 166 167 type TileEvidence<-input: Array,-depth: Nat,-state: Packing.Tiled()> is Data: 168 TileEvidence{-output: Array,cursor: TileCursor,equation: {state == restore_tiles(input,output,cursor) : Packing.Tiled()},balanced: Witness.Witness(Tile.Tile8,depth,output)} 169 170 def after_row_tile(-depth: Nat,-input: Array,position: U32,kernel: U32,channels: Nat,inner: Nat,-state: Packing.RowTiles(), 171 evidence: RowEvidence) -> TileEvidence: 172 RowEvidence{output,cursor,equation,balanced} = evidence 173 RowCursor{i,c,kr,kc,row,offset,ow,h,w,s,p,taps} = cursor 174 %Equal.sym(Packing.RowTiles(),state,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,ow,h,w,s,p,taps}),equation) : 175 TileEvidence 176 TileEvidence{output,TileCursor{i,(position + 8 : U32),h,w,kernel,s,p,ow,channels,taps,inner},{==},balanced} 177 178 def after_crossing(-depth: Nat,-input: Array,channels: Nat,taps: Nat,inner: Nat,-state: Packing.Crossings(),evidence: CrossingEvidence) -> 179 TileEvidence: 180 CrossingEvidence{output,cursor,equation,balanced} = evidence 181 CrossingCursor{i,q,position,h,w,k,s,p,ow} = cursor 182 %Equal.sym(Packing.Crossings(),state,restore_crossing(input,output,CrossingCursor{i,q,position,h,w,k,s,p,ow}),equation) : 183 TileEvidence 184 TileEvidence{output,TileCursor{i,(position + 8 : U32),h,w,k,s,p,ow,channels,taps,inner},{==},balanced} 185 186 def tile(+depth: Nat,within: Bool,wraps: Bool,-input: Array,-output: Array,+i: U32,+position: U32,+h: U32,+w: U32,+k: U32,+s: U32,+p: U32,+ow: U32, 187 +channels: Nat,+taps: Nat,+inner: Nat,+input_evidence: PackingReadWitness.family(input),balanced: Witness.Witness(Tile.Tile8,depth,output)) -> 188 TileEvidence: 189 match within wraps: 190 case True{} wraps: 191 +cursor = {RowCursor{i,0,0,0,(position / ow : U32),(position % ow : U32),ow,h,w,s,p,taps} : RowCursor} 192 after_row_tile(depth,input,position,k,channels,inner,Loop.run(~Packing.RowTiles(),~Packing.channel_step(~Packing.block),channels,restore_row(input,output,cursor)), 193 channel_loop(~Packing.block,~PackingReadWitness.row_block,depth,channels,input,restore_row(input,output,cursor),input_evidence,RowEvidence{output,cursor,{==},balanced})) 194 case False{} True{}: 195 +cursor = {RowCursor{i,0,0,0,(position / ow : U32),(position % ow : U32),ow,h,w,s,p,taps} : RowCursor} 196 after_row_tile(depth,input,position,k,channels,inner,Loop.run(~Packing.RowTiles(),~Packing.channel_step(~Packing.wrap_block),channels,restore_row(input,output,cursor)), 197 channel_loop(~Packing.wrap_block,~PackingReadWitness.wrap_block,depth,channels,input,restore_row(input,output,cursor),input_evidence,RowEvidence{output,cursor,{==},balanced})) 198 case False{} False{}: 199 +cursor = {CrossingCursor{i,0,position,h,w,k,s,p,ow} : CrossingCursor} 200 after_crossing(depth,input,channels,taps,inner,Loop.run(~Packing.Crossings(),~Packing.crossing_step,inner,restore_crossing(input,output,cursor)), 201 crossing_loop(depth,inner,input,restore_crossing(input,output,cursor),input_evidence,CrossingEvidence{output,cursor,{==},balanced})) 202 203 def tile_step(+depth: Nat,-input: Array,-state: Packing.Tiled(),+input_evidence: PackingReadWitness.family(input),evidence: TileEvidence) -> 204 TileEvidence: 205 TileEvidence{output,cursor,equation,balanced} = evidence 206 TileCursor{+i,+position,+h,+w,+k,+s,+p,+ow,+channels,+taps,+inner} = cursor 207 %Equal.sym(Packing.Tiled(),state,restore_tiles(input,output,TileCursor{i,position,h,w,k,s,p,ow,channels,taps,inner}),equation) : TileEvidence 208 tile(depth,U32.is_le((position % ow + 8 : U32),ow),U32.is_le((position % ow + 8 : U32),(ow + ow : U32)),input,output,i,position,h,w,k,s,p,ow,channels,taps,inner,input_evidence,balanced) 209 210 def tile_loop(+depth: Nat,count: Nat,-input: Array,-state: Packing.Tiled(),+input_evidence: PackingReadWitness.family(input),evidence: TileEvidence) -> 211 TileEvidence: 212 match count: 213 case 0n: evidence 214 case 1n+rest: tile_loop(depth,rest,input,Packing.tile_step(state),input_evidence,tile_step(depth,input,state,input_evidence,evidence)) 215 216 def finish(-depth: Nat,-input: Array,-state: Packing.Tiled(),evidence: TileEvidence) -> 217 TraversalWitness.ResultEvidence,Tile.Tile8,input,depth,Packing.finish(state)>: 218 TileEvidence{output,cursor,equation,balanced} = evidence 219 TileCursor{i,position,h,w,k,s,p,ow,channels,taps,inner} = cursor 220 %Equal.sym(Packing.Tiled(),state,restore_tiles(input,output,TileCursor{i,position,h,w,k,s,p,ow,channels,taps,inner}),equation) : 221 TraversalWitness.ResultEvidence,Tile.Tile8,input,depth,Packing.finish(_)> 222 TraversalWitness.ResultEvidence{output,{==},balanced} 223 224 def pack(+depth: Nat,+count: Nat,+start: U32,-input: Array,-output: Array,+h: U32,+w: U32,+k: U32,+s: U32,+p: U32,+ow: U32,+inner: U32, 225 +input_evidence: PackingReadWitness.family(input),balanced: Witness.Witness(Tile.Tile8,depth,output)) -> 226 TraversalWitness.ResultEvidence,Tile.Tile8,input,depth,Packing.pack(count,start,input,output,h,w,k,s,p,ow,inner)>: 227 +cursor = {TileCursor{0,start,h,w,k,s,p,ow,U32.to_nat((inner / (k * k) : U32)),U32.to_nat(k),U32.to_nat(inner)} : TileCursor} 228 finish(depth,input,Loop.run(~Packing.Tiled(),~Packing.tile_step,Nat.div(count,U32.to_nat(inner)),restore_tiles(input,output,cursor)), 229 tile_loop(depth,Nat.div(count,U32.to_nat(inner)),input,restore_tiles(input,output,cursor),input_evidence,TileEvidence{output,cursor,{==},balanced}))