numba-cuda 0.21.1__cp313-cp313-win_amd64.whl

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

@@ -0,0 +1,820 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda.cudadrv.driver import CudaAPIError
+import numpy as np
+import threading
+
+from numba.cuda.types import (
+    boolean,
+    float32,
+    float64,
+    int32,
+    int64,
+    uint32,
+    void,
+)
+from numba import cuda
+from numba.cuda import config, types
+from numba.cuda.core.errors import TypingError
+from numba.cuda.testing import (
+    cc_X_or_above,
+    skip_on_cudasim,
+    unittest,
+    CUDATestCase,
+)
+import math
+
+
+def add(x, y):
+    return x + y
+
+
+def add_kernel(r, x, y):
+    r[0] = x + y
+
+
+@skip_on_cudasim("Specialization not implemented in the simulator")
+class TestDispatcherSpecialization(CUDATestCase):
+    def _test_no_double_specialize(self, dispatcher, ty):
+        with self.assertRaises(RuntimeError) as e:
+            dispatcher.specialize(ty)
+
+        self.assertIn("Dispatcher already specialized", str(e.exception))
+
+    def test_no_double_specialize_sig_same_types(self):
+        # Attempting to specialize a kernel jitted with a signature is illegal,
+        # even for the same types the kernel is already specialized for.
+        @cuda.jit("void(float32[::1])")
+        def f(x):
+            pass
+
+        self._test_no_double_specialize(f, float32[::1])
+
+    def test_no_double_specialize_no_sig_same_types(self):
+        # Attempting to specialize an already-specialized kernel is illegal,
+        # even for the same types the kernel is already specialized for.
+        @cuda.jit
+        def f(x):
+            pass
+
+        f_specialized = f.specialize(float32[::1])
+        self._test_no_double_specialize(f_specialized, float32[::1])
+
+    def test_no_double_specialize_sig_diff_types(self):
+        # Attempting to specialize a kernel jitted with a signature is illegal.
+        @cuda.jit("void(int32[::1])")
+        def f(x):
+            pass
+
+        self._test_no_double_specialize(f, float32[::1])
+
+    def test_no_double_specialize_no_sig_diff_types(self):
+        # Attempting to specialize an already-specialized kernel is illegal.
+        @cuda.jit
+        def f(x):
+            pass
+
+        f_specialized = f.specialize(int32[::1])
+        self._test_no_double_specialize(f_specialized, float32[::1])
+
+    def test_specialize_cache_same(self):
+        # Ensure that the same dispatcher is returned for the same argument
+        # types, and that different dispatchers are returned for different
+        # argument types.
+        @cuda.jit
+        def f(x):
+            pass
+
+        self.assertEqual(len(f.specializations), 0)
+
+        f_float32 = f.specialize(float32[::1])
+        self.assertEqual(len(f.specializations), 1)
+
+        f_float32_2 = f.specialize(float32[::1])
+        self.assertEqual(len(f.specializations), 1)
+        self.assertIs(f_float32, f_float32_2)
+
+        f_int32 = f.specialize(int32[::1])
+        self.assertEqual(len(f.specializations), 2)
+        self.assertIsNot(f_int32, f_float32)
+
+    def test_specialize_cache_same_with_ordering(self):
+        # Ensure that the same dispatcher is returned for the same argument
+        # types, and that different dispatchers are returned for different
+        # argument types, taking into account array ordering and multiple
+        # arguments.
+        @cuda.jit
+        def f(x, y):
+            pass
+
+        self.assertEqual(len(f.specializations), 0)
+
+        # 'A' order specialization
+        f_f32a_f32a = f.specialize(float32[:], float32[:])
+        self.assertEqual(len(f.specializations), 1)
+
+        # 'C' order specialization
+        f_f32c_f32c = f.specialize(float32[::1], float32[::1])
+        self.assertEqual(len(f.specializations), 2)
+        self.assertIsNot(f_f32a_f32a, f_f32c_f32c)
+
+        # Reuse 'C' order specialization
+        f_f32c_f32c_2 = f.specialize(float32[::1], float32[::1])
+        self.assertEqual(len(f.specializations), 2)
+        self.assertIs(f_f32c_f32c, f_f32c_f32c_2)
+
+
+class TestDispatcher(CUDATestCase):
+    """Most tests based on those in numba.tests.test_dispatcher."""
+
+    def test_coerce_input_types(self):
+        # Do not allow unsafe conversions if we can still compile other
+        # specializations.
+        c_add = cuda.jit(add_kernel)
+
+        # Using a complex128 allows us to represent any result produced by the
+        # test
+        r = np.zeros(1, dtype=np.complex128)
+
+        c_add[1, 1](r, 123, 456)
+        self.assertEqual(r[0], add(123, 456))
+
+        c_add[1, 1](r, 12.3, 45.6)
+        self.assertEqual(r[0], add(12.3, 45.6))
+
+        c_add[1, 1](r, 12.3, 45.6j)
+        self.assertEqual(r[0], add(12.3, 45.6j))
+
+        c_add[1, 1](r, 12300000000, 456)
+        self.assertEqual(r[0], add(12300000000, 456))
+
+        # Now force compilation of only a single specialization
+        c_add = cuda.jit("(i4[::1], i4, i4)")(add_kernel)
+        r = np.zeros(1, dtype=np.int32)
+
+        c_add[1, 1](r, 123, 456)
+        self.assertPreciseEqual(r[0], add(123, 456))
+
+    @skip_on_cudasim("Simulator ignores signature")
+    @unittest.expectedFailure
+    def test_coerce_input_types_unsafe(self):
+        # Implicit (unsafe) conversion of float to int, originally from
+        # test_coerce_input_types. This test presently fails with the CUDA
+        # Dispatcher because argument preparation is done by
+        # _Kernel._prepare_args, which is currently inflexible with respect to
+        # the types it can accept when preparing.
+        #
+        # This test is marked as xfail until future changes enable this
+        # behavior.
+        c_add = cuda.jit("(i4[::1], i4, i4)")(add_kernel)
+        r = np.zeros(1, dtype=np.int32)
+
+        c_add[1, 1](r, 12.3, 45.6)
+        self.assertPreciseEqual(r[0], add(12, 45))
+
+    @skip_on_cudasim("Simulator ignores signature")
+    def test_coerce_input_types_unsafe_complex(self):
+        # Implicit conversion of complex to int disallowed
+        c_add = cuda.jit("(i4[::1], i4, i4)")(add_kernel)
+        r = np.zeros(1, dtype=np.int32)
+
+        with self.assertRaises(TypeError):
+            c_add[1, 1](r, 12.3, 45.6j)
+
+    @skip_on_cudasim("Simulator does not track overloads")
+    def test_ambiguous_new_version(self):
+        """Test compiling new version in an ambiguous case"""
+        c_add = cuda.jit(add_kernel)
+
+        r = np.zeros(1, dtype=np.float64)
+        INT = 1
+        FLT = 1.5
+
+        c_add[1, 1](r, INT, FLT)
+        self.assertAlmostEqual(r[0], INT + FLT)
+        self.assertEqual(len(c_add.overloads), 1)
+
+        c_add[1, 1](r, FLT, INT)
+        self.assertAlmostEqual(r[0], FLT + INT)
+        self.assertEqual(len(c_add.overloads), 2)
+
+        c_add[1, 1](r, FLT, FLT)
+        self.assertAlmostEqual(r[0], FLT + FLT)
+        self.assertEqual(len(c_add.overloads), 3)
+
+        # The following call is ambiguous because (int, int) can resolve
+        # to (float, int) or (int, float) with equal weight.
+        c_add[1, 1](r, 1, 1)
+        self.assertAlmostEqual(r[0], INT + INT)
+        self.assertEqual(
+            len(c_add.overloads), 4, "didn't compile a new version"
+        )
+
+    @skip_on_cudasim("Simulator doesn't support concurrent kernels")
+    def test_lock(self):
+        """
+        Test that (lazy) compiling from several threads at once doesn't
+        produce errors (see issue #908).
+        """
+        errors = []
+
+        @cuda.jit
+        def foo(r, x):
+            r[0] = x + 1
+
+        def wrapper():
+            try:
+                r = np.zeros(1, dtype=np.int64)
+                foo[1, 1](r, 1)
+                self.assertEqual(r[0], 2)
+            except Exception as e:
+                errors.append(e)
+
+        threads = [threading.Thread(target=wrapper) for i in range(16)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+        self.assertFalse(errors)
+
+    def _test_explicit_signatures(self, sigs):
+        f = cuda.jit(sigs)(add_kernel)
+
+        # Exact signature matches
+        r = np.zeros(1, dtype=np.int64)
+        f[1, 1](r, 1, 2)
+        self.assertPreciseEqual(r[0], 3)
+
+        r = np.zeros(1, dtype=np.float64)
+        f[1, 1](r, 1.5, 2.5)
+        self.assertPreciseEqual(r[0], 4.0)
+
+        if config.ENABLE_CUDASIM:
+            # Pass - we can't check for no conversion on the simulator.
+            return
+
+        # No conversion
+        with self.assertRaises(TypeError) as cm:
+            r = np.zeros(1, dtype=np.complex128)
+            f[1, 1](r, 1j, 1j)
+        self.assertIn("No matching definition", str(cm.exception))
+        self.assertEqual(len(f.overloads), 2, f.overloads)
+
+    def test_explicit_signatures_strings(self):
+        # Check with a list of strings for signatures
+        sigs = [
+            "(int64[::1], int64, int64)",
+            "(float64[::1], float64, float64)",
+        ]
+        self._test_explicit_signatures(sigs)
+
+    def test_explicit_signatures_tuples(self):
+        # Check with a list of tuples of argument types for signatures
+        sigs = [(int64[::1], int64, int64), (float64[::1], float64, float64)]
+        self._test_explicit_signatures(sigs)
+
+    def test_explicit_signatures_signatures(self):
+        # Check with a list of Signature objects for signatures
+        sigs = [
+            void(int64[::1], int64, int64),
+            void(float64[::1], float64, float64),
+        ]
+        self._test_explicit_signatures(sigs)
+
+    def test_explicit_signatures_mixed(self):
+        # Check when we mix types of signature objects in a list of signatures
+
+        # Tuple and string
+        sigs = [(int64[::1], int64, int64), "(float64[::1], float64, float64)"]
+        self._test_explicit_signatures(sigs)
+
+        # Tuple and Signature object
+        sigs = [
+            (int64[::1], int64, int64),
+            void(float64[::1], float64, float64),
+        ]
+        self._test_explicit_signatures(sigs)
+
+        # Signature object and string
+        sigs = [
+            void(int64[::1], int64, int64),
+            "(float64[::1], float64, float64)",
+        ]
+        self._test_explicit_signatures(sigs)
+
+    def test_explicit_signatures_same_type_class(self):
+        # A more interesting one...
+        # (Note that the type of r is deliberately float64 in both cases so
+        # that dispatch is differentiated on the types of x and y only, to
+        # closely preserve the intent of the original test from
+        # numba.tests.test_dispatcher)
+        sigs = [
+            "(float64[::1], float32, float32)",
+            "(float64[::1], float64, float64)",
+        ]
+        f = cuda.jit(sigs)(add_kernel)
+
+        r = np.zeros(1, dtype=np.float64)
+        f[1, 1](r, np.float32(1), np.float32(2**-25))
+        self.assertPreciseEqual(r[0], 1.0)
+
+        r = np.zeros(1, dtype=np.float64)
+        f[1, 1](r, 1, 2**-25)
+        self.assertPreciseEqual(r[0], 1.0000000298023224)
+
+    @skip_on_cudasim("No overload resolution in the simulator")
+    def test_explicit_signatures_ambiguous_resolution(self):
+        # Fail to resolve ambiguity between the two best overloads
+        # (Also deliberate float64[::1] for the first argument in all cases)
+        f = cuda.jit(
+            [
+                "(float64[::1], float32, float64)",
+                "(float64[::1], float64, float32)",
+                "(float64[::1], int64, int64)",
+            ]
+        )(add_kernel)
+        with self.assertRaises(TypeError) as cm:
+            r = np.zeros(1, dtype=np.float64)
+            f[1, 1](r, 1.0, 2.0)
+
+        # The two best matches are output in the error message, as well
+        # as the actual argument types.
+        self.assertRegex(
+            str(cm.exception),
+            r"Ambiguous overloading for <function add_kernel [^>]*> "
+            r"\(Array\(float64, 1, 'C', False, aligned=True\), float64,"
+            r" float64\):\n"
+            r"\(Array\(float64, 1, 'C', False, aligned=True\), float32,"
+            r" float64\) -> none\n"
+            r"\(Array\(float64, 1, 'C', False, aligned=True\), float64,"
+            r" float32\) -> none",
+        )
+        # The integer signature is not part of the best matches
+        self.assertNotIn("int64", str(cm.exception))
+
+    @skip_on_cudasim("Simulator does not use _prepare_args")
+    @unittest.expectedFailure
+    def test_explicit_signatures_unsafe(self):
+        # These tests are from test_explicit_signatures, but have to be xfail
+        # at present because _prepare_args in the CUDA target cannot handle
+        # unsafe conversions of arguments.
+        f = cuda.jit("(int64[::1], int64, int64)")(add_kernel)
+        r = np.zeros(1, dtype=np.int64)
+
+        # Approximate match (unsafe conversion)
+        f[1, 1](r, 1.5, 2.5)
+        self.assertPreciseEqual(r[0], 3)
+        self.assertEqual(len(f.overloads), 1, f.overloads)
+
+        sigs = [
+            "(int64[::1], int64, int64)",
+            "(float64[::1], float64, float64)",
+        ]
+        f = cuda.jit(sigs)(add_kernel)
+        r = np.zeros(1, dtype=np.float64)
+        # Approximate match (int32 -> float64 is a safe conversion)
+        f[1, 1](r, np.int32(1), 2.5)
+        self.assertPreciseEqual(r[0], 3.5)
+
+    def add_device_usecase(self, sigs):
+        # Generate a kernel that calls the add device function compiled with a
+        # given set of signatures
+        add_device = cuda.jit(sigs, device=True)(add)
+
+        @cuda.jit
+        def f(r, x, y):
+            r[0] = add_device(x, y)
+
+        return f
+
+    def test_explicit_signatures_device(self):
+        # Tests similar to test_explicit_signatures, but on a device function
+        # instead of a kernel
+        sigs = ["(int64, int64)", "(float64, float64)"]
+        f = self.add_device_usecase(sigs)
+
+        # Exact signature matches
+        r = np.zeros(1, dtype=np.int64)
+        f[1, 1](r, 1, 2)
+        self.assertPreciseEqual(r[0], 3)
+
+        r = np.zeros(1, dtype=np.float64)
+        f[1, 1](r, 1.5, 2.5)
+        self.assertPreciseEqual(r[0], 4.0)
+
+        if config.ENABLE_CUDASIM:
+            # Pass - we can't check for no conversion on the simulator.
+            return
+
+        # No conversion
+        with self.assertRaises(TypingError) as cm:
+            r = np.zeros(1, dtype=np.complex128)
+            f[1, 1](r, 1j, 1j)
+
+        msg = str(cm.exception)
+        self.assertIn("Invalid use of type", msg)
+        self.assertIn("with parameters (complex128, complex128)", msg)
+        self.assertEqual(len(f.overloads), 2, f.overloads)
+
+    def test_explicit_signatures_device_same_type_class(self):
+        # A more interesting one...
+        # (Note that the type of r is deliberately float64 in both cases so
+        # that dispatch is differentiated on the types of x and y only, to
+        # closely preserve the intent of the original test from
+        # numba.tests.test_dispatcher)
+        sigs = ["(float32, float32)", "(float64, float64)"]
+        f = self.add_device_usecase(sigs)
+
+        r = np.zeros(1, dtype=np.float64)
+        f[1, 1](r, np.float32(1), np.float32(2**-25))
+        self.assertPreciseEqual(r[0], 1.0)
+
+        r = np.zeros(1, dtype=np.float64)
+        f[1, 1](r, 1, 2**-25)
+        self.assertPreciseEqual(r[0], 1.0000000298023224)
+
+    def test_explicit_signatures_device_ambiguous(self):
+        # Ambiguity between the two best overloads resolves. This is somewhat
+        # surprising given that ambiguity is not permitted for dispatching
+        # overloads when launching a kernel, but seems to be the general
+        # behaviour of Numba (See Issue #8307:
+        # https://github.com/numba/numba/issues/8307).
+        sigs = ["(float32, float64)", "(float64, float32)", "(int64, int64)"]
+        f = self.add_device_usecase(sigs)
+
+        r = np.zeros(1, dtype=np.float64)
+        f[1, 1](r, 1.5, 2.5)
+        self.assertPreciseEqual(r[0], 4.0)
+
+    @skip_on_cudasim("CUDA Simulator does not force casting")
+    def test_explicit_signatures_device_unsafe(self):
+        # These tests are from test_explicit_signatures. The device function
+        # variant of these tests can succeed on CUDA because the compilation
+        # can handle unsafe casting (c.f. test_explicit_signatures_unsafe which
+        # has to xfail due to _prepare_args not supporting unsafe casting).
+        sigs = ["(int64, int64)"]
+        f = self.add_device_usecase(sigs)
+
+        # Approximate match (unsafe conversion)
+        r = np.zeros(1, dtype=np.int64)
+        f[1, 1](r, 1.5, 2.5)
+        self.assertPreciseEqual(r[0], 3)
+        self.assertEqual(len(f.overloads), 1, f.overloads)
+
+        sigs = ["(int64, int64)", "(float64, float64)"]
+        f = self.add_device_usecase(sigs)
+
+        # Approximate match (int32 -> float64 is a safe conversion)
+        r = np.zeros(1, dtype=np.float64)
+        f[1, 1](r, np.int32(1), 2.5)
+        self.assertPreciseEqual(r[0], 3.5)
+
+    def test_dispatcher_docstring(self):
+        # Ensure that CUDA-jitting a function preserves its docstring. See
+        # Issue #5902: https://github.com/numba/numba/issues/5902
+
+        @cuda.jit
+        def add_kernel(a, b):
+            """Add two integers, kernel version"""
+
+        @cuda.jit(device=True)
+        def add_device(a, b):
+            """Add two integers, device version"""
+
+        self.assertEqual("Add two integers, kernel version", add_kernel.__doc__)
+        self.assertEqual("Add two integers, device version", add_device.__doc__)
+
+    @skip_on_cudasim("Cudasim does not have device pointers")
+    def test_dispatcher_cpointer_arguments(self):
+        ptr = types.CPointer(types.int32)
+        sig = void(ptr, int32, ptr, ptr, uint32)
+
+        @cuda.jit(sig)
+        def axpy(r, a, x, y, n):
+            i = cuda.grid(1)
+            if i < n:
+                r[i] = a * x[i] + y[i]
+
+        N = 16
+        a = 5
+        hx = np.arange(10, dtype=np.int32)
+        hy = np.arange(10, dtype=np.int32) * 2
+        dx = cuda.to_device(hx)
+        dy = cuda.to_device(hy)
+        dr = cuda.device_array_like(dx)
+
+        r_ptr = dr.__cuda_array_interface__["data"][0]
+        x_ptr = dx.__cuda_array_interface__["data"][0]
+        y_ptr = dy.__cuda_array_interface__["data"][0]
+
+        axpy[1, 32](r_ptr, a, x_ptr, y_ptr, N)
+
+        expected = a * hx + hy
+        actual = dr.copy_to_host()
+        np.testing.assert_equal(expected, actual)
+
+
+@skip_on_cudasim("CUDA simulator doesn't implement kernel properties")
+class TestDispatcherKernelProperties(CUDATestCase):
+    def test_get_regs_per_thread_unspecialized(self):
+        # A kernel where the register usage per thread is likely to differ
+        # between different specializations
+        @cuda.jit
+        def pi_sin_array(x, n):
+            i = cuda.grid(1)
+            if i < n:
+                x[i] = 3.14 * math.sin(x[i])
+
+        # Call the kernel with different arguments to create two different
+        # definitions within the Dispatcher object
+        N = 10
+        arr_f32 = np.zeros(N, dtype=np.float32)
+        arr_f64 = np.zeros(N, dtype=np.float64)
+
+        pi_sin_array[1, N](arr_f32, N)
+        pi_sin_array[1, N](arr_f64, N)
+
+        # Check we get a positive integer for the two different variations
+        sig_f32 = void(float32[::1], int64)
+        sig_f64 = void(float64[::1], int64)
+        regs_per_thread_f32 = pi_sin_array.get_regs_per_thread(sig_f32)
+        regs_per_thread_f64 = pi_sin_array.get_regs_per_thread(sig_f64)
+
+        self.assertIsInstance(regs_per_thread_f32, int)
+        self.assertIsInstance(regs_per_thread_f64, int)
+
+        self.assertGreater(regs_per_thread_f32, 0)
+        self.assertGreater(regs_per_thread_f64, 0)
+
+        # Check that getting the registers per thread for all signatures
+        # provides the same values as getting the registers per thread for
+        # individual signatures.
+        regs_per_thread_all = pi_sin_array.get_regs_per_thread()
+        self.assertEqual(regs_per_thread_all[sig_f32.args], regs_per_thread_f32)
+        self.assertEqual(regs_per_thread_all[sig_f64.args], regs_per_thread_f64)
+
+        if regs_per_thread_f32 == regs_per_thread_f64:
+            # If the register usage is the same for both variants, there may be
+            # a bug, but this may also be an artifact of the compiler / driver
+            # / device combination, so produce an informational message only.
+            print("f32 and f64 variant thread usages are equal.")
+            print("This may warrant some investigation. Devices:")
+            cuda.detect()
+
+    def test_get_regs_per_thread_specialized(self):
+        @cuda.jit(void(float32[::1], int64))
+        def pi_sin_array(x, n):
+            i = cuda.grid(1)
+            if i < n:
+                x[i] = 3.14 * math.sin(x[i])
+
+        # Check we get a positive integer for the specialized variation
+        regs_per_thread = pi_sin_array.get_regs_per_thread()
+        self.assertIsInstance(regs_per_thread, int)
+        self.assertGreater(regs_per_thread, 0)
+
+    def test_get_const_mem_unspecialized(self):
+        @cuda.jit
+        def const_fmt_string(val, to_print):
+            # We guard the print with a conditional to prevent noise from the
+            # test suite
+            if to_print:
+                print(val)
+
+        # Call the kernel with different arguments to create two different
+        # definitions within the Dispatcher object
+        const_fmt_string[1, 1](1, False)
+        const_fmt_string[1, 1](1.0, False)
+
+        # Check we get a positive integer for the two different variations
+        sig_i64 = void(int64, boolean)
+        sig_f64 = void(float64, boolean)
+        const_mem_size_i64 = const_fmt_string.get_const_mem_size(sig_i64)
+        const_mem_size_f64 = const_fmt_string.get_const_mem_size(sig_f64)
+
+        self.assertIsInstance(const_mem_size_i64, int)
+        self.assertIsInstance(const_mem_size_f64, int)
+
+        # 6 bytes for the equivalent of b'%lld\n\0'
+        self.assertGreaterEqual(const_mem_size_i64, 6)
+        # 4 bytes for the equivalent of b'%f\n\0'
+        self.assertGreaterEqual(const_mem_size_f64, 4)
+
+        # Check that getting the const memory size for all signatures
+        # provides the same values as getting the const memory size for
+        # individual signatures.
+
+        const_mem_size_all = const_fmt_string.get_const_mem_size()
+        self.assertEqual(const_mem_size_all[sig_i64.args], const_mem_size_i64)
+        self.assertEqual(const_mem_size_all[sig_f64.args], const_mem_size_f64)
+
+    def test_get_const_mem_specialized(self):
+        arr = np.arange(32, dtype=np.int64)
+        sig = void(int64[::1])
+
+        @cuda.jit(sig)
+        def const_array_use(x):
+            C = cuda.const.array_like(arr)
+            i = cuda.grid(1)
+            x[i] = C[i]
+
+        const_mem_size = const_array_use.get_const_mem_size(sig)
+        self.assertIsInstance(const_mem_size, int)
+        self.assertGreaterEqual(const_mem_size, arr.nbytes)
+
+    def test_get_shared_mem_per_block_unspecialized(self):
+        N = 10
+
+        # A kernel where the shared memory per block is likely to differ
+        # between different specializations
+        @cuda.jit
+        def simple_smem(ary):
+            sm = cuda.shared.array(N, dtype=ary.dtype)
+            for j in range(N):
+                sm[j] = j
+            for j in range(N):
+                ary[j] = sm[j]
+
+        # Call the kernel with different arguments to create two different
+        # definitions within the Dispatcher object
+        arr_f32 = np.zeros(N, dtype=np.float32)
+        arr_f64 = np.zeros(N, dtype=np.float64)
+
+        simple_smem[1, 1](arr_f32)
+        simple_smem[1, 1](arr_f64)
+
+        sig_f32 = void(float32[::1])
+        sig_f64 = void(float64[::1])
+
+        sh_mem_f32 = simple_smem.get_shared_mem_per_block(sig_f32)
+        sh_mem_f64 = simple_smem.get_shared_mem_per_block(sig_f64)
+
+        self.assertIsInstance(sh_mem_f32, int)
+        self.assertIsInstance(sh_mem_f64, int)
+
+        self.assertEqual(sh_mem_f32, N * 4)
+        self.assertEqual(sh_mem_f64, N * 8)
+
+        # Check that getting the shared memory per block for all signatures
+        # provides the same values as getting the shared mem per block for
+        # individual signatures.
+        sh_mem_f32_all = simple_smem.get_shared_mem_per_block()
+        sh_mem_f64_all = simple_smem.get_shared_mem_per_block()
+        self.assertEqual(sh_mem_f32_all[sig_f32.args], sh_mem_f32)
+        self.assertEqual(sh_mem_f64_all[sig_f64.args], sh_mem_f64)
+
+    def test_get_shared_mem_per_block_specialized(self):
+        @cuda.jit(void(float32[::1]))
+        def simple_smem(ary):
+            sm = cuda.shared.array(100, dtype=float32)
+            i = cuda.grid(1)
+            if i == 0:
+                for j in range(100):
+                    sm[j] = j
+            cuda.syncthreads()
+            ary[i] = sm[i]
+
+        shared_mem_per_block = simple_smem.get_shared_mem_per_block()
+        self.assertIsInstance(shared_mem_per_block, int)
+        self.assertEqual(shared_mem_per_block, 400)
+
+    def test_get_max_threads_per_block_unspecialized(self):
+        N = 10
+
+        @cuda.jit
+        def simple_maxthreads(ary):
+            i = cuda.grid(1)
+            ary[i] = i
+
+        arr_f32 = np.zeros(N, dtype=np.float32)
+        simple_maxthreads[1, 1](arr_f32)
+        sig_f32 = void(float32[::1])
+        max_threads_f32 = simple_maxthreads.get_max_threads_per_block(sig_f32)
+
+        self.assertIsInstance(max_threads_f32, int)
+        self.assertGreater(max_threads_f32, 0)
+
+        max_threads_f32_all = simple_maxthreads.get_max_threads_per_block()
+        self.assertEqual(max_threads_f32_all[sig_f32.args], max_threads_f32)
+
+    def test_get_local_mem_per_thread_unspecialized(self):
+        # NOTE: A large amount of local memory must be allocated
+        # otherwise the compiler will optimize out the call to
+        # cuda.local.array and use local registers instead
+        N = 1000
+
+        @cuda.jit
+        def simple_lmem(ary):
+            lm = cuda.local.array(N, dtype=ary.dtype)
+            for j in range(N):
+                lm[j] = j
+            for j in range(N):
+                ary[j] = lm[j]
+
+        # Call the kernel with different arguments to create two different
+        # definitions within the Dispatcher object
+        arr_f32 = np.zeros(N, dtype=np.float32)
+        arr_f64 = np.zeros(N, dtype=np.float64)
+
+        simple_lmem[1, 1](arr_f32)
+        simple_lmem[1, 1](arr_f64)
+
+        sig_f32 = void(float32[::1])
+        sig_f64 = void(float64[::1])
+        local_mem_f32 = simple_lmem.get_local_mem_per_thread(sig_f32)
+        local_mem_f64 = simple_lmem.get_local_mem_per_thread(sig_f64)
+        self.assertIsInstance(local_mem_f32, int)
+        self.assertIsInstance(local_mem_f64, int)
+
+        self.assertGreaterEqual(local_mem_f32, N * 4)
+        self.assertGreaterEqual(local_mem_f64, N * 8)
+
+        # Check that getting the local memory per thread for all signatures
+        # provides the same values as getting the shared mem per block for
+        # individual signatures.
+        local_mem_all = simple_lmem.get_local_mem_per_thread()
+        self.assertEqual(local_mem_all[sig_f32.args], local_mem_f32)
+        self.assertEqual(local_mem_all[sig_f64.args], local_mem_f64)
+
+    def test_get_local_mem_per_thread_specialized(self):
+        # NOTE: A large amount of local memory must be allocated
+        # otherwise the compiler will optimize out the call to
+        # cuda.local.array and use local registers instead
+        N = 1000
+
+        @cuda.jit(void(float32[::1]))
+        def simple_lmem(ary):
+            lm = cuda.local.array(N, dtype=ary.dtype)
+            for j in range(N):
+                lm[j] = j
+            for j in range(N):
+                ary[j] = lm[j]
+
+        local_mem_per_thread = simple_lmem.get_local_mem_per_thread()
+        self.assertIsInstance(local_mem_per_thread, int)
+        self.assertGreaterEqual(local_mem_per_thread, N * 4)
+
+
+@skip_on_cudasim("Simulator does not support launch bounds")
+class TestLaunchBounds(CUDATestCase):
+    def _test_launch_bounds_common(self, launch_bounds):
+        @cuda.jit(launch_bounds=launch_bounds)
+        def f():
+            pass
+
+        # Test successful launch
+        f[1, 128]()
+
+        # Test launch bound exceeded
+        msg = "Call to cuLaunchKernel results in CUDA_ERROR_INVALID_VALUE"
+        with self.assertRaisesRegex(CudaAPIError, msg):
+            f[1, 256]()
+
+        sig = f.signatures[0]
+        ptx = f.inspect_asm(sig)
+        # Match either `.maxntid, 128, 1, 1` or `.maxntid 128` on a line by
+        # itself:
+        self.assertRegex(ptx, r".maxntid\s+128(?:,\s+1,\s+1)?\s*\n")
+
+        return ptx
+
+    def test_launch_bounds_scalar(self):
+        launch_bounds = 128
+        ptx = self._test_launch_bounds_common(launch_bounds)
+
+        self.assertNotIn(".minnctapersm", ptx)
+        self.assertNotIn(".maxclusterrank", ptx)
+
+    def test_launch_bounds_tuple(self):
+        launch_bounds = (128,)
+        ptx = self._test_launch_bounds_common(launch_bounds)
+
+        self.assertNotIn(".minnctapersm", ptx)
+        self.assertNotIn(".maxclusterrank", ptx)
+
+    def test_launch_bounds_with_min_cta(self):
+        launch_bounds = (128, 2)
+        ptx = self._test_launch_bounds_common(launch_bounds)
+
+        self.assertRegex(ptx, r".minnctapersm\s+2")
+        self.assertNotIn(".maxclusterrank", ptx)
+
+    @unittest.skipUnless(
+        cc_X_or_above(9, 0), "CC 9.0 needed for max cluster rank"
+    )
+    def test_launch_bounds_with_max_cluster_rank(self):
+        launch_bounds = (128, 2, 4)
+        ptx = self._test_launch_bounds_common(launch_bounds)
+
+        self.assertRegex(ptx, r".minnctapersm\s+2")
+        self.assertRegex(ptx, r".maxclusterrank\s+4")
+
+    def test_too_many_launch_bounds(self):
+        launch_bounds = (128, 2, 4, 8)
+        with self.assertRaisesRegex(ValueError, "Got 4 launch bounds:"):
+            cuda.jit("void()", launch_bounds=launch_bounds)(lambda: None)
+
+
+if __name__ == "__main__":
+    unittest.main()