kernel_packing_samples.bend

stelliferous@0.0.2.0 · 46 lines · 9 definitions

1# Pixel reads of an entry whose tile crosses an output row.
2import Base
3import stelliferous@0.0.2.0/kernel_tile.bend as Tile
4
5def masked(value: F32, inside: Bool) -> F32:
6  match inside:
7    case True{}:
8      value
9    case False{}:
10      0.0
11
12def mask_pixel(inside: Bool, pair: Array<F32> & F32) -> Array<F32> & F32:
13  (input_storage, value) = pair
14  (input_storage, masked(value, inside))
15
16def pixel(input_storage: Array<F32>, +output_position: U32, +reduction_index: U32, +input_height: U32, +input_width: U32, +kernel_size: U32, +stride: U32, +padding: U32, +output_width: U32) -> Array<F32> & F32:
17  +input_row = ((output_position / output_width) * stride + (reduction_index / kernel_size) % kernel_size - padding : U32)
18  +input_column = ((output_position % output_width) * stride + reduction_index % kernel_size - padding : U32)
19  at = ((reduction_index / (kernel_size * kernel_size)) * input_height * input_width + input_row * input_width + input_column : U32)
20  mask_pixel(U32.is_lt(input_row, input_height) && U32.is_lt(input_column, input_width), Array.get(F32, input_storage, at))
21
22type PixelBlock is Data:
23  PixelBlock{position: U32,kernel_row: U32,kernel_column: U32,channel_base: U32,height: U32,width: U32,stride: U32,padding: U32,output_width: U32}
24
25# Reduction coordinates are invariant across the whole block, at any width.
26
27def pixel_block(position: U32,+reduction: U32,+height: U32,+width: U32,+kernel: U32,stride: U32,padding: U32,output_width: U32) -> PixelBlock:
28  PixelBlock{position,(reduction / kernel % kernel : U32),(reduction % kernel : U32),
29    (reduction / (kernel * kernel) * height * width : U32),height,width,stride,padding,output_width}
30
31type PixelAddress is Data:
32  PixelAddress{index: U32,inside: Bool}
33
34def pixel_address(context: PixelBlock,index: Nat) -> PixelAddress:
35  PixelBlock{position,kernel_row,kernel_column,channel_base,+height,+width,+stride,+padding,+output_width} = context
36  +position = (position + U32.from_nat(index) : U32)
37  +row = (position / output_width * stride + kernel_row - padding : U32)
38  +column = (position % output_width * stride + kernel_column - padding : U32)
39  PixelAddress{(channel_base + row * width + column : U32),U32.is_lt(row,height) && U32.is_lt(column,width)}
40
41def read_pixel_address(input: Array<F32>,address: PixelAddress) -> Array<F32> & F32:
42  PixelAddress{index,inside} = address
43  mask_pixel(inside,Array.get(F32,input,index))
44
45def read_block_pixel(input: Array<F32>,context: PixelBlock,index: Nat) -> Array<F32> & F32:
46  read_pixel_address(input,pixel_address(context,index))