From c221592019ff31ca89464edde0afee937bfcd5e9 Mon Sep 17 00:00:00 2001 From: Kaeun Kim Date: Wed, 3 Dec 2025 00:33:57 +0000 Subject: [PATCH 1/2] fix: cache config enum --- numba_cuda/numba/cuda/cudadrv/driver.py | 2 +- .../cuda/tests/cudadrv/test_cuda_driver.py | 49 +++++++++++++++++++ 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/numba_cuda/numba/cuda/cudadrv/driver.py b/numba_cuda/numba/cuda/cudadrv/driver.py index ee9bac927..d40c62d0f 100644 --- a/numba_cuda/numba/cuda/cudadrv/driver.py +++ b/numba_cuda/numba/cuda/cudadrv/driver.py @@ -2393,7 +2393,7 @@ def cache_config( self, prefer_equal=False, prefer_cache=False, prefer_shared=False ): prefer_equal = prefer_equal or (prefer_cache and prefer_shared) - attr = binding.CUfunction_attribute + attr = binding.CUfunc_cache if prefer_equal: flag = attr.CU_FUNC_CACHE_PREFER_EQUAL elif prefer_cache: diff --git a/numba_cuda/numba/cuda/tests/cudadrv/test_cuda_driver.py b/numba_cuda/numba/cuda/tests/cudadrv/test_cuda_driver.py index d15dca6bd..6897d18e0 100644 --- a/numba_cuda/numba/cuda/tests/cudadrv/test_cuda_driver.py +++ b/numba_cuda/numba/cuda/tests/cudadrv/test_cuda_driver.py @@ -285,6 +285,55 @@ def b2d(bs): self.assertTrue(grid > 0) self.assertTrue(block > 0) + def test_cuda_jit_cache_config(self): + """Test cache_config with @cuda.jit decorated functions""" + from numba import types + + sig = (types.float32[::1], types.float32[::1]) + + @cuda.jit(sig) + def add_one(r, x): + i = cuda.grid(1) + if i < len(r): + r[i] = x[i] + 1 + + kernel = add_one.overloads[sig] + cufunc = kernel._codelibrary.get_cufunc() + + configs_to_test = [ + ("prefer_shared", dict(prefer_shared=True)), + ("prefer_cache", dict(prefer_cache=True)), + ("prefer_equal", dict(prefer_equal=True)), + ("default", dict()), + ] + + for name, kwargs in configs_to_test: + with self.subTest(config=name): + try: + cufunc.cache_config(**kwargs) + except Exception as e: + self.fail(f"cache_config({name}) failed: {e}") + + import numpy as np + + x = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float32) + r = np.zeros_like(x) + + d_x = cuda.to_device(x) + d_r = cuda.to_device(r) + + cufunc.cache_config(prefer_shared=True) + add_one[1, 5](d_r, d_x) + + result = d_r.copy_to_host() + expected = x + 1 + + np.testing.assert_array_almost_equal( + result, + expected, + err_msg="Kernel produced incorrect results after cache_config", + ) + class TestDevice(CUDATestCase): def test_device_get_uuid(self): From a23a7946c826195ce8e68eb41d0cf09513d451be Mon Sep 17 00:00:00 2001 From: Kaeun Kim Date: Wed, 3 Dec 2025 01:08:36 +0000 Subject: [PATCH 2/2] fix: test name --- numba_cuda/numba/cuda/tests/cudadrv/test_cuda_driver.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/numba_cuda/numba/cuda/tests/cudadrv/test_cuda_driver.py b/numba_cuda/numba/cuda/tests/cudadrv/test_cuda_driver.py index 6897d18e0..4a450d7d9 100644 --- a/numba_cuda/numba/cuda/tests/cudadrv/test_cuda_driver.py +++ b/numba_cuda/numba/cuda/tests/cudadrv/test_cuda_driver.py @@ -285,9 +285,9 @@ def b2d(bs): self.assertTrue(grid > 0) self.assertTrue(block > 0) - def test_cuda_jit_cache_config(self): - """Test cache_config with @cuda.jit decorated functions""" + def test_cuda_cache_config(self): from numba import types + import numpy as np sig = (types.float32[::1], types.float32[::1]) @@ -314,8 +314,6 @@ def add_one(r, x): except Exception as e: self.fail(f"cache_config({name}) failed: {e}") - import numpy as np - x = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float32) r = np.zeros_like(x)