math_activation.bend
stelliferous@0.0.3.0 · 191 lines · 0 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.3.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# The tanh approximation as x*logistic(2t), retaining argument rounding errors.
64def gelu_tanh_value(value: F32) -> F32:
65 Fp32.gelu_tanh(value)
66
67def identity() -> Activation:
68 Identity{}
69
70def relu() -> Activation:
71 ReLU{}
72
73def relu6() -> Activation:
74 ReLU6{}
75
76def leaky_relu(slope: F32) -> Activation:
77 LeakyReLU{slope}
78
79def sigmoid() -> Activation:
80 Sigmoid{}
81
82def silu() -> Activation:
83 SiLU{}
84
85def tanh() -> Activation:
86 Tanh{}
87
88def hard_sigmoid() -> Activation:
89 HardSigmoid{}
90
91def hard_swish() -> Activation:
92 HardSwish{}
93
94def gelu_tanh() -> Activation:
95 GELUTanh{}
96
97# The identity needs no pass at all.
98def is_identity(activation: Activation) -> Bool:
99 match activation:
100 case Identity{}: True{}
101 case activation: False{}
102
103# Operations per element, excluding the pass that applies them.
104def operations(activation: Activation) -> F32:
105 match activation:
106 case Identity{}: 0.0
107 # compare and select
108 case ReLU{}: 2.0
109 case ReLU6{}: 4.0
110 # compare, multiply and select
111 case LeakyReLU{slope}: 3.0
112 # 31 arithmetic, comparison, select and word instructions per eight lanes
113 case Sigmoid{}: 3.875
114 # 34 per eight lanes: scaled numerator and positive-tail selection
115 case SiLU{}: 4.25
116 # 54 per eight lanes
117 case Tanh{}: 6.75
118 # add, relu6 and divide; hard_swish also multiplies
119 case HardSigmoid{}: 6.0
120 case HardSwish{}: 7.0
121 # 68 per eight lanes: compensated argument, logistic and sign selection
122 case GELUTanh{}: 8.5
123 case Custom{operations}: operations
124
125# A custom function (F.conv2d_custom, F.apply) states its scalar operations per
126# element: one per arithmetic, comparison or select; F32.exp about 37 and
127# F32.tanh about 82 with their conversions. The count only steers parallel
128# planning; results never depend on it. When unsure, use this conservative
129# count: a call into the C math library and a few operations around it.
130def unknown_operations() -> F32:
131 90.0
132
133# Derivatives propagate NaNs independently of the ordered branch predicates.
134def derivative_nan(+x: F32,value: F32) -> F32:
135 Bool.pick(F32,F32.is_eq(x,x),value,x)
136
137def identity_derivative_value(value: F32) -> F32:
138 +value = value
139 derivative_nan(value,1.0)
140
141def relu_derivative_value(value: F32) -> F32:
142 +value = value
143 derivative_nan(value,Bool.pick(F32,F32.is_lt(0.0,value),1.0,0.0))
144
145def relu6_derivative_value(value: F32) -> F32:
146 +value = value
147 derivative_nan(value,Bool.pick(F32,F32.is_lt(0.0,value) && F32.is_lt(value,6.0),1.0,0.0))
148
149def leaky_relu_derivative_value(value: F32,slope: F32) -> F32:
150 +value = value
151 derivative_nan(value,Bool.pick(F32,F32.is_lt(0.0,value),1.0,slope))
152
153def hard_sigmoid_derivative_value(value: F32) -> F32:
154 +value = value
155 derivative_nan(value,Bool.pick(F32,F32.is_lt(F32.neg(3.0),value) && F32.is_lt(value,3.0),(1.0 / 6.0 : F32),0.0))
156
157# PyTorch's backward takes the outer branches at -3 and 3.
158def hard_swish_derivative_value(value: F32) -> F32:
159 +value = value
160 derivative_nan(value,Bool.pick(F32,F32.is_le(value,F32.neg(3.0)),0.0,
161 Bool.pick(F32,F32.is_lt(value,3.0),((value / 3.0 : F32) + 0.5 : F32),1.0)))
162
163# Packed arithmetic, comparisons, masks and selections per eight values,
164# divided by eight. NaN selection is included.
165def derivative_operations(activation: Activation) -> F32:
166 match activation:
167 case Identity{}: 0.25
168 case ReLU{}: 0.5
169 case ReLU6{}: 0.75
170 case LeakyReLU{slope}: 0.5
171 case HardSigmoid{}: 0.75
172 case HardSwish{}: 1.0
173 case Sigmoid{}: 5.375
174 case Tanh{}: 5.625
175 case SiLU{}: 8.375
176 case GELUTanh{}: 20.625
177 case activation: 0.0
178
179def sigmoid_derivative_value(value: F32) -> F32:
180 +value = value
181 derivative_nan(value,Fp32.sigmoid_derivative(value))
182
183def tanh_derivative_value(value: F32) -> F32:
184 +value = value
185 derivative_nan(value,Fp32.tanh_derivative(value))
186
187def silu_derivative_value(value: F32) -> F32:
188 Fp32.silu_derivative(value)
189
190def gelu_tanh_derivative_value(value: F32) -> F32:
191 Fp32.gelu_tanh_derivative(value)