kernel_spatial_packing.bend

stelliferous@0.0.2.0 · 175 lines · 33 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.2.0/traversal_loop.bend as Loop
11import stelliferous@0.0.2.0/kernel_packing_samples.bend as PackingSamples
12import stelliferous@0.0.2.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 RowState<ownership,-InputStorage: Kind(ownership),-OutputStorage: Kind(ownership)> 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
74def RowTiles() -> Type:
75  RowState<&1,Array<F32>,Array<Tile.Tile8>>
76
77def store_block(output: Array<Tile.Tile8>,+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<F32> & 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
82def 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
88def 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
92def 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
96def 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
100def 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.
105type CrossingState<ownership,-InputStorage: Kind(ownership),-OutputStorage: Kind(ownership)> 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
108def Crossings() -> Type:
109  CrossingState<&1,Array<F32>,Array<Tile.Tile8>>
110
111def store_crossing(output: Array<Tile.Tile8>,+index: U32,reduction: U32,position: U32,height: U32,width: U32,kernel: U32,stride: U32,padding: U32,output_width: U32,
112  pair: Array<F32> & 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
116def 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<F32>,~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.
124type TileState<ownership,-InputStorage: Kind(ownership),-OutputStorage: Kind(ownership)> 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
128def Tiled() -> Type:
129  TileState<&1,Array<F32>,Array<Tile.Tile8>>
130
131def 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
135def 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.
140def tile(within: Bool,wraps: Bool,input: Array<F32>,output: Array<Tile.Tile8>,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.
154def 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
159def finish(state: Tiled()) -> Array<F32> & Array<Tile.Tile8>:
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.
164def 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,
165  +inner: U32) -> Array<F32> & Array<Tile.Tile8>:
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.
170type 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
173def run(call: Call,input: Array<F32>,output: Array<Tile.Tile8>) -> Array<F32> & Array<Tile.Tile8>:
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)