Aluode/PerceptionLabPortable
0
1"""
2Provide math calls that uses intrinsics or libc math functions.
3"""
4
5import math
6import operator
7import sys
8import numpy as np
9
10import llvmlite.ir
11from llvmlite.ir import Constant
12
13from numba.core.imputils import Registry, impl_ret_untracked
14from numba import typeof
15from numba.core import types, utils, config, cgutils
16from numba.core.extending import overload
17from numba.core.typing import signature
18from numba.cpython.unsafe.numbers import trailing_zeros
19
20
21registry = Registry('mathimpl')
22lower = registry.lower
23
24
25# Helpers, shared with cmathimpl.
26_NP_FLT_FINFO = np.finfo(np.dtype('float32'))
27FLT_MAX = _NP_FLT_FINFO.max
28FLT_MIN = _NP_FLT_FINFO.tiny
29
30_NP_DBL_FINFO = np.finfo(np.dtype('float64'))
31DBL_MAX = _NP_DBL_FINFO.max
32DBL_MIN = _NP_DBL_FINFO.tiny
33
34FLOAT_ABS_MASK = 0x7fffffff
35FLOAT_SIGN_MASK = 0x80000000
36DOUBLE_ABS_MASK = 0x7fffffffffffffff
37DOUBLE_SIGN_MASK = 0x8000000000000000
38
39
40def is_nan(builder, val):
41 """
42 Return a condition testing whether *val* is a NaN.
43 """
44 return builder.fcmp_unordered('uno', val, val)
45
46def is_inf(builder, val):
47 """
48 Return a condition testing whether *val* is an infinite.
49 """
50 pos_inf = Constant(val.type, float("+inf"))
51 neg_inf = Constant(val.type, float("-inf"))
52 isposinf = builder.fcmp_ordered('==', val, pos_inf)
53 isneginf = builder.fcmp_ordered('==', val, neg_inf)
54 return builder.or_(isposinf, isneginf)
55
56def is_finite(builder, val):
57 """
58 Return a condition testing whether *val* is a finite.
59 """
60 # is_finite(x) <=> x - x != NaN
61 val_minus_val = builder.fsub(val, val)
62 return builder.fcmp_ordered('ord', val_minus_val, val_minus_val)
63
64def f64_as_int64(builder, val):
65 """
66 Bitcast a double into a 64-bit integer.
67 """
68 assert val.type == llvmlite.ir.DoubleType()
69 return builder.bitcast(val, llvmlite.ir.IntType(64))
70
71def int64_as_f64(builder, val):
72 """
73 Bitcast a 64-bit integer into a double.
74 """
75 assert val.type == llvmlite.ir.IntType(64)
76 return builder.bitcast(val, llvmlite.ir.DoubleType())
77
78def f32_as_int32(builder, val):
79 """
80 Bitcast a float into a 32-bit integer.
81 """
82 assert val.type == llvmlite.ir.FloatType()
83 return builder.bitcast(val, llvmlite.ir.IntType(32))
84
85def int32_as_f32(builder, val):
86 """
87 Bitcast a 32-bit integer into a float.
88 """
89 assert val.type == llvmlite.ir.IntType(32)
90 return builder.bitcast(val, llvmlite.ir.FloatType())
91
92def negate_real(builder, val):
93 """
94 Negate real number *val*, with proper handling of zeros.
95 """
96 # The negative zero forces LLVM to handle signed zeros properly.
97 return builder.fsub(Constant(val.type, -0.0), val)
98
99def call_fp_intrinsic(builder, name, args):
100 """
101 Call a LLVM intrinsic floating-point operation.
102 """
103 mod = builder.module
104 intr = mod.declare_intrinsic(name, [a.type for a in args])
105 return builder.call(intr, args)
106
107
108def _unary_int_input_wrapper_impl(wrapped_impl):
109 """
110 Return an implementation factory to convert the single integral input
111 argument to a float64, then defer to the *wrapped_impl*.
112 """
113 def implementer(context, builder, sig, args):
114 val, = args
115 input_type = sig.args[0]
116 fpval = context.cast(builder, val, input_type, types.float64)
117 inner_sig = signature(types.float64, types.float64)
118 res = wrapped_impl(context, builder, inner_sig, (fpval,))
119 return context.cast(builder, res, types.float64, sig.return_type)
120
121 return implementer
122
123def unary_math_int_impl(fn, float_impl):
124 impl = _unary_int_input_wrapper_impl(float_impl)
125 lower(fn, types.Integer)(impl)
126
127def unary_math_intr(fn, intrcode):
128 """
129 Implement the math function *fn* using the LLVM intrinsic *intrcode*.
130 """
131 @lower(fn, types.Float)
132 def float_impl(context, builder, sig, args):
133 res = call_fp_intrinsic(builder, intrcode, args)
134 return impl_ret_untracked(context, builder, sig.return_type, res)
135
136 unary_math_int_impl(fn, float_impl)
137 return float_impl
138
139def unary_math_extern(fn, f32extern, f64extern, int_restype=False):
140 """
141 Register implementations of Python function *fn* using the
142 external function named *f32extern* and *f64extern* (for float32
143 and float64 inputs, respectively).
144 If *int_restype* is true, then the function's return value should be
145 integral, otherwise floating-point.
146 """
147 f_restype = types.int64 if int_restype else None
148
149 def float_impl(context, builder, sig, args):
150 """
151 Implement *fn* for a types.Float input.
152 """
153 [val] = args
154 mod = builder.module
155 input_type = sig.args[0]
156 lty = context.get_value_type(input_type)
157 func_name = {
158 types.float32: f32extern,
159 types.float64: f64extern,
160 }[input_type]
161 fnty = llvmlite.ir.FunctionType(lty, [lty])
162 fn = cgutils.insert_pure_function(builder.module, fnty, name=func_name)
163 res = builder.call(fn, (val,))
164 res = context.cast(builder, res, input_type, sig.return_type)
165 return impl_ret_untracked(context, builder, sig.return_type, res)
166
167 lower(fn, types.Float)(float_impl)
168
169 # Implement wrapper for integer inputs
170 unary_math_int_impl(fn, float_impl)
171
172 return float_impl
173
174
175unary_math_intr(math.fabs, 'llvm.fabs')
176exp_impl = unary_math_intr(math.exp, 'llvm.exp')
177if sys.version_info >= (3, 11):
178 exp2_impl = unary_math_intr(math.exp2, 'llvm.exp2')
179log_impl = unary_math_intr(math.log, 'llvm.log')
180log10_impl = unary_math_intr(math.log10, 'llvm.log10')
181log2_impl = unary_math_intr(math.log2, 'llvm.log2')
182sin_impl = unary_math_intr(math.sin, 'llvm.sin')
183cos_impl = unary_math_intr(math.cos, 'llvm.cos')
184
185log1p_impl = unary_math_extern(math.log1p, "log1pf", "log1p")
186expm1_impl = unary_math_extern(math.expm1, "expm1f", "expm1")
187erf_impl = unary_math_extern(math.erf, "erff", "erf")
188erfc_impl = unary_math_extern(math.erfc, "erfcf", "erfc")
189
190tan_impl = unary_math_extern(math.tan, "tanf", "tan")
191asin_impl = unary_math_extern(math.asin, "asinf", "asin")
192acos_impl = unary_math_extern(math.acos, "acosf", "acos")
193atan_impl = unary_math_extern(math.atan, "atanf", "atan")
194
195asinh_impl = unary_math_extern(math.asinh, "asinhf", "asinh")
196acosh_impl = unary_math_extern(math.acosh, "acoshf", "acosh")
197atanh_impl = unary_math_extern(math.atanh, "atanhf", "atanh")
198sinh_impl = unary_math_extern(math.sinh, "sinhf", "sinh")
199cosh_impl = unary_math_extern(math.cosh, "coshf", "cosh")
200tanh_impl = unary_math_extern(math.tanh, "tanhf", "tanh")
201
202log2_impl = unary_math_extern(math.log2, "log2f", "log2")
203ceil_impl = unary_math_extern(math.ceil, "ceilf", "ceil", True)
204floor_impl = unary_math_extern(math.floor, "floorf", "floor", True)
205
206gamma_impl = unary_math_extern(math.gamma, "numba_gammaf", "numba_gamma") # work-around
207sqrt_impl = unary_math_extern(math.sqrt, "sqrtf", "sqrt")
208trunc_impl = unary_math_extern(math.trunc, "truncf", "trunc", True)
209lgamma_impl = unary_math_extern(math.lgamma, "lgammaf", "lgamma")
210
211
212@lower(math.isnan, types.Float)
213def isnan_float_impl(context, builder, sig, args):
214 [val] = args
215 res = is_nan(builder, val)
216 return impl_ret_untracked(context, builder, sig.return_type, res)
217
218@lower(math.isnan, types.Integer)
219def isnan_int_impl(context, builder, sig, args):
220 res = cgutils.false_bit
221 return impl_ret_untracked(context, builder, sig.return_type, res)
222
223
224@lower(math.isinf, types.Float)
225def isinf_float_impl(context, builder, sig, args):
226 [val] = args
227 res = is_inf(builder, val)
228 return impl_ret_untracked(context, builder, sig.return_type, res)
229
230@lower(math.isinf, types.Integer)
231def isinf_int_impl(context, builder, sig, args):
232 res = cgutils.false_bit
233 return impl_ret_untracked(context, builder, sig.return_type, res)
234
235
236@lower(math.isfinite, types.Float)
237def isfinite_float_impl(context, builder, sig, args):
238 [val] = args
239 res = is_finite(builder, val)
240 return impl_ret_untracked(context, builder, sig.return_type, res)
241
242
243@lower(math.isfinite, types.Integer)
244def isfinite_int_impl(context, builder, sig, args):
245 res = cgutils.true_bit
246 return impl_ret_untracked(context, builder, sig.return_type, res)
247
248
249@lower(math.copysign, types.Float, types.Float)
250def copysign_float_impl(context, builder, sig, args):
251 lty = args[0].type
252 mod = builder.module
253 fn = cgutils.get_or_insert_function(mod, llvmlite.ir.FunctionType(lty, (lty, lty)),
254 'llvm.copysign.%s' % lty.intrinsic_name)
255 res = builder.call(fn, args)
256 return impl_ret_untracked(context, builder, sig.return_type, res)
257
258
259# -----------------------------------------------------------------------------
260
261
262@lower(math.frexp, types.Float)
263def frexp_impl(context, builder, sig, args):
264 val, = args
265 fltty = context.get_data_type(sig.args[0])
266 intty = context.get_data_type(sig.return_type[1])
267 expptr = cgutils.alloca_once(builder, intty, name='exp')
268 fnty = llvmlite.ir.FunctionType(fltty, (fltty, llvmlite.ir.PointerType(intty)))
269 fname = {
270 "float": "numba_frexpf",
271 "double": "numba_frexp",
272 }[str(fltty)]
273 fn = cgutils.get_or_insert_function(builder.module, fnty, fname)
274 res = builder.call(fn, (val, expptr))
275 res = cgutils.make_anonymous_struct(builder, (res, builder.load(expptr)))
276 return impl_ret_untracked(context, builder, sig.return_type, res)
277
278
279@lower(math.ldexp, types.Float, types.intc)
280def ldexp_impl(context, builder, sig, args):
281 val, exp = args
282 fltty, intty = map(context.get_data_type, sig.args)
283 fnty = llvmlite.ir.FunctionType(fltty, (fltty, intty))
284 fname = {
285 "float": "numba_ldexpf",
286 "double": "numba_ldexp",
287 }[str(fltty)]
288 fn = cgutils.insert_pure_function(builder.module, fnty, name=fname)
289 res = builder.call(fn, (val, exp))
290 return impl_ret_untracked(context, builder, sig.return_type, res)
291
292
293# -----------------------------------------------------------------------------
294
295
296@lower(math.atan2, types.int64, types.int64)
297def atan2_s64_impl(context, builder, sig, args):
298 [y, x] = args
299 y = builder.sitofp(y, llvmlite.ir.DoubleType())
300 x = builder.sitofp(x, llvmlite.ir.DoubleType())
301 fsig = signature(types.float64, types.float64, types.float64)
302 return atan2_float_impl(context, builder, fsig, (y, x))
303
304@lower(math.atan2, types.uint64, types.uint64)
305def atan2_u64_impl(context, builder, sig, args):
306 [y, x] = args
307 y = builder.uitofp(y, llvmlite.ir.DoubleType())
308 x = builder.uitofp(x, llvmlite.ir.DoubleType())
309 fsig = signature(types.float64, types.float64, types.float64)
310 return atan2_float_impl(context, builder, fsig, (y, x))
311
312@lower(math.atan2, types.Float, types.Float)
313def atan2_float_impl(context, builder, sig, args):
314 assert len(args) == 2
315 mod = builder.module
316 ty = sig.args[0]
317 lty = context.get_value_type(ty)
318 func_name = {
319 types.float32: "atan2f",
320 types.float64: "atan2"
321 }[ty]
322 fnty = llvmlite.ir.FunctionType(lty, (lty, lty))
323 fn = cgutils.insert_pure_function(builder.module, fnty, name=func_name)
324 res = builder.call(fn, args)
325 return impl_ret_untracked(context, builder, sig.return_type, res)
326
327
328# -----------------------------------------------------------------------------
329
330
331@lower(math.hypot, types.int64, types.int64)
332def hypot_s64_impl(context, builder, sig, args):
333 [x, y] = args
334 y = builder.sitofp(y, llvmlite.ir.DoubleType())
335 x = builder.sitofp(x, llvmlite.ir.DoubleType())
336 fsig = signature(types.float64, types.float64, types.float64)
337 res = hypot_float_impl(context, builder, fsig, (x, y))
338 return impl_ret_untracked(context, builder, sig.return_type, res)
339
340
341@lower(math.hypot, types.uint64, types.uint64)
342def hypot_u64_impl(context, builder, sig, args):
343 [x, y] = args
344 y = builder.sitofp(y, llvmlite.ir.DoubleType())
345 x = builder.sitofp(x, llvmlite.ir.DoubleType())
346 fsig = signature(types.float64, types.float64, types.float64)
347 res = hypot_float_impl(context, builder, fsig, (x, y))
348 return impl_ret_untracked(context, builder, sig.return_type, res)
349
350
351@lower(math.hypot, types.Float, types.Float)
352def hypot_float_impl(context, builder, sig, args):
353 xty, yty = sig.args
354 assert xty == yty == sig.return_type
355 x, y = args
356
357 # Windows has alternate names for hypot/hypotf, see
358 # https://msdn.microsoft.com/fr-fr/library/a9yb3dbt%28v=vs.80%29.aspx
359 fname = {
360 types.float32: "_hypotf" if sys.platform == 'win32' else "hypotf",
361 types.float64: "_hypot" if sys.platform == 'win32' else "hypot",
362 }[xty]
363 plat_hypot = types.ExternalFunction(fname, sig)
364
365 if sys.platform == 'win32' and config.MACHINE_BITS == 32:
366 inf = xty(float('inf'))
367
368 def hypot_impl(x, y):
369 if math.isinf(x) or math.isinf(y):
370 return inf
371 return plat_hypot(x, y)
372 else:
373 def hypot_impl(x, y):
374 return plat_hypot(x, y)
375
376 res = context.compile_internal(builder, hypot_impl, sig, args)
377 return impl_ret_untracked(context, builder, sig.return_type, res)
378
379
380# -----------------------------------------------------------------------------
381
382@lower(math.radians, types.Float)
383def radians_float_impl(context, builder, sig, args):
384 [x] = args
385 coef = context.get_constant(sig.return_type, math.pi / 180)
386 res = builder.fmul(x, coef)
387 return impl_ret_untracked(context, builder, sig.return_type, res)
388
389unary_math_int_impl(math.radians, radians_float_impl)
390
391# -----------------------------------------------------------------------------
392
393@lower(math.degrees, types.Float)
394def degrees_float_impl(context, builder, sig, args):
395 [x] = args
396 coef = context.get_constant(sig.return_type, 180 / math.pi)
397 res = builder.fmul(x, coef)
398 return impl_ret_untracked(context, builder, sig.return_type, res)
399
400unary_math_int_impl(math.degrees, degrees_float_impl)
401
402# -----------------------------------------------------------------------------
403
404@lower(math.pow, types.Float, types.Float)
405@lower(math.pow, types.Float, types.Integer)
406def pow_impl(context, builder, sig, args):
407 impl = context.get_function(operator.pow, sig)
408 return impl(builder, args)
409
410# -----------------------------------------------------------------------------
411
412@lower(math.nextafter, types.Float, types.Float)
413def nextafter_impl(context, builder, sig, args):
414 assert len(args) == 2
415 ty = sig.args[0]
416 lty = context.get_value_type(ty)
417 func_name = {
418 types.float32: "nextafterf",
419 types.float64: "nextafter"
420 }[ty]
421 fnty = llvmlite.ir.FunctionType(lty, (lty, lty))
422 fn = cgutils.insert_pure_function(builder.module, fnty, name=func_name)
423 res = builder.call(fn, args)
424 return impl_ret_untracked(context, builder, sig.return_type, res)
425
426# -----------------------------------------------------------------------------
427
428def _unsigned(T):
429 """Convert integer to unsigned integer of equivalent width."""
430 pass
431
432@overload(_unsigned)
433def _unsigned_impl(T):
434 if T in types.unsigned_domain:
435 return lambda T: T
436 elif T in types.signed_domain:
437 newT = getattr(types, 'uint{}'.format(T.bitwidth))
438 return lambda T: newT(T)
439
440
441def gcd_impl(context, builder, sig, args):
442 xty, yty = sig.args
443 assert xty == yty == sig.return_type
444 x, y = args
445
446 def gcd(a, b):
447 """
448 Stein's algorithm, heavily cribbed from Julia implementation.
449 """
450 T = type(a)
451 if a == 0: return abs(b)
452 if b == 0: return abs(a)
453 za = trailing_zeros(a)
454 zb = trailing_zeros(b)
455 k = min(za, zb)
456 # Uses np.*_shift instead of operators due to return types
457 u = _unsigned(abs(np.right_shift(a, za)))
458 v = _unsigned(abs(np.right_shift(b, zb)))
459 while u != v:
460 if u > v:
461 u, v = v, u
462 v -= u
463 v = np.right_shift(v, trailing_zeros(v))
464 r = np.left_shift(T(u), k)
465 return r
466
467 res = context.compile_internal(builder, gcd, sig, args)
468 return impl_ret_untracked(context, builder, sig.return_type, res)
469
470
471lower(math.gcd, types.Integer, types.Integer)(gcd_impl)
472