kernel_spatial_packing.bend
stelliferous@0.0.3.0 · 189 lines · 0 definitions
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.
9import Base
10import stelliferous@0.0.3.0/traversal_loop.bend as Loop
11import stelliferous@0.0.3.0/kernel_packing_samples.bend as PackingSamples
12import stelliferous@0.0.3.0/kernel_tile.bend as Tile
13
14def masked(inside: Bool, pair: Array<F32> & F32) -> Array<F32> & F32:
15 (input_storage, value) = pair
16 (input_storage, PackingSamples.masked(value, inside))
17
18def pixel(input_storage: Array<F32>, +input_row: U32, +input_column: U32, +channel_base: U32, +input_height: U32, +input_width: U32) -> Array<F32> & 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
21type PixelBlock is Data:
22 PixelBlock{row: U32,column: U32,channel_base: U32,height: U32,width: U32,stride: U32}
23
24def read_block_pixel(input: Array<F32>,context: PixelBlock,index: Nat) -> Array<F32> & 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.
29def block(input: Array<F32>,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<F32> & Tile.Tile8:
30 +row = row
31 +column = column
32 +channel_base = channel_base
33 +height = height
34 +width = width
35 Tile.read_tail(~Array<F32>,~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`.
39type 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
42def wrap_pixel(same: Bool,input: Array<F32>,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> & 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
48def wrap_lane(input: Array<F32>,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> & 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
52def read_wrap_pixel(input: Array<F32>,context: WrapBlock,index: Nat) -> Array<F32> & 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
56def wrap_block(input: Array<F32>,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<F32> & Tile.Tile8:
57 +row = row
58 +column = column
59 +channel_base = channel_base
60 +height = height
61 +width = width
62 Tile.read_tail(~Array<F32>,~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
65def Block() -> Type:
66 Array<F32> -> U32 -> U32 -> U32 -> U32 -> U32 -> U32 -> U32 -> U32 -> U32 -> U32 -> U32 -> Array<F32> & 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.
70type 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
73def 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
76type RowState<ownership,-InputStorage: Kind(ownership),-OutputStorage: Kind(ownership)> 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
80def RowTiles() -> Type:
81 RowState<&1,Array<F32>,Array<Tile.Tile8>>
82
83def store_block(output: Array<Tile.Tile8>,+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<F32> & 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
88def 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
95def 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
99def 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
103def 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
107def 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.
112type CrossingState<ownership,-InputStorage: Kind(ownership),-OutputStorage: Kind(ownership)> is Kind(ownership):
113 Crossing{input: InputStorage,output: OutputStorage,index: U32,reduction: U32,position: U32,window: Stencil}
114
115def Crossings() -> Type:
116 CrossingState<&1,Array<F32>,Array<Tile.Tile8>>
117
118def store_crossing(output: Array<Tile.Tile8>,+index: U32,reduction: U32,position: U32,window: Stencil,
119 pair: Array<F32> & 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
123def 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<F32>,~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.
132type TileState<ownership,-InputStorage: Kind(ownership),-OutputStorage: Kind(ownership)> 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
135def Tiled() -> Type:
136 TileState<&1,Array<F32>,Array<Tile.Tile8>>
137
138def 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
142def 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.
147def tile(within: Bool,wraps: Bool,input: Array<F32>,output: Array<Tile.Tile8>,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.
163def 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
169def finish(state: Tiled()) -> Array<F32> & Array<Tile.Tile8>:
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.
174def pack_window(+window: Stencil,count: Nat,start: U32,input: Array<F32>,output: Array<Tile.Tile8>) -> Array<F32> & Array<Tile.Tile8>:
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
179def pack(count: Nat,start: U32,input: Array<F32>,output: Array<Tile.Tile8>,height: U32,width: U32,kernel: U32,stride: U32,+padding: U32,output_width: U32,
180 inner: U32) -> Array<F32> & Array<Tile.Tile8>:
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.
184type 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
187def run(call: Call,input: Array<F32>,output: Array<Tile.Tile8>) -> Array<F32> & Array<Tile.Tile8>:
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)