numba-cuda 0.8.1__py3-none-any.whl → 0.10.0__py3-none-any.whl

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

@@ -0,0 +1,257 @@
+import numba.cuda as cuda
+from numba.cuda.testing import unittest, CUDATestCase
+import numpy as np
+
+from numba import int16, int32, int64, uint16, uint32, uint64, float32, float64
+from numba.types import float16
+
+from numba.cuda.cuda_bf16 import (
+    nv_bfloat16,
+    htrunc,
+    hceil,
+    hfloor,
+    hrint,
+    hsqrt,
+    hrsqrt,
+    hrcp,
+    hlog,
+    hlog2,
+    hlog10,
+    hcos,
+    hsin,
+    hexp,
+    hexp2,
+    hexp10,
+)
+
+from numba.cuda.cudadrv.runtime import get_version
+
+cuda_version = get_version()
+
+dtypes = [int16, int32, int64, uint16, uint32, uint64, float32]
+
+
+@unittest.skipIf(
+    (cuda.get_current_device().compute_capability < (8, 0)),
+    "bfloat16 requires compute capability 8.0+",
+)
+class Bfloat16Test(CUDATestCase):
+    def test_ctor(self):
+        @cuda.jit
+        def simple_kernel():
+            a = nv_bfloat16(float64(1.0))  # noqa: F841
+            b = nv_bfloat16(float32(2.0))  # noqa: F841
+            c = nv_bfloat16(int16(3))  # noqa: F841
+            d = nv_bfloat16(int32(4))  # noqa: F841
+            e = nv_bfloat16(int64(5))  # noqa: F841
+            f = nv_bfloat16(uint16(6))  # noqa: F841
+            g = nv_bfloat16(uint32(7))  # noqa: F841
+            h = nv_bfloat16(uint64(8))  # noqa: F841
+
+        simple_kernel[1, 1]()
+
+        if cuda_version >= (12, 0):
+
+            @cuda.jit
+            def simple_kernel_fp16():
+                i = nv_bfloat16(float16(9))  # noqa: F841
+
+            simple_kernel_fp16[1, 1]()
+
+    def test_casts(self):
+        @cuda.jit
+        def simple_kernel(b, c, d, e, f, g, h):
+            a = nv_bfloat16(3.14)
+
+            b[0] = float32(a)
+            c[0] = int16(a)
+            d[0] = int32(a)
+            e[0] = int64(a)
+            f[0] = uint16(a)
+            g[0] = uint32(a)
+            h[0] = uint64(a)
+
+        b = np.zeros(1, dtype=np.float32)
+        c = np.zeros(1, dtype=np.int16)
+        d = np.zeros(1, dtype=np.int32)
+        e = np.zeros(1, dtype=np.int64)
+        f = np.zeros(1, dtype=np.uint16)
+        g = np.zeros(1, dtype=np.uint32)
+        h = np.zeros(1, dtype=np.uint64)
+
+        simple_kernel[1, 1](b, c, d, e, f, g, h)
+
+        np.testing.assert_allclose(b[0], 3.14, atol=1e-2)
+        assert c[0] == 3
+        assert d[0] == 3
+        assert e[0] == 3
+        assert f[0] == 3
+        assert g[0] == 3
+        assert h[0] == 3
+
+    def test_ctor_cast_loop(self):
+        for dtype in dtypes:
+            with self.subTest(dtype=dtype):
+
+                @cuda.jit
+                def simple_kernel(a):
+                    a[0] = dtype(nv_bfloat16(dtype(3.14)))
+
+                a = np.zeros(1, dtype=str(dtype))
+                simple_kernel[1, 1](a)
+
+                if np.dtype(str(dtype)).kind == "f":
+                    np.testing.assert_allclose(a[0], 3.14, atol=1e-2)
+                else:
+                    assert a[0] == 3
+
+    def test_arithmetic(self):
+        @cuda.jit
+        def simple_kernel(arith, logic):
+            # Binary Arithmetic Operators
+            a = nv_bfloat16(1.0)
+            b = nv_bfloat16(2.0)
+
+            arith[0] = float32(a + b)
+            arith[1] = float32(a - b)
+            arith[2] = float32(a * b)
+            arith[3] = float32(a / b)
+
+            # Arithmetic Assignment Operators
+            a = nv_bfloat16(1.0)
+            b = nv_bfloat16(2.0)
+
+            a += b
+            arith[4] = float32(a)
+            a -= b
+            arith[5] = float32(a)
+            a *= b
+            arith[6] = float32(a)
+            a /= b
+            arith[7] = float32(a)
+
+            # Unary Arithmetic Operators
+            a = nv_bfloat16(1.0)
+
+            arith[8] = float32(+a)
+            arith[9] = float32(-a)
+
+            # Comparison Operators
+            a = nv_bfloat16(1.0)
+            b = nv_bfloat16(2.0)
+
+            logic[0] = a == b
+            logic[1] = a != b
+            logic[2] = a > b
+            logic[3] = a < b
+            logic[4] = a >= b
+            logic[5] = a <= b
+
+        arith = np.zeros(10, dtype=np.float32)
+        logic = np.zeros(6, dtype=np.bool_)
+
+        simple_kernel[1, 1](arith, logic)
+
+        a = 1.0
+        b = 2.0
+        np.testing.assert_allclose(
+            arith,
+            [
+                a + b,
+                a - b,
+                a * b,
+                a / b,
+                a + b,
+                a + b - b,
+                (a + b - b) * b,
+                (a + b - b) * b / b,
+                +a,
+                -a,
+            ],
+            atol=1e-2,
+        )
+        np.testing.assert_equal(
+            logic, [a == b, a != b, a > b, a < b, a >= b, a <= b]
+        )
+
+    def test_math_func(self):
+        @cuda.jit
+        def simple_kernel(a):
+            x = nv_bfloat16(3.14)
+
+            a[0] = float32(htrunc(x))
+            a[1] = float32(hceil(x))
+            a[2] = float32(hfloor(x))
+            a[3] = float32(hrint(x))
+            a[4] = float32(hsqrt(x))
+            a[5] = float32(hrsqrt(x))
+            a[6] = float32(hrcp(x))
+            a[7] = float32(hlog(x))
+            a[8] = float32(hlog2(x))
+            a[9] = float32(hlog10(x))
+            a[10] = float32(hcos(x))
+            a[11] = float32(hsin(x))
+            a[12] = float32(hexp(x))
+            a[13] = float32(hexp2(x))
+            a[14] = float32(hexp10(x))
+
+        a = np.zeros(15, dtype=np.float32)
+        simple_kernel[1, 1](a)
+
+        x = 3.14
+        np.testing.assert_allclose(
+            a[:12],
+            [
+                np.trunc(x),
+                np.ceil(x),
+                np.floor(x),
+                np.rint(x),
+                np.sqrt(x),
+                1 / np.sqrt(x),
+                1 / x,
+                np.log(x),
+                np.log2(x),
+                np.log10(x),
+                np.cos(x),
+                np.sin(x),
+            ],
+            atol=1e-2,
+        )
+
+        np.testing.assert_allclose(
+            a[12:], [np.exp(x), np.exp2(x), np.power(10, x)], atol=1e2
+        )
+
+    def test_check_bfloat16_type(self):
+        @cuda.jit
+        def kernel(arr):
+            x = nv_bfloat16(3.14)
+            if isinstance(x, nv_bfloat16):
+                arr[0] = float32(x)
+            else:
+                arr[0] = float32(0.0)
+
+        arr = np.zeros(1, np.float32)
+        kernel[1, 1](arr)
+
+        np.testing.assert_allclose(arr, [3.14], atol=1e-2)
+
+    def test_use_within_device_func(self):
+        @cuda.jit(device=True)
+        def add_bf16(a, b):
+            return a + b
+
+        @cuda.jit
+        def kernel(arr):
+            a = nv_bfloat16(3.14)
+            b = nv_bfloat16(5)
+            arr[0] = float32(hfloor(add_bf16(a, b)))
+
+        arr = np.zeros(1, np.float32)
+        kernel[1, 1](arr)
+
+        np.testing.assert_allclose(arr, [8], atol=1e-2)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -17,13 +17,23 @@ RSQRT2PI = 0.39894228040143267793994605993438
 
 def cnd(d):
     K = 1.0 / (1.0 + 0.2316419 * np.abs(d))
