numba-cuda 0.0.0__py3-none-any.whl → 0.0.12__py3-none-any.whl

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

@@ -0,0 +1,318 @@
+import multiprocessing as mp
+import itertools
+import traceback
+import pickle
+
+import numpy as np
+
+from numba import cuda
+from numba.cuda.cudadrv import driver
+from numba.cuda.testing import (skip_on_arm, skip_on_cudasim,
+                                skip_under_cuda_memcheck,
+                                ContextResettingTestCase, ForeignArray)
+from numba.tests.support import linux_only, windows_only
+import unittest
+
+
+def core_ipc_handle_test(the_work, result_queue):
+    try:
+        arr = the_work()
+    # Catch anything going wrong in the worker function
+    except:  # noqa: E722
+        # FAILED. propagate the exception as a string
+        succ = False
+        out = traceback.format_exc()
+    else:
+        # OK. send the ndarray back
+        succ = True
+        out = arr
+    result_queue.put((succ, out))
+
+
+def base_ipc_handle_test(handle, size, result_queue):
+    def the_work():
+        dtype = np.dtype(np.intp)
+        with cuda.open_ipc_array(handle, shape=size // dtype.itemsize,
+                                 dtype=dtype) as darr:
+            # copy the data to host
+            return darr.copy_to_host()
+
+    core_ipc_handle_test(the_work, result_queue)
+
+
+def serialize_ipc_handle_test(handle, result_queue):
+    def the_work():
+        dtype = np.dtype(np.intp)
+        darr = handle.open_array(cuda.current_context(),
+                                 shape=handle.size // dtype.itemsize,
+                                 dtype=dtype)
+        # copy the data to host
+        arr = darr.copy_to_host()
+        handle.close()
+        return arr
+
+    core_ipc_handle_test(the_work, result_queue)
+
+
+def ipc_array_test(ipcarr, result_queue):
+    try:
+        with ipcarr as darr:
+            arr = darr.copy_to_host()
+            try:
+                # should fail to reopen
+                with ipcarr:
+                    pass
+            except ValueError as e:
+                if str(e) != 'IpcHandle is already opened':
+                    raise AssertionError('invalid exception message')
+            else:
+                raise AssertionError('did not raise on reopen')
+    # Catch any exception so we can propagate it
+    except:  # noqa: E722
+        # FAILED. propagate the exception as a string
+        succ = False
+        out = traceback.format_exc()
+    else:
+        # OK. send the ndarray back
+        succ = True
+        out = arr
+    result_queue.put((succ, out))
+
+
+@linux_only
+@skip_under_cuda_memcheck('Hangs cuda-memcheck')
+@skip_on_cudasim('Ipc not available in CUDASIM')
+@skip_on_arm('CUDA IPC not supported on ARM in Numba')
+class TestIpcMemory(ContextResettingTestCase):
+
+    def test_ipc_handle(self):
+        # prepare data for IPC
+        arr = np.arange(10, dtype=np.intp)
+        devarr = cuda.to_device(arr)
+
+        # create IPC handle
+        ctx = cuda.current_context()
+        ipch = ctx.get_ipc_handle(devarr.gpu_data)
+
+        # manually prepare for serialization as bytes
+        if driver.USE_NV_BINDING:
+            handle_bytes = ipch.handle.reserved
+        else:
+            handle_bytes = bytes(ipch.handle)
+        size = ipch.size
+
+        # spawn new process for testing
+        ctx = mp.get_context('spawn')
+        result_queue = ctx.Queue()
+        args = (handle_bytes, size, result_queue)
+        proc = ctx.Process(target=base_ipc_handle_test, args=args)
+        proc.start()
+        succ, out = result_queue.get()
+        if not succ:
+            self.fail(out)
+        else:
+            np.testing.assert_equal(arr, out)
+        proc.join(3)
+
+    def variants(self):
+        # Test with no slicing and various different slices
+        indices = (None, slice(3, None), slice(3, 8), slice(None, 8))
+        # Test with a Numba DeviceNDArray, or an array from elsewhere through
+        # the CUDA Array Interface
+        foreigns = (False, True)
+        return itertools.product(indices, foreigns)
+
+    def check_ipc_handle_serialization(self, index_arg=None, foreign=False):
+        # prepare data for IPC
+        arr = np.arange(10, dtype=np.intp)
+        devarr = cuda.to_device(arr)
+        if index_arg is not None:
+            devarr = devarr[index_arg]
+        if foreign:
+            devarr = cuda.as_cuda_array(ForeignArray(devarr))
+        expect = devarr.copy_to_host()
+
+        # create IPC handle
+        ctx = cuda.current_context()
+        ipch = ctx.get_ipc_handle(devarr.gpu_data)
+
+        # pickle
+        buf = pickle.dumps(ipch)
+        ipch_recon = pickle.loads(buf)
+        self.assertIs(ipch_recon.base, None)
+        self.assertEqual(ipch_recon.size, ipch.size)
+
+        if driver.USE_NV_BINDING:
+            self.assertEqual(ipch_recon.handle.reserved, ipch.handle.reserved)
+        else:
+            self.assertEqual(ipch_recon.handle.reserved[:],
+                             ipch.handle.reserved[:])
+
+        # spawn new process for testing
+        ctx = mp.get_context('spawn')
+        result_queue = ctx.Queue()
+        args = (ipch, result_queue)
+        proc = ctx.Process(target=serialize_ipc_handle_test, args=args)
+        proc.start()
+        succ, out = result_queue.get()
+        if not succ:
+            self.fail(out)
+        else:
+            np.testing.assert_equal(expect, out)
+        proc.join(3)
+
+    def test_ipc_handle_serialization(self):
+        for index, foreign, in self.variants():
+            with self.subTest(index=index, foreign=foreign):
+                self.check_ipc_handle_serialization(index, foreign)
+
+    def check_ipc_array(self, index_arg=None, foreign=False):
+        # prepare data for IPC
+        arr = np.arange(10, dtype=np.intp)
+        devarr = cuda.to_device(arr)
+        # Slice
+        if index_arg is not None:
+            devarr = devarr[index_arg]
+        if foreign:
+            devarr = cuda.as_cuda_array(ForeignArray(devarr))
+        expect = devarr.copy_to_host()
+        ipch = devarr.get_ipc_handle()
+
+        # spawn new process for testing
+        ctx = mp.get_context('spawn')
+        result_queue = ctx.Queue()
+        args = (ipch, result_queue)
+        proc = ctx.Process(target=ipc_array_test, args=args)
+        proc.start()
+        succ, out = result_queue.get()
+        if not succ:
+            self.fail(out)
+        else:
+            np.testing.assert_equal(expect, out)
+        proc.join(3)
+
+    def test_ipc_array(self):
+        for index, foreign, in self.variants():
+            with self.subTest(index=index, foreign=foreign):
+                self.check_ipc_array(index, foreign)
+
+
+def staged_ipc_handle_test(handle, device_num, result_queue):
+    def the_work():
+        with cuda.gpus[device_num]:
+            this_ctx = cuda.devices.get_context()
+            deviceptr = handle.open_staged(this_ctx)
+            arrsize = handle.size // np.dtype(np.intp).itemsize
+            hostarray = np.zeros(arrsize, dtype=np.intp)
+            cuda.driver.device_to_host(
+                hostarray, deviceptr, size=handle.size,
+            )
+            handle.close()
+        return hostarray
+
+    core_ipc_handle_test(the_work, result_queue)
+
+
+def staged_ipc_array_test(ipcarr, device_num, result_queue):
+    try:
+        with cuda.gpus[device_num]:
+            with ipcarr as darr:
+                arr = darr.copy_to_host()
+                try:
+                    # should fail to reopen
+                    with ipcarr:
+                        pass
+                except ValueError as e:
+                    if str(e) != 'IpcHandle is already opened':
+                        raise AssertionError('invalid exception message')
+                else:
+                    raise AssertionError('did not raise on reopen')
+    # Catch any exception so we can propagate it
+    except:  # noqa: E722
+        # FAILED. propagate the exception as a string
+        succ = False
+        out = traceback.format_exc()
+    else:
+        # OK. send the ndarray back
+        succ = True
+        out = arr
+    result_queue.put((succ, out))
+
+
+@linux_only
+@skip_under_cuda_memcheck('Hangs cuda-memcheck')
+@skip_on_cudasim('Ipc not available in CUDASIM')
+@skip_on_arm('CUDA IPC not supported on ARM in Numba')
+class TestIpcStaged(ContextResettingTestCase):
+    def test_staged(self):
+        # prepare data for IPC
+        arr = np.arange(10, dtype=np.intp)
+        devarr = cuda.to_device(arr)
+
+        # spawn new process for testing
+        mpctx = mp.get_context('spawn')
+        result_queue = mpctx.Queue()
+
+        # create IPC handle
+        ctx = cuda.current_context()
+        ipch = ctx.get_ipc_handle(devarr.gpu_data)
+        # pickle
+        buf = pickle.dumps(ipch)
+        ipch_recon = pickle.loads(buf)
+        self.assertIs(ipch_recon.base, None)
+        if driver.USE_NV_BINDING:
+            self.assertEqual(ipch_recon.handle.reserved, ipch.handle.reserved)
+        else:
+            self.assertEqual(
+                ipch_recon.handle.reserved[:],
+                ipch.handle.reserved[:]
+            )
+        self.assertEqual(ipch_recon.size, ipch.size)
+
+        # Test on every CUDA devices
+        for device_num in range(len(cuda.gpus)):
+            args = (ipch, device_num, result_queue)
+            proc = mpctx.Process(target=staged_ipc_handle_test, args=args)
+            proc.start()
+            succ, out = result_queue.get()
+            proc.join(3)
+            if not succ:
+                self.fail(out)
+            else:
+                np.testing.assert_equal(arr, out)
+
+    def test_ipc_array(self):
+        for device_num in range(len(cuda.gpus)):
+            # prepare data for IPC
+            arr = np.random.random(10)
+            devarr = cuda.to_device(arr)
+            ipch = devarr.get_ipc_handle()
+
+            # spawn new process for testing
+            ctx = mp.get_context('spawn')
+            result_queue = ctx.Queue()
+            args = (ipch, device_num, result_queue)
+            proc = ctx.Process(target=staged_ipc_array_test, args=args)
+            proc.start()
+            succ, out = result_queue.get()
+            proc.join(3)
+            if not succ:
+                self.fail(out)
+            else:
+                np.testing.assert_equal(arr, out)
+
+
+@windows_only
+@skip_on_cudasim('Ipc not available in CUDASIM')
+class TestIpcNotSupported(ContextResettingTestCase):
+    def test_unsupported(self):
+        arr = np.arange(10, dtype=np.intp)
+        devarr = cuda.to_device(arr)
+        with self.assertRaises(OSError) as raises:
+            devarr.get_ipc_handle()
+        errmsg = str(raises.exception)
+        self.assertIn('OS does not support CUDA IPC', errmsg)
+
+
+if __name__ == '__main__':
+    unittest.main()

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

