numba-cuda 0.0.1__py3-none-any.whl → 0.0.13__py3-none-any.whl

numba_cuda/numba/cuda/libdeviceimpl.py

@@ -0,0 +1,83 @@
+from llvmlite import ir
+from numba.core import cgutils, types
+from numba.core.imputils import Registry
+from numba.cuda import libdevice, libdevicefuncs
+
+registry = Registry()
+lower = registry.lower
+
+
+def libdevice_implement(func, retty, nbargs):
+    def core(context, builder, sig, args):
+        lmod = builder.module
+        fretty = context.get_value_type(retty)
+        fargtys = [context.get_value_type(arg.ty) for arg in nbargs]
+        fnty = ir.FunctionType(fretty, fargtys)
+        fn = cgutils.get_or_insert_function(lmod, fnty, func)
+        return builder.call(fn, args)
+
+    key = getattr(libdevice, func[5:])
+
+    argtys = [arg.ty for arg in args if not arg.is_ptr]
+    lower(key, *argtys)(core)
+
+
+def libdevice_implement_multiple_returns(func, retty, prototype_args):
+    sig = libdevicefuncs.create_signature(retty, prototype_args)
+    nb_retty = sig.return_type
+
+    def core(context, builder, sig, args):
+        lmod = builder.module
+
+        fargtys = []
+        for arg in prototype_args:
+            ty = context.get_value_type(arg.ty)
+            if arg.is_ptr:
+                ty = ty.as_pointer()
+            fargtys.append(ty)
+
+        fretty = context.get_value_type(retty)
+
+        fnty = ir.FunctionType(fretty, fargtys)
+        fn = cgutils.get_or_insert_function(lmod, fnty, func)
+
+        # For returned values that are returned through a pointer, we need to
+        # allocate variables on the stack and pass a pointer to them.
+        actual_args = []
+        virtual_args = []
+        arg_idx = 0
+        for arg in prototype_args:
+            if arg.is_ptr:
+                # Allocate space for return value and add to args
+                tmp_arg = cgutils.alloca_once(builder,
+                                              context.get_value_type(arg.ty))
+                actual_args.append(tmp_arg)
+                virtual_args.append(tmp_arg)
+            else:
+                actual_args.append(args[arg_idx])
+                arg_idx += 1
+
+        ret = builder.call(fn, actual_args)
+
+        # Following the call, we need to assemble the returned values into a
+        # tuple for returning back to the caller.
+        tuple_args = []
+        if retty != types.void:
+            tuple_args.append(ret)
+        for arg in virtual_args:
+            tuple_args.append(builder.load(arg))
+
+        if isinstance(nb_retty, types.UniTuple):
+            return cgutils.pack_array(builder, tuple_args)
+        else:
+            return cgutils.pack_struct(builder, tuple_args)
+
+    key = getattr(libdevice, func[5:])
+    lower(key, *sig.args)(core)
+
+
+for func, (retty, args) in libdevicefuncs.functions.items():
+    if any([arg.is_ptr for arg in args]):
+        libdevice_implement_multiple_returns(func, retty, args)
+    else:
+        libdevice_implement(func, retty, args)

numba_cuda/numba/cuda/mathimpl.py