-    ret_val = (RSQRT2PI * np.exp(-0.5 * d * d) *
-               (K * (A1 + K * (A2 + K * (A3 + K * (A4 + K * A5))))))
+    ret_val = (
+        RSQRT2PI
+        * np.exp(-0.5 * d * d)
+        * (K * (A1 + K * (A2 + K * (A3 + K * (A4 + K * A5)))))
+    )
     return np.where(d > 0, 1.0 - ret_val, ret_val)
 
 
-def black_scholes(callResult, putResult, stockPrice, optionStrike, optionYears,
-                  Riskfree, Volatility):
+def black_scholes(
+    callResult,
+    putResult,
+    stockPrice,
+    optionStrike,
+    optionYears,
+    Riskfree,
+    Volatility,
+):
     S = stockPrice
     X = optionStrike
     T = optionYears
@@ -35,9 +45,9 @@ def black_scholes(callResult, putResult, stockPrice, optionStrike, optionYears,
     cndd1 = cnd(d1)
     cndd2 = cnd(d2)
 
-    expRT = np.exp(- R * T)
-    callResult[:] = (S * cndd1 - X * expRT * cndd2)
-    putResult[:] = (X * expRT * (1.0 - cndd2) - S * (1.0 - cndd1))
+    expRT = np.exp(-R * T)
+    callResult[:] = S * cndd1 - X * expRT * cndd2
+    putResult[:] = X * expRT * (1.0 - cndd2) - S * (1.0 - cndd1)
 
 
 def randfloat(rand_var, low, high):
@@ -61,34 +71,54 @@ class TestBlackScholes(CUDATestCase):
 
         # numpy
         for i in range(iterations):
-            black_scholes(callResultNumpy, putResultNumpy, stockPrice,
-                          optionStrike, optionYears, RISKFREE, VOLATILITY)
-
-        @cuda.jit(double(double), device=True, inline=True)
+            black_scholes(
+                callResultNumpy,
+                putResultNumpy,
+                stockPrice,
+                optionStrike,
+                optionYears,
+                RISKFREE,
+                VOLATILITY,
+            )
+
+        @cuda.jit(double(double), device=True, inline="always")
         def cnd_cuda(d):
             K = 1.0 / (1.0 + 0.2316419 * math.fabs(d))
-            ret_val = (RSQRT2PI * math.exp(-0.5 * d * d) *
-                       (K * (A1 + K * (A2 + K * (A3 + K * (A4 + K * A5))))))
+            ret_val = (
+                RSQRT2PI
+                * math.exp(-0.5 * d * d)
+                * (K * (A1 + K * (A2 + K * (A3 + K * (A4 + K * A5)))))
+            )
             if d > 0:
                 ret_val = 1.0 - ret_val
             return ret_val
 
-        @cuda.jit(void(double[:], double[:], double[:], double[:], double[:],
-                       double, double))
+        @cuda.jit(
+            void(
+                double[:],
+                double[:],
+                double[:],
+                double[:],
+                double[:],
+                double,
+                double,
+            )
+        )
         def black_scholes_cuda(callResult, putResult, S, X, T, R, V):
             i = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
             if i >= S.shape[0]:
                 return
             sqrtT = math.sqrt(T[i])
-            d1 = ((math.log(S[i] / X[i]) + (R + 0.5 * V * V) * T[i])
-                  / (V * sqrtT))
+            d1 = (math.log(S[i] / X[i]) + (R + 0.5 * V * V) * T[i]) / (
+                V * sqrtT
+            )
             d2 = d1 - V * sqrtT
             cndd1 = cnd_cuda(d1)
             cndd2 = cnd_cuda(d2)
 
-            expRT = math.exp((-1. * R) * T[i])
-            callResult[i] = (S[i] * cndd1 - X[i] * expRT * cndd2)
-            putResult[i] = (X[i] * expRT * (1.0 - cndd2) - S[i] * (1.0 - cndd1))
+            expRT = math.exp((-1.0 * R) * T[i])
+            callResult[i] = S[i] * cndd1 - X[i] * expRT * cndd2
+            putResult[i] = X[i] * expRT * (1.0 - cndd2) - S[i] * (1.0 - cndd1)
 
         # numba
         blockdim = 512, 1
@@ -102,8 +132,14 @@ class TestBlackScholes(CUDATestCase):
 
         for i in range(iterations):
             black_scholes_cuda[griddim, blockdim, stream](
-                d_callResult, d_putResult, d_stockPrice, d_optionStrike,
-                d_optionYears, RISKFREE, VOLATILITY)
+                d_callResult,
+                d_putResult,
+                d_stockPrice,
+                d_optionStrike,
+                d_optionYears,
+                RISKFREE,
+                VOLATILITY,
+            )
         d_callResult.copy_to_host(callResultNumba, stream)
         d_putResult.copy_to_host(putResultNumba, stream)
         stream.synchronize()
