tensor_band_window.bend
stelliferous@0.0.3.0 · 149 lines · 0 definitions
1# Window geometry is independent of convolution arithmetic. Each output row
2# belongs to one existing core band; neighbours contribute only halo rows.
3import Base
4import stelliferous@0.0.3.0/tensor_band.bend as Bands
5import stelliferous@0.0.3.0/storage_buffer.bend as Storage
6import stelliferous@0.0.3.0/traversal_partition.bend as Partition
7
8type Geometry is Data:
9 Geometry{input: Bands.Geometry,kernel: U32,stride: U32,padding: U32,output_height: U32,output_width: U32}
10
11type Rows is Data:
12 Rows{first: U32,count: U32}
13
14# A leaf's output rows, its halo rows above and below the core, and the
15# image rows its window holds: from the stride-aligned first row its outputs
16# read to the last, so the window packs as a shorter image.
17type RowWindow is Data:
18 RowWindow{output: Rows,above: Rows,below: Rows,window: Rows}
19
20def ceiling(+value: U32,+stride: U32) -> U32:
21 (value / stride + Bool.pick(U32,U32.is_zero((value % stride : U32)),0,1) : U32)
22
23def subtract_clipped(+value: U32,+amount: U32) -> U32:
24 (U32.max(value,amount) - amount : U32)
25
26def clipped(+first: U32,+end: U32) -> Rows:
27 Rows{first,subtract_clipped(end,first)}
28
29def nonempty(active: Bool,window: RowWindow) -> RowWindow:
30 match active window:
31 case True{} window: window
32 case False{} RowWindow{output,Rows{above_first,above_count},Rows{below_first,below_count},Rows{window_first,window_count}}:
33 RowWindow{output,Rows{above_first,0},Rows{below_first,0},Rows{window_first,0}}
34
35# The first row a window holds: the first row its outputs read, moved down
36# to a stride multiple.
37def start(+height: U32,+output_first: U32,+stride: U32,+padding: U32) -> U32:
38 +needed = U32.min(height,subtract_clipped((output_first * stride : U32),padding))
39 subtract_clipped(needed,(needed % stride : U32))
40
41# Output partition boundaries follow the input cores, with the first and last
42# cores also owning any output rows introduced by padding. The halos hold
43# every window row outside the core, so a window is its image rows exactly.
44# Caller geometry checks establish positive stride/height and bounded products.
45def window(geometry: Geometry,+first: U32,+count: U32) -> RowWindow:
46 Geometry{Bands.Geometry{batch,channels,+height,width},+kernel,+stride,+padding,+output_height,output_width} = geometry
47 +image = (first / height : U32)
48 +core_first = (first % height : U32)
49 +core_end = (core_first + count : U32)
50 +out_first = U32.min(output_height,ceiling(core_first,stride))
51 +out_end = Bool.pick(U32,U32.is_eq(core_end,height),output_height,U32.min(output_height,ceiling(core_end,stride)))
52 +needed_first = start(height,out_first,stride,padding)
53 +needed_end = U32.min(height,subtract_clipped((subtract_clipped(out_end,1) * stride + kernel : U32),padding))
54 +above_first = (image * height + needed_first : U32)
55 +above_end = (image * height + U32.min(core_first,needed_end) : U32)
56 +below_first = (image * height + U32.max(core_end,needed_first) : U32)
57 +below_end = (image * height + needed_end : U32)
58 nonempty(U32.is_lt(out_first,out_end),RowWindow{Rows{(image * output_height + out_first : U32),subtract_clipped(out_end,out_first)},
59 clipped(above_first,above_end),clipped(below_first,below_end),clipped(above_first,below_end)})
60
61type Halo<-Element: Data> is Type:
62 Halo{above: Storage.Buffer<Element>,below: Storage.Buffer<Element>,window: RowWindow}
63
64def Halos(~Element: Data) -> Type:
65 Partition.Tree<&1,Halo<Element>,Unit>
66
67def extract(~Element: Data,+geometry: Bands.Geometry,fill: Element,range: Rows,bands: Bands.Bands(~Element)) -> Bands.Bands(~Element) & Storage.Buffer<Element>:
68 Rows{+first,+count} = range
69 Bands.range(~Element,geometry,first,count,bands,Bands.allocate(~Element,fill,Bands.local_count(geometry,count)))
70
71def halo(~Element: Data,+geometry: Bands.Geometry,+fill: Element,+first: U32,+count: U32,window: RowWindow,bands: Bands.Bands(~Element)) -> Bands.Bands(~Element) & Halos(~Element):
72 RowWindow{+output,+above,+below,+rows} = window
73 Storage.bind_pair(~Bands.Bands(~Element),~Storage.Buffer<Element>,~(Bands.Bands(~Element) & Halos(~Element)),
74 extract(~Element,geometry,fill,above,bands),bands => upper =>
75 Storage.bind_pair(~Bands.Bands(~Element),~Storage.Buffer<Element>,~(Bands.Bands(~Element) & Halos(~Element)),
76 extract(~Element,geometry,fill,below,bands),bands => lower =>
77 (bands,Partition.Leaf{Halo{upper,lower,RowWindow{output,above,below,rows}},first,count,Unit{}})))
78
79def copied_window(~Element: Data,owners: Bands.Bands(~Element) & Storage.Buffer<Element>) -> Storage.Buffer<Element>:
80 (sources,window) = owners
81 window
82
83# The window rows of every channel, stored contiguously. A window equal to
84# the core (for example a 1x1 stride-1 operation) is the core itself; any
85# other window is the shared range copy of its rows from the upper halo, the
86# core and the lower halo, one contiguous range per channel per source. The
87# three copies are direct calls: a leaf task must not reach a closure.
88def assembled(~Element: Data,exact: Bool,+fill: Element,+geometry: Bands.Geometry,+window: Rows,+core: Rows,+above: Rows,+below: Rows,
89 core_buffer: Storage.Buffer<Element>,above_buffer: Storage.Buffer<Element>,below_buffer: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
90 match exact:
91 case True{}: core_buffer
92 case False{}:
93 Rows{+window_first,+window_rows} = window
94 Rows{core_first,core_rows} = core
95 Rows{above_first,above_rows} = above
96 Rows{below_first,below_rows} = below
97 copied_window(~Element,Bands.range_leaf(~Element,geometry,window_first,window_rows,0,below_first,below_rows,below_buffer,
98 copied_window(~Element,Bands.range_leaf(~Element,geometry,window_first,window_rows,0,core_first,core_rows,core_buffer,
99 copied_window(~Element,Bands.range_leaf(~Element,geometry,window_first,window_rows,0,above_first,above_rows,above_buffer,
100 Bands.allocate(~Element,fill,Bands.local_count(geometry,window_rows))))))))
101
102def assemble(~Element: Data,+fill: Element,+geometry: Bands.Geometry,+window: Rows,+core: Rows,+above: Rows,+below: Rows,
103 core_buffer: Storage.Buffer<Element>,above_buffer: Storage.Buffer<Element>,below_buffer: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
104 Rows{+window_first,+window_rows} = window
105 Rows{+core_first,+core_rows} = core
106 assembled(~Element,U32.is_eq(window_first,core_first) && U32.is_eq(window_rows,core_rows),fill,geometry,window,core,above,below,
107 core_buffer,above_buffer,below_buffer)
108
109def input_geometry(geometry: Geometry) -> Bands.Geometry:
110 Geometry{input,kernel,stride,padding,output_height,output_width} = geometry
111 input
112
113# The full core tree is threaded through the preparation pass without cloning.
114# Each allocation is a requested neighbour range, never a core-sized window.
115# The resulting halo tree follows the same layout, ready to pair with cores.
116# What the halo of one band needs.
117type Padded<-Element: Data> is Data:
118 Padded{geometry: Geometry,fill: Element}
119
120def halo_band(~Element: Data,context: Padded<Element>,first: U32,count: U32,bands: Bands.Bands(~Element)) -> Bands.Bands(~Element) & Halos(~Element):
121 Padded{+geometry,fill} = context
122 +first = first
123 +count = count
124 halo(~Element,input_geometry(geometry),fill,first,count,window(geometry,first,count),bands)
125
126def prepare(~Element: Data,+geometry: Geometry,+fill: Element,layout: Bands.Plan,bands: Bands.Bands(~Element)) -> Bands.Bands(~Element) & Halos(~Element):
127 Bands.build(~Halo<Element>,~Bands.Bands(~Element),~Padded<Element>,~halo_band(~Element),layout,Padded{geometry,fill},bands)
128
129# A core with its halo goes through the leaf; a core without one is not
130# transformed.
131def window_band(~Element: Data,~Context: Data,~leaf: Context -> U32 -> U32 -> Storage.Buffer<Element> -> Halo<Element> -> Bands.Transformed(~Element),
132 context: Context,first: U32,rows: U32,pair: Bands.Paired<Storage.Buffer<Element>,Halo<Element>>,resource: Unit) -> Bands.Transformed(~Element):
133 match pair:
134 case Bands.Both{input,halo}: leaf(context,first,rows,input,halo)
135 case Bands.Only{input,own_first,own_rows}: Partition.Leaf{Bands.Checked{input,False{}},own_first,own_rows,Unit{}}
136
137def unwindowed(~Element: Data,tree: Partition.Tree<&1,Bands.Paired<Storage.Buffer<Element>,Halo<Element>>,Unit>) -> Bands.Transformed(~Element):
138 Bands.untransformed(~Element,Bands.left_of(~Storage.Buffer<Element>,~Halo<Element>,tree))
139
140# Every core leaf with its halo through its own checked call, in the shared
141# band recursion. The bands come back beside whether every call was usable;
142# trees of another shape are not transformed.
143def apply(~Element: Data,~Context: Data,~leaf: Context -> U32 -> U32 -> Storage.Buffer<Element> -> Halo<Element> -> Bands.Transformed(~Element),
144 layout: Bands.Plan,enabled: Bool,context: Context,bands: Bands.Bands(~Element),halos: Halos(~Element)) ->
145 Bands.Bands(~Element) & Bool:
146 Bands.reduce(~Bands.Checked<Element>,~Storage.Buffer<Element>,~Bool,~Unit,~Bands.checked_band(~Element),~Bool.and,Unit{},
147 Bands.each(~Bands.Paired<Storage.Buffer<Element>,Halo<Element>>,~Bands.Checked<Element>,~Unit,~Context,~Bands.unshared,
148 ~window_band(~Element,~Context,~leaf),~unwindowed(~Element),
149 layout,enabled,context,Bands.zip(~Storage.Buffer<Element>,~Halo<Element>,bands,halos),Unit{}))