numba-cuda 0.8.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

numba_cuda/numba/cuda/tests/cudadrv/test_detect.py

@@ -3,8 +3,12 @@ import sys
 import subprocess
 import threading
 from numba import cuda
-from numba.cuda.testing import (unittest, CUDATestCase, skip_on_cudasim,
-                                skip_under_cuda_memcheck)
+from numba.cuda.testing import (
+    unittest,
+    CUDATestCase,
+    skip_on_cudasim,
+    skip_under_cuda_memcheck,
+)
 from numba.tests.support import captured_stdout
 
 
@@ -14,21 +18,19 @@ class TestCudaDetect(CUDATestCase):
         with captured_stdout() as out:
             cuda.detect()
         output = out.getvalue()
-        self.assertIn('Found', output)
-        self.assertIn('CUDA devices', output)
+        self.assertIn("Found", output)
+        self.assertIn("CUDA devices", output)
 
 
-@skip_under_cuda_memcheck('Hangs cuda-memcheck')
+@skip_under_cuda_memcheck("Hangs cuda-memcheck")
 class TestCUDAFindLibs(CUDATestCase):
-
     def run_cmd(self, cmdline, env):
-        popen = subprocess.Popen(cmdline,
-                                 stdout=subprocess.PIPE,
-                                 stderr=subprocess.PIPE,
-                                 env=env)
+        popen = subprocess.Popen(
+            cmdline, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env
+        )
 
         # finish in 5 minutes or kill it
-        timeout = threading.Timer(5 * 60., popen.kill)
+        timeout = threading.Timer(5 * 60.0, popen.kill)
         try:
             timeout.start()
             out, err = popen.communicate()
@@ -51,8 +53,8 @@ class TestCUDAFindLibs(CUDATestCase):
         cmdline = [sys.executable, "-c", code]
         return self.run_cmd(cmdline, env_copy)
 
-    @skip_on_cudasim('Simulator does not hit device library search code path')
-    @unittest.skipIf(not sys.platform.startswith('linux'), "linux only")
+    @skip_on_cudasim("Simulator does not hit device library search code path")
+    @unittest.skipIf(not sys.platform.startswith("linux"), "linux only")
     def test_cuda_find_lib_errors(self):
         """
         This tests that the find_libs works as expected in the case of an
@@ -60,7 +62,7 @@ class TestCUDAFindLibs(CUDATestCase):
         """
         # one of these is likely to exist on linux, it's also unlikely that
         # someone has extracted the contents of libdevice into here!
-        locs = ['lib', 'lib64']
+        locs = ["lib", "lib64"]
 
         looking_for = None
         for l in locs:
@@ -71,11 +73,12 @@ class TestCUDAFindLibs(CUDATestCase):
         # This is the testing part, the test will only run if there's a valid
         # path in which to look
         if looking_for is not None:
-            out, err = self.run_test_in_separate_process("NUMBA_CUDA_DRIVER",
-                                                         looking_for)
+            out, err = self.run_test_in_separate_process(
+                "NUMBA_CUDA_DRIVER", looking_for
+            )
             self.assertTrue(out is not None)
             self.assertTrue(err is not None)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()

numba_cuda/numba/cuda/tests/cudadrv/test_emm_plugins.py

@@ -8,6 +8,7 @@ from numba.cuda.testing import unittest, CUDATestCase, skip_on_cudasim
 from numba.tests.support import linux_only
 
 if not config.ENABLE_CUDASIM:
+
     class DeviceOnlyEMMPlugin(cuda.HostOnlyCUDAMemoryManager):
         """
         Dummy EMM Plugin implementation for testing. It memorises which plugin
@@ -56,8 +57,9 @@ if not config.ENABLE_CUDASIM:
             # the reference count drops to zero.
             ctx = weakref.proxy(self.context)
             ptr = ctypes.c_void_p(alloc_count)
-            return cuda.cudadrv.driver.AutoFreePointer(ctx, ptr, size,
-                                                       finalizer=finalizer)
+            return cuda.cudadrv.driver.AutoFreePointer(
+                ctx, ptr, size, finalizer=finalizer
+            )
 
         def initialize(self):
             # No special initialization needed.
@@ -97,7 +99,7 @@ if not config.ENABLE_CUDASIM:
             return 2
 
 
-@skip_on_cudasim('EMM Plugins not supported on CUDA simulator')
+@skip_on_cudasim("EMM Plugins not supported on CUDA simulator")
 class TestDeviceOnlyEMMPlugin(CUDATestCase):
     """
     Tests that the API of an EMM Plugin that implements device allocations
