156 lines
4.7 KiB
Python
156 lines
4.7 KiB
Python
"""
|
|
Register external C functions necessary for Numba code generation.
|
|
"""
|
|
|
|
import sys
|
|
|
|
from llvmlite import ir
|
|
import llvmlite.binding as ll
|
|
|
|
from numba.core import utils, intrinsics
|
|
from numba import _helperlib
|
|
|
|
|
|
def _add_missing_symbol(symbol, addr):
|
|
"""Add missing symbol into LLVM internal symtab
|
|
"""
|
|
if not ll.address_of_symbol(symbol):
|
|
ll.add_symbol(symbol, addr)
|
|
|
|
|
|
def _get_msvcrt_symbol(symbol):
|
|
"""
|
|
Under Windows, look up a symbol inside the C runtime
|
|
and return the raw pointer value as an integer.
|
|
"""
|
|
from ctypes import cdll, cast, c_void_p
|
|
f = getattr(cdll.msvcrt, symbol)
|
|
return cast(f, c_void_p).value
|
|
|
|
|
|
def compile_multi3(context):
|
|
"""
|
|
Compile the multi3() helper function used by LLVM
|
|
for 128-bit multiplication on 32-bit platforms.
|
|
"""
|
|
codegen = context.codegen()
|
|
library = codegen.create_library("multi3")
|
|
|
|
ir_mod = library.create_ir_module("multi3")
|
|
|
|
i64 = ir.IntType(64)
|
|
i128 = ir.IntType(128)
|
|
lower_mask = ir.Constant(i64, 0xffffffff)
|
|
_32 = ir.Constant(i64, 32)
|
|
_64 = ir.Constant(i128, 64)
|
|
|
|
fn_type = ir.FunctionType(i128, [i128, i128])
|
|
fn = ir.Function(ir_mod, fn_type, name="multi3")
|
|
|
|
a, b = fn.args
|
|
bb = fn.append_basic_block()
|
|
builder = ir.IRBuilder(bb)
|
|
|
|
# This implementation mimics compiler-rt's.
|
|
al = builder.trunc(a, i64)
|
|
bl = builder.trunc(b, i64)
|
|
ah = builder.trunc(builder.ashr(a, _64), i64)
|
|
bh = builder.trunc(builder.ashr(b, _64), i64)
|
|
|
|
# Compute {rh, rl} = al * bl (unsigned 64-bit multiplication)
|
|
# rl = (al & 0xffffffff) * (bl & 0xffffffff)
|
|
rl = builder.mul(builder.and_(al, lower_mask), builder.and_(bl, lower_mask))
|
|
# t = rl >> 32
|
|
t = builder.lshr(rl, _32)
|
|
# rl &= 0xffffffff
|
|
rl = builder.and_(rl, lower_mask)
|
|
# t += (al >> 32) * (bl & 0xffffffff)
|
|
t = builder.add(t, builder.mul(builder.lshr(al, _32),
|
|
builder.and_(bl, lower_mask)))
|
|
# rl += t << 32
|
|
rl = builder.add(rl, builder.shl(t, _32))
|
|
# rh = t >> 32
|
|
rh = builder.lshr(t, _32)
|
|
# t = rl >> 32
|
|
t = builder.lshr(rl, _32)
|
|
# rl &= 0xffffffff
|
|
rl = builder.and_(rl, lower_mask)
|
|
# t += (bl >> 32) * (al & 0xffffffff)
|
|
t = builder.add(t, builder.mul(builder.lshr(bl, _32),
|
|
builder.and_(al, lower_mask)))
|
|
# rl += t << 32
|
|
rl = builder.add(rl, builder.shl(t, _32))
|
|
# rh += t >> 32
|
|
rh = builder.add(rh, builder.lshr(t, _32))
|
|
# rh += (al >> 32) * (bl >> 32)
|
|
rh = builder.add(rh, builder.mul(builder.lshr(al, _32),
|
|
builder.lshr(bl, _32)))
|
|
|
|
# rh += (bh * al) + (bl * ah)
|
|
rh = builder.add(rh, builder.mul(bh, al))
|
|
rh = builder.add(rh, builder.mul(bl, ah))
|
|
|
|
# r = rl + (rh << 64)
|
|
r = builder.zext(rl, i128)
|
|
r = builder.add(r, builder.shl(builder.zext(rh, i128), _64))
|
|
builder.ret(r)
|
|
|
|
library.add_ir_module(ir_mod)
|
|
library.finalize()
|
|
|
|
return library
|
|
|
|
|
|
class _Installer(object):
|
|
|
|
_installed = False
|
|
|
|
def install(self, context):
|
|
"""
|
|
Install the functions into LLVM. This only needs to be done once,
|
|
as the mappings are persistent during the process lifetime.
|
|
"""
|
|
if not self._installed:
|
|
self._do_install(context)
|
|
self._installed = True
|
|
|
|
|
|
class _ExternalMathFunctions(_Installer):
|
|
"""
|
|
Map the math functions from the C runtime library into the LLVM
|
|
execution environment.
|
|
"""
|
|
|
|
def _do_install(self, context):
|
|
is32bit = utils.MACHINE_BITS == 32
|
|
c_helpers = _helperlib.c_helpers
|
|
|
|
if sys.platform.startswith('win32') and is32bit:
|
|
# For Windows XP _ftol2 is not defined, we will just use
|
|
# _ftol as a replacement.
|
|
# On Windows 7, this is not necessary but will work anyway.
|
|
ftol = _get_msvcrt_symbol("_ftol")
|
|
_add_missing_symbol("_ftol2", ftol)
|
|
|
|
elif sys.platform.startswith('linux') and is32bit:
|
|
_add_missing_symbol("__fixunsdfdi", c_helpers["fptoui"])
|
|
_add_missing_symbol("__fixunssfdi", c_helpers["fptouif"])
|
|
|
|
if is32bit:
|
|
# Make the library immortal
|
|
self._multi3_lib = compile_multi3(context)
|
|
ptr = self._multi3_lib.get_pointer_to_function("multi3")
|
|
assert ptr
|
|
_add_missing_symbol("__multi3", ptr)
|
|
|
|
# List available C-math
|
|
for fname in intrinsics.INTR_MATH:
|
|
# Force binding from CPython's C runtime library.
|
|
# (under Windows, different versions of the C runtime can
|
|
# be loaded at the same time, for example msvcrt100 by
|
|
# CPython and msvcrt120 by LLVM)
|
|
ll.add_symbol(fname, c_helpers[fname])
|
|
|
|
|
|
c_math_functions = _ExternalMathFunctions()
|