tensor_operations.bend

stelliferous@0.0.2.0 · 286 lines · 36 definitions

1# Checked operations preserve all owners and their balanced-storage certificates.
2# Element functions enter hot loops only as direct templates.
3import Base
4import stelliferous@0.0.2.0/storage_buffer.bend as Storage
5import stelliferous@0.0.2.0/tensor_shape.bend as Shapes
6import stelliferous@0.0.2.0/tensor_view.bend as Views
7import stelliferous@0.0.2.0/traversal_array.bend as Traversal
8import stelliferous@0.0.2.0/tensor_access.bend as Access
9import stelliferous@0.0.2.0/proofs/storage_certificate.bend as Certificate
10import stelliferous@0.0.2.0/proofs/traversal_certificate.bend as TraversalCertificate
11import stelliferous@0.0.2.0/proofs/traversal_witness.bend as TraversalWitness
12import stelliferous@0.0.2.0/proofs/tensor_access_certificate.bend as AccessCertificate
13import stelliferous@0.0.2.0/traversal_lanes.bend as Lanes
14import stelliferous@0.0.2.0/proofs/traversal_lanes_certificate.bend as LanesCertificate
15
16type Outcome<-Value: Type> is Type:
17  Accepted{value: Value}
18  Rejected{value: Value}
19
20def valid_buffer(+length: Nat,capacity: U32,view: Views.View) -> Bool:
21  Nat.is_le(length,U32.to_nat(capacity)) && Views.fits(view,U32.from_nat(length))
22
23def pair_buffers(~Element: Data,input_length: Nat,output_length: Nat,result: Array<Element> & Array<Element>,
24  certificates: TraversalCertificate.Certificates<Array<Element>,Element,Certificate.family(~Element),result>) ->
25  Storage.Buffer<Element> & Storage.Buffer<Element>:
26  (input,output) = result
27  TraversalCertificate.Certificates{input_certificate,output_certificate} = certificates
28  (Storage.Buffer{input,input_length,input_certificate},Storage.Buffer{output,output_length,output_certificate})
29
30def copy_when(~Element: Data,valid: Bool,count: Maybe<&2,U32>,+source: Views.View,target: Views.View,
31  input: Storage.Buffer<Element>,output: Storage.Buffer<Element>) -> Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>:
32  match valid count:
33    case True{} Some{+count}:
34      Views.View{axes,+offset} = target
35      Storage.Buffer{input,input_length,input_certificate} = input
36      Storage.Buffer{output,output_length,output_certificate} = output
37      +contiguous = Views.is_contiguous(source)
38      Accepted{pair_buffers(~Element,input_length,output_length,
39        Access.copy_run(~Element,contiguous,U32.to_nat(count),source,offset,input,output),
40        AccessCertificate.copy(~Element,contiguous,U32.to_nat(count),source,offset,input,output,input_certificate,output_certificate))}
41    case valid count: Rejected{(input,output)}
42
43def copy(~Element: Data,input: Storage.Buffer<Element>,+source: Views.View,
44  output: Storage.Buffer<Element>,+target: Views.View) -> Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>:
45  Storage.Buffer{input,+input_length,input_certificate} = input
46  Certificate.Certificate{+input_depth,input_equation} = input_certificate
47  Storage.Buffer{output,+output_length,output_certificate} = output
48  Certificate.Certificate{+output_depth,output_equation} = output_certificate
49  copy_when(~Element,Shapes.same(Views.shape(source),Views.shape(target)) && Views.is_contiguous(target) &&
50    valid_buffer(input_length,U32.shln(1,input_depth),source) &&
51    valid_buffer(output_length,U32.shln(1,output_depth),target),Views.numel(target),source,target,
52    Storage.Buffer{input,input_length,Certificate.Certificate{input_depth,input_equation}},Storage.Buffer{output,output_length,Certificate.Certificate{output_depth,output_equation}})
53
54# output[address of element i of target] = input[offset of source + i]: a
55# contiguous source written through a strided target view, in row-major order.
56def scatter_when(~Element: Data,valid: Bool,count: Maybe<&2,U32>,source: Views.View,+target: Views.View,
57  input: Storage.Buffer<Element>,output: Storage.Buffer<Element>) -> Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>:
58  match valid count:
59    case True{} Some{+count}:
60      Views.View{axes,+offset} = source
61      Storage.Buffer{input,input_length,input_certificate} = input
62      Storage.Buffer{output,output_length,output_certificate} = output
63      Accepted{pair_buffers(~Element,input_length,output_length,
64        Access.scatter_run(~Element,U32.to_nat(count),target,offset,input,output),
65        AccessCertificate.scatter(~Element,U32.to_nat(count),target,offset,input,output,input_certificate,output_certificate))}
66    case valid count: Rejected{(input,output)}
67
68def scatter(~Element: Data,input: Storage.Buffer<Element>,+source: Views.View,
69  output: Storage.Buffer<Element>,+target: Views.View) -> Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>:
70  Storage.Buffer{input,+input_length,input_certificate} = input
71  Certificate.Certificate{+input_depth,input_equation} = input_certificate
72  Storage.Buffer{output,+output_length,output_certificate} = output
73  Certificate.Certificate{+output_depth,output_equation} = output_certificate
74  scatter_when(~Element,Shapes.same(Views.shape(source),Views.shape(target)) && Views.is_contiguous(source) &&
75    valid_buffer(input_length,U32.shln(1,input_depth),source) &&
76    valid_buffer(output_length,U32.shln(1,output_depth),target),Views.numel(target),source,target,
77    Storage.Buffer{input,input_length,Certificate.Certificate{input_depth,input_equation}},Storage.Buffer{output,output_length,Certificate.Certificate{output_depth,output_equation}})
78
79def combine_buffers(~Element: Data,left_length: Nat,right_length: Nat,result: Array<Element> & Array<Element>,
80  certificates: TraversalCertificate.Certificates<Array<Element>,Element,Certificate.family(~Element),result>) -> Storage.Buffer<Element> & Storage.Buffer<Element>:
81  (right,left) = result
82  TraversalCertificate.Certificates{right_certificate,left_certificate} = certificates
83  (Storage.Buffer{left,left_length,left_certificate},Storage.Buffer{right,right_length,right_certificate})
84
85def combine_kernel(~Element: Data,~operation: Element -> Element -> Element,+count: Nat,+right_view: Views.View,+offset: U32,
86  left: Storage.Buffer<Element>,right: Storage.Buffer<Element>) -> Storage.Buffer<Element> & Storage.Buffer<Element>:
87  Storage.Buffer{left,left_length,left_certificate} = left
88  Storage.Buffer{right,right_length,right_certificate} = right
89  +fast = Access.zero_contiguous(right_view) && U32.is_zero(offset)
90  combine_buffers(~Element,left_length,right_length,
91    Access.combine_run(~Element,~operation,fast,count,right_view,offset,right,left),
92    AccessCertificate.combine(~Element,~operation,fast,count,right_view,offset,right,left,right_certificate,left_certificate))
93
94def combine_when_with(~Element: Data,~kernel: @+count: Nat -> @+right_view: Views.View -> @+offset: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element> -> Storage.Buffer<Element> & Storage.Buffer<Element>,
95  valid: Bool,count: Maybe<&2,U32>,right_view: Views.View,target: Views.View,left: Storage.Buffer<Element>,right: Storage.Buffer<Element>) ->
96  Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>:
97  match valid count:
98    case True{} Some{count}:
99      Views.View{axes,offset} = target
100      Accepted{kernel(U32.to_nat(count),right_view,offset,left,right)}
101    case valid count: Rejected{(left,right)}
102
103def 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>,
104  left: Storage.Buffer<Element>,+left_view: Views.View,right: Storage.Buffer<Element>,+right_view: Views.View) ->
105  Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>:
106  Storage.Buffer{left,+left_length,left_certificate} = left
107  Certificate.Certificate{+left_depth,left_equation} = left_certificate
108  Storage.Buffer{right,+right_length,right_certificate} = right
109  Certificate.Certificate{+right_depth,right_equation} = right_certificate
110  combine_when_with(~Element,~kernel,Shapes.same(Views.shape(left_view),Views.shape(right_view)) && Views.is_contiguous(left_view) &&
111    valid_buffer(left_length,U32.shln(1,left_depth),left_view) &&
112    valid_buffer(right_length,U32.shln(1,right_depth),right_view),Views.numel(left_view),right_view,left_view,
113    Storage.Buffer{left,left_length,Certificate.Certificate{left_depth,left_equation}},Storage.Buffer{right,right_length,Certificate.Certificate{right_depth,right_equation}})
114
115def combine(~Element: Data,~operation: Element -> Element -> Element,
116  left: Storage.Buffer<Element>,+left_view: Views.View,right: Storage.Buffer<Element>,+right_view: Views.View) ->
117  Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>:
118  combine_with(~Element,~combine_kernel(~Element,~operation),left,left_view,right,right_view)
119
120# A mapper is the ordered map of count values from position on a certified
121# Buffer. mapped applies the operation to one value at a time.
122def mapped(~Element: Data,~operation: Element -> Element,+count: Nat,+position: U32,buffer: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
123  Storage.Buffer{storage,length,certificate} = buffer
124  Storage.Buffer{Traversal.map_in_place(~Element,~Unit,~Access.apply_element(~Element,~operation),count,position,Unit{},storage),length,
125    TraversalCertificate.map_in_place(~Element,~Unit,~Access.apply_element(~Element,~operation),count,position,Unit{},storage,certificate)}
126
127# The same map with the values of each round of eight computed together: an
128# operation written with FP32 arithmetic alone then compiles to vector code.
129# The certificate's depth decides whether the array holds the positions.
130def lanes(~Context: Data,~operation: Context -> F32 -> F32,+count: Nat,+position: U32,+context: Context,buffer: Storage.Buffer<F32>) -> Storage.Buffer<F32>:
131  Storage.Buffer{storage,length,Certificate.Certificate{+depth,equation}} = buffer
132  Storage.Buffer{Lanes.map_when(~Context,~operation,Lanes.covered(depth,position,count),count,position,context,storage),length,
133    LanesCertificate.map_when(~Context,~operation,Lanes.covered(depth,position,count),count,position,context,storage,Certificate.Certificate{depth,equation})}
134
135def mapped_lanes(~operation: F32 -> F32,+count: Nat,+position: U32,buffer: Storage.Buffer<F32>) -> Storage.Buffer<F32>:
136  lanes(~Unit,~Access.apply_element(~F32,~operation),count,position,Unit{},buffer)
137
138def apply_context_when(~Element: Data,~Context: Data,~mapper: @+count: Nat -> @+position: U32 -> @+context: Context -> Storage.Buffer<Element> -> Storage.Buffer<Element>,context: Context,valid: Bool,count: Maybe<&2,U32>,view: Views.View,
139  buffer: Storage.Buffer<Element>) -> Outcome<Storage.Buffer<Element>>:
140  match valid count:
141    case True{} Some{+count}:
142      Views.View{axes,+offset} = view
143      Accepted{mapper(U32.to_nat(count),offset,context,buffer)}
144    case valid count: Rejected{buffer}
145
146def apply_context(~Element: Data,~Context: Data,~mapper: @+count: Nat -> @+position: U32 -> @+context: Context -> Storage.Buffer<Element> -> Storage.Buffer<Element>,context: Context,buffer: Storage.Buffer<Element>,+view: Views.View) -> Outcome<Storage.Buffer<Element>>:
147  Storage.Buffer{storage,+length,certificate} = buffer
148  Certificate.Certificate{+depth,equation} = certificate
149  apply_context_when(~Element,~Context,~mapper,context,Views.is_contiguous(view) && valid_buffer(length,U32.shln(1,depth),view),Views.numel(view),view,
150    Storage.Buffer{storage,length,Certificate.Certificate{depth,equation}})
151
152# A mapper without a context, as the Unit-context case.
153def without_context(~Element: Data,~mapper: @+count: Nat -> @+position: U32 -> Storage.Buffer<Element> -> Storage.Buffer<Element>,+count: Nat,+position: U32,+context: Unit,buffer: Storage.Buffer<Element>) -> Storage.Buffer<Element>:
154  mapper(count,position,buffer)
155
156def apply_indexed_when(~Element: Data,~operation: U32 -> Element -> Element,valid: Bool,count: Maybe<&2,U32>,+view: Views.View,
157  input: Storage.Buffer<Element>,output: Storage.Buffer<Element>) -> Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>:
158  match valid count:
159    case True{} Some{+count}:
160      Storage.Buffer{input,input_length,input_certificate} = input
161      Storage.Buffer{output,output_length,output_certificate} = output
162      Accepted{pair_buffers(~Element,input_length,output_length,
163        Access.apply_indexed_run(~Element,~operation,U32.to_nat(count),view,input,output),
164        AccessCertificate.apply_indexed(~Element,~operation,U32.to_nat(count),view,input,output,input_certificate,output_certificate))}
165    case valid count: Rejected{(input,output)}
166
167def holds(count: Maybe<&2,U32>,+length: Nat,capacity: U32) -> Bool:
168  match count:
169    case None{}: False{}
170    case Some{count}: Nat.is_le(U32.to_nat(count),length) && Nat.is_le(length,U32.to_nat(capacity))
171
172# Writes f(index, element) for a contiguous input view into output[0..count).
173def apply_indexed(~Element: Data,~operation: U32 -> Element -> Element,input: Storage.Buffer<Element>,+view: Views.View,
174  output: Storage.Buffer<Element>) -> Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>:
175  Storage.Buffer{input,+input_length,input_certificate} = input
176  Certificate.Certificate{+input_depth,input_equation} = input_certificate
177  Storage.Buffer{output,+output_length,output_certificate} = output
178  Certificate.Certificate{+output_depth,output_equation} = output_certificate
179  apply_indexed_when(~Element,~operation,Views.is_contiguous(view) && valid_buffer(input_length,U32.shln(1,input_depth),view) &&
180    holds(Views.numel(view),output_length,U32.shln(1,output_depth)),Views.numel(view),view,
181    Storage.Buffer{input,input_length,Certificate.Certificate{input_depth,input_equation}},Storage.Buffer{output,output_length,Certificate.Certificate{output_depth,output_equation}})
182
183def small(value: U32) -> Bool:
184  U32.is_le(value,65535)
185
186# The window geometry is exact and every sample address stays in U32.
187def pool_geometry(+pool: Access.Pool) -> Bool:
188  Access.Pool{+height,+width,+kernel,+stride,+padding,+out_height,+out_width} = pool
189  small(height) && small(width) && small(kernel) && small(stride) && small(padding) && U32.is_lt(0,kernel) && U32.is_lt(0,stride) &&
190    U32.is_le(kernel,(height + 2 * padding : U32)) && U32.is_le(kernel,(width + 2 * padding : U32)) &&
191    U32.is_eq(out_height,((height + 2 * padding - kernel : U32) / stride + 1 : U32)) &&
192    U32.is_eq(out_width,((width + 2 * padding - kernel : U32) / stride + 1 : U32))
193
194# The unchecked sweep of two owners; the callers establish its bounds.
195def swept(~Element: Data,~operation: Element -> Element -> Element,+count: Nat,+context: Access.SweepContext<Element>,
196  input: Storage.Buffer<Element>,output: Storage.Buffer<Element>) -> Storage.Buffer<Element> & Storage.Buffer<Element>:
197  Storage.Buffer{input,input_length,input_certificate} = input
198  Storage.Buffer{output,output_length,output_certificate} = output
199  pair_buffers(~Element,input_length,output_length,
200    Access.sweep_run(~Element,~operation,count,context,input,output),
201    AccessCertificate.sweep(~Element,~operation,count,context,input,output,input_certificate,output_certificate))
202
203def sweep_when(~Element: Data,~operation: Element -> Element -> Element,valid: Bool,count: Maybe<&2,U32>,+context: Access.SweepContext<Element>,
204  input: Storage.Buffer<Element>,output: Storage.Buffer<Element>) -> Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>:
205  match valid count:
206    case True{} Some{count}: Accepted{swept(~Element,~operation,U32.to_nat(count),context,input,output)}
207    case valid count: Rejected{(input,output)}
208
209# The window geometry of a whole axis is exact and every sample address stays in U32.
210def sweep_geometry(+sweep: Access.Sweep) -> Bool:
211  Access.Sweep{+extent,+step,+kernel,+stride,+padding,+out_extent,+origin} = sweep
212  small(extent) && small(kernel) && small(stride) && small(padding) && U32.is_lt(0,kernel) && U32.is_lt(0,stride) && U32.is_lt(0,step) &&
213    U32.is_eq(origin,0) && U32.is_le(kernel,(extent + 2 * padding : U32)) &&
214    U32.is_eq(out_extent,((extent + 2 * padding - kernel : U32) / stride + 1 : U32))
215
216def sweep_input(+sweep: Access.Sweep,+outer: U32) -> Maybe<&2,U32>:
217  Access.Sweep{+extent,+step,kernel,stride,padding,out_extent,origin} = sweep
218  Shapes.product([outer,extent,step],Some{1})
219
220def sweep_output(+sweep: Access.Sweep,+outer: U32) -> Maybe<&2,U32>:
221  Access.Sweep{extent,+step,kernel,stride,padding,+out_extent,origin} = sweep
222  Shapes.product([outer,out_extent,step],Some{1})
223
224# Windows along one axis of outer contiguous blocks; see Access.Sweep.
225def sweep(~Element: Data,~operation: Element -> Element -> Element,+pad: Element,input: Storage.Buffer<Element>,+sweep: Access.Sweep,+outer: U32,
226  output: Storage.Buffer<Element>) -> Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>:
227  Storage.Buffer{input,+input_length,input_certificate} = input
228  Certificate.Certificate{+input_depth,input_equation} = input_certificate
229  Storage.Buffer{output,+output_length,output_certificate} = output
230  Certificate.Certificate{+output_depth,output_equation} = output_certificate
231  sweep_when(~Element,~operation,sweep_geometry(sweep) && holds(sweep_input(sweep,outer),input_length,U32.shln(1,input_depth)) &&
232    holds(sweep_output(sweep,outer),output_length,U32.shln(1,output_depth)),sweep_output(sweep,outer),Access.SweepContext{sweep,pad},
233    Storage.Buffer{input,input_length,Certificate.Certificate{input_depth,input_equation}},Storage.Buffer{output,output_length,Certificate.Certificate{output_depth,output_equation}})
234
235# The input and the output of a pool around its column-swept planes.
236def pool_finished(~Element: Data,input: Storage.Buffer<Element>,result: Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>) ->
237  Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>:
238  match result:
239    case Accepted{(middle,output)}: Accepted{(input,output)}
240    case Rejected{(middle,output)}: Rejected{(input,output)}
241
242def pool_rows(~Element: Data,~operation: Element -> Element -> Element,+pad: Element,result: Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>,
243  +rows: Access.Sweep,+channels: U32,output: Storage.Buffer<Element>) -> Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>:
244  match result:
245    case Accepted{(input,middle)}: pool_finished(~Element,input,sweep(~Element,~operation,pad,middle,rows,channels,output))
246    case Rejected{(input,middle)}: Rejected{(input,output)}
247
248def pool_lines(~Element: Data,~operation: Element -> Element -> Element,+pad: Element,lines: Maybe<&2,U32>,input: Storage.Buffer<Element>,+pool: Access.Pool,
249  +channels: U32,middle: Storage.Buffer<Element>,output: Storage.Buffer<Element>) -> Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>:
250  match lines:
251    case None{}: Rejected{(input,output)}
252    case Some{+lines}:
253      pool_rows(~Element,~operation,pad,sweep(~Element,~operation,pad,input,Access.pool_columns(pool),lines,middle),Access.pool_rows(pool,0),channels,output)
254
255# Windows over the last two axes of channels contiguous planes: a sweep along
256# the columns of every row into middle, [channels,height,out_width], then a
257# sweep along the rows into output. Each sweep checks its own geometry and
258# counts.
259def pool(~Element: Data,~operation: Element -> Element -> Element,+pad: Element,input: Storage.Buffer<Element>,+pool: Access.Pool,+channels: U32,
260  middle: Storage.Buffer<Element>,output: Storage.Buffer<Element>) -> Outcome<(Storage.Buffer<Element> & Storage.Buffer<Element>)>:
261  Access.Pool{+height,+width,+kernel,+stride,+padding,+out_height,+out_width} = pool
262  pool_lines(~Element,~operation,pad,Shapes.product([channels,height],Some{1}),input,Access.Pool{height,width,kernel,stride,padding,out_height,out_width},channels,middle,output)
263
264def reduce_buffer(~Element: Data,~Accumulator: Data,length: Nat,result: Array<Element> & Accumulator,
265  certificate: Certificate.Certificate<Element,TraversalWitness.first(~Array<Element>,~Accumulator,result)>) -> Storage.Buffer<Element> & Accumulator:
266  (storage,value) = result
267  (Storage.Buffer{storage,length,certificate},value)
268
269def reduce_when(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,
270  valid: Bool,count: Maybe<&2,U32>,+view: Views.View,buffer: Storage.Buffer<Element>,initial: Accumulator) ->
271  Outcome<(Storage.Buffer<Element> & Accumulator)>:
272  match valid count:
273    case True{} Some{+count}:
274      Storage.Buffer{storage,length,certificate} = buffer
275      +fast = Views.is_contiguous(view)
276      Accepted{reduce_buffer(~Element,~Accumulator,length,
277        Access.reduce_run(~Element,~Accumulator,~operation,fast,U32.to_nat(count),view,storage,initial),
278        AccessCertificate.reduce(~Element,~Accumulator,~operation,fast,U32.to_nat(count),view,storage,initial,certificate))}
279    case valid count: Rejected{(buffer,initial)}
280
281def reduce(~Element: Data,~Accumulator: Data,~operation: Accumulator -> Element -> Accumulator,
282  buffer: Storage.Buffer<Element>,+view: Views.View,initial: Accumulator) -> Outcome<(Storage.Buffer<Element> & Accumulator)>:
283  Storage.Buffer{storage,+length,certificate} = buffer
284  Certificate.Certificate{+depth,equation} = certificate
285  reduce_when(~Element,~Accumulator,~operation,valid_buffer(length,U32.shln(1,depth),view),Views.numel(view),view,
286    Storage.Buffer{storage,length,Certificate.Certificate{depth,equation}},initial)