@@ -175,7 +177,7 @@ class TestDeviceOnlyEMMPlugin(CUDATestCase):
         self.assertIn("Dummy IPC handle for alloc 1", ipch._ipc_handle)
 
 
-@skip_on_cudasim('EMM Plugins not supported on CUDA simulator')
+@skip_on_cudasim("EMM Plugins not supported on CUDA simulator")
 class TestBadEMMPluginVersion(CUDATestCase):
     """
     Ensure that Numba rejects EMM Plugins with incompatible version
@@ -185,8 +187,8 @@ class TestBadEMMPluginVersion(CUDATestCase):
     def test_bad_plugin_version(self):
         with self.assertRaises(RuntimeError) as raises:
             cuda.set_memory_manager(BadVersionEMMPlugin)
-        self.assertIn('version 1 required', str(raises.exception))
+        self.assertIn("version 1 required", str(raises.exception))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()

numba_cuda/numba/cuda/tests/cudadrv/test_events.py

@@ -34,5 +34,5 @@ class TestCudaEvent(CUDATestCase):
         evtstart.elapsed_time(evtend)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()

numba_cuda/numba/cuda/tests/cudadrv/test_host_alloc.py

@@ -10,10 +10,9 @@ class TestHostAlloc(ContextResettingTestCase):
         mem = cuda.current_context().memhostalloc(n, mapped=True)
 
         dtype = np.dtype(np.uint8)
-        ary = np.ndarray(shape=n // dtype.itemsize, dtype=dtype,
-                         buffer=mem)
+        ary = np.ndarray(shape=n // dtype.itemsize, dtype=dtype, buffer=mem)
 
-        magic = 0xab
+        magic = 0xAB
         driver.device_memset(mem, magic, n)
 
         self.assertTrue(np.all(ary == magic))
@@ -46,8 +45,10 @@ class TestHostAlloc(ContextResettingTestCase):
         self.assertTrue(sum(ary != 0) == 0)
 
     def test_host_operators(self):
-        for ary in [cuda.mapped_array(10, dtype=np.uint32),
-                    cuda.pinned_array(10, dtype=np.uint32)]:
+        for ary in [
+            cuda.mapped_array(10, dtype=np.uint32),
+            cuda.pinned_array(10, dtype=np.uint32),
+        ]:
             ary[:] = range(10)
             self.assertTrue(sum(ary + 1) == 55)
             self.assertTrue(sum((ary + 1) * 2 - 1) == 100)
@@ -55,11 +56,11 @@ class TestHostAlloc(ContextResettingTestCase):
             self.assertTrue(sum(ary <= 5) == 6)
             self.assertTrue(sum(ary > 6) == 3)
             self.assertTrue(sum(ary >= 6) == 4)
-            self.assertTrue(sum(ary ** 2) == 285)
+            self.assertTrue(sum(ary**2) == 285)
             self.assertTrue(sum(ary // 2) == 20)
             self.assertTrue(sum(ary / 2.0) == 22.5)
             self.assertTrue(sum(ary % 2) == 5)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()

numba_cuda/numba/cuda/tests/cudadrv/test_init.py

@@ -9,7 +9,7 @@ from numba.cuda.testing import skip_on_cudasim, unittest, CUDATestCase
 
 # A mock of cuInit that always raises a CudaAPIError
 def cuInit_raising(arg):
-    raise CudaAPIError(999, 'CUDA_ERROR_UNKNOWN')
+    raise CudaAPIError(999, "CUDA_ERROR_UNKNOWN")
 
 
 # Test code to run in a child that patches driver.cuInit to a variant that
@@ -82,45 +82,45 @@ def cuda_disabled_error_test(result_queue):
     result_queue.put((success, msg))
 
 
-@skip_on_cudasim('CUDA Simulator does not initialize driver')
+@skip_on_cudasim("CUDA Simulator does not initialize driver")
 class TestInit(CUDATestCase):
     def _test_init_failure(self, target, expected):
         # Run the initialization failure test in a separate subprocess
-        ctx = mp.get_context('spawn')
+        ctx = mp.get_context("spawn")
         result_queue = ctx.Queue()
         proc = ctx.Process(target=target, args=(result_queue,))
         proc.start()
-        proc.join(30) # should complete within 30s
+        proc.join(30)  # should complete within 30s
         success, msg = result_queue.get()
 
         # Ensure the child process raised an exception during initialization
         # before checking the message
         if not success:
-            self.fail('CudaSupportError not raised')
+            self.fail("CudaSupportError not raised")
 
         self.assertIn(expected, msg)
 
     def test_init_failure_raising(self):
-        expected = 'Error at driver init: CUDA_ERROR_UNKNOWN (999)'
+        expected = "Error at driver init: CUDA_ERROR_UNKNOWN (999)"
         self._test_init_failure(cuInit_raising_test, expected)
 
     def test_init_failure_error(self):
-        expected = 'CUDA_ERROR_UNKNOWN (999)'
+        expected = "CUDA_ERROR_UNKNOWN (999)"
         self._test_init_failure(initialization_error_test, expected)
 
     def _test_cuda_disabled(self, target):
         # Uses _test_init_failure to launch the test in a separate subprocess
         # with CUDA disabled.
-        cuda_disabled = os.environ.get('NUMBA_DISABLE_CUDA')
-        os.environ['NUMBA_DISABLE_CUDA'] = "1"
+        cuda_disabled = os.environ.get("NUMBA_DISABLE_CUDA")
+        os.environ["NUMBA_DISABLE_CUDA"] = "1"
         try:
-            expected = 'CUDA is disabled due to setting NUMBA_DISABLE_CUDA=1'
+            expected = "CUDA is disabled due to setting NUMBA_DISABLE_CUDA=1"
             self._test_init_failure(cuda_disabled_test, expected)
         finally:
             if cuda_disabled is not None:
-                os.environ['NUMBA_DISABLE_CUDA'] = cuda_disabled
+                os.environ["NUMBA_DISABLE_CUDA"] = cuda_disabled
             else:
-                os.environ.pop('NUMBA_DISABLE_CUDA')
+                os.environ.pop("NUMBA_DISABLE_CUDA")
 
     def test_cuda_disabled_raising(self):
         self._test_cuda_disabled(cuda_disabled_test)
@@ -135,5 +135,5 @@ class TestInit(CUDATestCase):
         self.assertIsNone(cuda.cuda_error())
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()

numba_cuda/numba/cuda/tests/cudadrv/test_inline_ptx.py

@@ -5,20 +5,23 @@ from numba.cuda.testing import unittest, ContextResettingTestCase
 from numba.cuda.testing import skip_on_cudasim
 
 
-@skip_on_cudasim('Inline PTX cannot be used in the simulator')
+@skip_on_cudasim("Inline PTX cannot be used in the simulator")
 class TestCudaInlineAsm(ContextResettingTestCase):
     def test_inline_rsqrt(self):
         mod = ir.Module(__name__)
-        mod.triple = 'nvptx64-nvidia-cuda'
+        mod.triple = "nvptx64-nvidia-cuda"
         nvvm.add_ir_version(mod)
         fnty = ir.FunctionType(ir.VoidType(), [ir.PointerType(ir.FloatType())])
-        fn = ir.Function(mod, fnty, 'cu_rsqrt')
-        bldr = ir.IRBuilder(fn.append_basic_block('entry'))
+        fn = ir.Function(mod, fnty, "cu_rsqrt")
+        bldr = ir.IRBuilder(fn.append_basic_block("entry"))
 
         rsqrt_approx_fnty = ir.FunctionType(ir.FloatType(), [ir.FloatType()])
-        inlineasm = ir.InlineAsm(rsqrt_approx_fnty,
-                                 'rsqrt.approx.f32 $0, $1;',
-                                 '=f,f', side_effect=True)
+        inlineasm = ir.InlineAsm(
+            rsqrt_approx_fnty,
+            "rsqrt.approx.f32 $0, $1;",
+            "=f,f",
+            side_effect=True,
+        )
         val = bldr.load(fn.args[0])
         res = bldr.call(inlineasm, [val])
 
@@ -30,8 +33,8 @@ class TestCudaInlineAsm(ContextResettingTestCase):
         nvvm.set_cuda_kernel(fn)
         nvvmir = str(mod)
         ptx = nvvm.compile_ir(nvvmir)
-        self.assertTrue('rsqrt.approx.f32' in str(ptx))
+        self.assertTrue("rsqrt.approx.f32" in str(ptx))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()

numba_cuda/numba/cuda/tests/cudadrv/test_linker.py

@@ -1,10 +1,9 @@
 import numpy as np
 import warnings
 from numba.cuda.testing import unittest
-from numba.cuda.testing import (skip_on_cudasim, skip_if_cuda_includes_missing)
+from numba.cuda.testing import skip_on_cudasim, skip_if_cuda_includes_missing
 from numba.cuda.testing import CUDATestCase, test_data_dir
-from numba.cuda.cudadrv.driver import (CudaAPIError, Linker,
-                                       LinkerError)
+from numba.cuda.cudadrv.driver import CudaAPIError, Linker, LinkerError
 from numba.cuda.cudadrv.error import NvrtcError
 from numba.cuda import require_context
 from numba.tests.support import ignore_internal_warnings
@@ -103,25 +102,24 @@ def simple_lmem(A, B, dty):
         B[i] = C[i]
 
 
-@skip_on_cudasim('Linking unsupported in the simulator')
+@skip_on_cudasim("Linking unsupported in the simulator")
 class TestLinker(CUDATestCase):
-    _NUMBA_NVIDIA_BINDING_0_ENV = {'NUMBA_CUDA_USE_NVIDIA_BINDING': '0'}
+    _NUMBA_NVIDIA_BINDING_0_ENV = {"NUMBA_CUDA_USE_NVIDIA_BINDING": "0"}
 
     @require_context
     def test_linker_basic(self):
-        '''Simply go through the constructor and destructor
-        '''
+        """Simply go through the constructor and destructor"""
         linker = Linker.new(cc=(5, 3))
         del linker
 
     def _test_linking(self, eager):
         global bar  # must be a global; other it is recognized as a freevar
-        bar = cuda.declare_device('bar', 'int32(int32)')
+        bar = cuda.declare_device("bar", "int32(int32)")
 
-        link = str(test_data_dir / 'jitlink.ptx')
+        link = str(test_data_dir / "jitlink.ptx")
 
         if eager:
-            args = ['void(int32[:], int32[:])']
+            args = ["void(int32[:], int32[:])"]
         else:
             args = []
 
@@ -144,9 +142,9 @@ class TestLinker(CUDATestCase):
         self._test_linking(eager=True)
 
     def test_linking_cu(self):
-        bar = cuda.declare_device('bar', 'int32(int32)')
+        bar = cuda.declare_device("bar", "int32(int32)")
 
-        link = str(test_data_dir / 'jitlink.cu')
+        link = str(test_data_dir / "jitlink.cu")
 
         @cuda.jit(link=[link])
         def kernel(r, x):
@@ -165,36 +163,37 @@ class TestLinker(CUDATestCase):
         np.testing.assert_array_equal(r, expected)
 
     def test_linking_cu_log_warning(self):
-        bar = cuda.declare_device('bar', 'int32(int32)')
+        bar = cuda.declare_device("bar", "int32(int32)")
 
-        link = str(test_data_dir / 'warn.cu')
+        link = str(test_data_dir / "warn.cu")
 
         with warnings.catch_warnings(record=True) as w:
             ignore_internal_warnings()
 
-            @cuda.jit('void(int32)', link=[link])
+            @cuda.jit("void(int32)", link=[link])
             def kernel(x):
                 bar(x)
 
-        self.assertEqual(len(w), 1, 'Expected warnings from NVRTC')
+        self.assertEqual(len(w), 1, "Expected warnings from NVRTC")
         # Check the warning refers to the log messages
-        self.assertIn('NVRTC log messages', str(w[0].message))
+        self.assertIn("NVRTC log messages", str(w[0].message))
         # Check the message pertaining to the unused variable is provided
-        self.assertIn('declared but never referenced', str(w[0].message))
+        self.assertIn("declared but never referenced", str(w[0].message))
 
     def test_linking_cu_error(self):
-        bar = cuda.declare_device('bar', 'int32(int32)')
+        bar = cuda.declare_device("bar", "int32(int32)")
 
-        link = str(test_data_dir / 'error.cu')
+        link = str(test_data_dir / "error.cu")
 
         with self.assertRaises(NvrtcError) as e:
-            @cuda.jit('void(int32)', link=[link])
+
+            @cuda.jit("void(int32)", link=[link])
             def kernel(x):
                 bar(x)
 
         msg = e.exception.args[0]
         # Check the error message refers to the NVRTC compile
-        self.assertIn('NVRTC Compilation failure', msg)
+        self.assertIn("NVRTC Compilation failure", msg)
         # Check the expected error in the CUDA source is reported
         self.assertIn('identifier "SYNTAX" is undefined', msg)
         # Check the filename is reported correctly
@@ -203,33 +202,37 @@ class TestLinker(CUDATestCase):
     def test_linking_unknown_filetype_error(self):
         expected_err = "Don't know how to link file with extension .cuh"
         with self.assertRaisesRegex(RuntimeError, expected_err):
-            @cuda.jit('void()', link=['header.cuh'])
+
+            @cuda.jit("void()", link=["header.cuh"])
             def kernel():
                 pass
 
     def test_linking_file_with_no_extension_error(self):
         expected_err = "Don't know how to link file with no extension"
         with self.assertRaisesRegex(RuntimeError, expected_err):
-            @cuda.jit('void()', link=['data'])
+
+            @cuda.jit("void()", link=["data"])
             def kernel():
                 pass
 
     @skip_if_cuda_includes_missing
     def test_linking_cu_cuda_include(self):
-        link = str(test_data_dir / 'cuda_include.cu')
+        link = str(test_data_dir / "cuda_include.cu")
 
         # An exception will be raised when linking this kernel due to the
         # compile failure if CUDA includes cannot be found by Nvrtc.
-        @cuda.jit('void()', link=[link])
+        @cuda.jit("void()", link=[link])
         def kernel():
             pass
 
     def test_try_to_link_nonexistent(self):
         with self.assertRaises(LinkerError) as e:
-            @cuda.jit('void(int32[::1])', link=['nonexistent.a'])
+
+            @cuda.jit("void(int32[::1])", link=["nonexistent.a"])
             def f(x):
                 x[0] = 0
-        self.assertIn('nonexistent.a not found', e.exception.args)
+
+        self.assertIn("nonexistent.a not found", e.exception.args)
 
     def test_set_registers_no_max(self):
         """Ensure that the jitted kernel used in the test_set_registers_* tests