@@ -116,5 +152,5 @@ class TestBlackScholes(CUDATestCase):
         self.assertTrue(max_abs_err < 1e-13)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()

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

@@ -12,13 +12,13 @@ def boolean_func(A, vertial):
 
 class TestCudaBoolean(CUDATestCase):
     def test_boolean(self):
-        func = cuda.jit('void(float64[:], bool_)')(boolean_func)
-        A = np.array([0], dtype='float64')
+        func = cuda.jit("void(float64[:], bool_)")(boolean_func)
+        A = np.array([0], dtype="float64")
         func[1, 1](A, True)
         self.assertTrue(A[0] == 123)
         func[1, 1](A, False)
         self.assertTrue(A[0] == 321)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()

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

@@ -8,15 +8,22 @@ import warnings
 
 from numba import cuda
 from numba.core.errors import NumbaWarning
-from numba.cuda.testing import (CUDATestCase, skip_on_cudasim,
-                                skip_unless_cc_60, skip_if_cudadevrt_missing,
-                                skip_if_mvc_enabled, test_data_dir)
+from numba.cuda.testing import (
+    CUDATestCase,
+    skip_on_cudasim,
+    skip_unless_cc_60,
+    skip_if_cudadevrt_missing,
+    skip_if_mvc_enabled,
+    test_data_dir,
+)
 from numba.tests.support import SerialMixin
