tensor_band_pool.bend
stelliferous@0.0.3.0 · 94 lines · 0 definitions
1# Each band leaf assembles the rows its outputs read into one contiguous
2# window and runs the two sweeps of the leaf pool on it: along the columns of
3# every window row, then along the rows, with output row zero's window
4# starting at the window row that the image row would. Sample order is the
5# row-major order.
6import Base
7import stelliferous@0.0.3.0/tensor_access.bend as Access
8import stelliferous@0.0.3.0/tensor_operations.bend as Operations
9import stelliferous@0.0.3.0/tensor_band.bend as Bands
10import stelliferous@0.0.3.0/tensor_band_window.bend as BandWindow
11import stelliferous@0.0.3.0/traversal_partition.bend as Partition
12import stelliferous@0.0.3.0/storage_buffer.bend as Storage
13import stelliferous@0.0.3.0/tensor.bend as Tensors
14import stelliferous@0.0.3.0/tensor_band_schedule.bend as BandSchedule
15import stelliferous@0.0.3.0/kernel_cost.bend as Cost
16import stelliferous@0.0.3.0/default_kernel_cost.bend as DefaultCost
17
18type Context<-Element: Data> is Data:
19 Context{geometry: BandWindow.Geometry,pad: Element,fill: Element,workers: U32}
20
21def window_size(window: BandWindow.RowWindow) -> BandSchedule.Size:
22 BandWindow.RowWindow{BandWindow.Rows{first,+count},above,below,rows} = window
23 BandSchedule.Size{1,count,count}
24
25def leaf_size(window: BandWindow.Geometry,first: U32,rows: U32) -> BandSchedule.Size:
26 window_size(BandWindow.window(window,first,rows))
27
28def plan_size(layout: Bands.Plan,+window: BandWindow.Geometry) -> BandSchedule.Size:
29 Bands.plan_fold(~BandSchedule.Size,~BandWindow.Geometry,~leaf_size,~Bands.plain(~BandSchedule.Size,~BandWindow.Geometry,~BandSchedule.merge),layout,window)
30
31# Bytes price loading, storing and the loop that moves them (copy_byte_ns is
32# measured on copy loops); operations count the rest. Per sample
33# (Access.read_line and the operation): an offset and a subtraction, one bound
34# check, a two-operation index, the padding select and the compare-and-select
35# maximum: 8. Each line adds ten for its origin (Access.line_at). An output is
36# one line along the rows; the lines along the columns are one for every
37# window row, stride for every output row.
38def parallel(window: BandWindow.Geometry,workers: U32,size: BandSchedule.Size) -> Bool:
39 BandWindow.Geometry{Bands.Geometry{batch,channels,ih,iw},kernel,stride,padding,oh,ow} = window
40 +config = DefaultCost.configuration()
41 +samples = U32.to_f32(kernel)
42 +lines = (U32.to_f32(stride) + 1.0 : F32)
43 +line_bytes = (samples * 4.0 + 8.0 : F32)
44 +line_operations = (samples * 8.0 + 10.0 : F32)
45 BandSchedule.sized(Bands.Geometry{batch,channels,oh,ow},workers,Cost.element_ns((lines * line_bytes : F32),(lines * line_operations : F32),config),config,size)
46
47# A band's pool as a checked band: its output beside whether both sweeps
48# accepted.
49def checked(~Element: Data,+first: U32,+rows: U32,result: Operations.Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>) -> Bands.Transformed(~Element):
50 match result:
51 case Operations.Accepted{(input,output)}: Partition.Leaf{Bands.Checked{output,True{}},first,rows,Unit{}}
52 case Operations.Rejected{(input,output)}: Partition.Leaf{Bands.Checked{output,False{}},first,rows,Unit{}}
53
54# The leaf pool of channels stored window rows, whose first output row's
55# window starts at origin. A band that holds no window start has no output
56# rows.
57def pooled(~Element: Data,~operation: Element -> Element -> Element,empty: Bool,+pool: Access.Pool,+pad: Element,+fill: Element,+origin: U32,+channels: U32,+first: U32,
58 input: Storage.Buffer<Element>) -> Bands.Transformed(~Element):
59 match empty:
60 case True{}: Partition.Leaf{Bands.Checked{Tensors.buffer(~Element,fill,0),True{}},first,0,Unit{}}
61 case False{}:
62 Access.Pool{+height,+width,+kernel,+stride,+padding,+out_height,+out_width} = pool
63 checked(~Element,first,out_height,Operations.pool(~Element,~operation,pad,input,Access.Pool{height,width,kernel,stride,padding,out_height,out_width},origin,channels,
64 Tensors.buffer(~Element,fill,(channels * height * out_width : U32)),Tensors.buffer(~Element,fill,(channels * out_height * out_width : U32))))
65
66def windowed(~Element: Data,~operation: Element -> Element -> Element,+geometry: BandWindow.Geometry,+pad: Element,+fill: Element,+first: U32,+rows: U32,
67 +rows_read: BandWindow.RowWindow,core: Storage.Buffer<Element>,above: Storage.Buffer<Element>,below: Storage.Buffer<Element>) -> Bands.Transformed(~Element):
68 BandWindow.Geometry{Bands.Geometry{batch,+channels,+height,+width},+kernel,+stride,+padding,+output_height,+output_width} = geometry
69 BandWindow.RowWindow{BandWindow.Rows{+output_first,+output_rows},+above_range,+below_range,+window} = rows_read
70 BandWindow.Rows{+window_first,+window_rows} = window
71 pooled(~Element,~operation,U32.is_eq(output_rows,0),Access.Pool{window_rows,width,kernel,stride,padding,output_rows,output_width},pad,fill,
72 ((output_first - first / height * output_height) * stride + first / height * height - window_first : U32),channels,output_first,
73 BandWindow.assemble(~Element,fill,BandWindow.input_geometry(geometry),window,BandWindow.Rows{first,rows},above_range,below_range,core,above,below))
74
75def leaf(~Element: Data,~operation: Element -> Element -> Element,context: Context<Element>,first: U32,rows: U32,core: Storage.Buffer<Element>,halo: BandWindow.Halo<Element>) -> Bands.Transformed(~Element):
76 Context{+geometry,+pad,+fill,workers} = context
77 BandWindow.Halo{above,below,+rows_read} = halo
78 +first = first
79 +rows = rows
80 windowed(~Element,~operation,geometry,pad,fill,first,rows,rows_read,core,above,below)
81
82def prepared(~Element: Data,~operation: Element -> Element -> Element,+context: Context<Element>,+layout: Bands.Plan,owners: Bands.Bands(~Element) & BandWindow.Halos(~Element)) ->
83 Bands.Bands(~Element) & Bool:
84 Context{+window,pad,fill,workers} = context
85 (bands,halos) = owners
86 BandWindow.apply(~Element,~Context<Element>,~leaf(~Element,~operation),layout,parallel(window,workers,plan_size(layout,window)),context,bands,halos)
87
88def planned(~Element: Data,~operation: Element -> Element -> Element,+context: Context<Element>,owners: Bands.Bands(~Element) & Bands.Plan) -> Bands.Bands(~Element) & Bool:
89 Context{window,pad,fill,workers} = context
90 (bands,+layout) = owners
91 prepared(~Element,~operation,context,layout,BandWindow.prepare(~Element,window,pad,layout,bands))
92
93def execute(~Element: Data,~operation: Element -> Element -> Element,context: Context<Element>,bands: Bands.Bands(~Element)) -> Bands.Bands(~Element) & Bool:
94 planned(~Element,~operation,context,Bands.plan(~Element,bands))