@@ -0,0 +1,99 @@
+from numba import cuda
+from numba.cuda.testing import unittest, CUDATestCase
+
+import numpy as np
+
+
+class TestIterators(CUDATestCase):
+
+    def test_enumerate(self):
+        @cuda.jit
+        def enumerator(x, error):
+            count = 0
+
+            for i, v in enumerate(x):
+                if count != i:
+                    error[0] = 1
+                if v != x[i]:
+                    error[0] = 2
+
+                count += 1
+
+            if count != len(x):
+                error[0] = 3
+
+        x = np.asarray((10, 9, 8, 7, 6))
+        error = np.zeros(1, dtype=np.int32)
+
+        enumerator[1, 1](x, error)
+        self.assertEqual(error[0], 0)
+
+    def _test_twoarg_function(self, f):
+        x = np.asarray((10, 9, 8, 7, 6))
+        y = np.asarray((1, 2, 3, 4, 5))
+        error = np.zeros(1, dtype=np.int32)
+
+        f[1, 1](x, y, error)
+        self.assertEqual(error[0], 0)
+
+    def test_zip(self):
+        @cuda.jit
+        def zipper(x, y, error):
+            i = 0
+
+            for xv, yv in zip(x, y):
+                if xv != x[i]:
+                    error[0] = 1
+                if yv != y[i]:
+                    error[0] = 2
+
+                i += 1
+
+            if i != len(x):
+                error[0] = 3
+
+        self._test_twoarg_function(zipper)
+
+    def test_enumerate_zip(self):
+        @cuda.jit
+        def enumerator_zipper(x, y, error):
+            count = 0
+
+            for i, (xv, yv) in enumerate(zip(x, y)):
+                if i != count:
+                    error[0] = 1
+                if xv != x[i]:
+                    error[0] = 2
+                if yv != y[i]:
+                    error[0] = 3
+
+                count += 1
+
+            if count != len(x):
+                error[0] = 4
+
+        self._test_twoarg_function(enumerator_zipper)
+
+    def test_zip_enumerate(self):
+        @cuda.jit
+        def zipper_enumerator(x, y, error):
+            count = 0
+
+            for (i, xv), yv in zip(enumerate(x), y):
+                if i != count:
+                    error[0] = 1
+                if xv != x[i]:
+                    error[0] = 2
+                if yv != y[i]:
+                    error[0] = 3
+
+                count += 1
+
+            if count != len(x):
+                error[0] = 4
+
+        self._test_twoarg_function(zipper_enumerator)
+
+
+if __name__ == '__main__':
+    unittest.main()

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

