kernel_band_pipeline.bend

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