@@ -276,7 +279,8 @@ class TestLinker(CUDATestCase):
     def test_get_shared_mem_per_specialized(self):
         compiled = cuda.jit(simple_smem)
         compiled_specialized = compiled.specialize(
-            np.zeros(100, dtype=np.int32), np.float64)
+            np.zeros(100, dtype=np.int32), np.float64
+        )
         shared_mem_size = compiled_specialized.get_shared_mem_per_block()
         self.assertEqual(shared_mem_size, 800)
 
@@ -307,11 +311,12 @@ class TestLinker(CUDATestCase):
         compiled_specialized = compiled.specialize(
             np.zeros(LMEM_SIZE, dtype=np.int32),
             np.zeros(LMEM_SIZE, dtype=np.int32),
-            np.float64)
+            np.float64,
+        )
         local_mem_size = compiled_specialized.get_local_mem_per_thread()
         calc_size = np.dtype(np.float64).itemsize * LMEM_SIZE
         self.assertGreaterEqual(local_mem_size, calc_size)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()

numba_cuda/numba/cuda/tests/cudadrv/test_managed_alloc.py

@@ -7,11 +7,10 @@ from numba.cuda.testing import skip_on_cudasim, skip_on_arm
 from numba.tests.support import linux_only
 
 
-@skip_on_cudasim('CUDA Driver API unsupported in the simulator')
+@skip_on_cudasim("CUDA Driver API unsupported in the simulator")
 @linux_only
