tensor_band.bend
stelliferous@0.0.3.0 · 875 lines · 0 definitions
1# Internal H-axis ownership. Offsets and spans are logical image rows, not
2# element positions: offset / height identifies the image. No leaf crosses an
3# image boundary. Its buffer stores [channels][span][width]. The common
4# partition tree owns all buffers; shape and element type remain parameters.
5import Base
6import stelliferous@0.0.3.0/storage_buffer.bend as Storage
7import stelliferous@0.0.3.0/storage_copy.bend as Copy
8import stelliferous@0.0.3.0/storage_clone.bend as Clone
9import stelliferous@0.0.3.0/traversal_partition.bend as Partition
10import stelliferous@0.0.3.0/traversal_array.bend as Traversal
11import stelliferous@0.0.3.0/proofs/traversal_certificate.bend as TraversalCertificate
12import stelliferous@0.0.3.0/tensor_operations.bend as Operations
13import stelliferous@0.0.3.0/storage_allocation.bend as Allocation
14import stelliferous@0.0.3.0/tensor_view.bend as Views
15import stelliferous@0.0.3.0/tensor_shape.bend as Shapes
16import stelliferous@0.0.3.0/tensor_indexed.bend as IndexedAccess
17import stelliferous@0.0.3.0/proofs/tensor_indexed_certificate.bend as IndexedCertificate
18import stelliferous@0.0.3.0/proofs/storage_certificate.bend as Certificate
19
20type Geometry is Data:
21 Geometry{batch: U32,channels: U32,height: U32,width: U32}
22
23def Bands(~Element: Data) -> Type:
24 Partition.Tree<&1,Storage.Buffer<Element>,Unit>
25
26# An existing tree with a state threaded through its leaves in order. The leaf
27# takes the state and gives it back beside what it leaves in the tree.
28def walk(~Input: Type,~Output: Type,~State: Type,~Context: Data,
29 ~leaf: Context -> U32 -> U32 -> Input -> State -> Partition.Tree<&1,Output,Unit> & State,
30 +context: Context,tree: Partition.Tree<&1,Input,Unit>,state: State) -> Partition.Tree<&1,Output,Unit> & State:
31 match tree:
32 case Partition.Leaf{input,first,rows,details}: leaf(context,first,rows,input,state)
33 case Partition.Fork{left,right}:
34 Storage.bind_pair(~Partition.Tree<&1,Output,Unit>,~State,~(Partition.Tree<&1,Output,Unit> & State),
35 walk(~Input,~Output,~State,~Context,~leaf,context,left,state),left => state =>
36 Storage.bind_pair(~Partition.Tree<&1,Output,Unit>,~State,~(Partition.Tree<&1,Output,Unit> & State),
37 walk(~Input,~Output,~State,~Context,~leaf,context,right,state),right => state => (Partition.Fork{left,right},state)))
38
39# Every leaf of a tree consumed into a state, in order.
40def fold(~Input: Type,~State: Type,~Context: Data,~leaf: Context -> U32 -> U32 -> Input -> State -> State,
41 +context: Context,tree: Partition.Tree<&1,Input,Unit>,state: State) -> State:
42 match tree:
43 case Partition.Leaf{input,first,rows,details}: leaf(context,first,rows,input,state)
44 case Partition.Fork{left,right}: fold(~Input,~State,~Context,~leaf,context,right,fold(~Input,~State,~Context,~leaf,context,left,state))
45
46def reduced(~Output: Type,~Value: Type,~join: Value -> Value -> Value,left: Partition.Tree<&1,Output,Unit> & Value,right: Partition.Tree<&1,Output,Unit> & Value) -> Partition.Tree<&1,Output,Unit> & Value:
47 (left,left_value) = left
48 (right,right_value) = right
49 (Partition.Fork{left,right},join(left_value,right_value))
50
51# A value from every leaf, joined where the tree forks. The leaf gives its
52# value beside what it leaves in the tree.
53def reduce(~Input: Type,~Output: Type,~Value: Type,~Context: Data,
54 ~leaf: Context -> U32 -> U32 -> Input -> Partition.Tree<&1,Output,Unit> & Value,~join: Value -> Value -> Value,
55 +context: Context,tree: Partition.Tree<&1,Input,Unit>) -> Partition.Tree<&1,Output,Unit> & Value:
56 match tree:
57 case Partition.Leaf{input,first,rows,details}: leaf(context,first,rows,input)
58 case Partition.Fork{left,right}:
59 reduced(~Output,~Value,~join,reduce(~Input,~Output,~Value,~Context,~leaf,~join,context,left),reduce(~Input,~Output,~Value,~Context,~leaf,~join,context,right))
60
61def elements(geometry: Geometry) -> U32:
62 Geometry{batch,channels,height,width} = geometry
63 (batch * channels * height * width : U32)
64
65def local_count(geometry: Geometry,rows: U32) -> U32:
66 Geometry{batch,channels,height,width} = geometry
67 (channels * rows * width : U32)
68
69# Channel-major local indices map to the same global NCHW index for every
70# partition. A leaf's first row is global across images; rows are image-local.
71def global_index(geometry: Geometry,+first: U32,+rows: U32,+index: U32) -> U32:
72 Geometry{batch,+channels,+height,+width} = geometry
73 +plane = (rows * width : U32)
74 (((first / height * channels + index / plane) * height + first % height) * width + index % plane : U32)
75
76def planes(~Element: Data,remaining: Nat,+channel: U32,+count: U32,+source_stride: U32,+destination_stride: U32,
77 +source_start: U32,+destination_start: U32,owners: Storage.Buffer<Element> & Storage.Buffer<Element>) ->
78 Storage.Buffer<Element> & Storage.Buffer<Element>:
79 match remaining:
80 case 0n: owners
81 case 1n+rest:
82 (source,destination) = owners
83 planes(~Element,rest,(channel + 1 : U32),count,source_stride,destination_stride,source_start,destination_start,
84 Copy.range(~Element,count,(source_start + channel * source_stride : U32),
85 (destination_start + channel * destination_stride : U32),source,destination))
86
87def destination(~Element: Data,owners: Storage.Buffer<Element> & Storage.Buffer<Element>) -> Storage.Buffer<Element>:
88 (source,destination) = owners
89 destination
90
91def gather_leaf(~Element: Data,geometry: Geometry,+first: U32,+rows: U32,input: Storage.Buffer<Element>,output: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
92 Geometry{batch,+channels,+height,+width} = geometry
93 destination(~Element,planes(~Element,U32.to_nat(channels),0,(rows * width : U32),(rows * width : U32),(height * width : U32),0,
94 ((first / height * channels * height + first % height) * width : U32),(input,output)))
95
96# One allocation is provided by the caller. Each value is copied once, using
97# one contiguous range per channel per band. No intermediate whole images.
98def gather_band(~Element: Data,geometry: Geometry,first: U32,rows: U32,input: Storage.Buffer<Element>,output: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
99 +first = first
100 +rows = rows
101 gather_leaf(~Element,geometry,first,rows,input,output)
102
103def gather(~Element: Data,+geometry: Geometry,bands: Bands(~Element),output: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
104 fold(~Storage.Buffer<Element>,~Storage.Buffer<Element>,~Geometry,~gather_band(~Element),geometry,bands,output)
105
106# A leaf that holds every row of the only image is that image in row-major
107# order.
108def whole_image(+geometry: Geometry,first: U32,rows: U32) -> Bool:
109 Geometry{batch,channels,+height,width} = geometry
110 U32.is_eq(batch,1) && U32.is_zero(first) && U32.is_eq(rows,height)
111
112def copied_leaf(~Element: Data,+first: U32,+rows: U32,owners: Storage.Buffer<Element> & Storage.Buffer<Element>) -> Bands(~Element) & Bands(~Element):
113 (left,right) = owners
114 (Partition.Leaf{left,first,rows,Unit{}},Partition.Leaf{right,first,rows,Unit{}})
115
116def copy_band(~Element: Data,context: Unit,first: U32,rows: U32,input: Storage.Buffer<Element>) -> Bands(~Element) & Bands(~Element):
117 +first = first
118 +rows = rows
119 copied_leaf(~Element,first,rows,Clone.clone(~Element,input))
120
121def fork(~Element: Data,left: Bands(~Element),right: Bands(~Element)) -> Bands(~Element):
122 Partition.Fork{left,right}
123
124def copy(~Element: Data,bands: Bands(~Element)) -> Bands(~Element) & Bands(~Element):
125 reduce(~Storage.Buffer<Element>,~Storage.Buffer<Element>,~Bands(~Element),~Unit,~copy_band(~Element),~fork(~Element),Unit{},bands)
126
127# A pure ownership plan contains no Buffer. It is a concrete Data type:
128# duplicating Tree<&2,...> through a computed type makes all constructors
129# shared in the official compiler, including the affine Buffer tree.
130type Plan is Data:
131 Rows{first: U32,count: U32}
132 Split{left: Plan,right: Plan}
133
134def plan_band(~Element: Data,context: Unit,first: U32,rows: U32,input: Storage.Buffer<Element>) -> Bands(~Element) & Plan:
135 +first = first
136 +rows = rows
137 (Partition.Leaf{input,first,rows,Unit{}},Rows{first,rows})
138
139def split_plan(left: Plan,right: Plan) -> Plan:
140 Split{left,right}
141
142def plan(~Element: Data,bands: Bands(~Element)) -> Bands(~Element) & Plan:
143 reduce(~Storage.Buffer<Element>,~Storage.Buffer<Element>,~Plan,~Unit,~plan_band(~Element),~split_plan,Unit{},bands)
144
145# A value from every leaf of a plan, joined where the plan splits. The join
146# also sees the two halves of the plan.
147def plan_fold(~Value: Type,~Context: Data,~leaf: Context -> U32 -> U32 -> Value,~join: Context -> Plan -> Plan -> Value -> Value -> Value,
148 layout: Plan,+context: Context) -> Value:
149 match layout:
150 case Rows{first,rows}: leaf(context,first,rows)
151 case Split{+left,+right}:
152 join(context,left,right,plan_fold(~Value,~Context,~leaf,~join,left,context),plan_fold(~Value,~Context,~leaf,~join,right,context))
153
154# A join that needs neither the context nor the halves of the plan.
155def plain(~Value: Type,~Context: Data,~join: Value -> Value -> Value,context: Context,left: Plan,right: Plan,first: Value,second: Value) -> Value:
156 join(first,second)
157
158def same_plan(left: Plan,right: Plan) -> Bool:
159 match left right:
160 case Rows{first,rows} Rows{other_first,other_rows}:
161 U32.is_eq(first,other_first) && U32.is_eq(rows,other_rows)
162 case Split{left,right} Split{other_left,other_right}:
163 same_plan(left,other_left) && same_plan(right,other_right)
164 case left right: False{}
165
166def same_geometry(left: Geometry,right: Geometry) -> Bool:
167 Geometry{batch,channels,height,width} = left
168 Geometry{other_batch,other_channels,other_height,other_width} = right
169 U32.is_eq(batch,other_batch) && U32.is_eq(channels,other_channels) && U32.is_eq(height,other_height) && U32.is_eq(width,other_width)
170
171# The geometry splits the image rows of an image shape, [channels, height,
172# width] or [batch, channels, height, width]. Operations that read neighbouring
173# rows or whole channels need this layout.
174def image_rows(shape: List<&2,U32>,geometry: Geometry) -> Bool:
175 match shape geometry:
176 case [c,h,w] Geometry{n,ci,hi,wi}: U32.is_eq(n,1) && U32.is_eq(c,ci) && U32.is_eq(h,hi) && U32.is_eq(w,wi)
177 case [n,c,h,w] Geometry{ni,ci,hi,wi}: U32.is_eq(n,ni) && U32.is_eq(c,ci) && U32.is_eq(h,hi) && U32.is_eq(w,wi)
178 case shape geometry: False{}
179
180def first(layout: Plan) -> U32:
181 match layout:
182 case Rows{first,count}: first
183 case Split{left,right}: first(left)
184
185def last(layout: Plan) -> U32:
186 match layout:
187 case Rows{first,count}: (first + count : U32)
188 case Split{left,right}: last(right)
189
190# A tree made from a plan, with a state threaded through the leaves in order.
191# The leaf takes the state and gives it back beside the band it made.
192def build(~Output: Type,~State: Type,~Context: Data,
193 ~leaf: Context -> U32 -> U32 -> State -> State & Partition.Tree<&1,Output,Unit>,
194 layout: Plan,+context: Context,state: State) -> State & Partition.Tree<&1,Output,Unit>:
195 match layout:
196 case Rows{first,rows}: leaf(context,first,rows,state)
197 case Split{left,right}:
198 Storage.bind_pair(~State,~Partition.Tree<&1,Output,Unit>,~(State & Partition.Tree<&1,Output,Unit>),
199 build(~Output,~State,~Context,~leaf,left,context,state),state => left =>
200 Storage.bind_pair(~State,~Partition.Tree<&1,Output,Unit>,~(State & Partition.Tree<&1,Output,Unit>),
201 build(~Output,~State,~Context,~leaf,right,context,state),state => right => (state,Partition.Fork{left,right})))
202
203# What a band made from scratch needs: its geometry and the value it starts with.
204type Filling<-Element: Data> is Data:
205 Filling{geometry: Geometry,fill: Element}
206
207# Two trees of one shape, joined leaf by leaf. A left leaf without a partner
208# keeps its own row range, so the left tree can be returned unchanged.
209type Paired<-Left: Type,-Right: Type> is Type:
210 Both{left: Left,right: Right}
211 Only{left: Left,first: U32,rows: U32}
212
213def only(~Left: Type,~Right: Type,tree: Partition.Tree<&1,Left,Unit>) -> Partition.Tree<&1,Paired<Left,Right>,Unit>:
214 match tree:
215 case Partition.Leaf{left,+first,+rows,details}: Partition.Leaf{Only{left,first,rows},first,rows,Unit{}}
216 case Partition.Fork{left,right}: Partition.Fork{only(~Left,~Right,left),only(~Left,~Right,right)}
217
218def zip(~Left: Type,~Right: Type,left: Partition.Tree<&1,Left,Unit>,right: Partition.Tree<&1,Right,Unit>) -> Partition.Tree<&1,Paired<Left,Right>,Unit>:
219 match left right:
220 case Partition.Leaf{left,first,rows,details} Partition.Leaf{right,other_first,other_rows,other_details}:
221 Partition.Leaf{Both{left,right},first,rows,Unit{}}
222 case Partition.Fork{left,right} Partition.Fork{other_left,other_right}:
223 Partition.Fork{zip(~Left,~Right,left,other_left),zip(~Left,~Right,right,other_right)}
224 case left right: only(~Left,~Right,left)
225
226def left_leaf(~Left: Type,~Right: Type,pair: Paired<Left,Right>,first: U32,rows: U32) -> Partition.Tree<&1,Left,Unit>:
227 match pair:
228 case Both{left,right}: Partition.Leaf{left,first,rows,Unit{}}
229 case Only{left,own_first,own_rows}: Partition.Leaf{left,own_first,own_rows,Unit{}}
230
231def left_of(~Left: Type,~Right: Type,tree: Partition.Tree<&1,Paired<Left,Right>,Unit>) -> Partition.Tree<&1,Left,Unit>:
232 match tree:
233 case Partition.Leaf{pair,first,rows,details}: left_leaf(~Left,~Right,pair,first,rows)
234 case Partition.Fork{left,right}: Partition.Fork{left_of(~Left,~Right,left),left_of(~Left,~Right,right)}
235
236def paired(~Left: Type,~Right: Type,~Resource: Type,~Context: Data,
237 ~leaf: Context -> U32 -> U32 -> Left -> Right -> Resource -> Partition.Tree<&1,Left,Unit>,
238 context: Context,first: U32,rows: U32,pair: Paired<Left,Right>,resource: Resource) -> Partition.Tree<&1,Left,Unit>:
239 match pair:
240 case Both{left,right}: leaf(context,first,rows,left,right,resource)
241 case Only{left,own_first,own_rows}: Partition.Leaf{left,own_first,own_rows,Unit{}}
242
243def unshared(resource: Unit) -> Unit & Unit:
244 (Unit{},Unit{})
245
246def same(~Block: Type,tree: Partition.Tree<&1,Block,Unit>) -> Partition.Tree<&1,Block,Unit>:
247 tree
248
249# What a split gives its halves: a state for each, or the whole state when it
250# cannot be divided.
251type Division<-State: Type> is Type:
252 Divided{left: State,right: State}
253 Whole{state: State}
254
255def divided(~State: Type,~Result: Type,division: Division<State>,whole: State -> Result,halves: State -> State -> Result) -> Result:
256 match division:
257 case Divided{left,right}: halves(left,right)
258 case Whole{state}: whole(state)
259
260# Every split may run its halves in parallel.
261def always(schedule: Unit,layout: Plan) -> Bool:
262 True{}
263
264# The one plan recursion. The plan gives every part its rows and decides the
265# structure. divide gives each half of a split its own state; a state it keeps
266# whole becomes one part through keep. A split whose schedule allows it runs
267# its halves in parallel. The last split calls its two leaves directly: a leaf
268# that neither forks nor calls a closure is then its own task, which the
269# runtime keeps for the work turn; a leaf reached through the forking
270# recursion would run in the grow turn on one worker per row.
271def each_part(~State: Type,~Output: Type,~Details: Data,~Context: Data,~Schedule: Data,
272 ~parallel: Schedule -> Plan -> Bool,
273 ~divide: Context -> Plan -> Plan -> State -> Division<State>,
274 ~leaf: Context -> U32 -> U32 -> State -> Partition.Tree<&1,Output,Details>,
275 ~keep: Context -> Plan -> State -> Partition.Tree<&1,Output,Details>,
276 layout: Plan,mode: Bool,+schedule: Schedule,+context: Context,state: State) -> Partition.Tree<&1,Output,Details>:
277 match layout mode:
278 case Rows{first,rows} mode: leaf(context,first,rows,state)
279 case Split{+left_plan,+right_plan} False{}:
280 divided(~State,~Partition.Tree<&1,Output,Details>,divide(context,left_plan,right_plan,state),whole => keep(context,Split{left_plan,right_plan},whole),
281 left_state => right_state =>
282 Partition.Fork{each_part(~State,~Output,~Details,~Context,~Schedule,~parallel,~divide,~leaf,~keep,left_plan,parallel(schedule,left_plan),schedule,context,left_state),
283 each_part(~State,~Output,~Details,~Context,~Schedule,~parallel,~divide,~leaf,~keep,right_plan,parallel(schedule,right_plan),schedule,context,right_state)})
284 case Split{Rows{left_first,left_rows},Rows{right_first,right_rows}} True{}:
285 divided(~State,~Partition.Tree<&1,Output,Details>,divide(context,Rows{left_first,left_rows},Rows{right_first,right_rows},state),
286 whole => keep(context,Split{Rows{left_first,left_rows},Rows{right_first,right_rows}},whole),
287 left_state => right_state =>
288 left right = leaf(context,left_first,left_rows,left_state) leaf(context,right_first,right_rows,right_state)
289 Partition.Fork{left,right})
290 case Split{Rows{left_first,left_rows},+right_plan} True{}:
291 divided(~State,~Partition.Tree<&1,Output,Details>,divide(context,Rows{left_first,left_rows},right_plan,state),
292 whole => keep(context,Split{Rows{left_first,left_rows},right_plan},whole),
293 left_state => right_state =>
294 left right = leaf(context,left_first,left_rows,left_state) each_part(~State,~Output,~Details,~Context,~Schedule,~parallel,~divide,~leaf,~keep,right_plan,parallel(schedule,right_plan),schedule,context,right_state)
295 Partition.Fork{left,right})
296 case Split{+left_plan,Rows{right_first,right_rows}} True{}:
297 divided(~State,~Partition.Tree<&1,Output,Details>,divide(context,left_plan,Rows{right_first,right_rows},state),
298 whole => keep(context,Split{left_plan,Rows{right_first,right_rows}},whole),
299 left_state => right_state =>
300 left right = each_part(~State,~Output,~Details,~Context,~Schedule,~parallel,~divide,~leaf,~keep,left_plan,parallel(schedule,left_plan),schedule,context,left_state) leaf(context,right_first,right_rows,right_state)
301 Partition.Fork{left,right})
302 case Split{+left_plan,+right_plan} True{}:
303 divided(~State,~Partition.Tree<&1,Output,Details>,divide(context,left_plan,right_plan,state),whole => keep(context,Split{left_plan,right_plan},whole),
304 left_state => right_state =>
305 left right = each_part(~State,~Output,~Details,~Context,~Schedule,~parallel,~divide,~leaf,~keep,left_plan,parallel(schedule,left_plan),schedule,context,left_state) each_part(~State,~Output,~Details,~Context,~Schedule,~parallel,~divide,~leaf,~keep,right_plan,parallel(schedule,right_plan),schedule,context,right_state)
306 Partition.Fork{left,right})
307
308# A band walk carries the tree beside the resource. A split divides the tree
309# where it forks; a tree of another shape is kept.
310def Banded(~Block: Type,~Resource: Type) -> Type:
311 Partition.Tree<&1,Block,Unit> & Resource
312
313# The planner decides once for an operation: every split of an enabled walk
314# may run its halves in parallel, inside an image or across images.
315def in_parallel(enabled: Bool,layout: Plan) -> Bool:
316 enabled
317
318def shared_halves(~Block: Type,~Resource: Type,left: Partition.Tree<&1,Block,Unit>,right: Partition.Tree<&1,Block,Unit>,resources: Resource & Resource) -> Division<Banded(~Block,~Resource)>:
319 (left_resource,right_resource) = resources
320 Divided{(left,left_resource),(right,right_resource)}
321
322def band_divide(~Block: Type,~Resource: Type,~Context: Data,~share: Resource -> Resource & Resource,
323 context: Context,left: Plan,right: Plan,state: Banded(~Block,~Resource)) -> Division<Banded(~Block,~Resource)>:
324 (tree,resource) = state
325 match tree:
326 case Partition.Leaf{block,first,rows,details}: Whole{(Partition.Leaf{block,first,rows,details},resource)}
327 case Partition.Fork{left_tree,right_tree}: shared_halves(~Block,~Resource,left_tree,right_tree,share(resource))
328
329def band_leaf(~Block: Type,~Output: Type,~Resource: Type,~Context: Data,
330 ~leaf: Context -> U32 -> U32 -> Block -> Resource -> Partition.Tree<&1,Output,Unit>,
331 ~keep: Partition.Tree<&1,Block,Unit> -> Partition.Tree<&1,Output,Unit>,
332 context: Context,first: U32,rows: U32,state: Banded(~Block,~Resource)) -> Partition.Tree<&1,Output,Unit>:
333 (tree,resource) = state
334 match tree:
335 case Partition.Leaf{block,offset,span,details}: leaf(context,first,rows,block,resource)
336 case Partition.Fork{left,right}: keep(Partition.Fork{left,right})
337
338def band_kept(~Block: Type,~Output: Type,~Resource: Type,~Context: Data,
339 ~keep: Partition.Tree<&1,Block,Unit> -> Partition.Tree<&1,Output,Unit>,
340 context: Context,layout: Plan,state: Banded(~Block,~Resource)) -> Partition.Tree<&1,Output,Unit>:
341 (tree,resource) = state
342 keep(tree)
343
344# The band recursion. share gives each half its own owner of a read-only
345# resource.
346def each(~Block: Type,~Output: Type,~Resource: Type,~Context: Data,
347 ~share: Resource -> Resource & Resource,
348 ~leaf: Context -> U32 -> U32 -> Block -> Resource -> Partition.Tree<&1,Output,Unit>,
349 ~keep: Partition.Tree<&1,Block,Unit> -> Partition.Tree<&1,Output,Unit>,
350 layout: Plan,+enabled: Bool,+context: Context,
351 tree: Partition.Tree<&1,Block,Unit>,resource: Resource) -> Partition.Tree<&1,Output,Unit>:
352 each_part(~Banded(~Block,~Resource),~Output,~Unit,~Context,~Bool,~in_parallel,~band_divide(~Block,~Resource,~Context,~share),
353 ~band_leaf(~Block,~Output,~Resource,~Context,~leaf,~keep),~band_kept(~Block,~Output,~Resource,~Context,~keep),
354 layout,enabled,enabled,context,(tree,resource))
355
356type Indexed is Data:
357 Indexed{geometry: Geometry,first: U32,rows: U32}
358
359def map_leaf(~Element: Data,~Context: Data,~operation: Context -> Element -> Element,+context: Context,buffer: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
360 Storage.Buffer{storage,+length,certificate} = buffer
361 Storage.Buffer{Traversal.map_in_place(~Element,~Context,~operation,length,0,context,storage),length,
362 TraversalCertificate.map_in_place(~Element,~Context,~operation,length,0,context,storage,certificate)}
363
364# A mapper over the whole of one band.
365def map_whole_context(~Element: Data,~Context: Data,~mapper: @+count: Nat -> @+position: U32 -> @+context: Context -> Storage.Buffer<Element> -> Storage.Buffer<Element>,+context: Context,buffer: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
366 Storage.Buffer{storage,+length,certificate} = buffer
367 mapper(length,0,context,Storage.Buffer{storage,length,certificate})
368
369def map_band(~Element: Data,~Context: Data,~leaf: @+context: Context -> Storage.Buffer<Element> -> Storage.Buffer<Element>,context: Context,first: U32,rows: U32,
370 buffer: Storage.Buffer<Element>,resource: Unit) -> Bands(~Element):
371 Partition.Leaf{leaf(context,buffer),first,rows,Unit{}}
372
373# Parallelism is a planner decision. Sequential execution keeps the bands and
374# their owners: a small operation need not pay either a fork or a gather.
375def map_with(~Element: Data,~Context: Data,~leaf: @+context: Context -> Storage.Buffer<Element> -> Storage.Buffer<Element>,+layout: Plan,+parallel: Bool,
376 context: Context,bands: Bands(~Element)) -> Bands(~Element):
377 each(~Storage.Buffer<Element>,~Storage.Buffer<Element>,~Unit,~Context,~unshared,~map_band(~Element,~Context,~leaf),~same(~Storage.Buffer<Element>),
378 layout,parallel,context,bands,Unit{})
379
380def scalar_value(~Element: Data,~operation: Element -> Element -> Element,value: Element,element: Element) -> Element:
381 operation(element,value)
382
383def scalar(~Element: Data,~operation: Element -> Element -> Element,layout: Plan,parallel: Bool,value: Element,bands: Bands(~Element)) -> Bands(~Element):
384 map_with(~Element,~Element,~map_leaf(~Element,~Element,~scalar_value(~Element,~operation)),layout,parallel,value,bands)
385
386def result(~Element: Data,owners: Storage.Buffer<Element> & Bands(~Element)) -> Bands(~Element):
387 (source,bands) = owners
388 bands
389
390def index_of(context: Indexed,index: U32) -> U32:
391 Indexed{geometry,first,rows} = context
392 global_index(geometry,first,rows,index)
393
394# Channel c of a leaf holds the positions from c * rows * width; their global
395# indices are consecutive from this base.
396def channel_base(context: Indexed,channel: U32) -> U32:
397 Indexed{Geometry{batch,channels,+height,width},+first,rows} = context
398 (((first / height * channels + channel) * height + first % height) * width : U32)
399
400def channels_of(context: Indexed) -> U32:
401 Indexed{Geometry{batch,channels,height,width},first,rows} = context
402 channels
403
404def plane_of(context: Indexed) -> U32:
405 Indexed{Geometry{batch,channels,height,width},first,rows} = context
406 (rows * width : U32)
407
408def indexed_buffers(~Element: Data,~operation: U32 -> Element -> Element,+count: U32,+context: Indexed,input: Storage.Buffer<Element>,output: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
409 Storage.Buffer{input,input_length,input_certificate} = input
410 Storage.Buffer{output,output_length,output_certificate} = output
411 destination(~Element,Operations.pair_buffers(~Element,input_length,output_length,
412 IndexedAccess.run_tiled(~Element,~Indexed,~index_of,~channel_base,~channels_of,~plane_of,~operation,count,context,input,output),
413 IndexedCertificate.tiled_certificates(~Element,~Indexed,~index_of,~channel_base,~channels_of,~plane_of,~operation,count,context,input,output,input_certificate,output_certificate)))
414
415def indexed_leaf(~Element: Data,~operation: U32 -> Element -> Element,
416 +geometry: Geometry,+first: U32,+rows: U32,fill: Element,input: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
417 +count = local_count(geometry,rows)
418 indexed_buffers(~Element,~operation,count,Indexed{geometry,first,rows},input,
419 Allocation.allocate(~Element,Storage.allocation_depth(U32.max(count,1)),fill,U32.to_nat(count)))
420
421type Filled<-Element: Data> is Data:
422 Filled{geometry: Geometry,fill: Element}
423
424def indexed_band(~Element: Data,~operation: U32 -> Element -> Element,context: Filled<Element>,first: U32,rows: U32,
425 input: Storage.Buffer<Element>,resource: Unit) -> Bands(~Element):
426 Filled{geometry,fill} = context
427 +first = first
428 +rows = rows
429 Partition.Leaf{indexed_leaf(~Element,~operation,geometry,first,rows,fill,input),first,rows,Unit{}}
430
431def apply_indexed(~Element: Data,~operation: U32 -> Element -> Element,
432 +layout: Plan,+parallel: Bool,+geometry: Geometry,fill: Element,bands: Bands(~Element)) -> Bands(~Element):
433 each(~Storage.Buffer<Element>,~Storage.Buffer<Element>,~Unit,~Filled<Element>,~unshared,~indexed_band(~Element,~operation),~same(~Storage.Buffer<Element>),
434 layout,parallel,Filled{geometry,fill},bands,Unit{})
435
436def range_result(~Element: Data,+first: U32,+rows: U32,owners: Storage.Buffer<Element> & Storage.Buffer<Element>) -> Bands(~Element) & Storage.Buffer<Element>:
437 (source,output) = owners
438 (Partition.Leaf{source,first,rows,Unit{}},output)
439
440# An empty intersection copies nothing, so it skips the per-channel loop.
441def range_planes(~Element: Data,empty: Bool,+channels: U32,+count: U32,+source_stride: U32,+destination_stride: U32,
442 +source_start: U32,+destination_start: U32,owners: Storage.Buffer<Element> & Storage.Buffer<Element>) ->
443 Storage.Buffer<Element> & Storage.Buffer<Element>:
444 match empty:
445 case True{}: owners
446 case False{}: planes(~Element,U32.to_nat(channels),0,count,source_stride,destination_stride,source_start,destination_start,owners)
447
448# Repartition and halo extraction share this owned range reader. Empty
449# intersections copy zero elements; each source owner is returned unchanged.
450def range_leaf(~Element: Data,geometry: Geometry,+target_first: U32,+target_rows: U32,target_offset: U32,+first: U32,+rows: U32,
451 source: Storage.Buffer<Element>,output: Storage.Buffer<Element>) -> Bands(~Element) & Storage.Buffer<Element>:
452 Geometry{batch,channels,height,+width} = geometry
453 +begin = U32.max(first,target_first)
454 +end = U32.max(begin,U32.min((first + rows : U32),(target_first + target_rows : U32)))
455 range_result(~Element,first,rows,range_planes(~Element,U32.is_eq(begin,end),channels,((end - begin) * width : U32),
456 (rows * width : U32),(target_rows * width : U32),((begin - first) * width : U32),(target_offset + (begin - target_first) * width : U32),(source,output)))
457
458# The rows a range reads and where it writes them.
459type Range is Data:
460 Range{geometry: Geometry,first: U32,rows: U32,offset: U32}
461
462def range_band(~Element: Data,context: Range,first: U32,rows: U32,input: Storage.Buffer<Element>,output: Storage.Buffer<Element>) -> Bands(~Element) & Storage.Buffer<Element>:
463 Range{geometry,+target_first,+target_rows,offset} = context
464 +first = first
465 +rows = rows
466 range_leaf(~Element,geometry,target_first,target_rows,offset,first,rows,input,output)
467
468def range_into(~Element: Data,+geometry: Geometry,+first: U32,+rows: U32,+target_offset: U32,bands: Bands(~Element),output: Storage.Buffer<Element>) -> Bands(~Element) & Storage.Buffer<Element>:
469 walk(~Storage.Buffer<Element>,~Storage.Buffer<Element>,~Storage.Buffer<Element>,~Range,~range_band(~Element),Range{geometry,first,rows,target_offset},bands,output)
470
471def range(~Element: Data,geometry: Geometry,first: U32,rows: U32,bands: Bands(~Element),output: Storage.Buffer<Element>) -> Bands(~Element) & Storage.Buffer<Element>:
472 range_into(~Element,geometry,first,rows,0,bands,output)
473
474def allocate(~Element: Data,fill: Element,+count: U32) -> Storage.Buffer<Element>:
475 Allocation.allocate(~Element,Storage.depth_for(count),fill,U32.to_nat(count))
476
477def repartition_band(~Element: Data,context: Filling<Element>,first: U32,rows: U32,bands: Bands(~Element)) -> Bands(~Element) & Bands(~Element):
478 Filling{+geometry,fill} = context
479 +first = first
480 +rows = rows
481 Storage.bind_pair(~Bands(~Element),~Storage.Buffer<Element>,~(Bands(~Element) & Bands(~Element)),
482 range(~Element,geometry,first,rows,bands,allocate(~Element,fill,local_count(geometry,rows))),source => output =>
483 (source,Partition.Leaf{output,first,rows,Unit{}}))
484
485def repartition(~Element: Data,geometry: Geometry,fill: Element,layout: Plan,bands: Bands(~Element)) -> Bands(~Element) & Bands(~Element):
486 build(~Storage.Buffer<Element>,~Bands(~Element),~Filling<Element>,~repartition_band(~Element),layout,Filling{geometry,fill},bands)
487
488def repartitioned(~Element: Data,owners: Bands(~Element) & Bands(~Element)) -> Bands(~Element):
489 (source,result) = owners
490 result
491
492def align_when(~Element: Data,same: Bool,geometry: Geometry,fill: Element,layout: Plan,bands: Bands(~Element)) -> Bands(~Element):
493 match same:
494 case True{}: bands
495 case False{}: repartitioned(~Element,repartition(~Element,geometry,fill,layout,bands))
496
497def align_planned(~Element: Data,geometry: Geometry,fill: Element,+layout: Plan,planned: Bands(~Element) & Plan) -> Bands(~Element):
498 (bands,other) = planned
499 align_when(~Element,same_plan(layout,other),geometry,fill,layout,bands)
500
501def align(~Element: Data,geometry: Geometry,fill: Element,layout: Plan,bands: Bands(~Element)) -> Bands(~Element):
502 align_planned(~Element,geometry,fill,layout,plan(~Element,bands))
503
504def whole_result(~Element: Data,+first: U32,+rows: U32,owners: Storage.Buffer<Element> & Storage.Buffer<Element>) -> Storage.Buffer<Element> & Bands(~Element):
505 (source,band) = owners
506 (source,Partition.Leaf{band,first,rows,Unit{}})
507
508# Rows [first, first + rows) of every plane of a whole buffer, in a buffer of
509# their own beside the source.
510def whole_rows(~Element: Data,geometry: Geometry,fill: Element,+first: U32,+rows: U32,input: Storage.Buffer<Element>) -> Storage.Buffer<Element> & Storage.Buffer<Element>:
511 Geometry{batch,+channels,+height,+width} = geometry
512 planes(~Element,U32.to_nat(channels),0,(rows * width : U32),(height * width : U32),(rows * width : U32),
513 ((first / height * channels * height + first % height) * width : U32),0,(input,allocate(~Element,fill,(channels * rows * width : U32))))
514
515def whole_leaf(~Element: Data,geometry: Geometry,fill: Element,+first: U32,+rows: U32,input: Storage.Buffer<Element>) -> Storage.Buffer<Element> & Bands(~Element):
516 whole_result(~Element,first,rows,whole_rows(~Element,geometry,fill,first,rows,input))
517
518def whole_band(~Element: Data,context: Filling<Element>,first: U32,rows: U32,input: Storage.Buffer<Element>) -> Storage.Buffer<Element> & Bands(~Element):
519 Filling{geometry,fill} = context
520 whole_leaf(~Element,geometry,fill,first,rows,input)
521
522def whole_ranges(~Element: Data,geometry: Geometry,fill: Element,layout: Plan,input: Storage.Buffer<Element>) -> Storage.Buffer<Element> & Bands(~Element):
523 build(~Storage.Buffer<Element>,~Storage.Buffer<Element>,~Filling<Element>,~whole_band(~Element),layout,Filling{geometry,fill},input)
524
525# A leaf that holds every row of the only image is that image in row-major
526# order: the whole buffer is the band, relabelled. Other rows are copied out.
527def from_rows(~Element: Data,whole: Bool,geometry: Geometry,fill: Element,+first: U32,+rows: U32,input: Storage.Buffer<Element>) -> Bands(~Element):
528 match whole:
529 case True{}: Partition.Leaf{input,first,rows,Unit{}}
530 case False{}: result(~Element,whole_ranges(~Element,geometry,fill,Rows{first,rows},input))
531
532# The bands of a row-major buffer cut by a plan.
533def from_whole(~Element: Data,+geometry: Geometry,fill: Element,layout: Plan,input: Storage.Buffer<Element>) -> Bands(~Element):
534 match layout:
535 case Rows{+first,+rows}: from_rows(~Element,whole_image(geometry,first,rows),geometry,fill,first,rows,input)
536 case Split{left,right}: result(~Element,whole_ranges(~Element,geometry,fill,Split{left,right},input))
537
538def row_plan(depth: Nat,split: Bool,+first: U32,+rows: U32) -> Plan:
539 match depth split:
540 case 0n split: Rows{first,rows}
541 case 1n+rest False{}: Rows{first,rows}
542 case 1n+ +rest True{}:
543 +half = (rows / 2 : U32)
544 Split{row_plan(rest,U32.is_lt(1,half),first,half),
545 row_plan(rest,U32.is_lt(1,(rows - half : U32)),(first + half : U32),(rows - half : U32))}
546
547def image_plan(remaining: Nat,+image: U32,+depth: Nat,+height: U32) -> Plan:
548 match remaining:
549 case 0n: Rows{0,0}
550 case 1n: row_plan(depth,U32.is_lt(1,height),(image * height : U32),height)
551 case 2n+rest:
552 Split{row_plan(depth,U32.is_lt(1,height),(image * height : U32),height),image_plan(1n+rest,(image + 1 : U32),depth,height)}
553
554# Every image of the geometry as one band, its rows unsplit.
555def unsplit(geometry: Geometry) -> Plan:
556 Geometry{batch,channels,+height,width} = geometry
557 image_plan(U32.to_nat(batch),0,0n,height)
558
559# Rows [first, first + count) halved depth times. The first half of every
560# split is a whole number of units.
561def unit_plan(depth: Nat,+unit: U32,+first: U32,+count: U32) -> Plan:
562 match depth:
563 case 0n: Rows{first,count}
564 case 1n+ +rest:
565 +half = (count / (unit * 2) * unit : U32)
566 Split{unit_plan(rest,unit,first,half),unit_plan(rest,unit,(first + half : U32),(count - half : U32))}
567
568def combined(~Element: Data,result: Operations.Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>) -> Storage.Buffer<Element>:
569 match result:
570 case Operations.Accepted{(left,right)}: left
571 case Operations.Rejected{(left,right)}: left
572
573def combine_leaf_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,+count: U32,left: Storage.Buffer<Element>,right: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
574 combined(~Element,Operations.combine_with(~Element,~kernel,left,Views.View{[Views.Axis{count,1}],0},right,Views.View{[Views.Axis{count,1}],0}))
575
576def combine_band_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,geometry: Geometry,first: U32,rows: U32,
577 left: Storage.Buffer<Element>,right: Storage.Buffer<Element>,resource: Unit) -> Bands(~Element):
578 +rows = rows
579 Partition.Leaf{combine_leaf_with(~Element,~kernel,local_count(geometry,rows),left,right),first,rows,Unit{}}
580
581def combine_planned_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,+geometry: Geometry,+parallel: Bool,planned: Bands(~Element) & Plan,right: Bands(~Element)) -> Bands(~Element):
582 (left,+layout) = planned
583 each(~Paired<Storage.Buffer<Element>,Storage.Buffer<Element>>,~Storage.Buffer<Element>,~Unit,~Geometry,~unshared,
584 ~paired(~Storage.Buffer<Element>,~Storage.Buffer<Element>,~Unit,~Geometry,~combine_band_with(~Element,~kernel)),
585 ~left_of(~Storage.Buffer<Element>,~Storage.Buffer<Element>),
586 layout,parallel,geometry,zip(~Storage.Buffer<Element>,~Storage.Buffer<Element>,left,right),Unit{})
587
588# The caller supplies the same plan on both sides, after shape and status
589# checks; the plan is read from the left. Trees of another shape keep the left owner.
590def combine_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,geometry: Geometry,parallel: Bool,left: Bands(~Element),right: Bands(~Element)) -> Bands(~Element):
591 combine_planned_with(~Element,~kernel,geometry,parallel,plan(~Element,left),right)
592
593# A source of the extents source read band by band as a tensor of the
594# extents shape under NumPy broadcasting, the bands those of the geometry.
595# Views are built again at every use: a View value used twice would make the
596# compiler share every View constructor.
597type Spread is Data:
598 Spread{geometry: Geometry,shape: List<&2,U32>,source: List<&2,U32>}
599
600def spread_geometry(context: Spread) -> Geometry:
601 Spread{geometry,shape,source} = context
602 geometry
603
604def stretched_view(view: Maybe<&2,Views.View>,+shape: List<&2,U32>) -> Maybe<&2,Views.View>:
605 match view:
606 case Some{view}: Views.stretch(view,shape)
607 case None{}: None{}
608
609def stretched_source(+source: List<&2,U32>,+shape: List<&2,U32>) -> Maybe<&2,Views.View>:
610 stretched_view(Views.contiguous(Shapes.Shape{source},0),shape)
611
612def sliced_rows(view: Maybe<&2,Views.View>,+start: U32,rows: U32) -> Maybe<&2,Views.View>:
613 match view:
614 case Some{view}: Views.slice(view,2n,start,rows,1)
615 case None{}: None{}
616
617def sliced_image(view: Maybe<&2,Views.View>,+image: U32) -> Maybe<&2,Views.View>:
618 match view:
619 case Some{view}: Views.slice(view,0n,image,1,1)
620 case None{}: None{}
621
622def sliced_plane_rows(view: Maybe<&2,Views.View>,+start: U32,rows: U32) -> Maybe<&2,Views.View>:
623 match view:
624 case Some{view}: Views.slice(view,1n,start,rows,1)
625 case None{}: None{}
626
627# The cells of a band of image rows: its image and its rows.
628def rows_view(+shape: List<&2,U32>,+height: U32,+first: U32,rows: U32,view: Maybe<&2,Views.View>) -> Maybe<&2,Views.View>:
629 match shape:
630 case [c,h,w]: sliced_plane_rows(view,(first % height : U32),rows)
631 case [n,c,h,w]: sliced_rows(sliced_image(view,(first / height : U32)),(first % height : U32),rows)
632 case shape: None{}
633
634def rows_extents(+shape: List<&2,U32>,rows: U32) -> Maybe<&2,List<&2,U32>>:
635 match shape:
636 case [c,h,w]: Some{[c,rows,w]}
637 case [n,c,h,w]: Some{[1,c,rows,w]}
638 case shape: None{}
639
640def band_view_when(whole: Bool,rows_shaped: Bool,+shape: List<&2,U32>,+height: U32,+first: U32,rows: U32,view: Maybe<&2,Views.View>) -> Maybe<&2,Views.View>:
641 match whole rows_shaped:
642 case True{} rows_shaped: view
643 case False{} True{}: rows_view(shape,height,first,rows,view)
644 case False{} False{}: None{}
645
646def band_extents_when(whole: Bool,rows_shaped: Bool,+shape: List<&2,U32>,rows: U32) -> Maybe<&2,List<&2,U32>>:
647 match whole rows_shaped:
648 case True{} rows_shaped: Some{shape}
649 case False{} True{}: rows_extents(shape,rows)
650 case False{} False{}: None{}
651
652# The view of the source a band reads. A band of a whole image reads the whole
653# stretched view, and a band of image rows of an image shape reads its image
654# and its rows. Other bands have none.
655def band_view(+context: Spread,+first: U32,+rows: U32) -> Maybe<&2,Views.View>:
656 Spread{+geometry,+shape,+source} = context
657 Geometry{batch,channels,+height,width} = geometry
658 band_view_when(whole_image(geometry,first,rows),image_rows(shape,geometry),shape,height,first,rows,stretched_source(source,shape))
659
660# The extents of the cells of that band, in the order the band stores them.
661def band_extents(+context: Spread,+first: U32,+rows: U32) -> Maybe<&2,List<&2,U32>>:
662 Spread{+geometry,+shape,source} = context
663 band_extents_when(whole_image(geometry,first,rows),image_rows(shape,geometry),shape,rows)
664
665def has_view(view: Maybe<&2,Views.View>) -> Bool:
666 match view:
667 case Some{view}: True{}
668 case None{}: False{}
669
670# Every band of the plan has a view.
671def spread_boxed(+context: Spread,layout: Plan) -> Bool:
672 match layout:
673 case Rows{first,rows}: has_view(band_view(context,first,rows))
674 case Split{left,right}: spread_boxed(context,left) && spread_boxed(context,right)
675
676def compact_axes(source: List<&2,U32>,axes: List<&2,Views.Axis>) -> List<&2,Views.Axis>:
677 match source axes:
678 case Con{+single,rest} Con{Views.Axis{extent,stride},others}: Views.Axis{Bool.pick(U32,U32.is_eq(single,1),1,extent),stride} <> compact_axes(rest,others)
679 case source axes: []
680
681# The elements of a band view that the source holds once: an axis the source
682# repeats keeps one element.
683def compact_of(+source: List<&2,U32>,view: Maybe<&2,Views.View>) -> Maybe<&2,Views.View>:
684 match view:
685 case Some{Views.View{+axes,offset}}: Some{Views.View{compact_axes(Shapes.with_rank(source,List.length(&2,Views.Axis,axes)),axes),offset}}
686 case None{}: None{}
687
688def compact_extents(source: List<&2,U32>,extents: List<&2,U32>) -> List<&2,U32>:
689 match source extents:
690 case Con{+single,rest} Con{extent,others}: Bool.pick(U32,U32.is_eq(single,1),1,extent) <> compact_extents(rest,others)
691 case source extents: []
692
693def compact_with(+source: List<&2,U32>,extents: Maybe<&2,List<&2,U32>>) -> Maybe<&2,List<&2,U32>>:
694 match extents:
695 case Some{+extents}: Some{compact_extents(Shapes.with_rank(source,List.length(&2,U32,extents)),extents)}
696 case None{}: None{}
697
698# The extents of the part of the source a band reads: those of its cells, an
699# axis the source repeats kept once.
700def band_compact(+context: Spread,+first: U32,+rows: U32) -> Maybe<&2,List<&2,U32>>:
701 Spread{geometry,shape,+source} = context
702 compact_with(source,band_extents(context,first,rows))
703
704def contiguous_of(extents: Maybe<&2,List<&2,U32>>) -> Maybe<&2,Views.View>:
705 match extents:
706 case Some{+extents}: Views.contiguous(Shapes.Shape{extents},0)
707 case None{}: None{}
708
709def count_of_extents(extents: Maybe<&2,List<&2,U32>>) -> Maybe<&2,U32>:
710 match extents:
711 case Some{extents}: Shapes.numel(Shapes.Shape{extents})
712 case None{}: None{}
713
714type Pieces<-Element: Data> is Data:
715 Pieces{spread: Spread,fill: Element}
716
717def piece_copied(~Element: Data,+first: U32,+rows: U32,result: Operations.Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>) -> Storage.Buffer<Element> & Bands(~Element):
718 match result:
719 case Operations.Accepted{(source,piece)}: (source,Partition.Leaf{piece,first,rows,Unit{}})
720 case Operations.Rejected{(source,piece)}: (source,Partition.Leaf{piece,first,rows,Unit{}})
721
722# A band view lies inside the source it was stretched from, so the band
723# copies through it without a check.
724def piece_views(~Element: Data,source_view: Maybe<&2,Views.View>,piece_view: Maybe<&2,Views.View>,count: Maybe<&2,U32>,+fill: Element,+first: U32,+rows: U32,
725 source: Storage.Buffer<Element>) -> Storage.Buffer<Element> & Bands(~Element):
726 match source_view piece_view count:
727 case Some{source_view} Some{piece_view} Some{+count}:
728 piece_copied(~Element,first,rows,Operations.copy_when(~Element,True{},Some{count},source_view,piece_view,source,allocate(~Element,fill,count)))
729 case source_view piece_view count: (source,Partition.Leaf{allocate(~Element,fill,0),first,rows,Unit{}})
730
731def source_of(context: Spread) -> List<&2,U32>:
732 Spread{geometry,shape,source} = context
733 source
734
735def piece_band(~Element: Data,context: Pieces<Element>,first: U32,rows: U32,source: Storage.Buffer<Element>) -> Storage.Buffer<Element> & Bands(~Element):
736 Pieces{+spread,+fill} = context
737 +first = first
738 +rows = rows
739 piece_views(~Element,compact_of(source_of(spread),band_view(spread,first,rows)),contiguous_of(band_compact(spread,first,rows)),
740 count_of_extents(band_compact(spread,first,rows)),fill,first,rows,source)
741
742def spread_leaf(~Element: Data,whole: Bool,+context: Spread,fill: Element,+first: U32,+rows: U32,source: Storage.Buffer<Element>) -> Bands(~Element):
743 match whole:
744 case True{}: Partition.Leaf{source,first,rows,Unit{}}
745 case False{}: result(~Element,piece_band(~Element,Pieces{context,fill},first,rows,source))
746
747# The part of a source each band of the plan reads, copied in plan order: the
748# elements of its band view, an axis that repeats one element kept once. One
749# band of a whole image reads all of the source, which it takes as is.
750def spread(~Element: Data,+context: Spread,fill: Element,layout: Plan,source: Storage.Buffer<Element>) -> Bands(~Element):
751 match layout:
752 case Rows{+first,+rows}: spread_leaf(~Element,whole_image(spread_geometry(context),first,rows),context,fill,first,rows,source)
753 case Split{left,right}:
754 result(~Element,build(~Storage.Buffer<Element>,~Storage.Buffer<Element>,~Pieces<Element>,~piece_band(~Element),Split{left,right},Pieces{context,fill},source))
755
756def stretched_to(view: Maybe<&2,Views.View>,target: Maybe<&2,List<&2,U32>>) -> Maybe<&2,Views.View>:
757 match view target:
758 case Some{view} Some{+target}: Views.stretch(view,target)
759 case view target: None{}
760
761def kernel_left(~Element: Data,owners: Storage.Buffer<Element> & Storage.Buffer<Element>) -> Storage.Buffer<Element>:
762 (left,right) = owners
763 left
764
765# The kernel on a band and its part: the views were checked when the
766# operation chose its layout, so the band runs the kernel itself.
767def spread_views(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,
768 +count: U32,right_view: Maybe<&2,Views.View>,left: Storage.Buffer<Element>,right: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
769 match right_view:
770 case Some{right_view}: kernel_left(~Element,kernel(U32.to_nat(count),right_view,0,left,right))
771 case None{}: left
772
773# A band combined with its part of the source: the band read with the
774# extents of its cells, the part stretched to them.
775def combine_spread_band_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,
776 context: Spread,first: U32,rows: U32,left: Storage.Buffer<Element>,right: Storage.Buffer<Element>,resource: Unit) -> Bands(~Element):
777 +context = context
778 +first = first
779 +rows = rows
780 Partition.Leaf{spread_views(~Element,~kernel,local_count(spread_geometry(context),rows),
781 stretched_to(contiguous_of(band_compact(context,first,rows)),band_extents(context,first,rows)),left,right),first,rows,Unit{}}
782
783def combine_spread_planned_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,
784 +context: Spread,+parallel: Bool,planned: Bands(~Element) & Plan,right: Bands(~Element)) -> Bands(~Element):
785 (left,+layout) = planned
786 each(~Paired<Storage.Buffer<Element>,Storage.Buffer<Element>>,~Storage.Buffer<Element>,~Unit,~Spread,~unshared,
787 ~paired(~Storage.Buffer<Element>,~Storage.Buffer<Element>,~Unit,~Spread,~combine_spread_band_with(~Element,~kernel)),
788 ~left_of(~Storage.Buffer<Element>,~Storage.Buffer<Element>),
789 layout,parallel,context,zip(~Storage.Buffer<Element>,~Storage.Buffer<Element>,left,right),Unit{})
790
791# The bands combined with the parts spread copies for them from a source.
792def combine_spread_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,
793 +context: Spread,parallel: Bool,left: Bands(~Element),right: Bands(~Element)) -> Bands(~Element):
794 combine_spread_planned_with(~Element,~kernel,context,parallel,plan(~Element,left),right)
795
796# A band made by a checked call, beside whether the call was usable.
797type Checked<-Element: Data> is Type:
798 Checked{output: Storage.Buffer<Element>,usable: Bool}
799
800def Transformed(~Element: Data) -> Type:
801 Partition.Tree<&1,Checked<Element>,Unit>
802
803def transform_band(~Element: Data,~Context: Data,~leaf: Context -> U32 -> U32 -> Storage.Buffer<Element> -> Transformed(~Element),
804 context: Context,first: U32,rows: U32,input: Storage.Buffer<Element>,resource: Unit) -> Transformed(~Element):
805 leaf(context,first,rows,input)
806
807# A tree the plan does not describe is not transformed.
808def untransformed(~Element: Data,bands: Bands(~Element)) -> Transformed(~Element):
809 match bands:
810 case Partition.Leaf{input,first,rows,details}: Partition.Leaf{Checked{input,False{}},first,rows,details}
811 case Partition.Fork{left,right}: Partition.Fork{untransformed(~Element,left),untransformed(~Element,right)}
812
813def checked_band(~Element: Data,context: Unit,first: U32,rows: U32,checked: Checked<Element>) -> Bands(~Element) & Bool:
814 Checked{output,usable} = checked
815 (Partition.Leaf{output,first,rows,Unit{}},usable)
816
817# Every band through its own checked call. The bands come back beside whether
818# every call was usable.
819def transform(~Element: Data,~Context: Data,~leaf: Context -> U32 -> U32 -> Storage.Buffer<Element> -> Transformed(~Element),
820 +layout: Plan,+parallel: Bool,context: Context,bands: Bands(~Element)) -> Bands(~Element) & Bool:
821 reduce(~Checked<Element>,~Storage.Buffer<Element>,~Bool,~Unit,~checked_band(~Element),~Bool.and,Unit{},
822 each(~Storage.Buffer<Element>,~Checked<Element>,~Unit,~Context,~unshared,~transform_band(~Element,~Context,~leaf),~untransformed(~Element),
823 layout,parallel,context,bands,Unit{}))
824
825def transform_pair(~Element: Data,~Context: Data,~leaf: Context -> U32 -> U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Transformed(~Element),
826 context: Context,first: U32,rows: U32,pair: Paired<Storage.Buffer<Element>,Storage.Buffer<Element>>,resource: Unit) -> Transformed(~Element):
827 match pair:
828 case Both{left,right}: leaf(context,first,rows,left,right)
829 case Only{left,own_first,own_rows}: Partition.Leaf{Checked{left,False{}},own_first,own_rows,Unit{}}
830
831def untransformed_pairs(~Element: Data,pairs: Partition.Tree<&1,Paired<Storage.Buffer<Element>,Storage.Buffer<Element>>,Unit>) -> Transformed(~Element):
832 untransformed(~Element,left_of(~Storage.Buffer<Element>,~Storage.Buffer<Element>,pairs))
833
834# Every band of the left through its own checked call with the band of the
835# right of the same rows; a left band without a partner is not usable.
836def transform_paired(~Element: Data,~Context: Data,~leaf: Context -> U32 -> U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Transformed(~Element),
837 +layout: Plan,+parallel: Bool,context: Context,pairs: Partition.Tree<&1,Paired<Storage.Buffer<Element>,Storage.Buffer<Element>>,Unit>) -> Bands(~Element) & Bool:
838 reduce(~Checked<Element>,~Storage.Buffer<Element>,~Bool,~Unit,~checked_band(~Element),~Bool.and,Unit{},
839 each(~Paired<Storage.Buffer<Element>,Storage.Buffer<Element>>,~Checked<Element>,~Unit,~Context,~unshared,~transform_pair(~Element,~Context,~leaf),~untransformed_pairs(~Element),
840 layout,parallel,context,pairs,Unit{}))
841
842type Inspection is Data:
843 Inspection{valid: Bool,first: U32,end: U32,leaves: U32}
844
845def join_inspection(left: Inspection,right: Inspection) -> Inspection:
846 Inspection{lv,lf,+le,ln} = left
847 Inspection{rv,+rf,re,rn} = right
848 Inspection{lv && rv && U32.is_eq(le,rf),lf,re,(ln + rn : U32)}
849
850def leaf_inspection(geometry: Geometry,+length: Nat,depth: Nat,+first: U32,+rows: U32) -> Inspection:
851 Geometry{batch,+channels,+height,+width} = geometry
852 Inspection{U32.is_le(rows,(height - first % height : U32)) &&
853 Nat.is_le(Nat.add(U32.to_nat(first),U32.to_nat(rows)),Nat.mul(U32.to_nat(batch),U32.to_nat(height))) &&
854 Nat.is_eq(length,Nat.mul(U32.to_nat(channels),Nat.mul(U32.to_nat(rows),U32.to_nat(width)))) &&
855 Nat.is_le(length,U32.to_nat(U32.shln(1,depth))),first,(first + rows : U32),1}
856
857def inspect_leaf(~Element: Data,geometry: Geometry,buffer: Storage.Buffer<Element>,+first: U32,+rows: U32) -> Bands(~Element) & Inspection:
858 Storage.Buffer{storage,+length,Certificate.Certificate{+depth,equation}} = buffer
859 (Partition.Leaf{Storage.Buffer{storage,length,Certificate.Certificate{depth,equation}},first,rows,Unit{}},
860 leaf_inspection(geometry,length,depth,first,rows))
861
862def inspect_band(~Element: Data,geometry: Geometry,first: U32,rows: U32,buffer: Storage.Buffer<Element>) -> Bands(~Element) & Inspection:
863 +first = first
864 +rows = rows
865 inspect_leaf(~Element,geometry,buffer,first,rows)
866
867def inspect(~Element: Data,+geometry: Geometry,bands: Bands(~Element)) -> Bands(~Element) & Inspection:
868 reduce(~Storage.Buffer<Element>,~Storage.Buffer<Element>,~Inspection,~Geometry,~inspect_band(~Element),~join_inspection,geometry,bands)
869
870# Leaves that fit the geometry and cover all of its rows in order.
871def covers(geometry: Geometry,inspection: Inspection) -> Bool:
872 Geometry{batch,channels,height,width} = geometry
873 Inspection{valid,first,end,leaves} = inspection
874 valid && U32.is_zero(first) && Nat.is_eq(U32.to_nat(end),Nat.mul(U32.to_nat(batch),U32.to_nat(height)))
875