-from numba.tests.test_caching import (DispatcherCacheUsecasesTest,
-                                      skip_bad_access)
+from numba.tests.test_caching import (
+    DispatcherCacheUsecasesTest,
+    skip_bad_access,
+)
 
 
-@skip_on_cudasim('Simulator does not implement caching')
+@skip_on_cudasim("Simulator does not implement caching")
 class CUDACachingTest(SerialMixin, DispatcherCacheUsecasesTest):
     here = os.path.dirname(__file__)
     usecases_file = os.path.join(here, "cache_usecases.py")
@@ -72,23 +79,23 @@ class CUDACachingTest(SerialMixin, DispatcherCacheUsecasesTest):
         mod = self.import_module()
         f = mod.many_locals
         f[1, 1]()
-        self.check_pycache(2) # 1 index, 1 data
+        self.check_pycache(2)  # 1 index, 1 data
 
     def test_closure(self):
         mod = self.import_module()
 
         with warnings.catch_warnings():
-            warnings.simplefilter('error', NumbaWarning)
+            warnings.simplefilter("error", NumbaWarning)
 
             f = mod.closure1
-            self.assertPreciseEqual(f(3), 6) # 3 + 3 = 6
+            self.assertPreciseEqual(f(3), 6)  # 3 + 3 = 6
             f = mod.closure2
-            self.assertPreciseEqual(f(3), 8) # 3 + 5 = 8
+            self.assertPreciseEqual(f(3), 8)  # 3 + 5 = 8
             f = mod.closure3
