kernel_band_pipeline.bend
stelliferous@0.0.3.0 · 101 lines · 0 definitions
1# Independent row bands keep their output owners. Fork distributes existing
2# core/halo owners; only read-only weights and bias need affine clones.
3import Base
4import stelliferous@0.0.3.0/tensor_band.bend as Bands
5import stelliferous@0.0.3.0/tensor_band_window.bend as BandWindow
6import stelliferous@0.0.3.0/kernel_window.bend as KernelWindow
7import stelliferous@0.0.3.0/kernel_pipeline.bend as Pipeline
8import stelliferous@0.0.3.0/kernel_tile.bend as Tile
9import stelliferous@0.0.3.0/storage_buffer.bend as Storage
10import stelliferous@0.0.3.0/storage_clone.bend as Clone
11import stelliferous@0.0.3.0/traversal_partition.bend as Partition
12import stelliferous@0.0.3.0/math_activation.bend as Activations
13
14type Context is Data:
15 Context{window: BandWindow.Geometry,channels: U32,inner: U32,parallel: Bool,activation: Activations.Activation}
16
17# A leaf assembles the rows its outputs read into one contiguous window whose
18# first row is a stride multiple, then runs the leaf on it: packing a
19# window is packing a shorter image with output rows shifted by first / stride.
20def packed_window(~custom: F32 -> F32,activation: Activations.Activation,+geometry: BandWindow.Geometry,+channels: U32,+inner: U32,+first: U32,+rows: U32,+output_first: U32,+output_rows: U32,
21 +window: BandWindow.Rows,+above_range: BandWindow.Rows,+below_range: BandWindow.Rows,
22 core: Storage.Buffer<F32>,above: Storage.Buffer<F32>,below: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>) -> Storage.Buffer<F32>:
23 BandWindow.Geometry{Bands.Geometry{batch,+ci,+height,+width},+kernel,+stride,+padding,+output_height,+output_width} = geometry
24 BandWindow.Rows{+window_first,+window_rows} = window
25 +positions = (output_rows * output_width : U32)
26 Pipeline.blocked(~custom,activation,
27 Pipeline.Work{channels,window_rows,width,kernel,stride,padding,output_width,positions,inner},
28 KernelWindow.start(((output_first - first / height * output_height) * output_width : U32),stride,output_width,(window_first - first / height * height : U32)),
29 BandWindow.assemble(~F32,0.0,BandWindow.input_geometry(geometry),window,BandWindow.Rows{first,rows},above_range,below_range,core,above,below),weights,bias)
30
31# A window is empty when the leaf has no outputs or its outputs read only
32# padding above the image. The leaf then packs the one row at the stride
33# position of its first output, which no output reads.
34def core_row(+geometry: BandWindow.Geometry,+first: U32,+output_first: U32) -> BandWindow.Rows:
35 BandWindow.Geometry{Bands.Geometry{batch,ci,+height,width},kernel,+stride,padding,+output_height,output_width} = geometry
36 BandWindow.Rows{(first / height * height + (output_first - first / height * output_height) * stride : U32),1}
37
38def source_rows(empty: Bool,+geometry: BandWindow.Geometry,+first: U32,+output_first: U32,window: BandWindow.Rows) -> BandWindow.Rows:
39 match empty:
40 case False{}: window
41 case True{}: core_row(geometry,first,output_first)
42
43def windowed(~custom: F32 -> F32,activation: Activations.Activation,+geometry: BandWindow.Geometry,+channels: U32,+inner: U32,+first: U32,+rows: U32,+rows_read: BandWindow.RowWindow,
44 core: Storage.Buffer<F32>,above: Storage.Buffer<F32>,below: Storage.Buffer<F32>,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>) -> Bands.Bands(~F32):
45 BandWindow.RowWindow{BandWindow.Rows{+output_first,+output_rows},+above_range,+below_range,BandWindow.Rows{+window_first,+window_rows}} = rows_read
46 Partition.Leaf{packed_window(~custom,activation,geometry,channels,inner,first,rows,output_first,output_rows,
47 source_rows(U32.is_zero(window_rows),geometry,first,output_first,BandWindow.Rows{window_first,window_rows}),
48 above_range,below_range,core,above,below,weights,bias),output_first,output_rows,Unit{}}
49
50def leaf(~custom: F32 -> F32,context: Context,+first: U32,+rows: U32,core: Storage.Buffer<F32>,halo: BandWindow.Halo<F32>,
51 weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>) -> Bands.Bands(~F32):
52 Context{+geometry,+channels,+inner,parallel,activation} = context
53 BandWindow.Halo{above,below,+rows_read} = halo
54 windowed(~custom,activation,geometry,channels,inner,first,rows,rows_read,core,above,below,weights,bias)
55
56def parallel(context: Context) -> Bool:
57 Context{window,channels,inner,parallel,activation} = context
58 parallel
59
60def Parameters() -> Type:
61 Storage.Buffer<Tile.Tile8> & Storage.Buffer<F32>
62
63def shared(weights: Storage.Buffer<Tile.Tile8> & Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32> & Storage.Buffer<F32>) -> Parameters() & Parameters():
64 (left_weights,right_weights) = weights
65 (left_bias,right_bias) = bias
66 ((left_weights,left_bias),(right_weights,right_bias))
67
68# Each half of a split owns its copy of the read-only weights and bias.
69def share(parameters: Parameters()) -> Parameters() & Parameters():
70 (weights,bias) = parameters
71 shared(Clone.clone(~Tile.Tile8,weights),Clone.clone(~F32,bias))
72
73def band(~custom: F32 -> F32,context: Context,first: U32,rows: U32,core: Storage.Buffer<F32>,halo: BandWindow.Halo<F32>,parameters: Parameters()) -> Bands.Bands(~F32):
74 (weights,bias) = parameters
75 +first = first
76 +rows = rows
77 leaf(~custom,context,first,rows,core,halo,weights,bias)
78
79# Every core leaf with its halo, through the shared band recursion. Trees of
80# another shape keep the cores.
81def run(~custom: F32 -> F32,plan: Bands.Plan,+context: Context,bands: Bands.Bands(~F32),halos: BandWindow.Halos(~F32),
82 weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>) -> Bands.Bands(~F32):
83 Bands.each(~Bands.Paired<Storage.Buffer<F32>,BandWindow.Halo<F32>>,~Storage.Buffer<F32>,~Parameters(),~Context,~share,
84 ~Bands.paired(~Storage.Buffer<F32>,~BandWindow.Halo<F32>,~Parameters(),~Context,~band(~custom)),
85 ~Bands.left_of(~Storage.Buffer<F32>,~BandWindow.Halo<F32>),
86 plan,parallel(context),context,Bands.zip(~Storage.Buffer<F32>,~BandWindow.Halo<F32>,bands,halos),(weights,bias))
87
88def prepared(~custom: F32 -> F32,+context: Context,+plan: Bands.Plan,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,
89 owners: Bands.Bands(~F32) & BandWindow.Halos(~F32)) -> Bands.Bands(~F32):
90 (bands,halos) = owners
91 run(~custom,plan,context,bands,halos,weights,bias)
92
93def planned(~custom: F32 -> F32,+context: Context,weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>,
94 owners: Bands.Bands(~F32) & Bands.Plan) -> Bands.Bands(~F32):
95 Context{window,channels,inner,parallel,activation} = context
96 (bands,+plan) = owners
97 prepared(~custom,context,plan,weights,bias,BandWindow.prepare(~F32,window,0.0,plan,bands))
98
99# The weights arrive in the order the kernel reads them.
100def execute(~custom: F32 -> F32,context: Context,bands: Bands.Bands(~F32),weights: Storage.Buffer<Tile.Tile8>,bias: Storage.Buffer<F32>) -> Bands.Bands(~F32):
101 planned(~custom,context,weights,bias,Bands.plan(~F32,bands))