@@ -0,0 +1,64 @@
+"""
+Test basic language features
+
+"""
+
+import numpy as np
+from numba import cuda, float64
+from numba.cuda.testing import unittest, CUDATestCase
+
+
+class TestLang(CUDATestCase):
+    def test_enumerate(self):
+        tup = (1., 2.5, 3.)
+
+        @cuda.jit("void(float64[:])")
+        def foo(a):
+            for i, v in enumerate(tup):
+                a[i] = v
+
+        a = np.zeros(len(tup))
+        foo[1, 1](a)
+        self.assertTrue(np.all(a == tup))
+
+    def test_zip(self):
+        t1 = (1, 2, 3)
+        t2 = (4.5, 5.6, 6.7)
+
+        @cuda.jit("void(float64[:])")
+        def foo(a):
+            c = 0
+            for i, j in zip(t1, t2):
+                c += i + j
+            a[0] = c
+
+        a = np.zeros(1)
+        foo[1, 1](a)
+        b = np.array(t1)
+        c = np.array(t2)
+        self.assertTrue(np.all(a == (b + c).sum()))
+
+    def test_issue_872(self):
+        '''
+        Ensure that typing and lowering of CUDA kernel API primitives works in
+        more than one block. Was originally to ensure that macro expansion works
+        for more than one block (issue #872), but macro expansion has been
+        replaced by a "proper" implementation of all kernel API functions.
+        '''
+
+        @cuda.jit("void(float64[:,:])")
+        def cuda_kernel_api_in_multiple_blocks(ary):
+            for i in range(2):
+                tx = cuda.threadIdx.x
+            for j in range(3):
+                ty = cuda.threadIdx.y
+            sm = cuda.shared.array((2, 3), float64)
+            sm[tx, ty] = 1.0
+            ary[tx, ty] = sm[tx, ty]
+
+        a = np.zeros((2, 3))
+        cuda_kernel_api_in_multiple_blocks[1, (2, 3)](a)
+
+
+if __name__ == '__main__':
+    unittest.main()

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

