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

numba_cuda/numba/cuda/tests/cudapy/test_casting.py

@@ -0,0 +1,257 @@
+import numpy as np
+
+from numba.cuda import compile_ptx
+from numba.core.types import f2, i1, i2, i4, i8, u1, u2, u4, u8
+from numba import cuda
+from numba.core import types
+from numba.cuda.testing import (CUDATestCase, skip_on_cudasim,
+                                skip_unless_cc_53)
+from numba.types import float16, float32
+import itertools
+import unittest
+
+
+def native_cast(x):
+    return float(x)
+
+
+def to_int8(x):
+    return np.int8(x)
+
+
+def to_int16(x):
+    return np.int16(x)
+
+
+def to_int32(x):
+    return np.int32(x)
+
+
+def to_int64(x):
+    return np.int64(x)
+
+
+def to_uint8(x):
+    return np.uint8(x)
+
+
+def to_uint16(x):
+    return np.uint16(x)
+
+
+def to_uint32(x):
+    return types.uint32(x)
+
+
+def to_uint64(x):
+    return types.uint64(x)
+
+
+def to_float16(x):
+    # When division and operators on float16 types are supported, this should
+    # be changed to match the implementation in to_float32.
+    return (np.float16(x) * np.float16(0.5))
+
+
+def to_float32(x):
+    return np.float32(x) / np.float32(2)
+
+
+def to_float64(x):
+    return np.float64(x) / np.float64(2)
+
+
+def to_complex64(x):
+    return np.complex64(x)
+
+
+def to_complex128(x):
+    return np.complex128(x)
+
+
+# Since multiplication of float16 is not supported via the operator * on
+# float16s yet, and the host does not implement cuda.fp16.*, we need two
+# versions of the following functions:
+#
+# - The device version uses cuda.fp16.hmul
+# - The host version uses the * operator
+
+def cuda_int_literal_to_float16(x):
+    # Note that we need to use `2` and not `np.float16(2)` to ensure that this
+    # types as a literal int and not a const float16.
+    return cuda.fp16.hmul(np.float16(x), 2)
+
+
+def reference_int_literal_to_float16(x):
+    return np.float16(x) * np.float16(2)
+
+
+def cuda_float_literal_to_float16(x):
+    # Note that `2.5` types as a const float64 and not a literal float, but
+    # this case is provided in case that changes in future.
+    return cuda.fp16.hmul(np.float16(x), 2.5)
+
+
+def reference_float_literal_to_float16(x):
+    return np.float16(x) * np.float16(2.5)
+
+
+class TestCasting(CUDATestCase):
+    def _create_wrapped(self, pyfunc, intype, outtype):
+        wrapped_func = cuda.jit(device=True)(pyfunc)
+
+        @cuda.jit
+        def cuda_wrapper_fn(arg, res):
+            res[0] = wrapped_func(arg[0])
+
+        def wrapper_fn(arg):
+            argarray = np.zeros(1, dtype=intype)
+            argarray[0] = arg
+            resarray = np.zeros(1, dtype=outtype)
+            cuda_wrapper_fn[1, 1](argarray, resarray)
+            return resarray[0]
+
+        return wrapper_fn
+
+    @skip_unless_cc_53
+    def test_float_to_int(self):
+        pyfuncs = (to_int8, to_int16, to_int32, to_int64)
+        totys = (np.int8, np.int16, np.int32, np.int64)
+        fromtys = (np.float16, np.float32, np.float64)
+
+        for pyfunc, toty in zip(pyfuncs, totys):
+            for fromty in fromtys:
+                with self.subTest(fromty=fromty, toty=toty):
+                    cfunc = self._create_wrapped(pyfunc, fromty, toty)
+                    self.assertEqual(cfunc(12.3), pyfunc(12.3))
+                    self.assertEqual(cfunc(12.3), int(12.3))
+                    self.assertEqual(cfunc(-12.3), pyfunc(-12.3))
+                    self.assertEqual(cfunc(-12.3), int(-12.3))
+
+    @skip_on_cudasim('Compilation unsupported in the simulator')
+    def test_float16_to_int_ptx(self):
+        pyfuncs = (to_int8, to_int16, to_int32, to_int64)
+        sizes = (8, 16, 32, 64)
+
+        for pyfunc, size in zip(pyfuncs, sizes):
+            ptx, _ = compile_ptx(pyfunc, (f2,), device=True)
+            self.assertIn(f"cvt.rni.s{size}.f16", ptx)
+
+    @skip_unless_cc_53
+    def test_float_to_uint(self):
+        pyfuncs = (to_int8, to_int16, to_int32, to_int64)
+        totys = (np.uint8, np.uint16, np.uint32, np.uint64)
+        fromtys = (np.float16, np.float32, np.float64)
+
+        for pyfunc, toty in zip(pyfuncs, totys):
+            for fromty in fromtys:
+                with self.subTest(fromty=fromty, toty=toty):
+                    cfunc = self._create_wrapped(pyfunc, fromty, toty)
+                    self.assertEqual(cfunc(12.3), pyfunc(12.3))
+                    self.assertEqual(cfunc(12.3), int(12.3))
+
+    @skip_on_cudasim('Compilation unsupported in the simulator')
+    def test_float16_to_uint_ptx(self):
+        pyfuncs = (to_uint8, to_uint16, to_uint32, to_uint64)
+        sizes = (8, 16, 32, 64)
+
+        for pyfunc, size in zip(pyfuncs, sizes):
+            ptx, _ = compile_ptx(pyfunc, (f2,), device=True)
+            self.assertIn(f"cvt.rni.u{size}.f16", ptx)
+
+    @skip_unless_cc_53
+    def test_int_to_float(self):
+        pyfuncs = (to_float16, to_float32, to_float64)
+        totys = (np.float16, np.float32, np.float64)
+
+        for pyfunc, toty in zip(pyfuncs, totys):
+            with self.subTest(toty=toty):
+                cfunc = self._create_wrapped(pyfunc, np.int64, toty)
+                self.assertEqual(cfunc(321), pyfunc(321))
+
+    @skip_unless_cc_53
+    def test_literal_to_float16(self):
+        cudafuncs = (cuda_int_literal_to_float16,
+                     cuda_float_literal_to_float16)
+        hostfuncs = (reference_int_literal_to_float16,
+                     reference_float_literal_to_float16)
+
+        for cudafunc, hostfunc in zip(cudafuncs, hostfuncs):
+            with self.subTest(func=cudafunc):
+                cfunc = self._create_wrapped(cudafunc, np.float16, np.float16)
+                self.assertEqual(cfunc(321), hostfunc(321))
+
+    @skip_on_cudasim('Compilation unsupported in the simulator')
+    def test_int_to_float16_ptx(self):
+        fromtys = (i1, i2, i4, i8)
+        sizes = (8, 16, 32, 64)
+
+        for ty, size in zip(fromtys, sizes):
+            ptx, _ = compile_ptx(to_float16, (ty,), device=True)
+            self.assertIn(f"cvt.rn.f16.s{size}", ptx)
+
+    @skip_on_cudasim('Compilation unsupported in the simulator')
+    def test_uint_to_float16_ptx(self):
+        fromtys = (u1, u2, u4, u8)
+        sizes = (8, 16, 32, 64)
+
+        for ty, size in zip(fromtys, sizes):
+            ptx, _ = compile_ptx(to_float16, (ty,), device=True)
+            self.assertIn(f"cvt.rn.f16.u{size}", ptx)
+
+    @skip_unless_cc_53
+    def test_float_to_float(self):
+        pyfuncs = (to_float16, to_float32, to_float64)
+        tys = (np.float16, np.float32, np.float64)
+
+        for (pyfunc, fromty), toty in itertools.product(zip(pyfuncs, tys), tys):
+            with self.subTest(fromty=fromty, toty=toty):
+                cfunc = self._create_wrapped(pyfunc, fromty, toty)
+                # For this test we cannot use the pyfunc for comparison because
+                # the CUDA target doesn't yet implement division (or operators)
+                # for float16 values, so we test by comparing with the computed
+                # expression instead.
+                np.testing.assert_allclose(cfunc(12.3),
+                                           toty(12.3) / toty(2), rtol=0.0003)
+                np.testing.assert_allclose(cfunc(-12.3),
+                                           toty(-12.3) / toty(2), rtol=0.0003)
+
+    @skip_on_cudasim('Compilation unsupported in the simulator')
+    def test_float16_to_float_ptx(self):
+        pyfuncs = (to_float32, to_float64)
+        postfixes = ("f32", "f64")
+
+        for pyfunc, postfix in zip(pyfuncs, postfixes):
+            ptx, _ = compile_ptx(pyfunc, (f2,), device=True)
+            self.assertIn(f"cvt.{postfix}.f16", ptx)
+
+    @skip_unless_cc_53
+    def test_float_to_complex(self):
+        pyfuncs = (to_complex64, to_complex128)
+        totys = (np.complex64, np.complex128)
+        fromtys = (np.float16, np.float32, np.float64)
+
+        for pyfunc, toty in zip(pyfuncs, totys):
+            for fromty in fromtys:
+                with self.subTest(fromty=fromty, toty=toty):
+                    cfunc = self._create_wrapped(pyfunc, fromty, toty)
+                    # Here we need to explicitly cast the input to the pyfunc
+                    # to match the casting that is automatically applied when
+                    # passing the input to the cfunc as part of wrapping it in
+                    # an array of type fromtype.
+                    np.testing.assert_allclose(cfunc(3.21),
+                                               pyfunc(fromty(3.21)))
+                    np.testing.assert_allclose(cfunc(-3.21),
+                                               pyfunc(fromty(-3.21)) + 0j)
+
+    @skip_on_cudasim('Compilation unsupported in the simulator')
+    def test_native_cast(self):
+        float32_ptx, _ = cuda.compile_ptx(native_cast, (float32,), device=True)
+        self.assertIn("st.f32", float32_ptx)
+
+        float16_ptx, _ = cuda.compile_ptx(native_cast, (float16,), device=True)
+        self.assertIn("st.u16", float16_ptx)
+
+
+if __name__ == '__main__':
+    unittest.main()

