tensor_f32.bend

stelliferous@0.0.2.0 · 405 lines · 102 definitions

1# The everyday F32 tensor module: import this one module and call F.add,
2# F.concat, F.check, ... without element-type arguments. Every function is a
3# direct specialization of the generic tensor modules; Base determines operand
4# order, rounding, NaN and signed-zero behavior.
5import Base
6import stelliferous@0.0.2.0/storage_buffer.bend as Buffers
7import stelliferous@0.0.2.0/math_activation.bend as Activations
8import stelliferous@0.0.2.0/math_fp32.bend as Arithmetic
9import stelliferous@0.0.2.0/math_compare.bend as Compare
10import stelliferous@0.0.2.0/tensor.bend as Tensors
11import stelliferous@0.0.2.0/tensor_operations.bend as Operations
12import stelliferous@0.0.2.0/tensor_lanes.bend as Lanes
13import stelliferous@0.0.2.0/tensor_select.bend as Select
14import stelliferous@0.0.2.0/tensor_layout.bend as Layout
15import stelliferous@0.0.2.0/tensor_window.bend as Window
16import stelliferous@0.0.2.0/tensor_convolution.bend as TensorConvolution
17import stelliferous@0.0.2.0/tensor_file.bend as Files
18import stelliferous@0.0.2.0/tensor_model.bend as Models
19
20def Tensor() -> Type:
21  Tensors.Tensor<F32>
22
23def Problem() -> Data:
24  Tensors.Problem
25
26# Storage exchanged with foreign I/O; see from_storage and into_storage.
27def Storage() -> Type:
28  Buffers.Buffer<F32>
29
30# Construction
31
32def full(+shape: List<&2,U32>,+value: F32) -> Tensor():
33  Tensors.full(~F32,shape,value)
34
35def zeros(+shape: List<&2,U32>) -> Tensor():
36  Tensors.full(~F32,shape,0.0)
37
38def ones(+shape: List<&2,U32>) -> Tensor():
39  full(shape,1.0)
40
41def from_list(+shape: List<&2,U32>,+values: List<&2,F32>) -> Tensor():
42  Tensors.from_list(~F32,shape,values,0.0)
43
44def from_storage(+shape: List<&2,U32>,storage: Storage()) -> Tensor():
45  Tensors.from_storage(~F32,shape,storage,0.0)
46
47# Raw little-endian FP32 files: the values of the shape in row-major order.
48# A file that cannot be read, or does not match the shape, gives a failed tensor.
49def load_raw(+path: String,+shape: List<&2,U32>) -> IO(Tensor()):
50  Files.load(path,shape)
51
52def save_raw(+path: String,tensor: Tensor()) -> IO(Tensor()):
53  Files.save(path,tensor)
54
55# NumPy .npy files of FP32 values; the file carries the shape.
56def load_npy(+path: String) -> IO(Tensor()):
57  Files.load_npy(path)
58
59def save_npy(+path: String,tensor: Tensor()) -> IO(Tensor()):
60  Files.save_npy(path,tensor)
61
62# Models: tensors by name, loaded once. In `do Model<T>:` a step takes a tensor
63# from the weights with take, or puts a copy there with keep for a later take.
64def Weights() -> Type:
65  Models.Store
66
67def Model(A: Type) -> Type:
68  Models.Model(A)
69
70def Model.pure(-A: Type,value: A) -> Model(A):
71  Models.Model.pure(A,value)
72
73def Model.bind(-A: Type,-B: Type,step: Model(A),next: A -> Model(B)) -> Model(B):
74  Models.Model.bind(A,B,step,next)
75
76def take(+name: String) -> Model(Tensor()):
77  Models.take(name)
78
79def keep(+name: String,tensor: Tensor()) -> Model(Tensor()):
80  Models.keep(name,tensor)
81
82# folder/names.txt lists one name per line; folder/<name>.npy holds its tensor.
83def load_weights(+folder: String) -> IO(Weights()):
84  Models.load(folder)
85
86def copy_weights(weights: Weights()) -> Weights() & Weights():
87  Models.copy(weights)
88
89def run(-A: Type,model: Model(A),weights: Weights()) -> A:
90  Models.result(A,model,weights)
91
92# The result with what remains of the weights, kept tensors included.
93def run_keeping(-A: Type,model: Model(A),weights: Weights()) -> A & Weights():
94  Models.run(A,model,weights)
95
96def save_weights(+folder: String,weights: Weights()) -> IO(Weights()):
97  Models.save(folder,weights)
98
99# Execution context: operations may use count workers, like a tensor's device.
100# Read the runtime's pool once with IO.thread_count(); results keep the count.
101def with_workers(tensor: Tensor(),+count: U32) -> Tensor():
102  Tensors.with_workers(~F32,tensor,count)
103
104def workers(tensor: Tensor()) -> Tensor() & U32:
105  Tensors.workers(~F32,tensor)
106
107# Observation and exits
108
109def shape(tensor: Tensor()) -> Tensor() & List<&2,U32>:
110  Tensors.shape(~F32,tensor)
111
112def check(tensor: Tensor()) -> Result<&1,&1,Tensor() & Problem(),Tensor()>:
113  Tensors.check(~F32,tensor)
114
115def describe(problem: Problem()) -> String:
116  Tensors.describe(problem)
117
118def expect(tensor: Tensor()) -> IO(Tensor()):
119  Tensors.expect(~F32,tensor)
120
121def to_list(tensor: Tensor()) -> Result<&2,&2,Problem(),List<&2,F32>>:
122  Tensors.to_list(~F32,tensor)
123
124def item(tensor: Tensor()) -> Result<&2,&2,Problem(),F32>:
125  Tensors.item(~F32,tensor)
126
127# Reads with a fallback for a failed tensor, for callers that do not inspect
128# the problem.
129def to_list_or(tensor: Tensor(),fallback: List<&2,F32>) -> List<&2,F32>:
130  Tensors.value_or(~List<&2,F32>,Tensors.to_list(~F32,tensor),fallback)
131
132def item_or(tensor: Tensor(),fallback: F32) -> F32:
133  Tensors.value_or(~F32,Tensors.item(~F32,tensor),fallback)
134
135def into_storage(tensor: Tensor()) -> Result<&2,&1,Problem(),Storage()>:
136  Tensors.into_storage(~F32,tensor)
137
138def expect_storage(tensor: Tensor()) -> IO(Storage()):
139  Tensors.expect_storage(~F32,tensor)
140
141# Operations consume their tensor arguments and return one result. Keep a
142# tensor by copying it first; copy uses then, and shape uses then_shape.
143
144def then(pair: Tensor() & Tensor(),body: Tensor() -> Tensor() -> Tensor()) -> Tensor():
145  Tensors.then(~Tensor(),~Tensor(),~Tensor(),pair,body)
146
147def then_shape(pair: Tensor() & List<&2,U32>,body: Tensor() -> List<&2,U32> -> Tensor()) -> Tensor():
148  Tensors.then(~Tensor(),~List<&2,U32>,~Tensor(),pair,body)
149
150def copy(tensor: Tensor()) -> Tensor() & Tensor():
151  Layout.copy(~F32,tensor)
152
153# Binary elementwise operations reuse the left tensor's storage. The number
154# after the name is the operation's scalar work for the band schedule.
155def add(left: Tensor(),right: Tensor()) -> Tensor():
156  Tensors.combine(~F32,~F32.add,"add",1.0,left,right)
157
158def sub(left: Tensor(),right: Tensor()) -> Tensor():
159  Tensors.combine(~F32,~F32.sub,"sub",1.0,left,right)
160
161def mul(left: Tensor(),right: Tensor()) -> Tensor():
162  Tensors.combine(~F32,~F32.mul,"mul",1.0,left,right)
163
164# Base division: a zero divisor gives an infinity or NaN, never a failure.
165def div(left: Tensor(),right: Tensor()) -> Tensor():
166  Tensors.combine(~F32,~F32.div,"div",1.0,left,right)
167
168def maximum(left: Tensor(),right: Tensor()) -> Tensor():
169  Tensors.combine(~F32,~F32.max,"maximum",2.0,left,right)
170
171def minimum(left: Tensor(),right: Tensor()) -> Tensor():
172  Tensors.combine(~F32,~F32.min,"minimum",2.0,left,right)
173
174def add_scalar(tensor: Tensor(),+value: F32) -> Tensor():
175  Tensors.combine_scalar(~F32,~F32.add,"add_scalar",1.0,tensor,value)
176
177def sub_scalar(tensor: Tensor(),+value: F32) -> Tensor():
178  Tensors.combine_scalar(~F32,~F32.sub,"sub_scalar",1.0,tensor,value)
179
180def mul_scalar(tensor: Tensor(),+value: F32) -> Tensor():
181  Tensors.combine_scalar(~F32,~F32.mul,"mul_scalar",1.0,tensor,value)
182
183def div_scalar(tensor: Tensor(),+value: F32) -> Tensor():
184  Tensors.combine_scalar(~F32,~F32.div,"div_scalar",1.0,tensor,value)
185
186def less(left: Tensor(),right: Tensor()) -> Tensor():
187  Tensors.combine_with(~F32,~Lanes.combine(~Compare.less_elements),"less",2.625,left,right)
188
189def less_equal(left: Tensor(),right: Tensor()) -> Tensor():
190  Tensors.combine_with(~F32,~Lanes.combine(~Compare.less_equal_elements),"less_equal",2.625,left,right)
191
192def greater(left: Tensor(),right: Tensor()) -> Tensor():
193  Tensors.combine_with(~F32,~Lanes.combine(~Compare.greater_elements),"greater",2.625,left,right)
194
195def greater_equal(left: Tensor(),right: Tensor()) -> Tensor():
196  Tensors.combine_with(~F32,~Lanes.combine(~Compare.greater_equal_elements),"greater_equal",2.625,left,right)
197
198def equal(left: Tensor(),right: Tensor()) -> Tensor():
199  Tensors.combine_with(~F32,~Lanes.combine(~Compare.equal_elements),"equal",1.375,left,right)
200
201def not_equal(left: Tensor(),right: Tensor()) -> Tensor():
202  Tensors.combine_with(~F32,~Lanes.combine(~Compare.not_equal_elements),"not_equal",1.375,left,right)
203
204# Scalar comparisons share the eight-lane map and preserve the tensor layout.
205# Costs count AVX2 arithmetic, comparison, selection and value-bit instructions
206# in one eight-element round, divided by eight; moves and addressing are excluded.
207def less_scalar(tensor: Tensor(),+value: F32) -> Tensor():
208  Tensors.apply_context(~F32,~F32,~Operations.lanes(~F32,~Compare.less),value,"less_scalar",2.75,tensor)
209
210def less_equal_scalar(tensor: Tensor(),+value: F32) -> Tensor():
211  Tensors.apply_context(~F32,~F32,~Operations.lanes(~F32,~Compare.less_equal),value,"less_equal_scalar",2.625,tensor)
212
213def greater_scalar(tensor: Tensor(),+value: F32) -> Tensor():
214  Tensors.apply_context(~F32,~F32,~Operations.lanes(~F32,~Compare.greater),value,"greater_scalar",2.625,tensor)
215
216def greater_equal_scalar(tensor: Tensor(),+value: F32) -> Tensor():
217  Tensors.apply_context(~F32,~F32,~Operations.lanes(~F32,~Compare.greater_equal),value,"greater_equal_scalar",2.75,tensor)
218
219def equal_scalar(tensor: Tensor(),+value: F32) -> Tensor():
220  Tensors.apply_context(~F32,~F32,~Operations.lanes(~F32,~Compare.equal),value,"equal_scalar",1.5,tensor)
221
222def not_equal_scalar(tensor: Tensor(),+value: F32) -> Tensor():
223  Tensors.apply_context(~F32,~F32,~Operations.lanes(~F32,~Compare.not_equal),value,"not_equal_scalar",1.5,tensor)
224
225def select(mask: Tensor(),yes: Tensor(),no: Tensor()) -> Tensor():
226  Select.select(mask,yes,no)
227
228def clamp(tensor: Tensor(),+low: F32,+high: F32) -> Tensor():
229  Tensors.apply_context(~F32,~Compare.Bounds,~Operations.lanes(~Compare.Bounds,~Compare.clamp),Compare.Bounds{low,high},"clamp",0.5,tensor)
230
231# Activations: Act.silu() and the other values of lib/math_activation.bend
232# carry their own planner cost. activate applies one to a tensor; relu, silu,
233# ... below are its shorthands. Custom carries only a cost; its function
234# reaches conv2d_custom or apply, and activate fails on it.
235def activate(tensor: Tensor(),+activation: Activations.Activation) -> Tensor():
236  match activation:
237    case Activations.Identity{}: tensor
238    case Activations.ReLU{}: Tensors.apply(~F32,~Operations.mapped(~F32,~Activations.relu_value),"relu",Activations.operations(Activations.ReLU{}),tensor)
239    case Activations.ReLU6{}: Tensors.apply(~F32,~Operations.mapped(~F32,~Activations.relu6_value),"relu6",Activations.operations(Activations.ReLU6{}),tensor)
240    case Activations.LeakyReLU{+slope}: Tensors.combine_scalar(~F32,~Activations.leaky_relu_value,"leaky_relu",Activations.operations(Activations.LeakyReLU{slope}),tensor,slope)
241    case Activations.Sigmoid{}: Tensors.apply(~F32,~Operations.mapped_lanes(~Activations.sigmoid_value),"sigmoid",Activations.operations(Activations.Sigmoid{}),tensor)
242    case Activations.SiLU{}: Tensors.apply(~F32,~Operations.mapped_lanes(~Activations.silu_value),"silu",Activations.operations(Activations.SiLU{}),tensor)
243    case Activations.Tanh{}: Tensors.apply(~F32,~Operations.mapped_lanes(~Activations.tanh_value),"tanh",Activations.operations(Activations.Tanh{}),tensor)
244    case Activations.HardSigmoid{}: Tensors.apply(~F32,~Operations.mapped(~F32,~Activations.hard_sigmoid_value),"hard_sigmoid",Activations.operations(Activations.HardSigmoid{}),tensor)
245    case Activations.HardSwish{}: Tensors.apply(~F32,~Operations.mapped(~F32,~Activations.hard_swish_value),"hard_swish",Activations.operations(Activations.HardSwish{}),tensor)
246    case Activations.GELUTanh{}: Tensors.apply(~F32,~Operations.mapped_lanes(~Activations.gelu_tanh_value),"gelu_tanh",Activations.operations(Activations.GELUTanh{}),tensor)
247    case Activations.Custom{count}: Tensors.rejected(~F32,"activate",tensor)
248
249def relu(tensor: Tensor()) -> Tensor():
250  activate(tensor,Activations.relu())
251
252def relu6(tensor: Tensor()) -> Tensor():
253  activate(tensor,Activations.relu6())
254
255def leaky_relu(tensor: Tensor(),slope: F32) -> Tensor():
256  activate(tensor,Activations.leaky_relu(slope))
257
258def sigmoid(tensor: Tensor()) -> Tensor():
259  activate(tensor,Activations.sigmoid())
260
261def silu(tensor: Tensor()) -> Tensor():
262  activate(tensor,Activations.silu())
263
264def tanh(tensor: Tensor()) -> Tensor():
265  activate(tensor,Activations.tanh())
266
267def hard_sigmoid(tensor: Tensor()) -> Tensor():
268  activate(tensor,Activations.hard_sigmoid())
269
270def hard_swish(tensor: Tensor()) -> Tensor():
271  activate(tensor,Activations.hard_swish())
272
273# PyTorch's gelu(approximate='tanh').
274def gelu_tanh(tensor: Tensor()) -> Tensor():
275  activate(tensor,Activations.gelu_tanh())
276
277# Ordered convolution over CHW or NCHW; the activation runs on each finished
278# leaf while it is in cache.
279def conv2d(input: Tensor(),weights: Tensor(),bias: Tensor(),stride: U32,padding: U32,activation: Activations.Activation) -> Tensor():
280  TensorConvolution.conv2d(~Activations.identity_value,activation,input,weights,bias,stride,padding)
281
282# conv2d with a custom activation of the given scalar operations per element.
283def conv2d_custom(~activation: F32 -> F32,operations: F32,input: Tensor(),weights: Tensor(),bias: Tensor(),stride: U32,padding: U32) -> Tensor():
284  TensorConvolution.conv2d(~activation,Activations.Custom{operations},input,weights,bias,stride,padding)
285
286# Custom elementwise functions are passed as direct templates, never records,
287# with their scalar operations per element.
288def apply(~operation: F32 -> F32,operations: F32,tensor: Tensor()) -> Tensor():
289  Tensors.apply(~F32,~Operations.mapped(~F32,~operation),"apply",operations,tensor)
290
291# The function also receives the element's row-major index.
292def apply_indexed(~operation: U32 -> F32 -> F32,operations: F32,tensor: Tensor()) -> Tensor():
293  Tensors.apply_indexed(~F32,~operation,operations,tensor)
294
295# start, start + 1, ... as a rank-1 tensor of count elements.
296def arange(+start: F32,+count: U32) -> Tensor():
297  add_scalar(Tensors.apply_indexed(~F32,~(index => value => U32.to_f32(index)),1.0,zeros([count])),start)
298
299# Reductions return a rank-0 tensor, so failures propagate like any result.
300
301def sum(tensor: Tensor()) -> Tensor():
302  Tensors.reduce_all(~F32,~F32.add,"sum",tensor,0.0)
303
304def max(tensor: Tensor()) -> Tensor():
305  Tensors.reduce_all(~F32,~F32.max,"max",tensor,F32.div(F32.neg(1.0),0.0))
306
307def min(tensor: Tensor()) -> Tensor():
308  Tensors.reduce_all(~F32,~F32.min,"min",tensor,F32.div(1.0,0.0))
309
310# sum divided by the number of elements; an empty tensor gives NaN.
311def mean(tensor: Tensor()) -> Tensor():
312  then_shape(shape(tensor),tensor => extents =>
313    Tensors.combine_scalar(~F32,~F32.div,"mean",1.0,Tensors.reduce_all(~F32,~F32.add,"mean",tensor,0.0),
314      U32.to_f32(Tensors.count_or_zero(extents))))
315
316# Layout: results are new dense tensors, except reshape, which relabels.
317
318def reshape(tensor: Tensor(),+shape: List<&2,U32>) -> Tensor():
319  Layout.reshape(~F32,tensor,shape)
320
321# The result's axis i is the argument's axis order[i].
322def transpose(tensor: Tensor(),order: List<&2,U32>) -> Tensor():
323  Layout.transpose(~F32,tensor,order)
324
325# count elements from start along axis, every step-th one.
326def slice(tensor: Tensor(),axis: U32,start: U32,count: U32,step: U32) -> Tensor():
327  Layout.slice(~F32,tensor,axis,start,count,step)
328
329# target with part written at start along axis.
330def set_slice(target: Tensor(),axis: U32,start: U32,part: Tensor()) -> Tensor():
331  Layout.set_slice(~F32,target,axis,start,part)
332
333def Parts() -> Type:
334  Tensors.Parts<F32>
335
336# Build the parts of concat: F.cons(a, F.cons(b, F.single(c))).
337def single(tensor: Tensor()) -> Parts():
338  Tensors.Single{tensor}
339
340def cons(tensor: Tensor(),rest: Parts()) -> Parts():
341  Tensors.Cons{tensor,rest}
342
343def concat(parts: Parts(),+axis: U32) -> Tensor():
344  Layout.concat(~F32,parts,axis)
345
346# Repeats each element of the last two axes scale times in each direction.
347def upsample_nearest(tensor: Tensor(),+scale: U32) -> Tensor():
348  Layout.upsample_nearest(~F32,tensor,scale)
349
350# Over the last two axes; positions outside the input count as -inf.
351def max_pool2d_with(tensor: Tensor(),+kernel: U32,+stride: U32,+padding: U32) -> Tensor():
352  Window.pool2d(~F32,~F32.max,"max_pool2d",F32.div(F32.neg(1.0),0.0),tensor,kernel,stride,padding)
353
354# Windows of kernel by kernel that do not overlap: stride kernel, padding 0.
355def max_pool2d(tensor: Tensor(),+kernel: U32) -> Tensor():
356  max_pool2d_with(tensor,kernel,kernel,0)
357
358# Broadcasting and axis operations.
359
360# The tensor repeated to shape under NumPy broadcasting; add, sub, mul, max
361# and min broadcast their operands in the same way.
362def broadcast_to(tensor: Tensor(),+shape: List<&2,U32>) -> Tensor():
363  Layout.broadcast_to(~F32,tensor,shape)
364
365# Reductions along one axis remove that axis. Each result element takes the
366# elements of the axis in increasing index, starting from 0, -inf or +inf.
367def sum_axis(tensor: Tensor(),+axis: U32) -> Tensor():
368  Layout.reduce_axis(~F32,~F32.add,"sum_axis",0.0,tensor,axis)
369
370def max_axis(tensor: Tensor(),+axis: U32) -> Tensor():
371  Layout.reduce_axis(~F32,~F32.max,"max_axis",F32.div(F32.neg(1.0),0.0),tensor,axis)
372
373def min_axis(tensor: Tensor(),+axis: U32) -> Tensor():
374  Layout.reduce_axis(~F32,~F32.min,"min_axis",F32.div(1.0,0.0),tensor,axis)
375
376# Elementwise functions share the eight-value traversal used by activations.
377def exp(tensor: Tensor()) -> Tensor():
378  Tensors.apply(~F32,~Operations.mapped_lanes(~Arithmetic.exp),"exp",4.5,tensor)
379
380def log(tensor: Tensor()) -> Tensor():
381  Tensors.apply(~F32,~Operations.mapped_lanes(~Arithmetic.log),"log",8.25,tensor)
382
383# Keep a singleton reduction axis for broadcasting. Both reductions use the
384# existing ordered fold; reduction gathers row bands to a dense owner.
385def softmax(tensor: Tensor(),+axis: U32) -> Tensor():
386  then_shape(shape(tensor),tensor => extents =>
387    +unit = Layout.with_extent(extents,U32.to_nat(axis),1)
388    then(copy(tensor),values => maximum_input =>
389      then(copy(exp(sub(values,reshape(Layout.reduce_axis(~F32,~F32.max,"softmax",F32.div(F32.neg(1.0),0.0),maximum_input,axis),unit)))),
390        values => sum_input => div(values,reshape(Layout.reduce_axis(~F32,~F32.add,"softmax",0.0,sum_input,axis),unit)))))
391
392# sum_axis divided by the extent of the axis.
393def mean_axis(tensor: Tensor(),+axis: U32) -> Tensor():
394  then_shape(shape(tensor),tensor => extents =>
395    Tensors.combine_scalar(~F32,~F32.div,"mean_axis",1.0,Layout.reduce_axis(~F32,~F32.add,"mean_axis",0.0,tensor,axis),
396      U32.to_f32(Layout.extent_at(extents,U32.to_nat(axis)))))
397
398# before[i] new elements in front of axis i and after[i] behind it, all equal
399# to value.
400def pad(tensor: Tensor(),+before: List<&2,U32>,+after: List<&2,U32>,+value: F32) -> Tensor():
401  Layout.pad(~F32,Layout.Constant{value},tensor,before,after)
402
403# The same, each new element repeating the nearest element of the tensor.
404def pad_edge(tensor: Tensor(),+before: List<&2,U32>,+after: List<&2,U32>) -> Tensor():
405  Layout.pad(~F32,Layout.Edge{},tensor,before,after)