@@ -0,0 +1,119 @@
+import numpy as np
+from numba import cuda, float64, void
+from numba.cuda.testing import unittest, CUDATestCase
+from numba.core import config
+
+# NOTE: CUDA kernel does not return any value
+
+if config.ENABLE_CUDASIM:
+    tpb = 4
+else:
+    tpb = 16
+SM_SIZE = tpb, tpb
+
+
+class TestCudaLaplace(CUDATestCase):
+    def test_laplace_small(self):
+
+        @cuda.jit(float64(float64, float64), device=True, inline=True)
+        def get_max(a, b):
+            if a > b:
+                return a
+            else:
+                return b
+
+        @cuda.jit(void(float64[:, :], float64[:, :], float64[:, :]))
+        def jocabi_relax_core(A, Anew, error):
+            err_sm = cuda.shared.array(SM_SIZE, dtype=float64)
+
+            ty = cuda.threadIdx.x
+            tx = cuda.threadIdx.y
+            bx = cuda.blockIdx.x
+            by = cuda.blockIdx.y
+
+            n = A.shape[0]
+            m = A.shape[1]
+
+            i, j = cuda.grid(2)
+
+            err_sm[ty, tx] = 0
+            if j >= 1 and j < n - 1 and i >= 1 and i < m - 1:
+                Anew[j, i] = 0.25 * ( A[j, i + 1] + A[j, i - 1]
+                                      + A[j - 1, i] + A[j + 1, i])
+                err_sm[ty, tx] = Anew[j, i] - A[j, i]
+
+            cuda.syncthreads()
+
+            # max-reduce err_sm vertically
+            t = tpb // 2
+            while t > 0:
+                if ty < t:
+                    err_sm[ty, tx] = get_max(err_sm[ty, tx], err_sm[ty + t, tx])
+                t //= 2
+                cuda.syncthreads()
+
+            # max-reduce err_sm horizontally
+            t = tpb // 2
+            while t > 0:
+                if tx < t and ty == 0:
+                    err_sm[ty, tx] = get_max(err_sm[ty, tx], err_sm[ty, tx + t])
+                t //= 2
+                cuda.syncthreads()
+
+            if tx == 0 and ty == 0:
+                error[by, bx] = err_sm[0, 0]
+
+        if config.ENABLE_CUDASIM:
+            NN, NM = 4, 4
+            iter_max = 20
+        else:
+            NN, NM = 256, 256
+            iter_max = 1000
+
+        A = np.zeros((NN, NM), dtype=np.float64)
+        Anew = np.zeros((NN, NM), dtype=np.float64)
+
+        n = NN
+
+        tol = 1.0e-6
+        error = 1.0
+
+        for j in range(n):
+            A[j, 0] = 1.0
+            Anew[j, 0] = 1.0
+
+        iter = 0
+
+        blockdim = (tpb, tpb)
+        griddim = (NN // blockdim[0], NM // blockdim[1])
+
+        error_grid = np.zeros(griddim)
+
+        stream = cuda.stream()
+
+        dA = cuda.to_device(A, stream)          # to device and don't come back
+        dAnew = cuda.to_device(Anew, stream)    # to device and don't come back
+        derror_grid = cuda.to_device(error_grid, stream)
+
+        while error > tol and iter < iter_max:
+            self.assertTrue(error_grid.dtype == np.float64)
+
+            jocabi_relax_core[griddim, blockdim, stream](dA, dAnew, derror_grid)
+
+            derror_grid.copy_to_host(error_grid, stream=stream)
+
+            # error_grid is available on host
+            stream.synchronize()
+
+            error = np.abs(error_grid).max()
+
+            # swap dA and dAnew
+            tmp = dA
+            dA = dAnew
+            dAnew = tmp
+
+            iter += 1
+
+
+if __name__ == '__main__':
+    unittest.main()