@@ -0,0 +1,448 @@
+import math
+import operator
+from llvmlite import ir
+from numba.core import types, typing, cgutils, targetconfig
+from numba.core.imputils import Registry
+from numba.types import float32, float64, int64, uint64
+from numba.cuda import libdevice
+from numba import cuda
+
+registry = Registry()
+lower = registry.lower
+
+
+booleans = []
+booleans += [('isnand', 'isnanf', math.isnan)]
+booleans += [('isinfd', 'isinff', math.isinf)]
+booleans += [('isfinited', 'finitef', math.isfinite)]
+
+unarys = []
+unarys += [('ceil', 'ceilf', math.ceil)]
+unarys += [('floor', 'floorf', math.floor)]
+unarys += [('fabs', 'fabsf', math.fabs)]
+unarys += [('exp', 'expf', math.exp)]
+unarys += [('expm1', 'expm1f', math.expm1)]
+unarys += [('erf', 'erff', math.erf)]
+unarys += [('erfc', 'erfcf', math.erfc)]
+unarys += [('tgamma', 'tgammaf', math.gamma)]
+unarys += [('lgamma', 'lgammaf', math.lgamma)]
+unarys += [('sqrt', 'sqrtf', math.sqrt)]
+unarys += [('log', 'logf', math.log)]
+unarys += [('log2', 'log2f', math.log2)]
+unarys += [('log10', 'log10f', math.log10)]
+unarys += [('log1p', 'log1pf', math.log1p)]
+unarys += [('acosh', 'acoshf', math.acosh)]
+unarys += [('acos', 'acosf', math.acos)]
+unarys += [('cos', 'cosf', math.cos)]
+unarys += [('cosh', 'coshf', math.cosh)]
+unarys += [('asinh', 'asinhf', math.asinh)]
+unarys += [('asin', 'asinf', math.asin)]
+unarys += [('sin', 'sinf', math.sin)]
+unarys += [('sinh', 'sinhf', math.sinh)]
+unarys += [('atan', 'atanf', math.atan)]
+unarys += [('atanh', 'atanhf', math.atanh)]
+unarys += [('tan', 'tanf', math.tan)]
+unarys += [('trunc', 'truncf', math.trunc)]
+
+unarys_fastmath = {}
+unarys_fastmath['cosf'] = 'fast_cosf'
+unarys_fastmath['sinf'] = 'fast_sinf'
+unarys_fastmath['tanf'] = 'fast_tanf'
+unarys_fastmath['expf'] = 'fast_expf'
+unarys_fastmath['log2f'] = 'fast_log2f'
+unarys_fastmath['log10f'] = 'fast_log10f'
+unarys_fastmath['logf'] = 'fast_logf'
+
+binarys = []
+binarys += [('copysign', 'copysignf', math.copysign)]
+binarys += [('atan2', 'atan2f', math.atan2)]
+binarys += [('pow', 'powf', math.pow)]
+binarys += [('fmod', 'fmodf', math.fmod)]
+binarys += [('hypot', 'hypotf', math.hypot)]
+binarys += [('remainder', 'remainderf', math.remainder)]
+
+binarys_fastmath = {}
+binarys_fastmath['powf'] = 'fast_powf'
+
+
+@lower(math.isinf, types.Integer)
+@lower(math.isnan, types.Integer)
+def math_isinf_isnan_int(context, builder, sig, args):
+    return context.get_constant(types.boolean, 0)
+
+
+@lower(operator.truediv, types.float32, types.float32)
+def maybe_fast_truediv(context, builder, sig, args):
+    if context.fastmath:
+        sig = typing.signature(float32, float32, float32)
+        impl = context.get_function(libdevice.fast_fdividef, sig)
+        return impl(builder, args)
+    else:
+        with cgutils.if_zero(builder, args[1]):
+            context.error_model.fp_zero_division(builder, ("division by zero",))
+        res = builder.fdiv(*args)
+        return res
+
+
+@lower(math.isfinite, types.Integer)
+def math_isfinite_int(context, builder, sig, args):
+    return context.get_constant(types.boolean, 1)
+
+
+@lower(math.sin, types.float16)
+def fp16_sin_impl(context, builder, sig, args):
+    def fp16_sin(x):
+        return cuda.fp16.hsin(x)
+
+    return context.compile_internal(builder, fp16_sin, sig, args)
+
+
+@lower(math.cos, types.float16)
+def fp16_cos_impl(context, builder, sig, args):
+    def fp16_cos(x):
+        return cuda.fp16.hcos(x)
+
+    return context.compile_internal(builder, fp16_cos, sig, args)
+
+
+@lower(math.log, types.float16)
+def fp16_log_impl(context, builder, sig, args):
+    def fp16_log(x):
+        return cuda.fp16.hlog(x)
+
+    return context.compile_internal(builder, fp16_log, sig, args)
+
+
+@lower(math.log10, types.float16)
+def fp16_log10_impl(context, builder, sig, args):
+    def fp16_log10(x):
+        return cuda.fp16.hlog10(x)
+
+    return context.compile_internal(builder, fp16_log10, sig, args)
+
+
+@lower(math.log2, types.float16)
+def fp16_log2_impl(context, builder, sig, args):
+    def fp16_log2(x):
+        return cuda.fp16.hlog2(x)
+
+    return context.compile_internal(builder, fp16_log2, sig, args)
+
+
+@lower(math.exp, types.float16)
+def fp16_exp_impl(context, builder, sig, args):
+    def fp16_exp(x):
+        return cuda.fp16.hexp(x)
+
+    return context.compile_internal(builder, fp16_exp, sig, args)
+
+
+@lower(math.floor, types.float16)
+def fp16_floor_impl(context, builder, sig, args):
+    def fp16_floor(x):
+        return cuda.fp16.hfloor(x)
+
+    return context.compile_internal(builder, fp16_floor, sig, args)
+
+
+@lower(math.ceil, types.float16)
+def fp16_ceil_impl(context, builder, sig, args):
+    def fp16_ceil(x):
+        return cuda.fp16.hceil(x)
+
+    return context.compile_internal(builder, fp16_ceil, sig, args)
+
+
+@lower(math.sqrt, types.float16)
+def fp16_sqrt_impl(context, builder, sig, args):
+    def fp16_sqrt(x):
+        return cuda.fp16.hsqrt(x)
+
+    return context.compile_internal(builder, fp16_sqrt, sig, args)
+
+
+@lower(math.fabs, types.float16)
+def fp16_fabs_impl(context, builder, sig, args):
+    def fp16_fabs(x):
+        return cuda.fp16.habs(x)
+
+    return context.compile_internal(builder, fp16_fabs, sig, args)
+
+
+@lower(math.trunc, types.float16)
+def fp16_trunc_impl(context, builder, sig, args):
+    def fp16_trunc(x):
+        return cuda.fp16.htrunc(x)
+
+    return context.compile_internal(builder, fp16_trunc, sig, args)
+
+
+def impl_boolean(key, ty, libfunc):
+    def lower_boolean_impl(context, builder, sig, args):
+        libfunc_impl = context.get_function(libfunc,
+                                            typing.signature(types.int32, ty))
+        result = libfunc_impl(builder, args)
+        return context.cast(builder, result, types.int32, types.boolean)
+
+    lower(key, ty)(lower_boolean_impl)
+
+
+def get_lower_unary_impl(key, ty, libfunc):
+    def lower_unary_impl(context, builder, sig, args):
+        actual_libfunc = libfunc
+        fast_replacement = None
+        if ty == float32 and context.fastmath:
+            fast_replacement = unarys_fastmath.get(libfunc.__name__)
+
+        if fast_replacement is not None:
+            actual_libfunc = getattr(libdevice, fast_replacement)
+
+        libfunc_impl = context.get_function(actual_libfunc,
+                                            typing.signature(ty, ty))
+        return libfunc_impl(builder, args)
+    return lower_unary_impl
+
+
+def get_unary_impl_for_fn_and_ty(fn, ty):
+    # tanh is a special case - because it is not registered like the other
+    # unary implementations, it does not appear in the unarys list. However,
+    # its implementation can be looked up by key like the other
+    # implementations, so we add it to the list we search here.
+    tanh_impls = ('tanh', 'tanhf', math.tanh)
+    for fname64, fname32, key in unarys + [tanh_impls]:
+        if fn == key:
+            if ty == float32:
+                impl = getattr(libdevice, fname32)
+            elif ty == float64:
+                impl = getattr(libdevice, fname64)
+
+            return get_lower_unary_impl(key, ty, impl)
+
+    raise RuntimeError(f"Implementation of {fn} for {ty} not found")
+
+
+def impl_unary(key, ty, libfunc):
+    lower_unary_impl = get_lower_unary_impl(key, ty, libfunc)
+    lower(key, ty)(lower_unary_impl)
+
+
+def impl_unary_int(key, ty, libfunc):
+    def lower_unary_int_impl(context, builder, sig, args):
+        if sig.args[0] == int64:
+            convert = builder.sitofp
+        elif sig.args[0] == uint64:
+            convert = builder.uitofp
+        else:
+            m = 'Only 64-bit integers are supported for generic unary int ops'
+            raise TypeError(m)
+
+        arg = convert(args[0], ir.DoubleType())
+        sig = typing.signature(float64, float64)
+        libfunc_impl = context.get_function(libfunc, sig)
+        return libfunc_impl(builder, [arg])
+
+    lower(key, ty)(lower_unary_int_impl)
+
+
+def get_lower_binary_impl(key, ty, libfunc):
+    def lower_binary_impl(context, builder, sig, args):
+        actual_libfunc = libfunc
+        fast_replacement = None
+        if ty == float32 and context.fastmath:
+            fast_replacement = binarys_fastmath.get(libfunc.__name__)
+
+        if fast_replacement is not None:
+            actual_libfunc = getattr(libdevice, fast_replacement)
+
+        libfunc_impl = context.get_function(actual_libfunc,
+                                            typing.signature(ty, ty, ty))
+        return libfunc_impl(builder, args)
+    return lower_binary_impl
+
+
+def get_binary_impl_for_fn_and_ty(fn, ty):
+    for fname64, fname32, key in binarys:
+        if fn == key:
+            if ty == float32:
+                impl = getattr(libdevice, fname32)
+            elif ty == float64:
+                impl = getattr(libdevice, fname64)
+
+            return get_lower_binary_impl(key, ty, impl)
+
+    raise RuntimeError(f"Implementation of {fn} for {ty} not found")
+
+
+def impl_binary(key, ty, libfunc):
+    lower_binary_impl = get_lower_binary_impl(key, ty, libfunc)
+    lower(key, ty, ty)(lower_binary_impl)
+
+
+def impl_binary_int(key, ty, libfunc):
+    def lower_binary_int_impl(context, builder, sig, args):
+        if sig.args[0] == int64:
+            convert = builder.sitofp
+        elif sig.args[0] == uint64:
+            convert = builder.uitofp
+        else:
+            m = 'Only 64-bit integers are supported for generic binary int ops'
+            raise TypeError(m)
+
+        args = [convert(arg, ir.DoubleType()) for arg in args]
+        sig = typing.signature(float64, float64, float64)
+        libfunc_impl = context.get_function(libfunc, sig)
+        return libfunc_impl(builder, args)
+
+    lower(key, ty, ty)(lower_binary_int_impl)
+
+
+for fname64, fname32, key in booleans:
+    impl32 = getattr(libdevice, fname32)
+    impl64 = getattr(libdevice, fname64)
+    impl_boolean(key, float32, impl32)
+    impl_boolean(key, float64, impl64)
+
+
+for fname64, fname32, key in unarys:
+    impl32 = getattr(libdevice, fname32)
+    impl64 = getattr(libdevice, fname64)
+    impl_unary(key, float32, impl32)
+    impl_unary(key, float64, impl64)
+    impl_unary_int(key, int64, impl64)
+    impl_unary_int(key, uint64, impl64)
+
+
+for fname64, fname32, key in binarys:
+    impl32 = getattr(libdevice, fname32)
+    impl64 = getattr(libdevice, fname64)
+    impl_binary(key, float32, impl32)
+    impl_binary(key, float64, impl64)
+    impl_binary_int(key, int64, impl64)
+    impl_binary_int(key, uint64, impl64)
+
+
+def impl_pow_int(ty, libfunc):
+    def lower_pow_impl_int(context, builder, sig, args):
+        powi_sig = typing.signature(ty, ty, types.int32)
+        libfunc_impl = context.get_function(libfunc, powi_sig)
+        return libfunc_impl(builder, args)
+
+    lower(math.pow, ty, types.int32)(lower_pow_impl_int)
+
+
+impl_pow_int(types.float32, libdevice.powif)
+impl_pow_int(types.float64, libdevice.powi)
+
+
+def impl_modf(ty, libfunc):
+    retty = types.UniTuple(ty, 2)
+
+    def lower_modf_impl(context, builder, sig, args):
+        modf_sig = typing.signature(retty, ty)
+        libfunc_impl = context.get_function(libfunc, modf_sig)
+        return libfunc_impl(builder, args)
+
+    lower(math.modf, ty)(lower_modf_impl)
+
+
+impl_modf(types.float32, libdevice.modff)
+impl_modf(types.float64, libdevice.modf)
+
+
+def impl_frexp(ty, libfunc):
+    retty = types.Tuple((ty, types.int32))
+
+    def lower_frexp_impl(context, builder, sig, args):
+        frexp_sig = typing.signature(retty, ty)
+        libfunc_impl = context.get_function(libfunc, frexp_sig)
+        return libfunc_impl(builder, args)
+
+    lower(math.frexp, ty)(lower_frexp_impl)
+
+
+impl_frexp(types.float32, libdevice.frexpf)
+impl_frexp(types.float64, libdevice.frexp)
+
+
+def impl_ldexp(ty, libfunc):
+    def lower_ldexp_impl(context, builder, sig, args):
+        ldexp_sig = typing.signature(ty, ty, types.int32)
+        libfunc_impl = context.get_function(libfunc, ldexp_sig)
+        return libfunc_impl(builder, args)
+
+    lower(math.ldexp, ty, types.int32)(lower_ldexp_impl)
+
+
+impl_ldexp(types.float32, libdevice.ldexpf)
+impl_ldexp(types.float64, libdevice.ldexp)
+
+
+def impl_tanh(ty, libfunc):
+    def lower_tanh_impl(context, builder, sig, args):
+        def get_compute_capability():
+            flags = targetconfig.ConfigStack().top()
+            return flags.compute_capability
+
+        def tanh_impl_libdevice():
+            tanh_sig = typing.signature(ty, ty)
+            libfunc_impl = context.get_function(libfunc, tanh_sig)
+            return libfunc_impl(builder, args)
+
+        def tanhf_impl_fastmath():
+            fnty = ir.FunctionType(ir.FloatType(), [ir.FloatType()])
+            asm = ir.InlineAsm(fnty, 'tanh.approx.f32 $0, $1;', '=f,f')
+            return builder.call(asm, args)
+
+        if ty == float32 and context.fastmath:
+            cc = get_compute_capability()
+            if cc >= (7,5):
+                return tanhf_impl_fastmath()
+
+        return tanh_impl_libdevice()
+
+    lower(math.tanh, ty)(lower_tanh_impl)
+
+
+impl_tanh(types.float32, libdevice.tanhf)
+impl_tanh(types.float64, libdevice.tanh)
+
+impl_unary_int(math.tanh, int64, libdevice.tanh)
+impl_unary_int(math.tanh, uint64, libdevice.tanh)
+
+# Complex power implementations - translations of _Py_c_pow from CPython
+# https://github.com/python/cpython/blob/a755410e054e1e2390de5830befc08fe80706c66/Objects/complexobject.c#L123-L151
+#
+# The complex64 variant casts all constants and some variables to ensure that
+# as much computation is done in single precision as possible. A small number
+# of operations are still done in 64-bit, but these come from libdevice code.
+
+
+def cpow_implement(fty, cty):
+    def core(context, builder, sig, args):
+        def cpow_internal(a, b):
+
+            if b.real == fty(0.0) and b.imag == fty(0.0):
+                return cty(1.0) + cty(0.0j)
+            elif a.real == fty(0.0) and b.real == fty(0.0):
+                return cty(0.0) + cty(0.0j)
+
+            vabs = math.hypot(a.real, a.imag)
+            len = math.pow(vabs, b.real)
+            at = math.atan2(a.imag, a.real)
+            phase = at * b.real
+            if b.imag != fty(0.0):
+                len /= math.exp(at * b.imag)
+                phase += b.imag * math.log(vabs)
+
+            return len * (cty(math.cos(phase)) +
+                          cty(math.sin(phase) * cty(1.0j)))
+
+        return context.compile_internal(builder, cpow_internal, sig, args)
+
+    lower(operator.pow, cty, cty)(core)
+    lower(operator.ipow, cty, cty)(core)
+    lower(pow, cty, cty)(core)
+
+
+cpow_implement(types.float32, types.complex64)
+cpow_implement(types.float64, types.complex128)

