1 # One packer for every output width. A tile holds eight consecutive output 2 # positions; its entries follow the reduction order: channel, kernel row, 3 # kernel column. Only a new tile divides: it finds the output row and column 4 # of its first position. A tile inside one or two output rows runs three 5 # nested loops whose counters give the input row and column of lane 0, and 6 # reads the lanes of each row at one stride. A tile that spans more rows (an 7 # output narrower than eight) locates every lane from its position, entry by 8 # entry. 9 import Base 10 import ./traversal_loop.bend as Loop 11 import ./kernel_packing_samples.bend as PackingSamples 12 import ./kernel_tile.bend as Tile 13 14 def masked(inside: Bool, pair: Array & F32) -> Array & F32: 15 (input_storage, value) = pair 16 (input_storage, PackingSamples.masked(value, inside)) 17 18 def pixel(input_storage: Array, +input_row: U32, +input_column: U32, +channel_base: U32, +input_height: U32, +input_width: U32) -> Array & F32: 19 masked(U32.is_lt(input_row, input_height) && U32.is_lt(input_column, input_width), Array.get(F32, input_storage, (channel_base + input_row * input_width + input_column : U32))) 20 21 type PixelBlock is Data: 22 PixelBlock{row: U32,column: U32,channel_base: U32,height: U32,width: U32,stride: U32} 23 24 def read_block_pixel(input: Array,context: PixelBlock,index: Nat) -> Array & F32: 25 PixelBlock{row,column,channel_base,height,width,+stride} = context 26 pixel(input,row,(column + U32.from_nat(index) * stride : U32),channel_base,height,width) 27 28 # Eight lanes of one output row. 29 def block(input: Array,row: U32,column: U32,channel_base: U32,height: U32,width: U32,stride: U32,next_row: U32,offset: U32,output_width: U32,kernel_column: U32,padding: U32) -> Array & Tile.Tile8: 30 +row = row 31 +column = column 32 +channel_base = channel_base 33 +height = height 34 +width = width 35 Tile.read_tail(~Array,~PixelBlock,~read_block_pixel,pixel(input,row,column,channel_base,height,width),PixelBlock{row,column,channel_base,height,width,stride}) 36 37 # Lane 0 is output column `offset`. A lane past the last output column lies 38 # in the next output row, whose input row is `next_row`. 39 type WrapBlock is Data: 40 WrapBlock{row: U32,column: U32,channel_base: U32,height: U32,width: U32,stride: U32,next_row: U32,offset: U32,output_width: U32,kernel_column: U32,padding: U32} 41 42 def wrap_pixel(same: Bool,input: Array,row: U32,column: U32,channel_base: U32,height: U32,width: U32,+stride: U32,next_row: U32,offset: U32,output_width: U32, 43 kernel_column: U32,padding: U32,+lane: U32) -> Array & F32: 44 match same: 45 case True{}: pixel(input,row,(column + lane * stride : U32),channel_base,height,width) 46 case False{}: pixel(input,next_row,((offset + lane - output_width) * stride + kernel_column - padding : U32),channel_base,height,width) 47 48 def wrap_lane(input: Array,row: U32,column: U32,channel_base: U32,height: U32,width: U32,stride: U32,next_row: U32,+offset: U32,+output_width: U32, 49 kernel_column: U32,padding: U32,+lane: U32) -> Array & F32: 50 wrap_pixel(U32.is_lt((offset + lane : U32),output_width),input,row,column,channel_base,height,width,stride,next_row,offset,output_width,kernel_column,padding,lane) 51 52 def read_wrap_pixel(input: Array,context: WrapBlock,index: Nat) -> Array & F32: 53 WrapBlock{row,column,channel_base,height,width,stride,next_row,offset,output_width,kernel_column,padding} = context 54 wrap_lane(input,row,column,channel_base,height,width,stride,next_row,offset,output_width,kernel_column,padding,U32.from_nat(index)) 55 56 def wrap_block(input: Array,row: U32,column: U32,channel_base: U32,height: U32,width: U32,stride: U32,next_row: U32,offset: U32,output_width: U32,kernel_column: U32,padding: U32) -> Array & Tile.Tile8: 57 +row = row 58 +column = column 59 +channel_base = channel_base 60 +height = height 61 +width = width 62 Tile.read_tail(~Array,~WrapBlock,~read_wrap_pixel,pixel(input,row,column,channel_base,height,width), 63 WrapBlock{row,column,channel_base,height,width,stride,next_row,offset,output_width,kernel_column,padding}) 64 65 def Block() -> Type: 66 Array -> U32 -> U32 -> U32 -> U32 -> U32 -> U32 -> U32 -> U32 -> U32 -> U32 -> U32 -> Array & Tile.Tile8 67 68 # The entries of a tile whose lane 0 is output column `offset` of output row 69 # `row` of an output `output_width` wide. block reads the eight lanes of an entry. 70 type Stencil is Data: 71 Stencil{base: U32,height: U32,width: U32,tap_rows: U32,tap_columns: U32,stride: U32,row_origin: U32,column_origin: U32,output_width: U32,inner: U32} 72 73 def square(height: U32,width: U32,+kernel: U32,stride: U32,+padding: U32,output_width: U32,inner: U32) -> Stencil: 74 Stencil{0,height,width,kernel,kernel,stride,(0 - padding : U32),(0 - padding : U32),output_width,inner} 75 76 type RowState is Kind(ownership): 77 RowTile{input: InputStorage,output: OutputStorage,index: U32,channel: U32,kernel_row: U32,kernel_column: U32,row: U32,offset: U32, 78 window: Stencil,tap_rows: Nat,tap_columns: Nat} 79 80 def RowTiles() -> Type: 81 RowState<&1,Array,Array> 82 83 def store_block(output: Array,+index: U32,channel: U32,kernel_row: U32,kernel_column: U32,row: U32,offset: U32,window: Stencil, 84 tap_rows: Nat,tap_columns: Nat,pair: Array & Tile.Tile8) -> RowTiles(): 85 (input,value) = pair 86 RowTile{input,Array.set(Tile.Tile8,output,index,value),(index + 1 : U32),channel,kernel_row,(kernel_column + 1 : U32),row,offset,window,tap_rows,tap_columns} 87 88 def column_step(~block: Block(),state: RowTiles()) -> RowTiles(): 89 RowTile{input,output,index,+channel,+kernel_row,+kernel_column,+row,+offset,+window,tap_rows,tap_columns} = state 90 Stencil{+base,+height,+width,tap_height,tap_width,+stride,+row_origin,+column_origin,+output_width,inner} = window 91 store_block(output,index,channel,kernel_row,kernel_column,row,offset,window,tap_rows,tap_columns, 92 block(input,(row * stride + kernel_row + row_origin : U32),(offset * stride + kernel_column + column_origin : U32),(base + channel * height * width : U32),height,width,stride, 93 ((row + 1) * stride + kernel_row + row_origin : U32),offset,output_width,kernel_column,(0 - column_origin : U32))) 94 95 def next_kernel_row(state: RowTiles()) -> RowTiles(): 96 RowTile{input,output,index,channel,kernel_row,kernel_column,row,offset,window,tap_rows,tap_columns} = state 97 RowTile{input,output,index,channel,(kernel_row + 1 : U32),0,row,offset,window,tap_rows,tap_columns} 98 99 def kernel_row_step(~block: Block(),state: RowTiles()) -> RowTiles(): 100 RowTile{input,output,index,channel,kernel_row,kernel_column,row,offset,window,tap_rows,+tap_columns} = state 101 next_kernel_row(Loop.run(~RowTiles(),~column_step(~block),tap_columns,RowTile{input,output,index,channel,kernel_row,kernel_column,row,offset,window,tap_rows,tap_columns})) 102 103 def next_channel(state: RowTiles()) -> RowTiles(): 104 RowTile{input,output,index,channel,kernel_row,kernel_column,row,offset,window,tap_rows,tap_columns} = state 105 RowTile{input,output,index,(channel + 1 : U32),0,0,row,offset,window,tap_rows,tap_columns} 106 107 def channel_step(~block: Block(),state: RowTiles()) -> RowTiles(): 108 RowTile{input,output,index,channel,kernel_row,kernel_column,row,offset,window,+tap_rows,tap_columns} = state 109 next_channel(Loop.run(~RowTiles(),~kernel_row_step(~block),tap_rows,RowTile{input,output,index,channel,kernel_row,kernel_column,row,offset,window,tap_rows,tap_columns})) 110 111 # The entries of a tile that crosses an output row, by reduction index. 112 type CrossingState is Kind(ownership): 113 Crossing{input: InputStorage,output: OutputStorage,index: U32,reduction: U32,position: U32,window: Stencil} 114 115 def Crossings() -> Type: 116 CrossingState<&1,Array,Array> 117 118 def store_crossing(output: Array,+index: U32,reduction: U32,position: U32,window: Stencil, 119 pair: Array & Tile.Tile8) -> Crossings(): 120 (input,value) = pair 121 Crossing{input,Array.set(Tile.Tile8,output,index,value),(index + 1 : U32),(reduction + 1 : U32),position,window} 122 123 def crossing_step(state: Crossings()) -> Crossings(): 124 Crossing{input,output,index,+reduction,+position,+window} = state 125 Stencil{+base,+height,+width,+tap_rows,+tap_columns,+stride,+row_origin,+column_origin,+output_width,inner} = window 126 store_crossing(output,index,reduction,position,window, 127 Tile.read_tail(~Array,~PackingSamples.WindowPixelBlock,~PackingSamples.read_window_pixel, 128 PackingSamples.window_pixel(input,position,reduction,base,height,width,tap_rows,tap_columns,stride,row_origin,column_origin,output_width), 129 PackingSamples.window_pixel_block(position,reduction,base,height,width,tap_rows,tap_columns,stride,row_origin,column_origin,output_width))) 130 131 # The tiles from output position `position`, eight positions each. 132 type TileState is Kind(ownership): 133 Tiles{input: InputStorage,output: OutputStorage,index: U32,position: U32,window: Stencil,channels: Nat,tap_rows: Nat,tap_columns: Nat,inner: Nat} 134 135 def Tiled() -> Type: 136 TileState<&1,Array,Array> 137 138 def after_row_tile(position: U32,channels: Nat,inner: Nat,state: RowTiles()) -> Tiled(): 139 RowTile{input,output,index,channel,kernel_row,kernel_column,row,offset,window,tap_rows,tap_columns} = state 140 Tiles{input,output,index,(position + 8 : U32),window,channels,tap_rows,tap_columns,inner} 141 142 def after_crossing(channels: Nat,tap_rows: Nat,tap_columns: Nat,inner: Nat,state: Crossings()) -> Tiled(): 143 Crossing{input,output,index,reduction,position,window} = state 144 Tiles{input,output,index,(position + 8 : U32),window,channels,tap_rows,tap_columns,inner} 145 146 # A tile inside one output row, one that ends in the next row, and any other. 147 def tile(within: Bool,wraps: Bool,input: Array,output: Array,index: U32,+position: U32,+window: Stencil, 148 +channels: Nat,tap_rows: Nat,tap_columns: Nat,+inner: Nat) -> Tiled(): 149 match within wraps: 150 case True{} wraps: 151 Stencil{base,height,width,tap_height,tap_width,stride,row_origin,column_origin,+output_width,inner_word} = window 152 after_row_tile(position,channels,inner,Loop.run(~RowTiles(),~channel_step(~block),channels, 153 RowTile{input,output,index,0,0,0,(position / output_width : U32),(position % output_width : U32),window,tap_rows,tap_columns})) 154 case False{} True{}: 155 Stencil{base,height,width,tap_height,tap_width,stride,row_origin,column_origin,+output_width,inner_word} = window 156 after_row_tile(position,channels,inner,Loop.run(~RowTiles(),~channel_step(~wrap_block),channels, 157 RowTile{input,output,index,0,0,0,(position / output_width : U32),(position % output_width : U32),window,tap_rows,tap_columns})) 158 case False{} False{}: 159 after_crossing(channels,tap_rows,tap_columns,inner,Loop.run(~Crossings(),~crossing_step,inner,Crossing{input,output,index,0,position,window})) 160 161 # Eight positions from output column c stay in one output row when c + 8 is at 162 # most the output width, and in two when it is at most twice the width. 163 def tile_step(state: Tiled()) -> Tiled(): 164 Tiles{input,output,index,+position,+window,channels,tap_rows,tap_columns,inner} = state 165 Stencil{base,height,width,tap_height,tap_width,stride,row_origin,column_origin,+output_width,inner_word} = window 166 tile(U32.is_le((position % output_width + 8 : U32),output_width),U32.is_le((position % output_width + 8 : U32),(output_width + output_width : U32)), 167 input,output,index,position,window,channels,tap_rows,tap_columns,inner) 168 169 def finish(state: Tiled()) -> Array & Array: 170 Tiles{input,output,index,position,window,channels,tap_rows,tap_columns,inner} = state 171 (input,output) 172 173 # count entries, inner per tile, from output position start. 174 def pack_window(+window: Stencil,count: Nat,start: U32,input: Array,output: Array) -> Array & Array: 175 Stencil{base,height,width,+tap_rows,+tap_columns,stride,row_origin,column_origin,output_width,+inner} = window 176 finish(Loop.run(~Tiled(),~tile_step,Nat.div(count,U32.to_nat(inner)), 177 Tiles{input,output,0,start,window,U32.to_nat((inner / (tap_rows * tap_columns) : U32)),U32.to_nat(tap_rows),U32.to_nat(tap_columns),U32.to_nat(inner)})) 178 179 def pack(count: Nat,start: U32,input: Array,output: Array,height: U32,width: U32,kernel: U32,stride: U32,+padding: U32,output_width: U32, 180 inner: U32) -> Array & Array: 181 pack_window(square(height,width,kernel,stride,padding,output_width,inner),count,start,input,output) 182 183 # One packing call shared by raw convolution and certified Buffer entry points. 184 type Call is Data: 185 Call{count: Nat,start: U32,height: U32,width: U32,kernel: U32,stride: U32,padding: U32,output_width: U32,inner: U32} 186 187 def run(call: Call,input: Array,output: Array) -> Array & Array: 188 Call{count,start,h,w,k,s,p,ow,inner} = call 189 pack(count,start,input,output,h,w,k,s,p,ow,inner)