math_activation.bend

stelliferous@0.0.2.0 · 135 lines · 27 definitions

1# FP32 activations, imported by callers as Act: Act.silu() is a value for
2# F.conv2d or F.activate, Act.silu_value the scalar function. Each built-in
3# states its operations per element for the planners, counted from the emitted
4# C and assembly: arithmetic, comparisons, selects and word operations;
5# addressing, moves and loop control are excluded. The sigmoid, SiLU, tanh and
6# GELU run eight lanes at a time, so their counts are the instructions per
7# eight lanes divided by eight; the others are scalar counts. The sigmoid and
8# tanh use ordered FP32 arithmetic and word operations (lib/math_fp32.bend).
9# Sigmoid and SiLU scale their negative tails before the final underflow;
10# accuracy is measured against MPFR rather than a global rounding guarantee.
11import Base
12import stelliferous@0.0.2.0/math_fp32.bend as Fp32
13
14# Custom is a caller-supplied function; only its operation count lives here.
15type Activation is Data:
16  Identity{}
17  ReLU{}
18  ReLU6{}
19  LeakyReLU{slope: F32}
20  Sigmoid{}
21  SiLU{}
22  Tanh{}
23  HardSigmoid{}
24  HardSwish{}
25  GELUTanh{}
26  Custom{operations: F32}
27
28def identity_value(value: F32) -> F32:
29  value
30
31def relu_value(value: F32) -> F32:
32  F32.max(value,0.0)
33
34def relu6_value(value: F32) -> F32:
35  F32.min(F32.max(value,0.0),6.0)
36
37def leaky(slope: F32,+value: F32) -> F32:
38  Bool.pick(F32,F32.is_lt(value,0.0),(value * slope : F32),value)
39
40# x for x >= 0, slope * x otherwise.
41def leaky_relu_value(value: F32,slope: F32) -> F32:
42  leaky(slope,value)
43
44def sigmoid_value(value: F32) -> F32:
45  Fp32.sigmoid(value)
46
47def silu_value(value: F32) -> F32:
48  Fp32.activate(value,True{})
49
50def tanh_value(value: F32) -> F32:
51  Fp32.tanh(value)
52
53# PyTorch's order: relu6(x + 3) / 6 and x * relu6(x + 3) / 6.
54def hard_sigmoid_value(value: F32) -> F32:
55  (relu6_value((value + 3.0 : F32)) / 6.0 : F32)
56
57def swished(+value: F32) -> F32:
58  ((value * relu6_value((value + 3.0 : F32)) : F32) / 6.0 : F32)
59
60def hard_swish_value(value: F32) -> F32:
61  swished(value)
62
63# PyTorch's gelu(approximate='tanh'): 0.5 x (1 + tanh(sqrt(2/pi) (x + 0.044715 x^3))).
64def gelu_tanh_of(+value: F32) -> F32:
65  +cube = ((value * value : F32) * value : F32)
66  ((0.5 * value : F32) * (1.0 + tanh_value((0.7978845608 * (value + (0.044715 * cube : F32) : F32) : F32)) : F32) : F32)
67
68def gelu_tanh_value(value: F32) -> F32:
69  gelu_tanh_of(value)
70
71def identity() -> Activation:
72  Identity{}
73
74def relu() -> Activation:
75  ReLU{}
76
77def relu6() -> Activation:
78  ReLU6{}
79
80def leaky_relu(slope: F32) -> Activation:
81  LeakyReLU{slope}
82
83def sigmoid() -> Activation:
84  Sigmoid{}
85
86def silu() -> Activation:
87  SiLU{}
88
89def tanh() -> Activation:
90  Tanh{}
91
92def hard_sigmoid() -> Activation:
93  HardSigmoid{}
94
95def hard_swish() -> Activation:
96  HardSwish{}
97
98def gelu_tanh() -> Activation:
99  GELUTanh{}
100
101# The identity needs no pass at all.
102def is_identity(activation: Activation) -> Bool:
103  match activation:
104    case Identity{}: True{}
105    case activation: False{}
106
107# Operations per element, excluding the pass that applies them.
108def operations(activation: Activation) -> F32:
109  match activation:
110    case Identity{}: 0.0
111    # compare and select
112    case ReLU{}: 2.0
113    case ReLU6{}: 4.0
114    # compare, multiply and select
115    case LeakyReLU{slope}: 3.0
116    # 31 arithmetic, comparison, select and word instructions per eight lanes
117    case Sigmoid{}: 3.875
118    # 34 per eight lanes: scaled numerator and positive-tail selection
119    case SiLU{}: 4.25
120    # 54 per eight lanes
121    case Tanh{}: 6.75
122    # add, relu6 and divide; hard_swish also multiplies
123    case HardSigmoid{}: 6.0
124    case HardSwish{}: 7.0
125    # 62 per eight lanes: the PyTorch expression around the tanh
126    case GELUTanh{}: 7.75
127    case Custom{operations}: operations
128
129# A custom function (F.conv2d_custom, F.apply) states its scalar operations per
130# element: one per arithmetic, comparison or select; F32.exp about 37 and
131# F32.tanh about 82 with their conversions. The count only steers parallel
132# planning; results never depend on it. When unsure, use this conservative
133# count: a call into the C math library and a few operations around it.
134def unknown_operations() -> F32:
135  90.0