numba_cuda/numba/cuda/models.py

@@ -0,0 +1,48 @@
+import functools
+
+from llvmlite import ir
+
+from numba.core.datamodel.registry import DataModelManager, register
+from numba.core.extending import models
+from numba.core import types
+from numba.cuda.types import Dim3, GridGroup, CUDADispatcher
+
+
+cuda_data_manager = DataModelManager()
+
+register_model = functools.partial(register, cuda_data_manager)
+
+
+@register_model(Dim3)
+class Dim3Model(models.StructModel):
+    def __init__(self, dmm, fe_type):
+        members = [
+            ('x', types.int32),
+            ('y', types.int32),
+            ('z', types.int32)
+        ]
+        super().__init__(dmm, fe_type, members)
+
+
+@register_model(GridGroup)
+class GridGroupModel(models.PrimitiveModel):
+    def __init__(self, dmm, fe_type):
+        be_type = ir.IntType(64)
+        super().__init__(dmm, fe_type, be_type)
+
+
+@register_model(types.Float)
+class FloatModel(models.PrimitiveModel):
+    def __init__(self, dmm, fe_type):
+        if fe_type == types.float16:
+            be_type = ir.IntType(16)
+        elif fe_type == types.float32:
+            be_type = ir.FloatType()
+        elif fe_type == types.float64:
+            be_type = ir.DoubleType()
+        else:
+            raise NotImplementedError(fe_type)
+        super(FloatModel, self).__init__(dmm, fe_type, be_type)
+
+
+register_model(CUDADispatcher)(models.OpaqueModel)