kernel_gather.bend

stelliferous@0.0.2.0 · 34 lines · 6 definitions

1import Base
2import stelliferous@0.0.2.0/traversal_array.bend as Traversal
3
4# Clip once per leaf. A channel is one contiguous copy; padding never becomes
5# a store into the next channel. Both owners stay affine across the channel loop.
6def gather_width_when(inside: Bool,offset: U32,output_positions: U32,span: U32) -> U32:
7  match inside:
8    case False{}: 0
9    case True{}: U32.min(span,(output_positions - offset : U32))
10
11def gather_width(+offset: U32,+output_positions: U32,span: U32) -> U32:
12  gather_width_when(U32.is_lt(offset,output_positions),offset,output_positions,span)
13
14def gather_channel(+channel: U32,+width: U32,+offset: U32,+output_positions: U32,+span: U32,
15  owners: Array<F32> & Array<F32>) -> Array<F32> & Array<F32>:
16  (input_storage,output_storage) = owners
17  Traversal.copy(~F32,U32.to_nat(width),(channel * span : U32),
18    (channel * output_positions + offset : U32),input_storage,output_storage)
19
20def gather_channels(remaining: Nat,+channel: U32,+width: U32,+offset: U32,+output_positions: U32,+span: U32,
21  owners: Array<F32> & Array<F32>) -> Array<F32> & Array<F32>:
22  match remaining:
23    case 0n: owners
24    case 1n+rest:
25      gather_channels(rest,(channel + 1 : U32),width,offset,output_positions,span,
26        gather_channel(channel,width,offset,output_positions,span,owners))
27
28def gathered_output(owners: Array<F32> & Array<F32>) -> Array<F32>:
29  (input_storage,output_storage) = owners
30  output_storage
31
32def gather_block(channels: U32,+offset: U32,+output_positions: U32,+span: U32,input_storage: Array<F32>,output_storage: Array<F32>) -> Array<F32>:
33  gathered_output(gather_channels(U32.to_nat(channels),0,gather_width(offset,output_positions,span),offset,output_positions,span,
34    (input_storage,output_storage)))