numba_cuda/numba/cuda/tests/cudapy/test_cffi.py

@@ -0,0 +1,33 @@
+import numpy as np
+
+from numba import cuda, types
+from numba.cuda.testing import (skip_on_cudasim, test_data_dir, unittest,
+                                CUDATestCase)
+from numba.tests.support import skip_unless_cffi
+
+
+@skip_unless_cffi
+@skip_on_cudasim('Simulator does not support linking')
+class TestCFFI(CUDATestCase):
+    def test_from_buffer(self):
+        import cffi
+        ffi = cffi.FFI()
+
+        link = str(test_data_dir / 'jitlink.ptx')
+        sig = types.void(types.CPointer(types.int32))
+        array_mutator = cuda.declare_device('array_mutator', sig)
+
+        @cuda.jit(link=[link])
+        def mutate_array(x):
+            x_ptr = ffi.from_buffer(x)
+            array_mutator(x_ptr)
+
+        x = np.arange(2).astype(np.int32)
+        mutate_array[1, 1](x)
+
+        # The foreign function should have copied element 1 to element 0
+        self.assertEqual(x[0], x[1])
+
+
+if __name__ == '__main__':
+    unittest.main()

numba_cuda/numba/cuda/tests/cudapy/test_compiler.py

@@ -0,0 +1,276 @@
+from math import sqrt
+from numba import cuda, float32, int16, int32, int64, uint32, void
+from numba.cuda import (compile, compile_for_current_device, compile_ptx,
+                        compile_ptx_for_current_device)
+from numba.cuda.cudadrv import runtime
+from numba.cuda.testing import skip_on_cudasim, unittest, CUDATestCase
+
+
+# A test function at the module scope to ensure we get the name right for the C
+# ABI whether a function is at module or local scope.
+def f_module(x, y):
+    return x + y
+
+
+@skip_on_cudasim('Compilation unsupported in the simulator')
+class TestCompile(unittest.TestCase):
+    def test_global_kernel(self):
+        def f(r, x, y):
+            i = cuda.grid(1)
+            if i < len(r):
+                r[i] = x[i] + y[i]
+
+        args = (float32[:], float32[:], float32[:])
+        ptx, resty = compile_ptx(f, args)
+
+        # Kernels should not have a func_retval parameter
+        self.assertNotIn('func_retval', ptx)
+        # .visible .func is used to denote a device function
+        self.assertNotIn('.visible .func', ptx)
+        # .visible .entry would denote the presence of a global function
+        self.assertIn('.visible .entry', ptx)
+        # Return type for kernels should always be void
+        self.assertEqual(resty, void)
+
+    def test_device_function(self):
+        def add(x, y):
+            return x + y
+
+        args = (float32, float32)
+        ptx, resty = compile_ptx(add, args, device=True)
+
+        # Device functions take a func_retval parameter for storing the
+        # returned value in by reference
+        self.assertIn('func_retval', ptx)
+        # .visible .func is used to denote a device function
+        self.assertIn('.visible .func', ptx)
+        # .visible .entry would denote the presence of a global function
+        self.assertNotIn('.visible .entry', ptx)
+        # Inferred return type as expected?
+        self.assertEqual(resty, float32)
+
+        # Check that function's output matches signature
+        sig_int32 = int32(int32, int32)
+        ptx, resty = compile_ptx(add, sig_int32, device=True)
+        self.assertEqual(resty, int32)
+
+        sig_int16 = int16(int16, int16)
+        ptx, resty = compile_ptx(add, sig_int16, device=True)
+        self.assertEqual(resty, int16)
+        # Using string as signature
+        sig_string = "uint32(uint32, uint32)"
+        ptx, resty = compile_ptx(add, sig_string, device=True)
+        self.assertEqual(resty, uint32)
+
+    def test_fastmath(self):
+        def f(x, y, z, d):
+            return sqrt((x * y + z) / d)
+
+        args = (float32, float32, float32, float32)
+        ptx, resty = compile_ptx(f, args, device=True)
+
+        # Without fastmath, fma contraction is enabled by default, but ftz and
+        # approximate div / sqrt is not.
+        self.assertIn('fma.rn.f32', ptx)
+        self.assertIn('div.rn.f32', ptx)
+        self.assertIn('sqrt.rn.f32', ptx)
+
+        ptx, resty = compile_ptx(f, args, device=True, fastmath=True)
+
+        # With fastmath, ftz and approximate div / sqrt are enabled
+        self.assertIn('fma.rn.ftz.f32', ptx)
+        self.assertIn('div.approx.ftz.f32', ptx)
+        self.assertIn('sqrt.approx.ftz.f32', ptx)
+
+    def check_debug_info(self, ptx):
+        # A debug_info section should exist in the PTX. Whitespace varies
+        # between CUDA toolkit versions.
+        self.assertRegex(ptx, '\\.section\\s+\\.debug_info')
+        # A .file directive should be produced and include the name of the
+        # source. The path and whitespace may vary, so we accept anything
+        # ending in the filename of this module.
+        self.assertRegex(ptx, '\\.file.*test_compiler.py"')
+
+    def test_device_function_with_debug(self):
+        # See Issue #6719 - this ensures that compilation with debug succeeds
+        # with CUDA 11.2 / NVVM 7.0 onwards. Previously it failed because NVVM
+        # IR version metadata was not added when compiling device functions,
+        # and NVVM assumed DBG version 1.0 if not specified, which is
+        # incompatible with the 3.0 IR we use. This was specified only for
+        # kernels.
+        def f():
+            pass
+
+        ptx, resty = compile_ptx(f, (), device=True, debug=True)
+        self.check_debug_info(ptx)
+
+    def test_kernel_with_debug(self):
+        # Inspired by (but not originally affected by) Issue #6719
+        def f():
+            pass
+
+        ptx, resty = compile_ptx(f, (), debug=True)
+        self.check_debug_info(ptx)
+
+    def check_line_info(self, ptx):
+        # A .file directive should be produced and include the name of the
+        # source. The path and whitespace may vary, so we accept anything
+        # ending in the filename of this module.
+        self.assertRegex(ptx, '\\.file.*test_compiler.py"')
+
+    def test_device_function_with_line_info(self):
+        def f():
+            pass
+
+        ptx, resty = compile_ptx(f, (), device=True, lineinfo=True)
+        self.check_line_info(ptx)
+
+    def test_kernel_with_line_info(self):
+        def f():
+            pass
+
+        ptx, resty = compile_ptx(f, (), lineinfo=True)
+        self.check_line_info(ptx)
+
+    def test_non_void_return_type(self):
+        def f(x, y):
+            return x[0] + y[0]
+
+        with self.assertRaisesRegex(TypeError, 'must have void return type'):
+            compile_ptx(f, (uint32[::1], uint32[::1]))
+
+    def test_c_abi_disallowed_for_kernel(self):
+        def f(x, y):
+            return x + y
+
+        with self.assertRaisesRegex(NotImplementedError,
+                                    "The C ABI is not supported for kernels"):
+            compile_ptx(f, (int32, int32), abi="c")
+
+    def test_unsupported_abi(self):
+        def f(x, y):
+            return x + y
+
+        with self.assertRaisesRegex(NotImplementedError,
+                                    "Unsupported ABI: fastcall"):
+            compile_ptx(f, (int32, int32), abi="fastcall")
+
+    def test_c_abi_device_function(self):
+        def f(x, y):
+            return x + y
+
+        ptx, resty = compile_ptx(f, int32(int32, int32), device=True, abi="c")
+        # There should be no more than two parameters
+        self.assertNotIn(ptx, "param_2")
+
+        # The function name should match the Python function name (not the
+        # qualname, which includes additional info), and its return value
+        # should be 32 bits
+        self.assertRegex(ptx, r"\.visible\s+\.func\s+\(\.param\s+\.b32\s+"
+                              r"func_retval0\)\s+f\(")
+
+        # If we compile for 64-bit integers, the return type should be 64 bits
+        # wide
+        ptx, resty = compile_ptx(f, int64(int64, int64), device=True, abi="c")
+        self.assertRegex(ptx, r"\.visible\s+\.func\s+\(\.param\s+\.b64")
+
+    def test_c_abi_device_function_module_scope(self):
+        ptx, resty = compile_ptx(f_module, int32(int32, int32), device=True,
+                                 abi="c")
+
+        # The function name should match the Python function name, and its
+        # return value should be 32 bits
+        self.assertRegex(ptx, r"\.visible\s+\.func\s+\(\.param\s+\.b32\s+"
+                              r"func_retval0\)\s+f_module\(")
+
+    def test_c_abi_with_abi_name(self):
+        abi_info = {'abi_name': '_Z4funcii'}
+        ptx, resty = compile_ptx(f_module, int32(int32, int32), device=True,
+                                 abi="c", abi_info=abi_info)
+
+        # The function name should match the one given in the ABI info, and its
+        # return value should be 32 bits
+        self.assertRegex(ptx, r"\.visible\s+\.func\s+\(\.param\s+\.b32\s+"
+                              r"func_retval0\)\s+_Z4funcii\(")
+
+    def test_compile_defaults_to_c_abi(self):
+        ptx, resty = compile(f_module, int32(int32, int32), device=True)
+
+        # The function name should match the Python function name, and its
+        # return value should be 32 bits
+        self.assertRegex(ptx, r"\.visible\s+\.func\s+\(\.param\s+\.b32\s+"
+                              r"func_retval0\)\s+f_module\(")
+
+    def test_compile_to_ltoir(self):
+        if runtime.get_version() < (11, 5):
+            self.skipTest("-gen-lto unavailable in this toolkit version")
+
+        ltoir, resty = compile(f_module, int32(int32, int32), device=True,
+                               output="ltoir")
+
+        # There are no tools to interpret the LTOIR output, but we can check
+        # that we appear to have obtained an LTOIR file. This magic number is
+        # not documented, but is expected to remain consistent.
+        LTOIR_MAGIC = 0x7F4E43ED
+        header = int.from_bytes(ltoir[:4], byteorder='little')
+        self.assertEqual(header, LTOIR_MAGIC)
+        self.assertEqual(resty, int32)
+
+    def test_compile_to_invalid_error(self):
+        illegal_output = "illegal"
+        msg = f"Unsupported output type: {illegal_output}"
+        with self.assertRaisesRegex(NotImplementedError, msg):
+            compile(f_module, int32(int32, int32), device=True,
+                    output=illegal_output)
+
+
+@skip_on_cudasim('Compilation unsupported in the simulator')
+class TestCompileForCurrentDevice(CUDATestCase):
+    def _check_ptx_for_current_device(self, compile_function):
+        def add(x, y):
+            return x + y
+
+        args = (float32, float32)
+        ptx, resty = compile_function(add, args, device=True)
+
+        # Check we target the current device's compute capability, or the
+        # closest compute capability supported by the current toolkit.
+        device_cc = cuda.get_current_device().compute_capability
+        cc = cuda.cudadrv.nvvm.find_closest_arch(device_cc)
+        target = f'.target sm_{cc[0]}{cc[1]}'
+        self.assertIn(target, ptx)
+
+    def test_compile_ptx_for_current_device(self):
+        self._check_ptx_for_current_device(compile_ptx_for_current_device)
+
+    def test_compile_for_current_device(self):
+        self._check_ptx_for_current_device(compile_for_current_device)
+
+
+@skip_on_cudasim('Compilation unsupported in the simulator')
+class TestCompileOnlyTests(unittest.TestCase):
+    '''For tests where we can only check correctness by examining the compiler
+    output rather than observing the effects of execution.'''
+
+    def test_nanosleep(self):
+        def use_nanosleep(x):
+            # Sleep for a constant time
+            cuda.nanosleep(32)
+            # Sleep for a variable time
+            cuda.nanosleep(x)
+
+        ptx, resty = compile_ptx(use_nanosleep, (uint32,), cc=(7, 0))
+
+        nanosleep_count = 0
+        for line in ptx.split('\n'):
+            if 'nanosleep.u32' in line:
+                nanosleep_count += 1
+
+        expected = 2
+        self.assertEqual(expected, nanosleep_count,
+                         (f'Got {nanosleep_count} nanosleep instructions, '
+                          f'expected {expected}'))
+
+
+if __name__ == '__main__':
+    unittest.main()