tensor.bend

stelliferous@0.0.2.0 · 601 lines · 99 definitions

1# Public tensors own dense row-major storage or image-row bands, their extents,
2# the value used to initialize allocations, a sticky status, and the number of
3# workers their operations may use. A failed operation leaves values untouched and
4# records why; later operations skip a failed tensor, and check is the only way
5# to observe or clear the failure.
6import Base
7import stelliferous@0.0.2.0/storage_buffer.bend as Storage
8import stelliferous@0.0.2.0/storage_allocation.bend as Allocation
9import stelliferous@0.0.2.0/tensor_shape.bend as Shapes
10import stelliferous@0.0.2.0/tensor_view.bend as Views
11import stelliferous@0.0.2.0/tensor_operations.bend as Operations
12import stelliferous@0.0.2.0/machine_limits.bend as Machine
13import stelliferous@0.0.2.0/default_machine_limits.bend as Default
14import stelliferous@0.0.2.0/tensor_band.bend as Bands
15import stelliferous@0.0.2.0/tensor_band_schedule.bend as BandSchedule
16
17type Reason is Data:
18  ShapeMismatch{left: List<&2,U32>,right: List<&2,U32>}
19  InvalidShape{shape: List<&2,U32>}
20  InvalidArgument{shape: List<&2,U32>}
21  FileFault{path: String,detail: String}
22  MissingName{name: String}
23
24type Problem is Data:
25  Problem{operation: String,reason: Reason}
26
27type Status is Data:
28  Usable{}
29  Failed{problem: Problem}
30
31# Like a PyTorch tensor's device, workers is execution context: the size of
32# the runtime's worker pool, read once with IO.thread_count and set with
33# with_workers. Pure constructors cannot read it and use one worker. Results
34# keep the largest worker count of their operands. It never changes a value.
35type Tensor<-Element: Data> is Type:
36  Tensor{storage: Storage.Buffer<Element>,shape: List<&2,U32>,fill: Element,status: Status,workers: U32}
37  Banded{bands: Bands.Bands(~Element),geometry: Bands.Geometry,shape: List<&2,U32>,fill: Element,status: Status,workers: U32}
38
39# A nonempty sequence of tensors. Base List is not used for tensors: the
40# compiler shares every constructor of a list type that is duplicated anywhere,
41# and that sharing would reach the certificate inside each tensor's storage.
42type Parts<-Element: Data> is Type:
43  Single{tensor: Tensor<Element>}
44  Cons{tensor: Tensor<Element>,rest: Parts<Element>}
45
46# Element count of a shape that the default machine can allocate.
47def within(fits: Bool,count: U32) -> Maybe<&2,U32>:
48  match fits:
49    case True{}: Some{count}
50    case False{}: None{}
51
52def allocatable(count: Maybe<&2,U32>) -> Maybe<&2,U32>:
53  match count:
54    case None{}: None{}
55    case Some{+count}: within(U32.is_le(count,Machine.capacity(Default.limits())),count)
56
57def count_of(+shape: List<&2,U32>) -> Maybe<&2,U32>:
58  allocatable(Shapes.numel(Shapes.Shape{shape}))
59
60def known_count(count: Maybe<&2,U32>) -> U32:
61  match count:
62    case Some{count}: count
63    case None{}: 0
64
65# The element count as a number; a shape too large to allocate counts 0.
66def count_or_zero(+shape: List<&2,U32>) -> U32:
67  known_count(count_of(shape))
68
69def view_of(+shape: List<&2,U32>) -> Maybe<&2,Views.View>:
70  Views.contiguous(Shapes.Shape{shape},0)
71
72# Bits of value, counted by tail recursion. Storage.bit_length is not tail
73# recursive, and a non-tail recursive callee turns the calling chain into
74# worklist dispatch: measured allocation of 2^20 floats took twice as long.
75def significant_bits(remaining: Nat,+value: U32,+bits: Nat,zero: Bool) -> Nat:
76  Storage.significant_bits(remaining,value,bits,zero)
77
78# The smallest depth whose capacity holds count elements (at least one).
79def depth_for(+count: U32) -> Nat:
80  Storage.depth_for(count)
81
82def buffer(~Element: Data,+fill: Element,+count: U32) -> Storage.Buffer<Element>:
83  Allocation.allocate(~Element,depth_for(count),fill,U32.to_nat(count))
84
85def with_dense(~Element: Data,~Result: Type,tensor: Tensor<Element>,
86  body: Storage.Buffer<Element> -> List<&2,U32> -> Element -> Status -> U32 -> Result) -> Result:
87  match tensor:
88    case Tensor{storage,shape,fill,status,workers}: body(storage,shape,fill,status,workers)
89    case Banded{bands,+geometry,shape,+fill,status,workers}:
90      body(Bands.gather(~Element,geometry,bands,buffer(~Element,fill,Bands.elements(geometry))),shape,fill,status,workers)
91
92# Metadata inspection retains the original layout and affine storage.
93def with_metadata(~Element: Data,~Result: Type,tensor: Tensor<Element>,
94  body: Tensor<Element> -> List<&2,U32> -> Element -> Status -> U32 -> Result) -> Result:
95  match tensor:
96    case Tensor{storage,+shape,+fill,+status,+workers}: body(Tensor{storage,shape,fill,status,workers},shape,fill,status,workers)
97    case Banded{bands,geometry,+shape,+fill,+status,+workers}: body(Banded{bands,geometry,shape,fill,status,workers},shape,fill,status,workers)
98
99def failed(~Element: Data,operation: String,reason: Reason,shape: List<&2,U32>,+fill: Element,workers: U32) -> Tensor<Element>:
100  Tensor{buffer(~Element,fill,0),shape,fill,Failed{Problem{operation,reason}},workers}
101
102# A usable tensor fails with an invalid argument of the operation; a failed one keeps its problem.
103def rejected(~Element: Data,name: String,tensor: Tensor<Element>) -> Tensor<Element>:
104  match tensor:
105    case Tensor{storage,+shape,fill,Usable{},workers}: Tensor{storage,shape,fill,Failed{Problem{name,InvalidArgument{shape}}},workers}
106    case Tensor{storage,shape,fill,Failed{problem},workers}: Tensor{storage,shape,fill,Failed{problem},workers}
107    case Banded{bands,geometry,+shape,fill,Usable{},workers}: Banded{bands,geometry,shape,fill,Failed{Problem{name,InvalidArgument{shape}}},workers}
108    case Banded{bands,geometry,shape,fill,Failed{problem},workers}: Banded{bands,geometry,shape,fill,Failed{problem},workers}
109
110def mismatch(same: Bool,left: List<&2,U32>,right: List<&2,U32>) -> Reason:
111  match same:
112    case True{}: InvalidShape{left}
113    case False{}: ShapeMismatch{left,right}
114
115def same_or_invalid(+left: List<&2,U32>,+right: List<&2,U32>) -> Reason:
116  mismatch(Shapes.same_extents(left,right),left,right)
117
118# Construction
119
120def full_count(~Element: Data,+shape: List<&2,U32>,count: Maybe<&2,U32>,+value: Element) -> Tensor<Element>:
121  match count:
122    case Some{count}: Tensor{buffer(~Element,value,count),shape,value,Usable{},1}
123    case None{}: failed(~Element,"full",InvalidShape{shape},shape,value,1)
124
125def full(~Element: Data,+shape: List<&2,U32>,+value: Element) -> Tensor<Element>:
126  full_count(~Element,shape,count_of(shape),value)
127
128def listed_count(~Element: Data,matches: Bool,+shape: List<&2,U32>,+count: U32,values: List<&2,Element>,+fill: Element) -> Tensor<Element>:
129  match matches:
130    case True{}:
131      Tensor{Storage.from_list(~Element,values,depth_for(count),fill,U32.to_nat(count)),shape,fill,Usable{},1}
132    case False{}: failed(~Element,"from_list",InvalidArgument{shape},shape,fill,1)
133
134def listed(~Element: Data,+shape: List<&2,U32>,count: Maybe<&2,U32>,+values: List<&2,Element>,+fill: Element) -> Tensor<Element>:
135  match count:
136    case Some{+count}: listed_count(~Element,Nat.is_eq(List.length(&2,Element,values),U32.to_nat(count)),shape,count,values,fill)
137    case None{}: failed(~Element,"from_list",InvalidShape{shape},shape,fill,1)
138
139# fill initializes storage this tensor or its results allocate; it is never
140# observable in a successful result.
141def from_list(~Element: Data,+shape: List<&2,U32>,+values: List<&2,Element>,+fill: Element) -> Tensor<Element>:
142  listed(~Element,shape,count_of(shape),values,fill)
143
144def stored(~Element: Data,matches: Bool,storage: Storage.Buffer<Element>,+shape: List<&2,U32>,fill: Element) -> Tensor<Element>:
145  match matches:
146    case True{}: Tensor{storage,shape,fill,Usable{},1}
147    case False{}: Tensor{storage,shape,fill,Failed{Problem{"from_storage",InvalidShape{shape}}},1}
148
149def matches_count(count: Maybe<&2,U32>,length: Nat) -> Bool:
150  match count:
151    case Some{count}: Nat.is_eq(length,U32.to_nat(count))
152    case None{}: False{}
153
154def storage_of(~Element: Data,+shape: List<&2,U32>,count: Maybe<&2,U32>,storage: Storage.Buffer<Element>,fill: Element) -> Tensor<Element>:
155  Storage.Buffer{storage,+length,certificate} = storage
156  stored(~Element,matches_count(count,length),Storage.Buffer{storage,length,certificate},shape,fill)
157
158# Storage handed over by foreign I/O keeps its documented trust boundary.
159def from_storage(~Element: Data,+shape: List<&2,U32>,storage: Storage.Buffer<Element>,+fill: Element) -> Tensor<Element>:
160  storage_of(~Element,shape,count_of(shape),storage,fill)
161
162# Execution context
163
164# Sets the worker count, at least one; values and status are unchanged.
165def with_workers(~Element: Data,tensor: Tensor<Element>,+count: U32) -> Tensor<Element>:
166  match tensor:
167    case Tensor{storage,shape,fill,status,workers}: Tensor{storage,shape,fill,status,U32.max(1,count)}
168    case Banded{bands,geometry,shape,fill,status,workers}: Banded{bands,geometry,shape,fill,status,U32.max(1,count)}
169
170def workers(~Element: Data,tensor: Tensor<Element>) -> Tensor<Element> & U32:
171  match tensor:
172    case Tensor{storage,shape,fill,status,+workers}: (Tensor{storage,shape,fill,status,workers},workers)
173    case Banded{bands,geometry,shape,fill,status,+workers}: (Banded{bands,geometry,shape,fill,status,workers},workers)
174
175# Observation
176
177def shape(~Element: Data,tensor: Tensor<Element>) -> Tensor<Element> & List<&2,U32>:
178  match tensor:
179    case Tensor{storage,+shape,fill,status,workers}: (Tensor{storage,shape,fill,status,workers},shape)
180    case Banded{bands,geometry,+shape,fill,status,workers}: (Banded{bands,geometry,shape,fill,status,workers},shape)
181
182def checked(~Element: Data,status: Status,storage: Storage.Buffer<Element>,shape: List<&2,U32>,fill: Element,workers: U32) ->
183  Result<&1,&1,Tensor<Element> & Problem,Tensor<Element>>:
184  match status:
185    case Usable{}: Done{Tensor{storage,shape,fill,Usable{},workers}}
186    case Failed{problem}: Fail{(Tensor{storage,shape,fill,Usable{},workers},problem)}
187
188def checked_bands(~Element: Data,status: Status,bands: Bands.Bands(~Element),geometry: Bands.Geometry,shape: List<&2,U32>,fill: Element,workers: U32) ->
189  Result<&1,&1,Tensor<Element> & Problem,Tensor<Element>>:
190  match status:
191    case Usable{}: Done{Banded{bands,geometry,shape,fill,Usable{},workers}}
192    case Failed{problem}: Fail{(Banded{bands,geometry,shape,fill,Usable{},workers},problem)}
193
194def check(~Element: Data,tensor: Tensor<Element>) -> Result<&1,&1,Tensor<Element> & Problem,Tensor<Element>>:
195  match tensor:
196    case Tensor{storage,shape,fill,status,workers}: checked(~Element,status,storage,shape,fill,workers)
197    case Banded{bands,geometry,shape,fill,status,workers}: checked_bands(~Element,status,bands,geometry,shape,fill,workers)
198
199def released(~Element: Data,status: Status,storage: Storage.Buffer<Element>) -> Result<&2,&1,Problem,Storage.Buffer<Element>>:
200  match status:
201    case Usable{}: Done{storage}
202    case Failed{problem}: Fail{problem}
203
204def into_storage(~Element: Data,tensor: Tensor<Element>) -> Result<&2,&1,Problem,Storage.Buffer<Element>>:
205  with_dense(~Element,~Result<&2,&1,Problem,Storage.Buffer<Element>>,tensor,storage => shape => fill => status => workers => released(~Element,status,storage))
206
207def listed_values(~Element: Data,status: Status,storage: Array<Element>,length: Nat) -> Result<&2,&2,Problem,List<&2,Element>>:
208  match status:
209    case Usable{}: Done{Storage.read_list(~Element,length,0,storage)}
210    case Failed{problem}: Fail{problem}
211
212def listed_buffer(~Element: Data,status: Status,buffer: Storage.Buffer<Element>) -> Result<&2,&2,Problem,List<&2,Element>>:
213  Storage.Buffer{storage,length,certificate} = buffer
214  listed_values(~Element,status,storage,length)
215
216def to_list(~Element: Data,tensor: Tensor<Element>) -> Result<&2,&2,Problem,List<&2,Element>>:
217  with_dense(~Element,~Result<&2,&2,Problem,List<&2,Element>>,tensor,buffer => shape => fill => status => workers =>
218    listed_buffer(~Element,status,buffer))
219
220def only_value(~Element: Data,+shape: List<&2,U32>,values: List<&2,Element>) -> Result<&2,&2,Problem,Element>:
221  match values:
222    case Con{value,Nil{}}: Done{value}
223    case values: Fail{Problem{"item",InvalidShape{shape}}}
224
225def single(~Element: Data,+shape: List<&2,U32>,values: Result<&2,&2,Problem,List<&2,Element>>) -> Result<&2,&2,Problem,Element>:
226  match values:
227    case Done{values}: only_value(~Element,shape,values)
228    case Fail{problem}: Fail{problem}
229
230def item(~Element: Data,tensor: Tensor<Element>) -> Result<&2,&2,Problem,Element>:
231  with_dense(~Element,~Result<&2,&2,Problem,Element>,tensor,storage => shape => fill => status => workers =>
232    +shape = shape
233    single(~Element,shape,to_list(~Element,Tensor{storage,shape,fill,status,workers})))
234
235# The value read, or fallback when the tensor carries a failure.
236def value_or(~Value: Data,result: Result<&2,&2,Problem,Value>,fallback: Value) -> Value:
237  match result:
238    case Done{value}: value
239    case Fail{problem}: fallback
240
241# Operations consume their tensor arguments and return one result; callers
242# keep a tensor by copying it first. Binary elementwise operands follow NumPy
243# broadcasting; the result reuses the left storage when it has the left
244# extents. The result carries the first failure of its operands.
245
246def first_failure(left: Status,right: Status) -> Status:
247  match left:
248    case Usable{}: right
249    case Failed{problem}: Failed{problem}
250
251def combined(~Element: Data,operation: String,+left_shape: List<&2,U32>,+right_shape: List<&2,U32>,+fill: Element,+workers: U32,
252  result: Operations.Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>) -> Tensor<Element>:
253  match result:
254    case Operations.Accepted{(left,right)}: Tensor{left,left_shape,fill,Usable{},workers}
255    case Operations.Rejected{(left,right)}: Tensor{left,left_shape,fill,Failed{Problem{operation,same_or_invalid(left_shape,right_shape)}},workers}
256
257def combine_views_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,left_view: Maybe<&2,Views.View>,right_view: Maybe<&2,Views.View>,
258  left: Storage.Buffer<Element>,+left_shape: List<&2,U32>,+fill: Element,right: Storage.Buffer<Element>,+right_shape: List<&2,U32>,+workers: U32) -> Tensor<Element>:
259  match left_view right_view:
260    case Some{left_view} Some{right_view}:
261      combined(~Element,name,left_shape,right_shape,fill,workers,Operations.combine_with(~Element,~kernel,left,left_view,right,right_view))
262    case left_view right_view: Tensor{left,left_shape,fill,Failed{Problem{name,same_or_invalid(left_shape,right_shape)}},workers}
263
264# A view of an operand read with the extents target; see Views.stretch.
265def stretched(view: Maybe<&2,Views.View>,+target: List<&2,U32>) -> Maybe<&2,Views.View>:
266  match view:
267    case Some{view}: Views.stretch(view,target)
268    case None{}: None{}
269
270# The left operand was copied through its stretched view into new storage of
271# the common extents; the right operand is then combined into that storage.
272# The view of that storage is built once for each of the two steps: using one
273# view value twice would make the compiler share every View constructor.
274def combined_fresh_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,+shape: List<&2,U32>,+fill: Element,+workers: U32,
275  right: Storage.Buffer<Element>,right_view: Views.View,target: Views.View,
276  copied: Operations.Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>) -> Tensor<Element>:
277  match copied:
278    case Operations.Accepted{(left,output)}:
279      combined(~Element,name,shape,shape,fill,workers,Operations.combine_with(~Element,~kernel,output,target,right,right_view))
280    case Operations.Rejected{(left,output)}: Tensor{output,shape,fill,Failed{Problem{name,InvalidShape{shape}}},workers}
281
282def combine_fresh_views_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,left_view: Maybe<&2,Views.View>,right_view: Maybe<&2,Views.View>,
283  target: Maybe<&2,Views.View>,output: Maybe<&2,Views.View>,count: Maybe<&2,U32>,left: Storage.Buffer<Element>,right: Storage.Buffer<Element>,
284  +shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensor<Element>:
285  match left_view right_view target output count:
286    case Some{left_view} Some{right_view} Some{target} Some{output} Some{count}:
287      combined_fresh_with(~Element,~kernel,name,shape,fill,workers,right,right_view,output,
288        Operations.copy(~Element,left,left_view,buffer(~Element,fill,count),target))
289    case left_view right_view target output count: failed(~Element,name,InvalidShape{shape},shape,fill,workers)
290
291# Operands of different extents with the common extents shape. When shape is
292# the left extents, only the right operand is stretched and the result reuses
293# the left storage. Otherwise the left operand is first copied through its
294# stretched view into new storage, which the right operand is combined into;
295# both steps run the innermost axis as one loop, which the three-buffer
296# combination with two cursors does not.
297def 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>,name: String,in_place: Bool,
298  left: Storage.Buffer<Element>,+left_shape: List<&2,U32>,+fill: Element,right: Storage.Buffer<Element>,+right_shape: List<&2,U32>,+shape: List<&2,U32>,+workers: U32) -> Tensor<Element>:
299  match in_place:
300    case True{}:
301      combine_views_with(~Element,~kernel,name,view_of(left_shape),stretched(view_of(right_shape),left_shape),left,left_shape,fill,right,right_shape,workers)
302    case False{}:
303      combine_fresh_views_with(~Element,~kernel,name,stretched(view_of(left_shape),shape),stretched(view_of(right_shape),shape),view_of(shape),view_of(shape),count_of(shape),
304        left,right,shape,fill,workers)
305
306def 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,shape: Maybe<&2,List<&2,U32>>,
307  left: Storage.Buffer<Element>,+left_shape: List<&2,U32>,+fill: Element,right: Storage.Buffer<Element>,+right_shape: List<&2,U32>,+workers: U32) -> Tensor<Element>:
308  match shape:
309    case Some{+shape}:
310      combine_stretched_with(~Element,~kernel,name,Shapes.same_extents(shape,left_shape),left,left_shape,fill,right,right_shape,shape,workers)
311    case None{}: Tensor{left,left_shape,fill,Failed{Problem{name,ShapeMismatch{left_shape,right_shape}}},workers}
312
313# Equal extents take the direct path. Different extents follow NumPy
314# broadcasting (Shapes.broadcast) and fail only when they are incompatible.
315def combine_shapes_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,same: Bool,
316  left: Storage.Buffer<Element>,+left_shape: List<&2,U32>,+fill: Element,right: Storage.Buffer<Element>,+right_shape: List<&2,U32>,+workers: U32) -> Tensor<Element>:
317  match same:
318    case True{}: combine_views_with(~Element,~kernel,name,view_of(left_shape),view_of(right_shape),left,left_shape,fill,right,right_shape,workers)
319    case False{}: combine_broadcast_with(~Element,~kernel,name,Shapes.broadcast(left_shape,right_shape),left,left_shape,fill,right,right_shape,workers)
320
321def combine_status_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,status: Status,
322  left: Storage.Buffer<Element>,+left_shape: List<&2,U32>,+fill: Element,right: Storage.Buffer<Element>,+right_shape: List<&2,U32>,+workers: U32) -> Tensor<Element>:
323  match status:
324    case Usable{}: combine_shapes_with(~Element,~kernel,name,Shapes.same_extents(left_shape,right_shape),left,left_shape,fill,right,right_shape,workers)
325    case Failed{problem}: Tensor{left,left_shape,fill,Failed{problem},workers}
326
327# Band layout is kept only for equal extents: operands that need broadcasting
328# are gathered and combined on dense storage.
329def combine_dense_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,left: Tensor<Element>,right: Tensor<Element>) -> Tensor<Element>:
330  with_dense(~Element,~Tensor<Element>,left,left => left_shape => fill => left_status => left_workers =>
331    with_dense(~Element,~Tensor<Element>,right,right => right_shape => right_fill => right_status => right_workers =>
332      combine_status_with(~Element,~kernel,name,first_failure(left_status,right_status),left,left_shape,fill,right,right_shape,U32.max(left_workers,right_workers))))
333
334def with_status(~Element: Data,status: Status,tensor: Tensor<Element>) -> Tensor<Element>:
335  match tensor:
336    case Tensor{storage,shape,fill,previous,workers}: Tensor{storage,shape,fill,status,workers}
337    case Banded{bands,geometry,shape,fill,previous,workers}: Banded{bands,geometry,shape,fill,status,workers}
338
339def aligned(~Element: Data,geometry: Bands.Geometry,layout: Bands.Plan,tensor: Tensor<Element>) -> Bands.Bands(~Element):
340  match tensor:
341    case Tensor{storage,shape,fill,status,workers}: Bands.from_dense(~Element,geometry,fill,layout,storage)
342    case Banded{bands,previous,shape,fill,status,workers}: Bands.align(~Element,geometry,fill,layout,bands)
343
344def combine_bands_checked_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,status: Status,same: Bool,+geometry: Bands.Geometry,
345  +layout: Bands.Plan,+left_shape: List<&2,U32>,+right_shape: List<&2,U32>,fill: Element,+workers: U32,left: Tensor<Element>,right: Tensor<Element>) -> Tensor<Element>:
346  match status same:
347    case Usable{} True{}:
348      Banded{Bands.combine_with(~Element,~kernel,geometry,BandSchedule.parallel(geometry,workers,12.0,operations,layout),aligned(~Element,geometry,layout,left),aligned(~Element,geometry,layout,right)),
349        geometry,left_shape,fill,Usable{},workers}
350    case Usable{} False{}: combine_dense_with(~Element,~kernel,name,left,right)
351    case Failed{problem} same: with_status(~Element,Failed{problem},left)
352
353def combine_bands_plan_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,status: Status,same: Bool,+geometry: Bands.Geometry,
354  +left_shape: List<&2,U32>,+right_shape: List<&2,U32>,+fill: Element,+workers: U32,right: Tensor<Element>,planned: Bands.Bands(~Element) & Bands.Plan) -> Tensor<Element>:
355  (left,layout) = planned
356  combine_bands_checked_with(~Element,~kernel,name,operations,status,same,geometry,layout,left_shape,right_shape,fill,workers,
357    Banded{left,geometry,left_shape,fill,Usable{},workers},right)
358
359def combine_right_plan_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,status: Status,same: Bool,+geometry: Bands.Geometry,
360  +left_shape: List<&2,U32>,+right_shape: List<&2,U32>,left_fill: Element,+right_fill: Element,+workers: U32,left: Storage.Buffer<Element>,
361  planned: Bands.Bands(~Element) & Bands.Plan) -> Tensor<Element>:
362  (right,layout) = planned
363  +left_fill = left_fill
364  combine_bands_checked_with(~Element,~kernel,name,operations,status,same,geometry,layout,left_shape,right_shape,left_fill,workers,
365    Tensor{left,left_shape,left_fill,Usable{},workers},Banded{right,geometry,right_shape,right_fill,Usable{},workers})
366
367def combine_pair_plans_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,status: Status,same: Bool,+geometry: Bands.Geometry,right_geometry: Bands.Geometry,
368  +left_shape: List<&2,U32>,+right_shape: List<&2,U32>,+left_fill: Element,right_fill: Element,+workers: U32,
369  left: Bands.Bands(~Element) & Bands.Plan,right: Bands.Bands(~Element) & Bands.Plan) -> Tensor<Element>:
370  (left,left_plan) = left
371  (right,right_plan) = right
372  combine_bands_checked_with(~Element,~kernel,name,operations,status,same,geometry,BandSchedule.alignment(geometry,left_plan,right_plan),left_shape,right_shape,left_fill,workers,
373    Banded{left,geometry,left_shape,left_fill,Usable{},workers},Banded{right,right_geometry,right_shape,right_fill,Usable{},workers})
374
375# operation(left element, right element), retaining band layout when either
376# operand is banded. Alignment is done once, before any parallel leaf work.
377# operations counts the operation's scalar work for the band schedule.
378def 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>:
379  match left right:
380    case Tensor{left,left_shape,left_fill,left_status,left_workers} Tensor{right,right_shape,right_fill,right_status,right_workers}:
381      combine_status_with(~Element,~kernel,name,first_failure(left_status,right_status),left,left_shape,left_fill,right,right_shape,U32.max(left_workers,right_workers))
382    case Banded{left,geometry,+left_shape,left_fill,left_status,left_workers} Tensor{right,+right_shape,right_fill,+right_status,+right_workers}:
383      combine_bands_plan_with(~Element,~kernel,name,operations,first_failure(left_status,right_status),Shapes.same_extents(left_shape,right_shape),geometry,
384        left_shape,right_shape,left_fill,U32.max(left_workers,right_workers),Tensor{right,right_shape,right_fill,right_status,right_workers},Bands.plan(~Element,left))
385    case Banded{left,geometry,+left_shape,left_fill,left_status,left_workers} Banded{right,right_geometry,+right_shape,right_fill,+right_status,+right_workers}:
386      combine_pair_plans_with(~Element,~kernel,name,operations,first_failure(left_status,right_status),Shapes.same_extents(left_shape,right_shape),geometry,right_geometry,
387        left_shape,right_shape,left_fill,right_fill,U32.max(left_workers,right_workers),Bands.plan(~Element,left),Bands.plan(~Element,right))
388    case Tensor{left,+left_shape,left_fill,left_status,left_workers} Banded{right,geometry,+right_shape,right_fill,right_status,right_workers}:
389      combine_right_plan_with(~Element,~kernel,name,operations,first_failure(left_status,right_status),Shapes.same_extents(left_shape,right_shape),geometry,
390        left_shape,right_shape,left_fill,right_fill,U32.max(left_workers,right_workers),left,Bands.plan(~Element,right))
391
392# The elementwise form: the kernel applies operation to each pair in order.
393def combine(~Element: Data,~operation: Element -> Element -> Element,name: String,operations: F32,left: Tensor<Element>,right: Tensor<Element>) -> Tensor<Element>:
394  combine_with(~Element,~Operations.combine_kernel(~Element,~operation),name,operations,left,right)
395
396def applied(~Element: Data,name: String,+shape: List<&2,U32>,+fill: Element,+workers: U32,result: Operations.Outcome<Storage.Buffer<Element>>) -> Tensor<Element>:
397  match result:
398    case Operations.Accepted{storage}: Tensor{storage,shape,fill,Usable{},workers}
399    case Operations.Rejected{storage}: Tensor{storage,shape,fill,Failed{Problem{name,InvalidShape{shape}}},workers}
400
401def apply_view_context(~Element: Data,~Context: Data,~mapper: @+count: Nat -> @+position: U32 -> @+context: Context -> Storage.Buffer<Element> -> Storage.Buffer<Element>,context: Context,name: String,view: Maybe<&2,Views.View>,storage: Storage.Buffer<Element>,+shape: List<&2,U32>,+fill: Element,
402  +workers: U32) -> Tensor<Element>:
403  match view:
404    case Some{view}: applied(~Element,name,shape,fill,workers,Operations.apply_context(~Element,~Context,~mapper,context,storage,view))
405    case None{}: Tensor{storage,shape,fill,Failed{Problem{name,InvalidShape{shape}}},workers}
406
407def apply_status_context(~Element: Data,~Context: Data,~mapper: @+count: Nat -> @+position: U32 -> @+context: Context -> Storage.Buffer<Element> -> Storage.Buffer<Element>,context: Context,name: String,status: Status,storage: Storage.Buffer<Element>,+shape: List<&2,U32>,+fill: Element,
408  +workers: U32) -> Tensor<Element>:
409  match status:
410    case Usable{}: apply_view_context(~Element,~Context,~mapper,context,name,view_of(shape),storage,shape,fill,workers)
411    case Failed{problem}: Tensor{storage,shape,fill,Failed{problem},workers}
412
413def 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,
414  shape: List<&2,U32>,fill: Element,+workers: U32) -> Tensor<Element>:
415  match status:
416    case Usable{}: Banded{BandSchedule.map_context(~Element,~Context,~mapper,context,geometry,workers,operations,bands),geometry,shape,fill,Usable{},workers}
417    case Failed{problem}: Banded{bands,geometry,shape,fill,Failed{problem},workers}
418
419def apply_context(~Element: Data,~Context: Data,~mapper: @+count: Nat -> @+position: U32 -> @+context: Context -> Storage.Buffer<Element> -> Storage.Buffer<Element>,context: Context,name: String,operations: F32,tensor: Tensor<Element>) -> Tensor<Element>:
420  match tensor:
421    case Tensor{storage,shape,fill,status,workers}: apply_status_context(~Element,~Context,~mapper,context,name,status,storage,shape,fill,workers)
422    case Banded{bands,geometry,shape,fill,status,workers}: apply_bands_context(~Element,~Context,~mapper,context,operations,status,bands,geometry,shape,fill,workers)
423
424def apply(~Element: Data,~mapper: @+count: Nat -> @+position: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element>,name: String,operations: F32,tensor: Tensor<Element>) -> Tensor<Element>:
425  apply_context(~Element,~Unit,~Operations.without_context(~Element,~mapper),Unit{},name,operations,tensor)
426
427
428# The function also receives each element's row-major index.
429def indexed(~Element: Data,+shape: List<&2,U32>,+fill: Element,+workers: U32,result: Operations.Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>) -> Tensor<Element>:
430  match result:
431    case Operations.Accepted{(input,output)}: Tensor{output,shape,fill,Usable{},workers}
432    case Operations.Rejected{(input,output)}: Tensor{input,shape,fill,Failed{Problem{"apply_indexed",InvalidShape{shape}}},workers}
433
434def apply_indexed_count(~Element: Data,~operation: U32 -> Element -> Element,count: Maybe<&2,U32>,view: Maybe<&2,Views.View>,
435  storage: Storage.Buffer<Element>,+shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensor<Element>:
436  match count view:
437    case Some{count} Some{view}:
438      indexed(~Element,shape,fill,workers,Operations.apply_indexed(~Element,~operation,storage,view,buffer(~Element,fill,count)))
439    case count view: Tensor{storage,shape,fill,Failed{Problem{"apply_indexed",InvalidShape{shape}}},workers}
440
441def apply_indexed_status(~Element: Data,~operation: U32 -> Element -> Element,status: Status,storage: Storage.Buffer<Element>,
442  +shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensor<Element>:
443  match status:
444    case Usable{}: apply_indexed_count(~Element,~operation,count_of(shape),view_of(shape),storage,shape,fill,workers)
445    case Failed{problem}: Tensor{storage,shape,fill,Failed{problem},workers}
446
447def indexed_width_positive(geometry: Bands.Geometry) -> Bool:
448  Bands.Geometry{batch,channels,height,width} = geometry
449  Nat.is_lt(0n,U32.to_nat(width))
450
451# Empty widths use the existing dense indexed operation after one gather.
452# The choice is made once, before the band scheduler or element loop.
453def apply_indexed_positive(~Element: Data,~operation: U32 -> Element -> Element,positive: Bool,operations: F32,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,
454  +shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensor<Element>:
455  match positive:
456    case True{}: Banded{BandSchedule.indexed(~Element,~operation,geometry,workers,operations,fill,bands),geometry,shape,fill,Usable{},workers}
457    case False{}:
458      with_dense(~Element,~Tensor<Element>,Banded{bands,geometry,shape,fill,Usable{},workers},
459        storage => shape => fill => status => workers =>
460          apply_indexed_status(~Element,~operation,status,storage,shape,fill,workers))
461
462def apply_indexed_bands(~Element: Data,~operation: U32 -> Element -> Element,operations: F32,status: Status,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,
463  shape: List<&2,U32>,+fill: Element,+workers: U32) -> Tensor<Element>:
464  match status:
465    case Usable{}: apply_indexed_positive(~Element,~operation,indexed_width_positive(geometry),operations,bands,geometry,shape,fill,workers)
466    case Failed{problem}: Banded{bands,geometry,shape,fill,Failed{problem},workers}
467
468def apply_indexed(~Element: Data,~operation: U32 -> Element -> Element,operations: F32,tensor: Tensor<Element>) -> Tensor<Element>:
469  match tensor:
470    case Tensor{storage,shape,fill,status,workers}: apply_indexed_status(~Element,~operation,status,storage,shape,fill,workers)
471    case Banded{bands,geometry,shape,fill,status,workers}: apply_indexed_bands(~Element,~operation,operations,status,bands,geometry,shape,fill,workers)
472
473# Reductions visit elements in row-major order.
474def reduced(~Element: Data,~Accumulator: Data,name: String,+shape: List<&2,U32>,result: Operations.Outcome<(Storage.Buffer<Element> & Accumulator)>) ->
475  Result<&2,&2,Problem,Accumulator>:
476  match result:
477    case Operations.Accepted{(storage,value)}: Done{value}
478    case Operations.Rejected{(storage,value)}: Fail{Problem{name,InvalidShape{shape}}}
479
480def reduce_view(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,name: String,view: Maybe<&2,Views.View>,
481  storage: Storage.Buffer<Element>,+shape: List<&2,U32>,initial: Accumulator) -> Result<&2,&2,Problem,Accumulator>:
482  match view:
483    case Some{view}: reduced(~Element,~Accumulator,name,shape,Operations.reduce(~Element,~Accumulator,~operation,storage,view,initial))
484    case None{}: Fail{Problem{name,InvalidShape{shape}}}
485
486def reduce_status(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,name: String,status: Status,
487  storage: Storage.Buffer<Element>,+shape: List<&2,U32>,initial: Accumulator) -> Result<&2,&2,Problem,Accumulator>:
488  match status:
489    case Usable{}: reduce_view(~Element,~Accumulator,~operation,name,view_of(shape),storage,shape,initial)
490    case Failed{problem}: Fail{problem}
491
492def reduce(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,name: String,tensor: Tensor<Element>,initial: Accumulator) ->
493  Result<&2,&2,Problem,Accumulator>:
494  with_dense(~Element,~Result<&2,&2,Problem,Accumulator>,tensor,storage => shape => fill => status => workers =>
495    reduce_status(~Element,~Accumulator,~operation,name,status,storage,shape,initial))
496
497# A reduction as a rank-0 tensor, so a failure propagates like any result.
498def scalar_result(~Element: Data,+fill: Element,+count: U32,result: Result<&2,&2,Problem,Element>) -> Tensor<Element>:
499  match result:
500    case Done{value}: with_workers(~Element,full(~Element,[],value),count)
501    case Fail{problem}: Tensor{buffer(~Element,fill,0),[],fill,Failed{problem},count}
502
503def reduce_counted(~Element: Data,~operation: Element -> Element -> Element,name: String,pair: Tensor<Element> & U32,+initial: Element) -> Tensor<Element>:
504  (tensor,+count) = pair
505  scalar_result(~Element,initial,count,reduce(~Element,~Element,~operation,name,tensor,initial))
506
507def reduce_all(~Element: Data,~operation: Element -> Element -> Element,name: String,tensor: Tensor<Element>,+initial: Element) -> Tensor<Element>:
508  reduce_counted(~Element,~operation,name,workers(~Element,tensor),initial)
509
510# Scalar operands are read through a stride-zero view of one element.
511def broadcast_axes(axes: List<&2,Views.Axis>) -> List<&2,Views.Axis>:
512  match axes:
513    case []: []
514    case Views.Axis{extent,stride} <> rest: Views.Axis{extent,0} <> broadcast_axes(rest)
515
516def broadcast(view: Views.View) -> Views.View:
517  Views.View{axes,offset} = view
518  Views.View{broadcast_axes(axes),0}
519
520def scalar_applied(~Element: Data,name: String,+shape: List<&2,U32>,+fill: Element,+workers: U32,
521  result: Operations.Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>) -> Tensor<Element>:
522  match result:
523    case Operations.Accepted{(storage,scalar)}: Tensor{storage,shape,fill,Usable{},workers}
524    case Operations.Rejected{(storage,scalar)}: Tensor{storage,shape,fill,Failed{Problem{name,InvalidShape{shape}}},workers}
525
526def scalar_view(~Element: Data,~operation: Element -> Element -> Element,name: String,view: Maybe<&2,Views.View>,
527  storage: Storage.Buffer<Element>,+shape: List<&2,U32>,+fill: Element,+value: Element,+workers: U32) -> Tensor<Element>:
528  match view:
529    case Some{+view}:
530      scalar_applied(~Element,name,shape,fill,workers,Operations.combine(~Element,~operation,storage,view,buffer(~Element,value,1),broadcast(view)))
531    case None{}: Tensor{storage,shape,fill,Failed{Problem{name,InvalidShape{shape}}},workers}
532
533def scalar_status(~Element: Data,~operation: Element -> Element -> Element,name: String,status: Status,
534  storage: Storage.Buffer<Element>,+shape: List<&2,U32>,+fill: Element,+value: Element,+workers: U32) -> Tensor<Element>:
535  match status:
536    case Usable{}: scalar_view(~Element,~operation,name,view_of(shape),storage,shape,fill,value,workers)
537    case Failed{problem}: Tensor{storage,shape,fill,Failed{problem},workers}
538
539# Updates every element x to operation(x, value).
540def scalar_bands(~Element: Data,~operation: Element -> Element -> Element,operations: F32,status: Status,bands: Bands.Bands(~Element),+geometry: Bands.Geometry,
541  shape: List<&2,U32>,fill: Element,+workers: U32,value: Element) -> Tensor<Element>:
542  match status:
543    case Usable{}: Banded{BandSchedule.scalar(~Element,~operation,geometry,workers,operations,value,bands),geometry,shape,fill,Usable{},workers}
544    case Failed{problem}: Banded{bands,geometry,shape,fill,Failed{problem},workers}
545
546def combine_scalar(~Element: Data,~operation: Element -> Element -> Element,name: String,operations: F32,tensor: Tensor<Element>,+value: Element) -> Tensor<Element>:
547  match tensor:
548    case Tensor{storage,shape,fill,status,workers}: scalar_status(~Element,~operation,name,status,storage,shape,fill,value,workers)
549    case Banded{bands,geometry,shape,fill,status,workers}:
550      scalar_bands(~Element,~operation,operations,status,bands,geometry,shape,fill,workers,value)
551
552# Messages
553
554def show_extents(extents: List<&2,U32>) -> String:
555  match extents:
556    case []: ""
557    case [extent]: U32.show(extent)
558    case extent <> rest: String.append(U32.show(extent),String.append(", ",show_extents(rest)))
559
560def show_shape(extents: List<&2,U32>) -> String:
561  String.append("[",String.append(show_extents(extents),"]"))
562
563def explain(reason: Reason) -> String:
564  match reason:
565    case ShapeMismatch{left,right}:
566      String.append("shapes ",String.append(show_shape(left),String.append(" and ",String.append(show_shape(right)," do not match"))))
567    case InvalidShape{shape}:
568      String.append("shape ",String.append(show_shape(shape)," is invalid or exceeds the machine capacity"))
569    case InvalidArgument{shape}:
570      String.append("invalid argument for shape ",show_shape(shape))
571    case FileFault{path,detail}:
572      String.append("file ",String.append(path,String.append(" ",detail)))
573    case MissingName{name}:
574      String.append("no tensor named ",name)
575
576def describe(problem: Problem) -> String:
577  Problem{operation,reason} = problem
578  String.append(operation,String.append(": ",explain(reason)))
579
580def expected(~Element: Data,result: Result<&1,&1,Tensor<Element> & Problem,Tensor<Element>>) -> IO(Tensor<Element>):
581  match result:
582    case Done{tensor}: IO.pure(Tensor<Element>,tensor)
583    case Fail{(tensor,problem)}: IO.die(Tensor<Element>,1,describe(problem))
584
585# Stops the program with the recorded problem, for scripts and examples.
586def expect(~Element: Data,tensor: Tensor<Element>) -> IO(Tensor<Element>):
587  expected(~Element,check(~Element,tensor))
588
589def expected_storage(~Element: Data,result: Result<&2,&1,Problem,Storage.Buffer<Element>>) -> IO(Storage.Buffer<Element>):
590  match result:
591    case Done{storage}: IO.pure(Storage.Buffer<Element>,storage)
592    case Fail{problem}: IO.die(Storage.Buffer<Element>,1,describe(problem))
593
594def expect_storage(~Element: Data,tensor: Tensor<Element>) -> IO(Storage.Buffer<Element>):
595  expected_storage(~Element,into_storage(~Element,tensor))
596
597# Read-only operations return their tensor in a pair; Bend destructures only
598# parameters, so a continuation receives both halves.
599def then(~First: Type,~Second: Type,~Result: Type,pair: First & Second,body: First -> Second -> Result) -> Result:
600  (first,second) = pair
601  body(first,second)