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,window: Packing.Stencil,tap_rows: Nat,tap_columns: Nat} 15 16 def restore_row(input: Array,output: Array,cursor: RowCursor) -> Packing.RowTiles(): 17 RowCursor{i,c,kr,kc,row,offset,window,tap_rows,tap_columns} = cursor 18 Packing.RowTile{input,output,i,c,kr,kc,row,offset,window,tap_rows,tap_columns} 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,window,tap_rows,tap_columns} = cursor 25 Packing.store_block(output,i,c,kr,kc,row,offset,window,tap_rows,tap_columns,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,window,tap_rows,tap_columns} = 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,window,tap_rows,tap_columns},{==},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,+window,+tap_rows,+tap_columns} = cursor 43 Packing.Stencil{+base,+h,+w,tr,tc,+s,+oy,+ox,+ow,inner} = window 44 %Equal.sym(Packing.RowTiles(),state,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,window,tap_rows,tap_columns}),equation) : RowEvidence 45 +iy = (row * s + kr + oy : U32) 46 +ix = (offset * s + kc + ox : U32) 47 +cb = (base + c * h * w : U32) 48 +p = (0 - ox : U32) 49 store_row(depth,input,output,RowCursor{i,c,kr,kc,row,offset,window,tap_rows,tap_columns},block(input,iy,ix,cb,h,w,s,((row + 1) * s + kr + oy : U32),offset,ow,kc,p),balanced, 50 sampled(input,iy,ix,cb,h,w,s,((row + 1) * s + kr + oy : U32),offset,ow,kc,p,input_evidence)) 51 52 def column_loop(~block: Packing.Block(), 53 ~sampled: @-input: Array -> @+row: U32 -> @+column: U32 -> @+channel: U32 -> @+height: U32 -> @+width: U32 -> @+stride: U32 -> @+next_row: U32 -> @+offset: U32 -> @+output_width: U32 -> 54 @+kernel_column: U32 -> @+padding: U32 -> 55 @+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)>, 56 +depth: Nat,count: Nat,-input: Array,-state: Packing.RowTiles(),+input_evidence: PackingReadWitness.family(input),evidence: RowEvidence) -> 57 RowEvidence: 58 match count: 59 case 0n: evidence 60 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)) 61 62 def after_columns(-depth: Nat,-input: Array,-state: Packing.RowTiles(),evidence: RowEvidence) -> 63 RowEvidence: 64 RowEvidence{output,cursor,equation,balanced} = evidence 65 RowCursor{i,c,kr,kc,row,offset,window,tap_rows,tap_columns} = cursor 66 %Equal.sym(Packing.RowTiles(),state,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,window,tap_rows,tap_columns}),equation) : RowEvidence 67 RowEvidence{output,RowCursor{i,c,(kr + 1 : U32),0,row,offset,window,tap_rows,tap_columns},{==},balanced} 68 69 def kernel_row(~block: Packing.Block(), 70 ~sampled: @-input: Array -> @+row: U32 -> @+column: U32 -> @+channel: U32 -> @+height: U32 -> @+width: U32 -> @+stride: U32 -> @+next_row: U32 -> @+offset: U32 -> @+output_width: U32 -> 71 @+kernel_column: U32 -> @+padding: U32 -> 72 @+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)>, 73 +depth: Nat,-input: Array,-state: Packing.RowTiles(),+input_evidence: PackingReadWitness.family(input),evidence: RowEvidence) -> 74 RowEvidence: 75 RowEvidence{output,cursor,equation,balanced} = evidence 76 RowCursor{+i,+c,+kr,+kc,+row,+offset,+window,+tap_rows,+tap_columns} = cursor 77 %Equal.sym(Packing.RowTiles(),state,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,window,tap_rows,tap_columns}),equation) : RowEvidence 78 after_columns(depth,input,Loop.run(~Packing.RowTiles(),~Packing.column_step(~block),tap_columns,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,window,tap_rows,tap_columns})), 79 column_loop(~block,~sampled,depth,tap_columns,input,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,window,tap_rows,tap_columns}),input_evidence, 80 RowEvidence{output,RowCursor{i,c,kr,kc,row,offset,window,tap_rows,tap_columns},{==},balanced})) 81 82 def kernel_row_loop(~block: Packing.Block(), 83 ~sampled: @-input: Array -> @+row: U32 -> @+column: U32 -> @+channel: U32 -> @+height: U32 -> @+width: U32 -> @+stride: U32 -> @+next_row: U32 -> @+offset: U32 -> @+output_width: U32 -> 84 @+kernel_column: U32 -> @+padding: U32 -> 85 @+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)>, 86 +depth: Nat,count: Nat,-input: Array,-state: Packing.RowTiles(),+input_evidence: PackingReadWitness.family(input),evidence: RowEvidence) -> 87 RowEvidence: 88 match count: 89 case 0n: evidence 90 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)) 91 92 def after_kernel_rows(-depth: Nat,-input: Array,-state: Packing.RowTiles(),evidence: RowEvidence) -> 93 RowEvidence: 94 RowEvidence{output,cursor,equation,balanced} = evidence 95 RowCursor{i,c,kr,kc,row,offset,window,tap_rows,tap_columns} = cursor 96 %Equal.sym(Packing.RowTiles(),state,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,window,tap_rows,tap_columns}),equation) : RowEvidence 97 RowEvidence{output,RowCursor{i,(c + 1 : U32),0,0,row,offset,window,tap_rows,tap_columns},{==},balanced} 98 99 def channel(~block: Packing.Block(), 100 ~sampled: @-input: Array -> @+row: U32 -> @+column: U32 -> @+channel: U32 -> @+height: U32 -> @+width: U32 -> @+stride: U32 -> @+next_row: U32 -> @+offset: U32 -> @+output_width: U32 -> 101 @+kernel_column: U32 -> @+padding: U32 -> 102 @+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)>, 103 +depth: Nat,-input: Array,-state: Packing.RowTiles(),+input_evidence: PackingReadWitness.family(input),evidence: RowEvidence) -> 104 RowEvidence: 105 RowEvidence{output,cursor,equation,balanced} = evidence 106 RowCursor{+i,+c,+kr,+kc,+row,+offset,+window,+tap_rows,+tap_columns} = cursor 107 %Equal.sym(Packing.RowTiles(),state,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,window,tap_rows,tap_columns}),equation) : RowEvidence 108 after_kernel_rows(depth,input,Loop.run(~Packing.RowTiles(),~Packing.kernel_row_step(~block),tap_rows,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,window,tap_rows,tap_columns})), 109 kernel_row_loop(~block,~sampled,depth,tap_rows,input,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,window,tap_rows,tap_columns}),input_evidence, 110 RowEvidence{output,RowCursor{i,c,kr,kc,row,offset,window,tap_rows,tap_columns},{==},balanced})) 111 112 def channel_loop(~block: Packing.Block(), 113 ~sampled: @-input: Array -> @+row: U32 -> @+column: U32 -> @+channel: U32 -> @+height: U32 -> @+width: U32 -> @+stride: U32 -> @+next_row: U32 -> @+offset: U32 -> @+output_width: U32 -> 114 @+kernel_column: U32 -> @+padding: U32 -> 115 @+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)>, 116 +depth: Nat,count: Nat,-input: Array,-state: Packing.RowTiles(),+input_evidence: PackingReadWitness.family(input),evidence: RowEvidence) -> 117 RowEvidence: 118 match count: 119 case 0n: evidence 120 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)) 121 122 type CrossingCursor is Data: 123 CrossingCursor{index: U32,reduction: U32,position: U32,window: Packing.Stencil} 124 125 def restore_crossing(input: Array,output: Array,cursor: CrossingCursor) -> Packing.Crossings(): 126 CrossingCursor{i,q,position,window} = cursor 127 Packing.Crossing{input,output,i,q,position,window} 128 129 type CrossingEvidence<-input: Array,-depth: Nat,-state: Packing.Crossings()> is Data: 130 CrossingEvidence{-output: Array,cursor: CrossingCursor,equation: {state == restore_crossing(input,output,cursor) : Packing.Crossings()}, 131 balanced: Witness.Witness(Tile.Tile8,depth,output)} 132 133 def stored_crossing(output: Array,cursor: CrossingCursor,result: Array & Tile.Tile8) -> Packing.Crossings(): 134 CrossingCursor{i,q,position,window} = cursor 135 Packing.store_crossing(output,i,q,position,window,result) 136 137 def store_crossing(+depth: Nat,-input: Array,-output: Array,cursor: CrossingCursor,-result: Array & Tile.Tile8, 138 balanced: Witness.Witness(Tile.Tile8,depth,output),sample: TraversalWitness.Sample,Tile.Tile8,input,result>) -> 139 CrossingEvidence: 140 CrossingCursor{+i,q,position,window} = cursor 141 TraversalWitness.Sample{value,equation} = sample 142 %Equal.sym(Array & Tile.Tile8,result,(input,value),equation) : CrossingEvidence 143 CrossingEvidence{Array.set(Tile.Tile8,output,i,value),CrossingCursor{(i + 1 : U32),(q + 1 : U32),position,window},{==},Witness.write(Tile.Tile8,depth,output,i,value,balanced)} 144 145 def crossing(+depth: Nat,-input: Array,-state: Packing.Crossings(),+input_evidence: PackingReadWitness.family(input),evidence: CrossingEvidence) -> 146 CrossingEvidence: 147 CrossingEvidence{output,cursor,equation,balanced} = evidence 148 CrossingCursor{+i,+q,+position,+window} = cursor 149 Packing.Stencil{+base,+h,+w,+tr,+tc,+s,+oy,+ox,+ow,inner} = window 150 %Equal.sym(Packing.Crossings(),state,restore_crossing(input,output,CrossingCursor{i,q,position,window}),equation) : CrossingEvidence 151 +context = PackingSamples.window_pixel_block(position,q,base,h,w,tr,tc,s,oy,ox,ow) 152 -first = PackingSamples.window_pixel(input,position,q,base,h,w,tr,tc,s,oy,ox,ow) 153 store_crossing(depth,input,output,CrossingCursor{i,q,position,window},Tile.read_tail(~Array,~PackingSamples.WindowPixelBlock,~PackingSamples.read_window_pixel,first,context),balanced, 154 TileWitness.tail(~Array,~PackingSamples.WindowPixelBlock,~PackingReadWitness.family,~PackingSamples.read_window_pixel,~PackingReadWitness.window_read,input,context,first,input_evidence, 155 PackingReadWitness.window_first(input,position,q,base,h,w,tr,tc,s,oy,ox,ow,input_evidence))) 156 157 def crossing_loop(+depth: Nat,count: Nat,-input: Array,-state: Packing.Crossings(),+input_evidence: PackingReadWitness.family(input),evidence: CrossingEvidence) -> 158 CrossingEvidence: 159 match count: 160 case 0n: evidence 161 case 1n+rest: crossing_loop(depth,rest,input,Packing.crossing_step(state),input_evidence,crossing(depth,input,state,input_evidence,evidence)) 162 163 type TileCursor is Data: 164 TileCursor{index: U32,position: U32,window: Packing.Stencil,channels: Nat,tap_rows: Nat,tap_columns: Nat,inner: Nat} 165 166 def restore_tiles(input: Array,output: Array,cursor: TileCursor) -> Packing.Tiled(): 167 TileCursor{i,position,window,channels,tap_rows,tap_columns,inner} = cursor 168 Packing.Tiles{input,output,i,position,window,channels,tap_rows,tap_columns,inner} 169 170 type TileEvidence<-input: Array,-depth: Nat,-state: Packing.Tiled()> is Data: 171 TileEvidence{-output: Array,cursor: TileCursor,equation: {state == restore_tiles(input,output,cursor) : Packing.Tiled()},balanced: Witness.Witness(Tile.Tile8,depth,output)} 172 173 def after_row_tile(-depth: Nat,-input: Array,position: U32,channels: Nat,inner: Nat,-state: Packing.RowTiles(), 174 evidence: RowEvidence) -> TileEvidence: 175 RowEvidence{output,cursor,equation,balanced} = evidence 176 RowCursor{i,c,kr,kc,row,offset,window,tap_rows,tap_columns} = cursor 177 %Equal.sym(Packing.RowTiles(),state,restore_row(input,output,RowCursor{i,c,kr,kc,row,offset,window,tap_rows,tap_columns}),equation) : 178 TileEvidence 179 TileEvidence{output,TileCursor{i,(position + 8 : U32),window,channels,tap_rows,tap_columns,inner},{==},balanced} 180 181 def after_crossing(-depth: Nat,-input: Array,channels: Nat,tap_rows: Nat,tap_columns: Nat,inner: Nat,-state: Packing.Crossings(),evidence: CrossingEvidence) -> 182 TileEvidence: 183 CrossingEvidence{output,cursor,equation,balanced} = evidence 184 CrossingCursor{i,q,position,window} = cursor 185 %Equal.sym(Packing.Crossings(),state,restore_crossing(input,output,CrossingCursor{i,q,position,window}),equation) : 186 TileEvidence 187 TileEvidence{output,TileCursor{i,(position + 8 : U32),window,channels,tap_rows,tap_columns,inner},{==},balanced} 188 189 def tile(+depth: Nat,within: Bool,wraps: Bool,-input: Array,-output: Array,+i: U32,+position: U32,+window: Packing.Stencil, 190 +channels: Nat,+tap_rows: Nat,+tap_columns: Nat,+inner: Nat,+input_evidence: PackingReadWitness.family(input),balanced: Witness.Witness(Tile.Tile8,depth,output)) -> 191 TileEvidence: 192 match within wraps: 193 case True{} wraps: 194 Packing.Stencil{base,h,w,tr,tc,s,oy,ox,+ow,inner_word} = window 195 +cursor = {RowCursor{i,0,0,0,(position / ow : U32),(position % ow : U32),window,tap_rows,tap_columns} : RowCursor} 196 after_row_tile(depth,input,position,channels,inner,Loop.run(~Packing.RowTiles(),~Packing.channel_step(~Packing.block),channels,restore_row(input,output,cursor)), 197 channel_loop(~Packing.block,~PackingReadWitness.row_block,depth,channels,input,restore_row(input,output,cursor),input_evidence,RowEvidence{output,cursor,{==},balanced})) 198 case False{} True{}: 199 Packing.Stencil{base,h,w,tr,tc,s,oy,ox,+ow,inner_word} = window 200 +cursor = {RowCursor{i,0,0,0,(position / ow : U32),(position % ow : U32),window,tap_rows,tap_columns} : RowCursor} 201 after_row_tile(depth,input,position,channels,inner,Loop.run(~Packing.RowTiles(),~Packing.channel_step(~Packing.wrap_block),channels,restore_row(input,output,cursor)), 202 channel_loop(~Packing.wrap_block,~PackingReadWitness.wrap_block,depth,channels,input,restore_row(input,output,cursor),input_evidence,RowEvidence{output,cursor,{==},balanced})) 203 case False{} False{}: 204 +cursor = {CrossingCursor{i,0,position,window} : CrossingCursor} 205 after_crossing(depth,input,channels,tap_rows,tap_columns,inner,Loop.run(~Packing.Crossings(),~Packing.crossing_step,inner,restore_crossing(input,output,cursor)), 206 crossing_loop(depth,inner,input,restore_crossing(input,output,cursor),input_evidence,CrossingEvidence{output,cursor,{==},balanced})) 207 208 def tile_step(+depth: Nat,-input: Array,-state: Packing.Tiled(),+input_evidence: PackingReadWitness.family(input),evidence: TileEvidence) -> 209 TileEvidence: 210 TileEvidence{output,cursor,equation,balanced} = evidence 211 TileCursor{+i,+position,+window,+channels,+tap_rows,+tap_columns,+inner} = cursor 212 Packing.Stencil{base,h,w,tr,tc,s,oy,ox,+ow,inner_word} = window 213 %Equal.sym(Packing.Tiled(),state,restore_tiles(input,output,TileCursor{i,position,window,channels,tap_rows,tap_columns,inner}),equation) : TileEvidence 214 tile(depth,U32.is_le((position % ow + 8 : U32),ow),U32.is_le((position % ow + 8 : U32),(ow + ow : U32)),input,output,i,position,window,channels,tap_rows,tap_columns,inner,input_evidence,balanced) 215 216 def tile_loop(+depth: Nat,count: Nat,-input: Array,-state: Packing.Tiled(),+input_evidence: PackingReadWitness.family(input),evidence: TileEvidence) -> 217 TileEvidence: 218 match count: 219 case 0n: evidence 220 case 1n+rest: tile_loop(depth,rest,input,Packing.tile_step(state),input_evidence,tile_step(depth,input,state,input_evidence,evidence)) 221 222 def finish(-depth: Nat,-input: Array,-state: Packing.Tiled(),evidence: TileEvidence) -> 223 TraversalWitness.ResultEvidence,Tile.Tile8,input,depth,Packing.finish(state)>: 224 TileEvidence{output,cursor,equation,balanced} = evidence 225 TileCursor{i,position,window,channels,tap_rows,tap_columns,inner} = cursor 226 %Equal.sym(Packing.Tiled(),state,restore_tiles(input,output,TileCursor{i,position,window,channels,tap_rows,tap_columns,inner}),equation) : 227 TraversalWitness.ResultEvidence,Tile.Tile8,input,depth,Packing.finish(_)> 228 TraversalWitness.ResultEvidence{output,{==},balanced} 229 230 def pack_window(+depth: Nat,+window: Packing.Stencil,+count: Nat,+start: U32,-input: Array,-output: Array, 231 +input_evidence: PackingReadWitness.family(input),balanced: Witness.Witness(Tile.Tile8,depth,output)) -> 232 TraversalWitness.ResultEvidence,Tile.Tile8,input,depth,Packing.pack_window(window,count,start,input,output)>: 233 Packing.Stencil{base,h,w,+tr,+tc,s,oy,ox,ow,+inner} = window 234 +cursor = {TileCursor{0,start,window,U32.to_nat((inner / (tr * tc) : U32)),U32.to_nat(tr),U32.to_nat(tc),U32.to_nat(inner)} : TileCursor} 235 finish(depth,input,Loop.run(~Packing.Tiled(),~Packing.tile_step,Nat.div(count,U32.to_nat(inner)),restore_tiles(input,output,cursor)), 236 tile_loop(depth,Nat.div(count,U32.to_nat(inner)),input,restore_tiles(input,output,cursor),input_evidence,TileEvidence{output,cursor,{==},balanced})) 237 238 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, 239 +input_evidence: PackingReadWitness.family(input),balanced: Witness.Witness(Tile.Tile8,depth,output)) -> 240 TraversalWitness.ResultEvidence,Tile.Tile8,input,depth,Packing.pack(count,start,input,output,h,w,k,s,p,ow,inner)>: 241 pack_window(depth,Packing.square(h,w,k,s,p,ow,inner),count,start,input,output,input_evidence,balanced)