1	# Ordered FP32 arithmetic on the shared sequence model. No separate container.
2	import Base
3	import ./math_sequence.bend as Sequence
4	
5	def multiply_add(accumulator: F32,input: F32,weight: F32) -> F32:
6	  (accumulator + input * weight : F32)
7	
8	def weighted_add(weight: F32,accumulator: F32,input: F32) -> F32:
9	  multiply_add(accumulator,input,weight)
10	
11	def zip_mac(length: Nat,accumulator: Sequence.Values(F32,length),input: Sequence.Values(F32,length),weight: F32) -> Sequence.Values(F32,length):
12	  Sequence.zip_map(~F32,~F32,~weighted_add,length,weight,accumulator,input)
13	
14	# The Cephes expf polynomial after its constant and linear terms.
15	def exponential_polynomial(+r: F32) -> F32:
16	  (((((0.00019875691500 * r + 0.0013981999507 : F32) * r + 0.0083334519073 : F32) * r + 0.041665795894 : F32) * r + 0.16666665459 : F32) * r + 0.50000001201 : F32)
17	
18	# The 32-bit word of a U32, read as an F32.
19	def from_bits(value: U32) -> F32:
20	  U32{word} = value
21	  F32{word}
22	
23	# exp(r)-1-r near |r| <= ln(2)/2. Chebyshev coefficients, rounded to F32;
24	# paired evaluation shortens the dependency chain without contracting products.
25	def exponential_remainder(+r: F32) -> F32:
26	  +z = (r * r : F32)
27	  (z * ((0.5 + 0.16666577756404877 * r : F32) + z * ((0.041666556149721146 + 0.008363173343241215 * r : F32) + z * 0.0013926175888627768 : F32) : F32) : F32)
28	
29	# b=RN(1+q). Recover the small term lost in that sum before adding the
30	# reduced exponential. Writing (1-b)+q keeps the eight lanes together under
31	# the default vectorizer; q-(b-1) makes it mix unrelated addition trees.
32	def logistic_denominator(q: F32,+base: F32,r: F32,correction: F32) -> F32:
33	  (base + (r + (correction + ((1.0 - base : F32) + q : F32) : F32) : F32) : F32)
34	
35	def selected(+mask: U32,yes: F32,+no: F32) -> F32:
36	  from_bits(U32.xor(F32.bits(no),U32.and(U32.xor(F32.bits(yes),F32.bits(no)),mask)))
37	
38	# Above 18 SiLU rounds to x. Pass x through, including positive infinity;
39	# the negative tail uses the clamped input before the final underflow.
40	def silu_finish(+x: F32,value: F32) -> F32:
41	  +word = F32.bits(x)
42	  +positive = U32.not((0 - U32.shrn(word,31n) : U32))
43	  +large = U32.and(positive,(0 - U32.shrn((1099956224 - word : U32),31n) : U32))
44	  selected(large,x,value)
45	
46	def logistic_scaled(silu: Bool,x: F32,value: F32,power: F32,denominator: F32) -> F32:
47	  match silu:
48	    case True{}:
49	      silu_finish(x,(((value * power : F32) / denominator : F32) * 0.00000000000000000005421010862427522 : F32))
50	    case False{}:
51	      ((power / denominator : F32) * 0.00000000000000000005421010862427522 : F32)
52	
53	# u=-value, n approximates u/ln(2) to the nearest integer, r=u-n*ln(2).
54	# The exponent word makes
55	# H=2^(64-n), so sigmoid(x)=H/(exp(r)+2^-n)*2^-64. H remains normal;
56	# only the last multiplication underflows. The tiny denominator addend q
57	# may be floored to the least normal F32 without changing its rounded sum.
58	def sigmoid_reduced(silu: Bool,x: F32,+value: F32,+biased: F32) -> F32:
59	  +n = (biased - 12582912.0 : F32)
60	  +r = ((n * F32.neg(0.693359375) - value : F32) + n * 0.00021219444005469058 : F32)
61	  +power = from_bits((1602224128 - U32.shln(F32.bits(biased),23n) : U32))
62	  +q = from_bits((U32.max(F32.bits(power),545259520) - 536870912 : U32))
63	  logistic_scaled(silu,x,value,power,logistic_denominator(q,(1.0 + q : F32),r,exponential_remainder(r)))
64	
65	def sigmoid_input(silu: Bool,x: F32,+value: F32) -> F32:
66	  sigmoid_reduced(silu,x,value,(value * F32.neg(1.44269504088896341) + 12582912.0 : F32))
67	
68	# At x>=18 sigmoid rounds to 1. Below -110 both sigmoid and SiLU round
69	# to zero; multiplying the clamped x before the final scaling preserves -0
70	# for SiLU, including at -infinity. NaNs pass the clamp and propagate.
71	def sigmoid(+x: F32) -> F32:
72	  sigmoid_input(False{},x,F32.max(F32.min(18.0,x),F32.neg(110.0)))
73	
74	def activate(+x: F32,enabled: Bool) -> F32:
75	  match enabled:
76	    case True{}:
77	      sigmoid_input(True{},x,F32.max(F32.min(18.0,x),F32.neg(110.0)))
78	    case False{}:
79	      x
80	
81	# Sign restoration and the tiny interval are word selections. For |x|<2^-12,
82	# |x-tanh(x)|<|x|^3/3 is below half an F32 ulp, so return x, including -0.
83	def tanh_finish(+x: F32,+value: F32) -> F32:
84	  +word = F32.bits(x)
85	  +small = (0 - U32.shrn((U32.and(word,2147483647) - 964689920 : U32),31n) : U32)
86	  +keep = U32.or(small,2147483648)
87	  from_bits(U32.or(U32.and(word,keep),U32.and(F32.bits(value),U32.not(keep))))
88	
89	# At |x|>=10, 1-tanh(|x|)<2*exp(-20)<2^-25. Clamping before arithmetic
90	# also keeps infinities away from intermediate products; NaNs propagate.
91	def tanh_magnitude(+x: F32) -> F32:
92	  F32.max(F32.min(10.0,from_bits(U32.and(F32.bits(x),2147483647))),0.000244140625)
93	
94	# Chebyshev interpolation of (tanh(sqrt(z))/sqrt(z)-1)/z, |t|<=0.75.
95	# The coefficients are F32 values.
96	def tanh_small(+t: F32,+z: F32) -> F32:
97	  (t + (t * z : F32) * (((((0.001909503829665482 * z + F32.neg(0.007933800108730793) : F32) * z + 0.021615320816636086 : F32) * z + F32.neg(0.05393586680293083) : F32) * z + 0.1333317905664444 : F32) * z + F32.neg(0.3333333134651184) : F32) : F32)
98	
99	# One additional term for the expm1 correction in the tanh quotient.
100	def tanh_remainder(+r: F32) -> F32:
101	  ((((((0.00019890980911441147 * r + 0.0013933641603216529 : F32) * r + 0.00833331048488617 : F32) * r + 0.04166646674275398 : F32) * r + 0.1666666716337204 : F32) * r + 0.5 : F32) * (r * r : F32) : F32)
102	
103	def tanh_exponential(+a: F32,+biased: F32,x: F32) -> F32:
104	  +n = (biased - 12582912.0 : F32)
105	  +r = (((a * 2.0 : F32) - n * 0.693359375 : F32) - n * F32.neg(0.00021219444005469058) : F32)
106	  +q = from_bits((1065353216 - U32.shln(F32.bits(biased),23n) : U32))
107	  +t = F32.min(0.75,a)
108	  +small = tanh_small(t,(t * t : F32))
109	  +large = (1.0 - 2.0 * (q / ((1.0 + q : F32) + (r + tanh_remainder(r) : F32) : F32) : F32) : F32)
110	  +mask = (0 - U32.shrn((F32.bits(a) - 1061158912 : U32),31n) : U32)
111	  tanh_finish(x,selected(mask,small,large))
112	
113	def tanh_input(+a: F32,x: F32) -> F32:
114	  tanh_exponential(a,((a * 2.0 : F32) * 1.44269504088896341 + 12582912.0 : F32),x)
115	
116	def tanh(+x: F32) -> F32:
117	  tanh_input(tanh_magnitude(x),x)
118	
119	# Split 2^n at +64 for nonnegative inputs and -64 for negative inputs.
120	# The first product is always normal and exact; only the last product can
121	# underflow. The exponent words come directly from the biased integer n.
122	def exponential_scale(biased: F32,x: F32,p: F32) -> F32:
123	  +sign = U32.shrn(U32.and(F32.bits(x),2147483648),1n)
124	  ((p * from_bits((1602224128 - sign : U32)) : F32) *
125	    from_bits((528482304 + U32.shln(F32.bits(biased),23n) + sign : U32)) : F32)
126	
127	# FastTwoSum retains the low bits of 1+r before adding the quadratic tail.
128	def exponential_precise(+r: F32) -> F32:
129	  +sum = (1.0 + r : F32)
130	  +error = (r - (sum - 1.0 : F32) : F32)
131	  (sum + ((exponential_polynomial(r) * (r * r : F32) : F32) + error : F32) : F32)
132	
133	def exponential_reduced(+x: F32,+biased: F32) -> F32:
134	  +n = (biased - 12582912.0 : F32)
135	  +high = (x - n * 0.693359375 : F32)
136	  +low = (n * F32.neg(0.00021219444005469058) : F32)
137	  +r = (high - low : F32)
138	  exponential_scale(biased,x,exponential_precise(r))
139	
140	def exponential_input(+x: F32) -> F32:
141	  exponential_reduced(x,(x * 1.44269504088896341 + 12582912.0 : F32))
142	
143	def exp(x: F32) -> F32:
144	  exponential_input(F32.max(F32.min(89.0,x),F32.neg(104.0)))
145	
146	# For s near r/(2+r), d = r-(2+r)*s gives the exact real identity
147	# log(1+r) = 2*atanh(s) + log(1+d/(1+s)). The odd series ends at s^11.
148	# Split r and s into 12-bit heads so their head product is exact; retain
149	# the low products in d. (1-s)*(1+s^2+s^4) approximates 1/(1+s).
150	# FastTwoSum retains the low bits when adding the exponent's ln(2) part.
151	def logarithm_series(+r: F32,+n: F32) -> F32:
152	  +s = (r / (2.0 + r : F32) : F32)
153	  +square = (s * s : F32)
154	  +r_high = from_bits(U32.and(F32.bits(r),4294963200))
155	  +r_low = (r - r_high : F32)
156	  +s_high = from_bits(U32.and(F32.bits(s),4294963200))
157	  +s_low = (s - s_high : F32)
158	  +product_low = ((r_high * s_low + r_low * s_high : F32) + r_low * s_low : F32)
159	  +correction = (((r - 2.0 * s : F32) - r_high * s_high : F32) - product_low : F32)
160	  +p = ((((0.09090909090909091 * square + 0.1111111111111111 : F32) * square + 0.14285714285714285 : F32) * square + 0.2 : F32) * square + 0.3333333333333333 : F32)
161	  +main = (2.0 * s : F32)
162	  +rest = ((main * square : F32) * p + correction * ((1.0 - s : F32) * (1.0 + square * (1.0 + square : F32) : F32) : F32) : F32)
163	  +high = (n * 0.693359375 : F32)
164	  +sum = (high + main : F32)
165	  +tail = (main - (sum - high : F32) : F32)
166	  (sum + (rest + (n * F32.neg(0.000212194440) + tail : F32) : F32) : F32)
167	
168	# The exponent carry at mantissa word 0x3504f4 selects the same reduced
169	# interval as halving m above sqrt(2), without a floating-point selection.
170	def logarithm_normalized(x: F32,adjustment: F32) -> F32:
171	  +word = F32.bits(x)
172	  +adjusted = (word + 4913932 : U32)
173	  logarithm_series((from_bits((word - U32.and(adjusted,4286578688) + 1065353216 : U32)) - 1.0 : F32),
174	    ((from_bits(U32.or(U32.and(U32.shrn(adjusted,23n),255),1258291200)) - 8388735.0 : F32) - adjustment : F32))
175	
176	# Scaling a subnormal by 2^23 is exact. Zero, negatives, infinities and
177	# NaNs are selected explicitly, after normalizing the finite positive input.
178	def logarithm_finish(+x: F32,value: F32) -> F32:
179	  +word = F32.bits(x)
180	  +magnitude = U32.and(word,2147483647)
181	  +zero = (0 - U32.shrn((magnitude - 1 : U32),31n) : U32)
182	  +negative = (0 - U32.shrn(word,31n) : U32)
183	  +special = U32.not((0 - U32.shrn((magnitude - 2139095040 : U32),31n) : U32))
184	  selected(zero,from_bits(4286578688),selected(negative,from_bits(2143289344),selected(special,(x + x : F32),value)))
185	
186	def logarithm_input(+x: F32) -> F32:
187	  +small = (0 - U32.shrn((U32.and(F32.bits(x),2147483647) - 8388608 : U32),31n) : U32)
188	  logarithm_finish(x,logarithm_normalized(selected(small,(x * 8388608.0 : F32),x),from_bits(U32.and(small,1102577664))))
189	
190	def log(x: F32) -> F32:
191	  logarithm_input(x)