tensor.bend
stelliferous@0.0.3.0 · 676 lines · 0 definitions
1# A public tensor owns bands of row-major storage and the geometry they split,
2# its extents, the value used to initialize allocations, a sticky status, and
3# the number of workers its operations may use. A whole tensor is one band that
4# holds every element. A failed operation leaves values untouched and records
5# why; later operations skip a failed tensor, and check is the only way to
6# observe or clear the failure.
7import Base
8import stelliferous@0.0.3.0/storage_buffer.bend as Storage
9import stelliferous@0.0.3.0/storage_allocation.bend as Allocation
10import stelliferous@0.0.3.0/tensor_shape.bend as Shapes
11import stelliferous@0.0.3.0/tensor_view.bend as Views
12import stelliferous@0.0.3.0/tensor_operations.bend as Operations
13import stelliferous@0.0.3.0/machine_limits.bend as Machine
14import stelliferous@0.0.3.0/default_machine_limits.bend as Default
15import stelliferous@0.0.3.0/tensor_band.bend as Bands
16import stelliferous@0.0.3.0/tensor_band_schedule.bend as BandSchedule
17import stelliferous@0.0.3.0/tensor_band_view.bend as BandViews
18import stelliferous@0.0.3.0/traversal_partition.bend as Partition
19
20type Reason is Data:
21 ShapeMismatch{left: List<&2,U32>,right: List<&2,U32>}
22 InvalidShape{shape: List<&2,U32>}
23 InvalidArgument{shape: List<&2,U32>}
24 FileFault{path: String,detail: String}
25 MissingName{name: String}
26
27type Problem is Data:
28 Problem{operation: String,reason: Reason}
29
30type Status is Data:
31 Usable{}
32 Failed{problem: Problem}
33
34# Like a PyTorch tensor's device, workers is execution context: the size of
35# the runtime's worker pool, read once with IO.thread_count and set with
36# with_workers. Pure constructors cannot read it and use one worker. Results
37# keep the largest worker count of their operands. It never changes a value.
38type Tensor<-Element: Data> is Type:
39 Tensor{bands: Bands.Bands(~Element),geometry: Bands.Geometry,shape: List<&2,U32>,fill: Element,status: Status,workers: U32}
40
41# A nonempty sequence of tensors. Base List is not used for tensors: the
42# compiler shares every constructor of a list type that is duplicated anywhere,
43# and that sharing would reach the certificate inside each tensor's storage.
44type Parts<-Element: Data> is Type:
45 Single{tensor: Tensor<Element>}
46 Cons{tensor: Tensor<Element>,rest: Parts<Element>}
47
48# Element count of a shape that the default machine can allocate.
49def within(fits: Bool,count: U32) -> Maybe<&2,U32>:
50 match fits:
51 case True{}: Some{count}
52 case False{}: None{}
53
54def allocatable(count: Maybe<&2,U32>) -> Maybe<&2,U32>:
55 match count:
56 case None{}: None{}
57 case Some{+count}: within(U32.is_le(count,Machine.capacity(Default.limits())),count)
58
59def count_of(+shape: List<&2,U32>) -> Maybe<&2,U32>:
60 allocatable(Shapes.numel(Shapes.Shape{shape}))
61
62def known_count(count: Maybe<&2,U32>) -> U32:
63 match count:
64 case Some{count}: count
65 case None{}: 0
66
67# The element count as a number; a shape too large to allocate counts 0.
68def count_or_zero(+shape: List<&2,U32>) -> U32:
69 known_count(count_of(shape))
70
71def view_of(+shape: List<&2,U32>) -> Maybe<&2,Views.View>:
72 Views.contiguous(Shapes.Shape{shape},0)
73
74# Bits of value, counted by tail recursion. Storage.bit_length is not tail
75# recursive, and a non-tail recursive callee turns the calling chain into
76# worklist dispatch: measured allocation of 2^20 floats took twice as long.
77def significant_bits(remaining: Nat,+value: U32,+bits: Nat,zero: Bool) -> Nat:
78 Storage.significant_bits(remaining,value,bits,zero)
79
80# The smallest depth whose capacity holds count elements (at least one).
81def depth_for(+count: U32) -> Nat:
82 Storage.depth_for(count)
83
84def buffer(~Element: Data,+fill: Element,+count: U32) -> Storage.Buffer<Element>:
85 Allocation.allocate(~Element,depth_for(count),fill,U32.to_nat(count))
86
87# Row-major storage as a tensor: one band of one row whose width is every
88# element of the shape.
89def whole(~Element: Data,storage: Storage.Buffer<Element>,+shape: List<&2,U32>,fill: Element,status: Status,workers: U32) -> Tensor<Element>:
90 Tensor{Partition.Leaf{storage,0,1,Unit{}},Bands.Geometry{1,1,1,count_or_zero(shape)},shape,fill,status,workers}
91
92# The whole storage of bands. A leaf that holds the whole image already is it;
93# other bands are gathered into one allocation.
94def leaf_whole_storage(~Element: Data,whole: Bool,+geometry: Bands.Geometry,+fill: Element,data: Storage.Buffer<Element>,first: U32,rows: U32) -> Storage.Buffer<Element>:
95 match whole:
96 case True{}: data
97 case False{}: Bands.gather(~Element,geometry,Partition.Leaf{data,first,rows,Unit{}},buffer(~Element,fill,Bands.elements(geometry)))
98
99def whole_storage(~Element: Data,+geometry: Bands.Geometry,+fill: Element,bands: Bands.Bands(~Element)) -> Storage.Buffer<Element>:
100 match bands:
101 case Partition.Leaf{data,+first,+rows,details}: leaf_whole_storage(~Element,Bands.whole_image(geometry,first,rows),geometry,fill,data,first,rows)
102 case Partition.Fork{left,right}: Bands.gather(~Element,geometry,Partition.Fork{left,right},buffer(~Element,fill,Bands.elements(geometry)))
103
104def arrive_when(~Element: Data,same: Bool,geometry: Bands.Geometry,+fill: Element,layout: Bands.Plan,other: Bands.Geometry,bands: Bands.Bands(~Element)) -> Bands.Bands(~Element):
105 match same:
106 case True{}: Bands.align(~Element,geometry,fill,layout,bands)
107 case False{}: Bands.from_whole(~Element,geometry,fill,layout,whole_storage(~Element,other,fill,bands))
108
109# The bands of a tensor in the layout an operation runs in. Bands of the same
110# geometry are aligned to the plan. Bands of another geometry are read in
111# row-major order and cut by the plan, which relabels a whole image as is.
112def arrive(~Element: Data,+geometry: Bands.Geometry,layout: Bands.Plan,tensor: Tensor<Element>) -> Bands.Bands(~Element):
113 Tensor{bands,+other,shape,fill,status,workers} = tensor
114 arrive_when(~Element,Bands.same_geometry(geometry,other),geometry,fill,layout,other,bands)
115
116# The storage of a tensor arrived in the geometry of whole: the one leaf of
117# its unsplit layout.
118def with_whole(~Element: Data,~Result: Type,tensor: Tensor<Element>,
119 body: Storage.Buffer<Element> -> List<&2,U32> -> Element -> Status -> U32 -> Result) -> Result:
120 Tensor{bands,geometry,+shape,+fill,+status,+workers} = tensor
121 +whole = {Bands.Geometry{1,1,1,count_or_zero(shape)} : Bands.Geometry}
122 body(whole_storage(~Element,whole,fill,arrive(~Element,whole,Bands.unsplit(whole),Tensor{bands,geometry,shape,fill,status,workers})),shape,fill,status,workers)
123
124# Metadata inspection retains the original layout and affine storage.
125def with_metadata(~Element: Data,~Result: Type,tensor: Tensor<Element>,
126 body: Tensor<Element> -> List<&2,U32> -> Element -> Status -> U32 -> Result) -> Result:
127 Tensor{bands,geometry,+shape,+fill,+status,+workers} = tensor
128 body(Tensor{bands,geometry,shape,fill,status,workers},shape,fill,status,workers)
129
130def usable(status: Status) -> Bool:
131 match status:
132 case Usable{}: True{}
133 case Failed{problem}: False{}
134
135# The band a checked call made, as whole storage. It counts as usable only when
136# the call was usable.
137def checked_band(~Element: Data,+first: U32,+rows: U32,result: Tensor<Element>) -> Bands.Transformed(~Element):
138 Tensor{bands,+geometry,shape,+fill,status,workers} = result
139 Partition.Leaf{Bands.Checked{whole_storage(~Element,geometry,fill,bands),usable(status)},first,rows,Unit{}}
140
141def failed(~Element: Data,operation: String,reason: Reason,shape: List<&2,U32>,+fill: Element,workers: U32) -> Tensor<Element>:
142 whole(~Element,buffer(~Element,fill,0),shape,fill,Failed{Problem{operation,reason}},workers)
143
144def transformed_accepted(~Element: Data,accepted: Bool,operation: String,reason: Reason,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,+shape: List<&2,U32>,
145 fill: Element,workers: U32) -> Tensor<Element>:
146 match accepted:
147 case True{}: Tensor{bands,geometry,shape,fill,Usable{},workers}
148 case False{}: failed(~Element,operation,reason,shape,fill,workers)
149
150def transformed_inspected(~Element: Data,usable: Bool,operation: String,reason: Reason,+geometry: Bands.Geometry,+shape: List<&2,U32>,fill: Element,workers: U32,
151 inspected: Bands.Bands(~Element) & Bands.Inspection) -> Tensor<Element>:
152 (bands,inspection) = inspected
153 transformed_accepted(~Element,usable && Bands.covers(geometry,inspection),operation,reason,bands,geometry,shape,fill,workers)
154
155# Bands made by checked calls are the result when every call was usable and
156# the bands cover the geometry; otherwise the operation fails for reason.
157def transformed(~Element: Data,operation: String,reason: Reason,+geometry: Bands.Geometry,+shape: List<&2,U32>,fill: Element,workers: U32,
158 result: Bands.Bands(~Element) & Bool) -> Tensor<Element>:
159 (bands,usable) = result
160 transformed_inspected(~Element,usable,operation,reason,geometry,shape,fill,workers,Bands.inspect(~Element,geometry,bands))
161
162# A usable tensor fails with an invalid argument of the operation; a failed one keeps its problem.
163def rejected(~Element: Data,name: String,tensor: Tensor<Element>) -> Tensor<Element>:
164 match tensor:
165 case Tensor{bands,geometry,+shape,fill,Usable{},workers}: Tensor{bands,geometry,shape,fill,Failed{Problem{name,InvalidArgument{shape}}},workers}
166 case Tensor{bands,geometry,shape,fill,Failed{problem},workers}: Tensor{bands,geometry,shape,fill,Failed{problem},workers}
167
168def mismatch(same: Bool,left: List<&2,U32>,right: List<&2,U32>) -> Reason:
169 match same:
170 case True{}: InvalidShape{left}
171 case False{}: ShapeMismatch{left,right}
172
173def same_or_invalid(+left: List<&2,U32>,+right: List<&2,U32>) -> Reason:
174 mismatch(Shapes.same_extents(left,right),left,right)
175
176# Construction
177
178def full_count(~Element: Data,+shape: List<&2,U32>,count: Maybe<&2,U32>,+value: Element) -> Tensor<Element>:
179 match count:
180 case Some{count}: whole(~Element,buffer(~Element,value,count),shape,value,Usable{},1)
181 case None{}: failed(~Element,"full",InvalidShape{shape},shape,value,1)
182
183def full(~Element: Data,+shape: List<&2,U32>,+value: Element) -> Tensor<Element>:
184 full_count(~Element,shape,count_of(shape),value)
185
186def listed_count(~Element: Data,matches: Bool,+shape: List<&2,U32>,+count: U32,values: List<&2,Element>,+fill: Element) -> Tensor<Element>:
187 match matches:
188 case True{}:
189 whole(~Element,Storage.from_list(~Element,values,depth_for(count),fill,U32.to_nat(count)),shape,fill,Usable{},1)
190 case False{}: failed(~Element,"from_list",InvalidArgument{shape},shape,fill,1)
191
192def listed(~Element: Data,+shape: List<&2,U32>,count: Maybe<&2,U32>,+values: List<&2,Element>,+fill: Element) -> Tensor<Element>:
193 match count:
194 case Some{+count}: listed_count(~Element,Nat.is_eq(List.length(&2,Element,values),U32.to_nat(count)),shape,count,values,fill)
195 case None{}: failed(~Element,"from_list",InvalidShape{shape},shape,fill,1)
196
197# fill initializes storage this tensor or its results allocate; it is never
198# observable in a successful result.
199def from_list(~Element: Data,+shape: List<&2,U32>,+values: List<&2,Element>,+fill: Element) -> Tensor<Element>:
200 listed(~Element,shape,count_of(shape),values,fill)
201
202def stored(~Element: Data,matches: Bool,storage: Storage.Buffer<Element>,+shape: List<&2,U32>,fill: Element) -> Tensor<Element>:
203 match matches:
204 case True{}: whole(~Element,storage,shape,fill,Usable{},1)
205 case False{}: whole(~Element,storage,shape,fill,Failed{Problem{"from_storage",InvalidShape{shape}}},1)
206
207def matches_count(count: Maybe<&2,U32>,length: Nat) -> Bool:
208 match count:
209 case Some{count}: Nat.is_eq(length,U32.to_nat(count))
210 case None{}: False{}
211
212def storage_of(~Element: Data,+shape: List<&2,U32>,count: Maybe<&2,U32>,storage: Storage.Buffer<Element>,fill: Element) -> Tensor<Element>:
213 Storage.Buffer{storage,+length,certificate} = storage
214 stored(~Element,matches_count(count,length),Storage.Buffer{storage,length,certificate},shape,fill)
215
216# Storage handed over by foreign I/O keeps its documented trust boundary.
217def from_storage(~Element: Data,+shape: List<&2,U32>,storage: Storage.Buffer<Element>,+fill: Element) -> Tensor<Element>:
218 storage_of(~Element,shape,count_of(shape),storage,fill)
219
220# Execution context
221
222# Sets the worker count, at least one; values and status are unchanged.
223def with_workers(~Element: Data,tensor: Tensor<Element>,+count: U32) -> Tensor<Element>:
224 Tensor{bands,geometry,shape,fill,status,workers} = tensor
225 Tensor{bands,geometry,shape,fill,status,U32.max(1,count)}
226
227def workers(~Element: Data,tensor: Tensor<Element>) -> Tensor<Element> & U32:
228 Tensor{bands,geometry,shape,fill,status,+workers} = tensor
229 (Tensor{bands,geometry,shape,fill,status,workers},workers)
230
231# Observation
232
233def shape(~Element: Data,tensor: Tensor<Element>) -> Tensor<Element> & List<&2,U32>:
234 Tensor{bands,geometry,+shape,fill,status,workers} = tensor
235 (Tensor{bands,geometry,shape,fill,status,workers},shape)
236
237def checked(~Element: Data,status: Status,bands: Bands.Bands(~Element),geometry: Bands.Geometry,shape: List<&2,U32>,fill: Element,workers: U32) ->
238 Result<&1,&1,Tensor<Element> & Problem,Tensor<Element>>:
239 match status:
240 case Usable{}: Done{Tensor{bands,geometry,shape,fill,Usable{},workers}}
241 case Failed{problem}: Fail{(Tensor{bands,geometry,shape,fill,Usable{},workers},problem)}
242
243def check(~Element: Data,tensor: Tensor<Element>) -> Result<&1,&1,Tensor<Element> & Problem,Tensor<Element>>:
244 Tensor{bands,geometry,shape,fill,status,workers} = tensor
245 checked(~Element,status,bands,geometry,shape,fill,workers)
246
247def released(~Element: Data,status: Status,storage: Storage.Buffer<Element>) -> Result<&2,&1,Problem,Storage.Buffer<Element>>:
248 match status:
249 case Usable{}: Done{storage}
250 case Failed{problem}: Fail{problem}
251
252def into_storage(~Element: Data,tensor: Tensor<Element>) -> Result<&2,&1,Problem,Storage.Buffer<Element>>:
253 with_whole(~Element,~Result<&2,&1,Problem,Storage.Buffer<Element>>,tensor,storage => shape => fill => status => workers => released(~Element,status,storage))
254
255def listed_values(~Element: Data,status: Status,storage: Array<Element>,length: Nat) -> Result<&2,&2,Problem,List<&2,Element>>:
256 match status:
257 case Usable{}: Done{Storage.read_list(~Element,length,0,storage)}
258 case Failed{problem}: Fail{problem}
259
260def listed_buffer(~Element: Data,status: Status,buffer: Storage.Buffer<Element>) -> Result<&2,&2,Problem,List<&2,Element>>:
261 Storage.Buffer{storage,length,certificate} = buffer
262 listed_values(~Element,status,storage,length)
263
264def to_list(~Element: Data,tensor: Tensor<Element>) -> Result<&2,&2,Problem,List<&2,Element>>:
265 with_whole(~Element,~Result<&2,&2,Problem,List<&2,Element>>,tensor,buffer => shape => fill => status => workers =>
266 listed_buffer(~Element,status,buffer))
267
268def only_value(~Element: Data,+shape: List<&2,U32>,values: List<&2,Element>) -> Result<&2,&2,Problem,Element>:
269 match values:
270 case Con{value,Nil{}}: Done{value}
271 case values: Fail{Problem{"item",InvalidShape{shape}}}
272
273def single(~Element: Data,+shape: List<&2,U32>,values: Result<&2,&2,Problem,List<&2,Element>>) -> Result<&2,&2,Problem,Element>:
274 match values:
275 case Done{values}: only_value(~Element,shape,values)
276 case Fail{problem}: Fail{problem}
277
278def item(~Element: Data,tensor: Tensor<Element>) -> Result<&2,&2,Problem,Element>:
279 with_whole(~Element,~Result<&2,&2,Problem,Element>,tensor,storage => shape => fill => status => workers =>
280 +shape = shape
281 single(~Element,shape,listed_buffer(~Element,status,storage)))
282
283# The value read, or fallback when the tensor carries a failure.
284def value_or(~Value: Data,result: Result<&2,&2,Problem,Value>,fallback: Value) -> Value:
285 match result:
286 case Done{value}: value
287 case Fail{problem}: fallback
288
289# Operations consume their tensor arguments and return one result; callers
290# keep a tensor by copying it first. Binary elementwise operands follow NumPy
291# broadcasting; the result reuses the left storage when it has the left
292# extents. The result carries the first failure of its operands.
293
294def first_failure(left: Status,right: Status) -> Status:
295 match left:
296 case Usable{}: right
297 case Failed{problem}: Failed{problem}
298
299# A view of an operand read with the extents target; see Views.stretch.
300def stretched(view: Maybe<&2,Views.View>,+target: List<&2,U32>) -> Maybe<&2,Views.View>:
301 match view:
302 case Some{view}: Views.stretch(view,target)
303 case None{}: None{}
304
305def with_status(~Element: Data,status: Status,tensor: Tensor<Element>) -> Tensor<Element>:
306 Tensor{bands,geometry,shape,fill,previous,workers} = tensor
307 Tensor{bands,geometry,shape,fill,status,workers}
308
309def offered(~Element: Data,+geometry: Bands.Geometry,shape: List<&2,U32>,fill: Element,status: Status,workers: U32,
310 owners: Bands.Bands(~Element) & Bands.Plan) -> Tensor<Element> & BandSchedule.Target:
311 (bands,+layout) = owners
312 (Tensor{bands,geometry,shape,fill,status,workers},BandSchedule.Target{geometry,layout,BandSchedule.splits(layout)})
313
314# The layout a tensor offers to an operation: its own geometry and plan.
315def offer(~Element: Data,tensor: Tensor<Element>) -> Tensor<Element> & BandSchedule.Target:
316 Tensor{bands,+geometry,shape,fill,status,workers} = tensor
317 offered(~Element,geometry,shape,fill,status,workers,Bands.plan(~Element,bands))
318
319# Read-only operations return their tensor in a pair; Bend destructures only
320# parameters, so a continuation receives both halves.
321def then(~First: Type,~Second: Type,~Result: Type,pair: First & Second,body: First -> Second -> Result) -> Result:
322 (first,second) = pair
323 body(first,second)
324
325# The two operands arrive in the layout they met in and one band kernel
326# combines them.
327def combine_met_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,operations: F32,
328 target: BandSchedule.Target,shape: List<&2,U32>,fill: Element,+workers: U32,left: Tensor<Element>,right: Tensor<Element>) -> Tensor<Element>:
329 BandSchedule.Target{+geometry,+layout,split} = target
330 Tensor{Bands.combine_with(~Element,~kernel,geometry,BandSchedule.parallel(geometry,workers,12.0,operations,layout),arrive(~Element,geometry,layout,left),arrive(~Element,geometry,layout,right)),
331 geometry,shape,fill,Usable{},workers}
332
333def combine_offered_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,operations: F32,
334 +workers: U32,+shape: List<&2,U32>,+fill: Element,left: Tensor<Element>,right: Tensor<Element>) -> Tensor<Element>:
335 then(~Tensor<Element>,~BandSchedule.Target,~Tensor<Element>,offer(~Element,left),left => left_target =>
336 then(~Tensor<Element>,~BandSchedule.Target,~Tensor<Element>,offer(~Element,right),right => right_target =>
337 combine_met_with(~Element,~kernel,operations,BandSchedule.meet(left_target,right_target),shape,fill,workers,left,right)))
338
339# Usable operands of equal extents meet in one layout. A failed operand fails
340# the call, which keeps the left bands, geometry, extents and fill.
341def combine_equal_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,operations: F32,
342 status: Status,+workers: U32,+shape: List<&2,U32>,left: Bands.Bands(~Element),+geometry: Bands.Geometry,+left_fill: Element,
343 right: Bands.Bands(~Element),+right_geometry: Bands.Geometry,right_fill: Element) -> Tensor<Element>:
344 match status:
345 case Usable{}:
346 combine_offered_with(~Element,~kernel,operations,workers,shape,left_fill,Tensor{left,geometry,shape,left_fill,Usable{},workers},Tensor{right,right_geometry,shape,right_fill,Usable{},workers})
347 case Failed{problem}: Tensor{left,geometry,shape,left_fill,Failed{problem},workers}
348
349# A checked copy of the elements a view selects into a new whole tensor:
350# (result, source).
351def copied(~Element: Data,name: String,+source_shape: List<&2,U32>,+target_shape: List<&2,U32>,+fill: Element,+workers: U32,
352 result: Operations.Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>) -> Tensor<Element> & Tensor<Element>:
353 match result:
354 case Operations.Accepted{(input,output)}:
355 (whole(~Element,output,target_shape,fill,Usable{},workers),whole(~Element,input,source_shape,fill,Usable{},workers))
356 case Operations.Rejected{(input,output)}:
357 (whole(~Element,output,target_shape,fill,Failed{Problem{name,InvalidArgument{source_shape}}},workers),
358 whole(~Element,input,source_shape,fill,Usable{},workers))
359
360def gather_views(~Element: Data,name: String,source_view: Maybe<&2,Views.View>,+reversed: List<&2,Bool>,target_view: Maybe<&2,Views.View>,count: Maybe<&2,U32>,
361 input: Storage.Buffer<Element>,+source_shape: List<&2,U32>,+target_shape: List<&2,U32>,+fill: Element,+workers: U32) ->
362 Tensor<Element> & Tensor<Element>:
363 match source_view target_view count:
364 case Some{source_view} Some{target_view} Some{count}:
365 copied(~Element,name,source_shape,target_shape,fill,workers,Operations.copy(~Element,input,source_view,reversed,buffer(~Element,fill,count),target_view))
366 case source_view target_view count:
367 (failed(~Element,name,InvalidArgument{source_shape},target_shape,fill,workers),whole(~Element,input,source_shape,fill,Usable{},workers))
368
369# Copy the elements selected by source_view, in its row-major order with the
370# flagged axes in reverse, into a new tensor whose storage is laid out as
371# view_shape and labelled target_shape.
372def gather_status(~Element: Data,name: String,status: Status,source_view: Maybe<&2,Views.View>,+reversed: List<&2,Bool>,+view_shape: List<&2,U32>,
373 +target_shape: List<&2,U32>,input: Storage.Buffer<Element>,+source_shape: List<&2,U32>,+fill: Element,+workers: U32) ->
374 Tensor<Element> & Tensor<Element>:
375 match status:
376 case Usable{}:
377 gather_views(~Element,name,source_view,reversed,view_of(view_shape),count_of(view_shape),input,source_shape,target_shape,fill,workers)
378 case Failed{+problem}:
379 (whole(~Element,buffer(~Element,fill,0),target_shape,fill,Failed{problem},workers),
380 whole(~Element,input,source_shape,fill,Failed{problem},workers))
381
382def result_of(~Element: Data,pair: Tensor<Element> & Tensor<Element>) -> Tensor<Element>:
383 (result,source) = pair
384 result
385
386# What every band of a gather shares: the operation, the banded view of the
387# input geometry, the flags of the local view and the fill of fresh storage.
388type Gathered<-Element: Data> is Data:
389 Gathered{name: String,banded: BandViews.Banded,geometry: Bands.Geometry,reversed: List<&2,Bool>,fill: Element}
390
391# A band gathers the cells its local view selects into the output band of the
392# same rows, each row repeated as the view repeats it. A band of no rows has
393# no cells to read.
394def gathered_rows(~Element: Data,empty: Bool,name: String,+banded: BandViews.Banded,+batch: U32,+channels: U32,+height: U32,+width: U32,+reversed: List<&2,Bool>,+fill: Element,
395 +first: U32,+rows: U32,input: Storage.Buffer<Element>) -> Bands.Transformed(~Element):
396 match empty:
397 case True{}:
398 +scale = BandViews.scale(banded)
399 checked_band(~Element,(first * scale : U32),(rows * scale : U32),whole(~Element,buffer(~Element,fill,0),BandViews.extents(banded,rows),fill,Usable{},1))
400 case False{}:
401 +geometry = {Bands.Geometry{batch,channels,height,width} : Bands.Geometry}
402 +scale = BandViews.scale(banded)
403 +extents = BandViews.extents(banded,rows)
404 checked_band(~Element,(first * scale : U32),(rows * scale : U32),
405 result_of(~Element,gather_status(~Element,name,Usable{},Some{BandViews.local(banded,geometry,rows)},reversed,extents,extents,input,[channels,rows,width],fill,1)))
406
407def gathered_band(~Element: Data,context: Gathered<Element>,first: U32,rows: U32,input: Storage.Buffer<Element>) -> Bands.Transformed(~Element):
408 Gathered{name,+banded,Bands.Geometry{+batch,+channels,+height,+width},+reversed,+fill} = context
409 +first = first
410 +rows = rows
411 gathered_rows(~Element,U32.is_zero(rows),name,banded,batch,channels,height,width,reversed,fill,first,rows,input)
412
413def gather_banded(~Element: Data,+name: String,+banded: BandViews.Banded,+geometry: Bands.Geometry,+reversed: List<&2,Bool>,+shape: List<&2,U32>,+target: List<&2,U32>,
414 +fill: Element,+workers: U32,bands: Bands.Bands(~Element)) -> Tensor<Element>:
415 transformed(~Element,name,InvalidArgument{shape},BandViews.output(banded,geometry),target,fill,workers,
416 BandSchedule.transform(~Element,~Gathered<Element>,~gathered_band(~Element),Gathered{name,banded,geometry,reversed,fill},geometry,workers,12.0,0.0,bands))
417
418def gather_whole_local(~Element: Data,+name: String,+banded: BandViews.Banded,local: Maybe<&2,List<&2,Bool>>,+whole: Bands.Geometry,+shape: List<&2,U32>,+target: List<&2,U32>,
419 +fill: Element,+workers: U32,tensor: Tensor<Element>) -> Tensor<Element>:
420 match local:
421 case Some{+flags}: gather_banded(~Element,name,banded,whole,flags,shape,target,fill,workers,arrive(~Element,whole,Bands.unsplit(whole),tensor))
422 case None{}: failed(~Element,name,InvalidArgument{shape},target,fill,workers)
423
424def gather_whole(~Element: Data,+name: String,found: Maybe<&2,BandViews.Banded>,+reversed: List<&2,Bool>,+whole: Bands.Geometry,+shape: List<&2,U32>,+target: List<&2,U32>,
425 +fill: Element,+workers: U32,tensor: Tensor<Element>) -> Tensor<Element>:
426 match found:
427 case Some{+banded}: gather_whole_local(~Element,name,banded,BandViews.local_flags(banded,whole,reversed),whole,shape,target,fill,workers,tensor)
428 case None{}: failed(~Element,name,InvalidArgument{shape},target,fill,workers)
429
430# Other bands first arrive as the one band of one row of the whole geometry.
431def gather_arrived(~Element: Data,+name: String,+reversed: List<&2,Bool>,+axes: List<&2,Views.Axis>,+offset: U32,bands: Bands.Bands(~Element),
432 +geometry: Bands.Geometry,+shape: List<&2,U32>,+target: List<&2,U32>,+fill: Element,+workers: U32) -> Tensor<Element>:
433 +whole = {Bands.Geometry{1,1,1,count_or_zero(shape)} : Bands.Geometry}
434 gather_whole(~Element,name,BandViews.classify(axes,offset,whole),reversed,whole,shape,target,fill,workers,Tensor{bands,geometry,shape,fill,Usable{},workers})
435
436def gather_local(~Element: Data,+name: String,+banded: BandViews.Banded,local: Maybe<&2,List<&2,Bool>>,+reversed: List<&2,Bool>,+axes: List<&2,Views.Axis>,+offset: U32,
437 bands: Bands.Bands(~Element),+geometry: Bands.Geometry,+shape: List<&2,U32>,+target: List<&2,U32>,+fill: Element,+workers: U32) -> Tensor<Element>:
438 match local:
439 case Some{+flags}: gather_banded(~Element,name,banded,geometry,flags,shape,target,fill,workers,bands)
440 case None{}: gather_arrived(~Element,name,reversed,axes,offset,bands,geometry,shape,target,fill,workers)
441
442# Bands whose images and rows the view keeps in place and reads forward
443# gather their own cells; other bands first arrive as one band.
444def gather_found(~Element: Data,+name: String,found: Maybe<&2,BandViews.Banded>,+reversed: List<&2,Bool>,+axes: List<&2,Views.Axis>,+offset: U32,bands: Bands.Bands(~Element),
445 +geometry: Bands.Geometry,
446 +shape: List<&2,U32>,+target: List<&2,U32>,+fill: Element,+workers: U32) -> Tensor<Element>:
447 match found:
448 case Some{+banded}: gather_local(~Element,name,banded,BandViews.local_flags(banded,geometry,reversed),reversed,axes,offset,bands,geometry,shape,target,fill,workers)
449 case None{}: gather_arrived(~Element,name,reversed,axes,offset,bands,geometry,shape,target,fill,workers)
450
451def gather_counted(~Element: Data,+name: String,empty: Bool,+axes: List<&2,Views.Axis>,+offset: U32,+reversed: List<&2,Bool>,bands: Bands.Bands(~Element),
452 +geometry: Bands.Geometry,+shape: List<&2,U32>,+target: List<&2,U32>,+fill: Element,+workers: U32) -> Tensor<Element>:
453 match empty:
454 case True{}: whole(~Element,buffer(~Element,fill,0),target,fill,Usable{},workers)
455 case False{}: gather_found(~Element,name,BandViews.classify(axes,offset,geometry),reversed,axes,offset,bands,geometry,shape,target,fill,workers)
456
457# The result must have a count. The view comes taken apart, so that no View
458# value is used twice: that would make every View constructor shared.
459def gather_axes(~Element: Data,+name: String,count: Maybe<&2,U32>,+axes: List<&2,Views.Axis>,+offset: U32,+reversed: List<&2,Bool>,bands: Bands.Bands(~Element),
460 +geometry: Bands.Geometry,+shape: List<&2,U32>,+target: List<&2,U32>,+fill: Element,+workers: U32) -> Tensor<Element>:
461 match count:
462 case Some{+count}: gather_counted(~Element,name,U32.is_zero(count),axes,offset,reversed,bands,geometry,shape,target,fill,workers)
463 case None{}: failed(~Element,name,InvalidArgument{shape},target,fill,workers)
464
465# The elements a view of the tensor's row-major elements selects, in the
466# view's row-major order with the flagged axes in reverse, as a tensor of
467# extents target. An axis without a flag is read forward.
468def gather(~Element: Data,name: String,status: Status,source: Maybe<&2,Views.View>,+reversed: List<&2,Bool>,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,
469 +shape: List<&2,U32>,+target: List<&2,U32>,+fill: Element,+workers: U32) -> Tensor<Element>:
470 match status source:
471 case Usable{} Some{Views.View{+axes,+offset}}: gather_axes(~Element,name,count_of(target),axes,offset,reversed,bands,geometry,shape,target,fill,workers)
472 case Usable{} None{}: failed(~Element,name,InvalidArgument{shape},target,fill,workers)
473 case Failed{+problem} source: whole(~Element,buffer(~Element,fill,0),target,fill,Failed{problem},workers)
474
475# The same gather labelled with the extents of the view.
476def gather_extents(~Element: Data,name: String,status: Status,source: Maybe<&2,Views.View>,+reversed: List<&2,Bool>,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,
477 +shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensor<Element>:
478 match source:
479 case Some{Views.View{+axes,offset}}: gather(~Element,name,status,Some{Views.View{axes,offset}},reversed,bands,geometry,shape,Views.axis_extents(axes),fill,workers)
480 case None{}: gather(~Element,name,status,None{},reversed,bands,geometry,shape,[],fill,workers)
481
482# The tensor repeated to the extents target under NumPy broadcasting: one
483# gather through the stretched view.
484def broadcast_view(~Element: Data,status: Status,view: Maybe<&2,Views.View>,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,+shape: List<&2,U32>,
485 +target: List<&2,U32>,+fill: Element,+workers: U32) -> Tensor<Element>:
486 match status view:
487 case Usable{} None{}: failed(~Element,"broadcast_to",ShapeMismatch{shape,target},target,fill,workers)
488 case status view: gather(~Element,"broadcast_to",status,view,[],bands,geometry,shape,target,fill,workers)
489
490def broadcast_to(~Element: Data,tensor: Tensor<Element>,+target: List<&2,U32>) -> Tensor<Element>:
491 Tensor{bands,+geometry,+shape,+fill,status,+workers} = tensor
492 broadcast_view(~Element,status,stretched(view_of(shape),target),bands,geometry,shape,target,fill,workers)
493
494# Each band of the left takes the part of the right source it reads through
495# the stretched view. A layout whose bands have no view of their cells reads
496# the right repeated to the common extents.
497def combine_spread_when(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,operations: F32,
498 boxed: Bool,+shape: List<&2,U32>,+workers: U32,+geometry: Bands.Geometry,+layout: Bands.Plan,left: Bands.Bands(~Element),+fill: Element,
499 source: Storage.Buffer<Element>,+right_shape: List<&2,U32>,right_fill: Element) -> Tensor<Element>:
500 match boxed:
501 case True{}:
502 Tensor{Bands.combine_spread_with(~Element,~kernel,Bands.Spread{geometry,shape,right_shape},BandSchedule.parallel(geometry,workers,12.0,operations,layout),left,
503 Bands.spread(~Element,Bands.Spread{geometry,shape,right_shape},right_fill,layout,source)),geometry,shape,fill,Usable{},workers}
504 case False{}:
505 combine_offered_with(~Element,~kernel,operations,workers,shape,fill,Tensor{left,geometry,shape,fill,Usable{},workers},
506 broadcast_to(~Element,whole(~Element,source,right_shape,right_fill,Usable{},workers),shape))
507
508def combine_viewed_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,operations: F32,
509 +shape: List<&2,U32>,+workers: U32,+geometry: Bands.Geometry,+fill: Element,planned: Bands.Bands(~Element) & Bands.Plan,
510 source: Storage.Buffer<Element>,+right_shape: List<&2,U32>,right_fill: Element) -> Tensor<Element>:
511 (left,+layout) = planned
512 combine_spread_when(~Element,~kernel,operations,Bands.spread_boxed(Bands.Spread{geometry,shape,right_shape},layout) && Bool.not(Shapes.has_zero(shape)),
513 shape,workers,geometry,layout,left,fill,source,right_shape,right_fill)
514
515# A left operand of the common extents keeps its layout; the right operand is
516# read through its view stretched to those extents.
517def combine_stretched_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,operations: F32,
518 +shape: List<&2,U32>,+workers: U32,left: Tensor<Element>,right: Tensor<Element>) -> Tensor<Element>:
519 Tensor{bands,+geometry,left_shape,+fill,status,left_workers} = left
520 with_whole(~Element,~Tensor<Element>,right,source => right_shape => right_fill => right_status => right_workers =>
521 +right_shape = right_shape
522 combine_viewed_with(~Element,~kernel,operations,shape,workers,geometry,fill,Bands.plan(~Element,bands),source,right_shape,right_fill))
523
524def combine_repeated_when(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,operations: F32,
525 same: Bool,+shape: List<&2,U32>,+workers: U32,+fill: Element,left: Tensor<Element>,right: Tensor<Element>) -> Tensor<Element>:
526 match same:
527 case True{}: combine_offered_with(~Element,~kernel,operations,workers,shape,fill,left,right)
528 case False{}: combine_stretched_with(~Element,~kernel,operations,shape,workers,left,right)
529
530# A repeated left operand meets a right operand of the common extents as an
531# operand of equal extents.
532def combine_repeated_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,operations: F32,
533 +shape: List<&2,U32>,+workers: U32,+fill: Element,left: Tensor<Element>,right: Tensor<Element>) -> Tensor<Element>:
534 Tensor{right_bands,right_geometry,+right_shape,right_fill,right_status,right_workers} = right
535 combine_repeated_when(~Element,~kernel,operations,Shapes.same_extents(right_shape,shape),shape,workers,fill,left,
536 Tensor{right_bands,right_geometry,right_shape,right_fill,right_status,right_workers})
537
538# A left operand of other extents is first repeated to the common extents.
539def combine_common_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,name: String,operations: F32,
540 same: Bool,count: Maybe<&2,U32>,+shape: List<&2,U32>,+workers: U32,+fill: Element,left: Tensor<Element>,right: Tensor<Element>) -> Tensor<Element>:
541 match same count:
542 case True{} count: combine_stretched_with(~Element,~kernel,operations,shape,workers,left,right)
543 case False{} Some{count}: combine_repeated_with(~Element,~kernel,operations,shape,workers,fill,broadcast_to(~Element,left,shape),right)
544 case False{} None{}: failed(~Element,name,InvalidShape{shape},shape,fill,workers)
545
546# Operands of other extents follow NumPy broadcasting (Shapes.broadcast) and
547# fail only when they are incompatible.
548def combine_broadcast_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,name: String,operations: F32,
549 status: Status,common: Maybe<&2,List<&2,U32>>,+workers: U32,left: Bands.Bands(~Element),+geometry: Bands.Geometry,+left_shape: List<&2,U32>,+left_fill: Element,
550 right: Bands.Bands(~Element),+right_geometry: Bands.Geometry,+right_shape: List<&2,U32>,right_fill: Element) -> Tensor<Element>:
551 match status common:
552 case Usable{} Some{+shape}:
553 combine_common_with(~Element,~kernel,name,operations,Shapes.same_extents(left_shape,shape),count_of(shape),shape,workers,left_fill,
554 Tensor{left,geometry,left_shape,left_fill,Usable{},workers},Tensor{right,right_geometry,right_shape,right_fill,Usable{},workers})
555 case Usable{} None{}: Tensor{left,geometry,left_shape,left_fill,Failed{Problem{name,ShapeMismatch{left_shape,right_shape}}},workers}
556 case Failed{problem} common: Tensor{left,geometry,left_shape,left_fill,Failed{problem},workers}
557
558# Operands of equal extents run in the layout they meet in; operands of other
559# extents broadcast.
560def combine_shapes(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,name: String,operations: F32,same: Bool,
561 left: Bands.Bands(~Element),+geometry: Bands.Geometry,+left_shape: List<&2,U32>,+left_fill: Element,left_status: Status,left_workers: U32,
562 right: Bands.Bands(~Element),+right_geometry: Bands.Geometry,+right_shape: List<&2,U32>,right_fill: Element,right_status: Status,right_workers: U32) -> Tensor<Element>:
563 match same:
564 case True{}:
565 combine_equal_with(~Element,~kernel,operations,first_failure(left_status,right_status),U32.max(left_workers,right_workers),left_shape,left,geometry,left_fill,right,right_geometry,right_fill)
566 case False{}:
567 combine_broadcast_with(~Element,~kernel,name,operations,first_failure(left_status,right_status),Shapes.broadcast(left_shape,right_shape),U32.max(left_workers,right_workers),
568 left,geometry,left_shape,left_fill,right,right_geometry,right_shape,right_fill)
569
570# operation(left element, right element), retaining the band layout of the
571# operands. Alignment is done once, before any parallel leaf work.
572# operations counts the operation's scalar work for the band schedule.
573def 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>,name: String,operations: F32,left: Tensor<Element>,right: Tensor<Element>) -> Tensor<Element>:
574 match left right:
575 case Tensor{left,geometry,+left_shape,left_fill,left_status,left_workers} Tensor{right,right_geometry,+right_shape,right_fill,right_status,right_workers}:
576 combine_shapes(~Element,~kernel,name,operations,Shapes.same_extents(left_shape,right_shape),
577 left,geometry,left_shape,left_fill,left_status,left_workers,right,right_geometry,right_shape,right_fill,right_status,right_workers)
578
579# The elementwise form: the kernel applies operation to each pair in order.
580def combine(~Element: Data,~operation: Element -> Element -> Element,name: String,operations: F32,left: Tensor<Element>,right: Tensor<Element>) -> Tensor<Element>:
581 combine_with(~Element,~Operations.combine_kernel(~Element,~operation),name,operations,left,right)
582
583def apply_bands_context(~Element: Data,~Context: Data,~mapper: @+count: Nat -> @+position: U32 -> @+context: Context -> Storage.Buffer<Element> -> Storage.Buffer<Element>,context: Context,operations: F32,status: Status,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,
584 shape: List<&2,U32>,fill: Element,+workers: U32) -> Tensor<Element>:
585 match status:
586 case Usable{}: Tensor{BandSchedule.map_context(~Element,~Context,~mapper,context,geometry,workers,operations,bands),geometry,shape,fill,Usable{},workers}
587 case Failed{problem}: Tensor{bands,geometry,shape,fill,Failed{problem},workers}
588
589def apply_context(~Element: Data,~Context: Data,~mapper: @+count: Nat -> @+position: U32 -> @+context: Context -> Storage.Buffer<Element> -> Storage.Buffer<Element>,context: Context,operations: F32,tensor: Tensor<Element>) -> Tensor<Element>:
590 Tensor{bands,geometry,shape,fill,status,workers} = tensor
591 apply_bands_context(~Element,~Context,~mapper,context,operations,status,bands,geometry,shape,fill,workers)
592
593def apply(~Element: Data,~mapper: @+count: Nat -> @+position: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element>,operations: F32,tensor: Tensor<Element>) -> Tensor<Element>:
594 apply_context(~Element,~Unit,~Operations.without_context(~Element,~mapper),Unit{},operations,tensor)
595
596def indexed_width_positive(geometry: Bands.Geometry) -> Bool:
597 Bands.Geometry{batch,channels,height,width} = geometry
598 Nat.is_lt(0n,U32.to_nat(width))
599
600# An empty width holds no element, so the bands stay as they are. The choice
601# is made once, before the band scheduler or element loop.
602def apply_indexed_positive(~Element: Data,~operation: U32 -> Element -> Element,positive: Bool,operations: F32,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,
603 +shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensor<Element>:
604 match positive:
605 case True{}: Tensor{BandSchedule.indexed(~Element,~operation,geometry,workers,operations,fill,bands),geometry,shape,fill,Usable{},workers}
606 case False{}: Tensor{bands,geometry,shape,fill,Usable{},workers}
607
608def apply_indexed_bands(~Element: Data,~operation: U32 -> Element -> Element,operations: F32,status: Status,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,
609 shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensor<Element>:
610 match status:
611 case Usable{}: apply_indexed_positive(~Element,~operation,indexed_width_positive(geometry),operations,bands,geometry,shape,fill,workers)
612 case Failed{problem}: Tensor{bands,geometry,shape,fill,Failed{problem},workers}
613
614# The function also receives each element's row-major index.
615def apply_indexed(~Element: Data,~operation: U32 -> Element -> Element,operations: F32,tensor: Tensor<Element>) -> Tensor<Element>:
616 Tensor{bands,geometry,shape,fill,status,workers} = tensor
617 apply_indexed_bands(~Element,~operation,operations,status,bands,geometry,shape,fill,workers)
618
619def same(~Element: Data,value: Element) -> Element:
620 value
621
622# Updates every element x to operation(x, value).
623def scalar_bands(~Element: Data,~operation: Element -> Element -> Element,operations: F32,status: Status,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,
624 shape: List<&2,U32>,fill: Element,+workers: U32,value: Element) -> Tensor<Element>:
625 match status:
626 case Usable{}: Tensor{BandSchedule.scalar(~Element,~operation,geometry,workers,operations,value,bands),geometry,shape,fill,Usable{},workers}
627 case Failed{problem}: Tensor{bands,geometry,shape,fill,Failed{problem},workers}
628
629def combine_scalar(~Element: Data,~operation: Element -> Element -> Element,operations: F32,tensor: Tensor<Element>,value: Element) -> Tensor<Element>:
630 Tensor{bands,geometry,shape,fill,status,workers} = tensor
631 scalar_bands(~Element,~operation,operations,status,bands,geometry,shape,fill,workers,value)
632
633# Messages
634
635def show_extents(extents: List<&2,U32>) -> String:
636 match extents:
637 case []: ""
638 case [extent]: U32.show(extent)
639 case extent <> rest: String.append(U32.show(extent),String.append(", ",show_extents(rest)))
640
641def show_shape(extents: List<&2,U32>) -> String:
642 String.append("[",String.append(show_extents(extents),"]"))
643
644def explain(reason: Reason) -> String:
645 match reason:
646 case ShapeMismatch{left,right}:
647 String.append("shapes ",String.append(show_shape(left),String.append(" and ",String.append(show_shape(right)," do not match"))))
648 case InvalidShape{shape}:
649 String.append("shape ",String.append(show_shape(shape)," is invalid or exceeds the machine capacity"))
650 case InvalidArgument{shape}:
651 String.append("invalid argument for shape ",show_shape(shape))
652 case FileFault{path,detail}:
653 String.append("file ",String.append(path,String.append(" ",detail)))
654 case MissingName{name}:
655 String.append("no tensor named ",name)
656
657def describe(problem: Problem) -> String:
658 Problem{operation,reason} = problem
659 String.append(operation,String.append(": ",explain(reason)))
660
661def expected(~Element: Data,result: Result<&1,&1,Tensor<Element> & Problem,Tensor<Element>>) -> IO(Tensor<Element>):
662 match result:
663 case Done{tensor}: IO.pure(Tensor<Element>,tensor)
664 case Fail{(tensor,problem)}: IO.die(Tensor<Element>,1,describe(problem))
665
666# Stops the program with the recorded problem, for scripts and examples.
667def expect(~Element: Data,tensor: Tensor<Element>) -> IO(Tensor<Element>):
668 expected(~Element,check(~Element,tensor))
669
670def expected_storage(~Element: Data,result: Result<&2,&1,Problem,Storage.Buffer<Element>>) -> IO(Storage.Buffer<Element>):
671 match result:
672 case Done{storage}: IO.pure(Storage.Buffer<Element>,storage)
673 case Fail{problem}: IO.die(Storage.Buffer<Element>,1,describe(problem))
674
675def expect_storage(~Element: Data,tensor: Tensor<Element>) -> IO(Storage.Buffer<Element>):
676 expected_storage(~Element,into_storage(~Element,tensor))