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.
11	import Base
12	import ./math_fp32.bend as Fp32
13	
14	# Custom is a caller-supplied function; only its operation count lives here.
15	type 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	
28	def identity_value(value: F32) -> F32:
29	  value
30	
31	def relu_value(value: F32) -> F32:
32	  F32.max(value,0.0)
33	
34	def relu6_value(value: F32) -> F32:
35	  F32.min(F32.max(value,0.0),6.0)
36	
37	def 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.
41	def leaky_relu_value(value: F32,slope: F32) -> F32:
42	  leaky(slope,value)
43	
44	def sigmoid_value(value: F32) -> F32:
45	  Fp32.sigmoid(value)
46	
47	def silu_value(value: F32) -> F32:
48	  Fp32.activate(value,True{})
49	
50	def tanh_value(value: F32) -> F32:
51	  Fp32.tanh(value)
52	
53	# PyTorch's order: relu6(x + 3) / 6 and x * relu6(x + 3) / 6.
54	def hard_sigmoid_value(value: F32) -> F32:
55	  (relu6_value((value + 3.0 : F32)) / 6.0 : F32)
56	
57	def swished(+value: F32) -> F32:
58	  ((value * relu6_value((value + 3.0 : F32)) : F32) / 6.0 : F32)
59	
60	def hard_swish_value(value: F32) -> F32:
61	  swished(value)
62	
63	# The tanh approximation as x*logistic(2t), retaining argument rounding errors.
64	def gelu_tanh_value(value: F32) -> F32:
65	  Fp32.gelu_tanh(value)
66	
67	def identity() -> Activation:
68	  Identity{}
69	
70	def relu() -> Activation:
71	  ReLU{}
72	
73	def relu6() -> Activation:
74	  ReLU6{}
75	
76	def leaky_relu(slope: F32) -> Activation:
77	  LeakyReLU{slope}
78	
79	def sigmoid() -> Activation:
80	  Sigmoid{}
81	
82	def silu() -> Activation:
83	  SiLU{}
84	
85	def tanh() -> Activation:
86	  Tanh{}
87	
88	def hard_sigmoid() -> Activation:
89	  HardSigmoid{}
90	
91	def hard_swish() -> Activation:
92	  HardSwish{}
93	
94	def gelu_tanh() -> Activation:
95	  GELUTanh{}
96	
97	# The identity needs no pass at all.
98	def 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.
104	def 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.
130	def unknown_operations() -> F32:
131	  90.0
132	
133	# Derivatives propagate NaNs independently of the ordered branch predicates.
134	def derivative_nan(+x: F32,value: F32) -> F32:
135	  Bool.pick(F32,F32.is_eq(x,x),value,x)
136	
137	def identity_derivative_value(value: F32) -> F32:
138	  +value = value
139	  derivative_nan(value,1.0)
140	
141	def 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	
145	def 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	
149	def 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	
153	def 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.
158	def 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.
165	def 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	
179	def sigmoid_derivative_value(value: F32) -> F32:
180	  +value = value
181	  derivative_nan(value,Fp32.sigmoid_derivative(value))
182	
183	def tanh_derivative_value(value: F32) -> F32:
184	  +value = value
185	  derivative_nan(value,Fp32.tanh_derivative(value))
186	
187	def silu_derivative_value(value: F32) -> F32:
188	  Fp32.silu_derivative(value)
189	
190	def gelu_tanh_derivative_value(value: F32) -> F32:
191	  Fp32.gelu_tanh_derivative(value)