-@skip_on_arm('Managed Alloc support is experimental/untested on ARM')
+@skip_on_arm("Managed Alloc support is experimental/untested on ARM")
 class TestManagedAlloc(ContextResettingTestCase):
-
     def get_total_gpu_memory(self):
         # We use a driver function to directly get the total GPU memory because
         # an EMM plugin may report something different (or not implement
@@ -48,7 +47,9 @@ class TestManagedAlloc(ContextResettingTestCase):
     def test_managed_alloc_driver_undersubscribe(self):
         msg = "Managed memory unsupported prior to CC 3.0"
         self.skip_if_cc_major_lt(3, msg)
-        self._test_managed_alloc_driver(0.5)
+        # We keep the allocation small so that it doesn't hang on GPUs
+        # with large memory (H100)
+        self._test_managed_alloc_driver(0.1)
 
     # This test is skipped by default because it is easy to hang the machine
     # for a very long time or get OOM killed if the GPU memory size is >50% of
@@ -85,7 +86,7 @@ class TestManagedAlloc(ContextResettingTestCase):
         n_elems = n_bytes // dtype.itemsize
         ary = np.ndarray(shape=n_elems, dtype=dtype, buffer=mem)
 
-        magic = 0xab
+        magic = 0xAB
         device_memset(mem, magic, n_bytes)
         ctx.synchronize()
 
@@ -102,7 +103,7 @@ class TestManagedAlloc(ContextResettingTestCase):
         ary.fill(123.456)
         self.assertTrue(all(ary == 123.456))
 
-        @cuda.jit('void(double[:])')
+        @cuda.jit("void(double[:])")
         def kernel(x):
             i = cuda.grid(1)
             if i < x.shape[0]:
@@ -123,5 +124,5 @@ class TestManagedAlloc(ContextResettingTestCase):
         self._test_managed_array(attach_global=False)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()