kernel_packing_samples.bend

stelliferous@0.0.3.0 · 47 lines · 0 definitions

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