-            self.assertPreciseEqual(f(3), 10) # 3 + 7 = 10
+            self.assertPreciseEqual(f(3), 10)  # 3 + 7 = 10
             f = mod.closure4
-            self.assertPreciseEqual(f(3), 12) # 3 + 9 = 12
-            self.check_pycache(5) # 1 nbi, 4 nbc
+            self.assertPreciseEqual(f(3), 12)  # 3 + 9 = 12
+            self.check_pycache(5)  # 1 nbi, 4 nbc
 
     def test_cache_reuse(self):
         mod = self.import_module()
@@ -158,7 +165,7 @@ class CUDACachingTest(SerialMixin, DispatcherCacheUsecasesTest):
 
     @skip_unless_cc_60
     @skip_if_cudadevrt_missing
-    @skip_if_mvc_enabled('CG not supported with MVC')
+    @skip_if_mvc_enabled("CG not supported with MVC")
     def test_cache_cg(self):
         # Functions using cooperative groups should be cacheable. See Issue
         # #8888: https://github.com/numba/numba/issues/8888
@@ -174,7 +181,7 @@ class CUDACachingTest(SerialMixin, DispatcherCacheUsecasesTest):
 
     @skip_unless_cc_60
     @skip_if_cudadevrt_missing
-    @skip_if_mvc_enabled('CG not supported with MVC')
+    @skip_if_mvc_enabled("CG not supported with MVC")
     def test_cache_cg_clean_run(self):
         # See Issue #9432: https://github.com/numba/numba/issues/9432
         # If a cached function using CG sync was the first thing to compile,
