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 RowState is Kind(ownership): 71 RowTile{input: InputStorage,output: OutputStorage,index: U32,channel: U32,kernel_row: U32,kernel_column: U32,row: U32,offset: U32,output_width: U32, 72 height: U32,width: U32,stride: U32,padding: U32,taps: Nat} 73 74 def RowTiles() -> Type: 75 RowState<&1,Array,Array> 76 77 def store_block(output: Array,+index: U32,channel: U32,kernel_row: U32,kernel_column: U32,row: U32,offset: U32,output_width: U32, 78 height: U32,width: U32,stride: U32,padding: U32,taps: Nat,pair: Array & Tile.Tile8) -> RowTiles(): 79 (input,value) = pair 80 RowTile{input,Array.set(Tile.Tile8,output,index,value),(index + 1 : U32),channel,kernel_row,(kernel_column + 1 : U32),row,offset,output_width,height,width,stride,padding,taps} 81 82 def column_step(~block: Block(),state: RowTiles()) -> RowTiles(): 83 RowTile{input,output,index,+channel,+kernel_row,+kernel_column,+row,+offset,+output_width,+height,+width,+stride,+padding,taps} = state 84 store_block(output,index,channel,kernel_row,kernel_column,row,offset,output_width,height,width,stride,padding,taps, 85 block(input,(row * stride + kernel_row - padding : U32),(offset * stride + kernel_column - padding : U32),(channel * height * width : U32),height,width,stride, 86 ((row + 1) * stride + kernel_row - padding : U32),offset,output_width,kernel_column,padding)) 87 88 def next_kernel_row(state: RowTiles()) -> RowTiles(): 89 RowTile{input,output,index,channel,kernel_row,kernel_column,row,offset,output_width,height,width,stride,padding,taps} = state 90 RowTile{input,output,index,channel,(kernel_row + 1 : U32),0,row,offset,output_width,height,width,stride,padding,taps} 91 92 def kernel_row_step(~block: Block(),state: RowTiles()) -> RowTiles(): 93 RowTile{input,output,index,channel,kernel_row,kernel_column,row,offset,output_width,height,width,stride,padding,+taps} = state 94 next_kernel_row(Loop.run(~RowTiles(),~column_step(~block),taps,RowTile{input,output,index,channel,kernel_row,kernel_column,row,offset,output_width,height,width,stride,padding,taps})) 95 96 def next_channel(state: RowTiles()) -> RowTiles(): 97 RowTile{input,output,index,channel,kernel_row,kernel_column,row,offset,output_width,height,width,stride,padding,taps} = state 98 RowTile{input,output,index,(channel + 1 : U32),0,0,row,offset,output_width,height,width,stride,padding,taps} 99 100 def channel_step(~block: Block(),state: RowTiles()) -> RowTiles(): 101 RowTile{input,output,index,channel,kernel_row,kernel_column,row,offset,output_width,height,width,stride,padding,+taps} = state 102 next_channel(Loop.run(~RowTiles(),~kernel_row_step(~block),taps,RowTile{input,output,index,channel,kernel_row,kernel_column,row,offset,output_width,height,width,stride,padding,taps})) 103 104 # The entries of a tile that crosses an output row, by reduction index. 105 type CrossingState is Kind(ownership): 106 Crossing{input: InputStorage,output: OutputStorage,index: U32,reduction: U32,position: U32,height: U32,width: U32,kernel: U32,stride: U32,padding: U32,output_width: U32} 107 108 def Crossings() -> Type: 109 CrossingState<&1,Array,Array> 110 111 def store_crossing(output: Array,+index: U32,reduction: U32,position: U32,height: U32,width: U32,kernel: U32,stride: U32,padding: U32,output_width: U32, 112 pair: Array & Tile.Tile8) -> Crossings(): 113 (input,value) = pair 114 Crossing{input,Array.set(Tile.Tile8,output,index,value),(index + 1 : U32),(reduction + 1 : U32),position,height,width,kernel,stride,padding,output_width} 115 116 def crossing_step(state: Crossings()) -> Crossings(): 117 Crossing{input,output,index,+reduction,+position,+height,+width,+kernel,+stride,+padding,+output_width} = state 118 store_crossing(output,index,reduction,position,height,width,kernel,stride,padding,output_width, 119 Tile.read_tail(~Array,~PackingSamples.PixelBlock,~PackingSamples.read_block_pixel, 120 PackingSamples.pixel(input,position,reduction,height,width,kernel,stride,padding,output_width), 121 PackingSamples.pixel_block(position,reduction,height,width,kernel,stride,padding,output_width))) 122 123 # The tiles from output position `position`, eight positions each. 124 type TileState is Kind(ownership): 125 Tiles{input: InputStorage,output: OutputStorage,index: U32,position: U32,height: U32,width: U32,kernel: U32,stride: U32,padding: U32,output_width: U32, 126 channels: Nat,taps: Nat,inner: Nat} 127 128 def Tiled() -> Type: 129 TileState<&1,Array,Array> 130 131 def after_row_tile(position: U32,kernel: U32,channels: Nat,inner: Nat,state: RowTiles()) -> Tiled(): 132 RowTile{input,output,index,channel,kernel_row,kernel_column,row,offset,output_width,height,width,stride,padding,taps} = state 133 Tiles{input,output,index,(position + 8 : U32),height,width,kernel,stride,padding,output_width,channels,taps,inner} 134 135 def after_crossing(channels: Nat,taps: Nat,inner: Nat,state: Crossings()) -> Tiled(): 136 Crossing{input,output,index,reduction,position,height,width,kernel,stride,padding,output_width} = state 137 Tiles{input,output,index,(position + 8 : U32),height,width,kernel,stride,padding,output_width,channels,taps,inner} 138 139 # A tile inside one output row, one that ends in the next row, and any other. 140 def tile(within: Bool,wraps: Bool,input: Array,output: Array,index: U32,+position: U32,height: U32,width: U32,kernel: U32,stride: U32,padding: U32,+output_width: U32, 141 +channels: Nat,taps: Nat,+inner: Nat) -> Tiled(): 142 match within wraps: 143 case True{} wraps: 144 after_row_tile(position,kernel,channels,inner,Loop.run(~RowTiles(),~channel_step(~block),channels, 145 RowTile{input,output,index,0,0,0,(position / output_width : U32),(position % output_width : U32),output_width,height,width,stride,padding,taps})) 146 case False{} True{}: 147 after_row_tile(position,kernel,channels,inner,Loop.run(~RowTiles(),~channel_step(~wrap_block),channels, 148 RowTile{input,output,index,0,0,0,(position / output_width : U32),(position % output_width : U32),output_width,height,width,stride,padding,taps})) 149 case False{} False{}: 150 after_crossing(channels,taps,inner,Loop.run(~Crossings(),~crossing_step,inner,Crossing{input,output,index,0,position,height,width,kernel,stride,padding,output_width})) 151 152 # Eight positions from output column c stay in one output row when c + 8 is at 153 # most the output width, and in two when it is at most twice the width. 154 def tile_step(state: Tiled()) -> Tiled(): 155 Tiles{input,output,index,+position,height,width,kernel,stride,padding,+output_width,channels,taps,inner} = state 156 tile(U32.is_le((position % output_width + 8 : U32),output_width),U32.is_le((position % output_width + 8 : U32),(output_width + output_width : U32)), 157 input,output,index,position,height,width,kernel,stride,padding,output_width,channels,taps,inner) 158 159 def finish(state: Tiled()) -> Array & Array: 160 Tiles{input,output,index,position,height,width,kernel,stride,padding,output_width,channels,taps,inner} = state 161 (input,output) 162 163 # count entries, inner per tile, from output position start. 164 def pack(count: Nat,start: U32,input: Array,output: Array,height: U32,width: U32,+kernel: U32,stride: U32,padding: U32,output_width: U32, 165 +inner: U32) -> Array & Array: 166 finish(Loop.run(~Tiled(),~tile_step,Nat.div(count,U32.to_nat(inner)), 167 Tiles{input,output,0,start,height,width,kernel,stride,padding,output_width,U32.to_nat((inner / (kernel * kernel) : U32)),U32.to_nat(kernel),U32.to_nat(inner)})) 168 169 # One packing call shared by raw convolution and certified Buffer entry points. 170 type Call is Data: 171 Call{count: Nat,start: U32,height: U32,width: U32,kernel: U32,stride: U32,padding: U32,output_width: U32,inner: U32} 172 173 def run(call: Call,input: Array,output: Array) -> Array & Array: 174 Call{count,start,h,w,k,s,p,ow,inner} = call 175 pack(count,start,input,output,h,w,k,s,p,ow,inner)