@@ -191,9 +198,11 @@ class CUDACachingTest(SerialMixin, DispatcherCacheUsecasesTest):
             mod.cg_usecase(0)
             """ % dict(tempdir=self.tempdir, modname=self.modname)
 
-        popen = subprocess.Popen([sys.executable, "-c", code],
-                                 stdout=subprocess.PIPE,
-                                 stderr=subprocess.PIPE)
+        popen = subprocess.Popen(
+            [sys.executable, "-c", code],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
         out, err = popen.communicate(timeout=60)
         if popen.returncode != 0:
             raise AssertionError(
@@ -212,8 +221,9 @@ class CUDACachingTest(SerialMixin, DispatcherCacheUsecasesTest):
         f = mod.add_usecase
         # Remove this function's cache files at the end, to avoid accumulation
         # across test calls.
-        self.addCleanup(shutil.rmtree, f.func.stats.cache_path,
-                        ignore_errors=True)
+        self.addCleanup(
+            shutil.rmtree, f.func.stats.cache_path, ignore_errors=True
+        )
 
         self.assertPreciseEqual(f(2, 3), 6)
         # It's a cache miss since the file was copied to a new temp location
@@ -230,8 +240,9 @@ class CUDACachingTest(SerialMixin, DispatcherCacheUsecasesTest):
         self.check_pycache(0)
 
     @skip_bad_access
-    @unittest.skipIf(os.name == "nt",
-                     "cannot easily make a directory read-only on Windows")
+    @unittest.skipIf(
+        os.name == "nt", "cannot easily make a directory read-only on Windows"
+    )
     def test_non_creatable_pycache(self):
         # Make it impossible to create the __pycache__ directory
         old_perms = os.stat(self.tempdir).st_mode
@@ -241,11 +252,12 @@ class CUDACachingTest(SerialMixin, DispatcherCacheUsecasesTest):
         self._test_pycache_fallback()
 
     @skip_bad_access
-    @unittest.skipIf(os.name == "nt",
-                     "cannot easily make a directory read-only on Windows")
+    @unittest.skipIf(
+        os.name == "nt", "cannot easily make a directory read-only on Windows"
+    )
     def test_non_writable_pycache(self):
         # Make it impossible to write to the __pycache__ directory
-        pycache = os.path.join(self.tempdir, '__pycache__')
+        pycache = os.path.join(self.tempdir, "__pycache__")
         os.mkdir(pycache)
         old_perms = os.stat(pycache).st_mode
         os.chmod(pycache, 0o500)
@@ -254,15 +266,16 @@ class CUDACachingTest(SerialMixin, DispatcherCacheUsecasesTest):
         self._test_pycache_fallback()
 
     def test_cannot_cache_linking_libraries(self):
-        link = str(test_data_dir / 'jitlink.ptx')
-        msg = 'Cannot pickle CUDACodeLibrary with linking files'
+        link = str(test_data_dir / "jitlink.ptx")
+        msg = "Cannot pickle CUDACodeLibrary with linking files"
         with self.assertRaisesRegex(RuntimeError, msg):
-            @cuda.jit('void()', cache=True, link=[link])
+
+            @cuda.jit("void()", cache=True, link=[link])
             def f():
                 pass
 
 
-@skip_on_cudasim('Simulator does not implement caching')
+@skip_on_cudasim("Simulator does not implement caching")
 class CUDAAndCPUCachingTest(SerialMixin, DispatcherCacheUsecasesTest):
     here = os.path.dirname(__file__)
     usecases_file = os.path.join(here, "cache_with_cpu_usecases.py")
@@ -353,7 +366,7 @@ def get_different_cc_gpus():
     return None
 
 
-@skip_on_cudasim('Simulator does not implement caching')
+@skip_on_cudasim("Simulator does not implement caching")
 class TestMultiCCCaching(SerialMixin, DispatcherCacheUsecasesTest):
     here = os.path.dirname(__file__)
     usecases_file = os.path.join(here, "cache_usecases.py")
@@ -370,7 +383,7 @@ class TestMultiCCCaching(SerialMixin, DispatcherCacheUsecasesTest):
     def test_cache(self):
         gpus = get_different_cc_gpus()
         if not gpus:
-            self.skipTest('Need two different CCs for multi-CC cache test')
+            self.skipTest("Need two different CCs for multi-CC cache test")
 
         self.check_pycache(0)
         mod = self.import_module()
@@ -482,13 +495,13 @@ def child_initializer():
     # Disable occupancy and implicit copy warnings in processes in a
     # multiprocessing pool.
     from numba.core import config
+
     config.CUDA_LOW_OCCUPANCY_WARNINGS = 0
     config.CUDA_WARN_ON_IMPLICIT_COPY = 0
 
 
-@skip_on_cudasim('Simulator does not implement caching')
+@skip_on_cudasim("Simulator does not implement caching")
 class TestMultiprocessCache(SerialMixin, DispatcherCacheUsecasesTest):
-
     # Nested multiprocessing.Pool raises AssertionError:
     # "daemonic processes are not allowed to have children"
     _numba_parallel_test_ = False
@@ -513,7 +526,7 @@ class TestMultiprocessCache(SerialMixin, DispatcherCacheUsecasesTest):
         f = mod.simple_usecase_caller
         n = 3
         try:
-            ctx = multiprocessing.get_context('spawn')
+            ctx = multiprocessing.get_context("spawn")
         except AttributeError:
             ctx = multiprocessing
 
@@ -526,7 +539,7 @@ class TestMultiprocessCache(SerialMixin, DispatcherCacheUsecasesTest):
         self.assertEqual(res, n * (n - 1) // 2)
 
 
-@skip_on_cudasim('Simulator does not implement the CUDACodeLibrary')
+@skip_on_cudasim("Simulator does not implement the CUDACodeLibrary")
 class TestCUDACodeLibrary(CUDATestCase):
     # For tests of miscellaneous CUDACodeLibrary behaviour that we wish to
     # explicitly check
@@ -539,7 +552,7 @@ class TestCUDACodeLibrary(CUDATestCase):
         # Usually a CodeLibrary requires a real CodeGen, but since we don't
         # interact with it, anything will do
         codegen = object()
-        name = 'library'
+        name = "library"
         cl = CUDACodeLibrary(codegen, name)
-        with self.assertRaisesRegex(RuntimeError, 'Cannot pickle unfinalized'):
+        with self.assertRaisesRegex(RuntimeError, "Cannot pickle unfinalized"):
             cl._reduce_states()