From dc93dac0be8644d38c99366f02bb3ce75213a026 Mon Sep 17 00:00:00 2001 From: Atmn Patel Date: Thu, 30 Oct 2025 12:09:07 -0700 Subject: [PATCH 1/4] [Refactor][NFC] Vendor-in compiler_lock for future CUDA-specific changes --- numba_cuda/numba/cuda/compiler.py | 2 +- numba_cuda/numba/cuda/core/base.py | 2 +- numba_cuda/numba/cuda/core/compiler_lock.py | 59 +++++++++++++++++++ .../numba/cuda/core/compiler_machinery.py | 2 +- numba_cuda/numba/cuda/dispatcher.py | 2 +- numba_cuda/numba/cuda/target.py | 2 +- 6 files changed, 64 insertions(+), 5 deletions(-) create mode 100644 numba_cuda/numba/cuda/core/compiler_lock.py diff --git a/numba_cuda/numba/cuda/compiler.py b/numba_cuda/numba/cuda/compiler.py index a1653c192..13cc39eed 100644 --- a/numba_cuda/numba/cuda/compiler.py +++ b/numba_cuda/numba/cuda/compiler.py @@ -10,7 +10,7 @@ from numba.cuda.core import bytecode from numba.cuda import types from numba.cuda.core.options import ParallelOptions -from numba.core.compiler_lock import global_compiler_lock +from numba.cuda.core.compiler_lock import global_compiler_lock from numba.cuda.core.errors import NumbaWarning, NumbaInvalidConfigWarning from numba.cuda.core.interpreter import Interpreter diff --git a/numba_cuda/numba/cuda/core/base.py b/numba_cuda/numba/cuda/core/base.py index c20b50e30..13059d952 100644 --- a/numba_cuda/numba/cuda/core/base.py +++ b/numba_cuda/numba/cuda/core/base.py @@ -15,7 +15,7 @@ from numba.cuda.core import imputils, targetconfig, funcdesc from numba.cuda import cgutils, debuginfo, types, utils, datamodel, config from numba.cuda.core import errors -from numba.core.compiler_lock import global_compiler_lock +from numba.cuda.core.compiler_lock import global_compiler_lock from numba.cuda.core.pythonapi import PythonAPI from numba.cuda.core.imputils import ( user_function, diff --git a/numba_cuda/numba/cuda/core/compiler_lock.py b/numba_cuda/numba/cuda/core/compiler_lock.py new file mode 100644 index 000000000..f3da53f58 --- /dev/null +++ b/numba_cuda/numba/cuda/core/compiler_lock.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-2-Clause + +import threading +import functools +import numba.cuda.core.event as ev + + +# Lock for the preventing multiple compiler execution +class _CompilerLock(object): + def __init__(self): + self._lock = threading.RLock() + + def acquire(self): + ev.start_event("numba.cuda:compiler_lock") + self._lock.acquire() + + def release(self): + self._lock.release() + ev.end_event("numba.cuda:compiler_lock") + + def __enter__(self): + self.acquire() + + def __exit__(self, exc_val, exc_type, traceback): + self.release() + + def is_locked(self): + is_owned = getattr(self._lock, "_is_owned") + if not callable(is_owned): + is_owned = self._is_owned + return is_owned() + + def __call__(self, func): + @functools.wraps(func) + def _acquire_compile_lock(*args, **kwargs): + with self: + return func(*args, **kwargs) + + return _acquire_compile_lock + + def _is_owned(self): + # This method is borrowed from threading.Condition. + # Return True if lock is owned by current_thread. + # This method is called only if _lock doesn't have _is_owned(). + if self._lock.acquire(0): + self._lock.release() + return False + else: + return True + + +global_compiler_lock = _CompilerLock() + + +def require_global_compiler_lock(): + """Sentry that checks the global_compiler_lock is acquired.""" + # Use assert to allow turning off this check + assert global_compiler_lock.is_locked() diff --git a/numba_cuda/numba/cuda/core/compiler_machinery.py b/numba_cuda/numba/cuda/core/compiler_machinery.py index cb91f057e..cd92eaf99 100644 --- a/numba_cuda/numba/cuda/core/compiler_machinery.py +++ b/numba_cuda/numba/cuda/core/compiler_machinery.py @@ -7,8 +7,8 @@ import inspect -from numba.core.compiler_lock import global_compiler_lock from numba.cuda.core import errors +from numba.cuda.core.compiler_lock import global_compiler_lock from numba.cuda.core import config from numba.cuda import utils from numba.cuda.core import transforms diff --git a/numba_cuda/numba/cuda/dispatcher.py b/numba_cuda/numba/cuda/dispatcher.py index 50e2c60dd..4cc726237 100644 --- a/numba_cuda/numba/cuda/dispatcher.py +++ b/numba_cuda/numba/cuda/dispatcher.py @@ -19,7 +19,7 @@ from numba.cuda import serialize, utils from numba import cuda -from numba.core.compiler_lock import global_compiler_lock +from numba.cuda.core.compiler_lock import global_compiler_lock from numba.cuda.typeconv.rules import default_type_manager from numba.cuda.typing.templates import fold_arguments from numba.cuda.typing.typeof import Purpose, typeof diff --git a/numba_cuda/numba/cuda/target.py b/numba_cuda/numba/cuda/target.py index 3425bd181..dde6ab760 100644 --- a/numba_cuda/numba/cuda/target.py +++ b/numba_cuda/numba/cuda/target.py @@ -10,7 +10,7 @@ from numba.cuda import types from numba.cuda import HAS_NUMBA -from numba.core.compiler_lock import global_compiler_lock +from numba.cuda.core.compiler_lock import global_compiler_lock from numba.cuda.core.errors import NumbaWarning from numba.cuda.core.base import BaseContext from numba.cuda.typing import cmathdecl From b8e5b64bbc7fa3e0f22c05cbf95203c777d804ce Mon Sep 17 00:00:00 2001 From: Atmn Patel Date: Fri, 31 Oct 2025 14:52:05 -0700 Subject: [PATCH 2/4] set the numba cuda global compiler lock to also control the lock of numba --- numba_cuda/numba/cuda/core/compiler_lock.py | 58 ++++++++++++++++++++- 1 file changed, 57 insertions(+), 1 deletion(-) diff --git a/numba_cuda/numba/cuda/core/compiler_lock.py b/numba_cuda/numba/cuda/core/compiler_lock.py index f3da53f58..57dcb9759 100644 --- a/numba_cuda/numba/cuda/core/compiler_lock.py +++ b/numba_cuda/numba/cuda/core/compiler_lock.py @@ -4,6 +4,14 @@ import threading import functools import numba.cuda.core.event as ev +from numba.cuda import HAS_NUMBA + +if HAS_NUMBA: + from numba.core.compiler_lock import ( + global_compiler_lock as _numba_compiler_lock, + ) +else: + _numba_compiler_lock = None # Lock for the preventing multiple compiler execution @@ -50,7 +58,55 @@ def _is_owned(self): return True -global_compiler_lock = _CompilerLock() +_numba_cuda_compiler_lock = _CompilerLock() + + +# Wrapper that coordinates both numba and numba-cuda compiler locks +class _DualCompilerLock(object): + """Wrapper that coordinates both the numba-cuda and upstream numba compiler locks.""" + + def __init__(self, cuda_lock, numba_lock): + self._cuda_lock = cuda_lock + self._numba_lock = numba_lock + + def acquire(self): + if self._numba_lock: + self._numba_lock.acquire() + self._cuda_lock.acquire() + + def release(self): + self._cuda_lock.release() + if self._numba_lock: + self._numba_lock.release() + + def __enter__(self): + self.acquire() + + def __exit__(self, exc_val, exc_type, traceback): + self.release() + + def is_locked(self): + cuda_locked = self._cuda_lock.is_locked() + if self._numba_lock: + return cuda_locked and self._numba_lock.is_locked() + return cuda_locked + + def __call__(self, func): + @functools.wraps(func) + def _acquire_compile_lock(*args, **kwargs): + with self: + return func(*args, **kwargs) + + return _acquire_compile_lock + + +# Create the global compiler lock, wrapping both locks if numba is available +if HAS_NUMBA: + global_compiler_lock = _DualCompilerLock( + _numba_cuda_compiler_lock, _numba_compiler_lock + ) +else: + global_compiler_lock = _numba_cuda_compiler_lock def require_global_compiler_lock(): From 46d66bba0cf14f4593fccd4c9ea0be7abddd2f3c Mon Sep 17 00:00:00 2001 From: Atmn Patel Date: Mon, 3 Nov 2025 03:52:48 -0800 Subject: [PATCH 3/4] remove deadcode, change event logging prefix --- numba_cuda/numba/cuda/core/compiler_lock.py | 21 +++++---------------- 1 file changed, 5 insertions(+), 16 deletions(-) diff --git a/numba_cuda/numba/cuda/core/compiler_lock.py b/numba_cuda/numba/cuda/core/compiler_lock.py index 57dcb9759..51ba1e1a5 100644 --- a/numba_cuda/numba/cuda/core/compiler_lock.py +++ b/numba_cuda/numba/cuda/core/compiler_lock.py @@ -20,12 +20,12 @@ def __init__(self): self._lock = threading.RLock() def acquire(self): - ev.start_event("numba.cuda:compiler_lock") + ev.start_event("numba-cuda:compiler_lock") self._lock.acquire() def release(self): self._lock.release() - ev.end_event("numba.cuda:compiler_lock") + ev.end_event("numba-cuda:compiler_lock") def __enter__(self): self.acquire() @@ -70,14 +70,12 @@ def __init__(self, cuda_lock, numba_lock): self._numba_lock = numba_lock def acquire(self): - if self._numba_lock: - self._numba_lock.acquire() + self._numba_lock.acquire() self._cuda_lock.acquire() def release(self): self._cuda_lock.release() - if self._numba_lock: - self._numba_lock.release() + self._numba_lock.release() def __enter__(self): self.acquire() @@ -86,10 +84,7 @@ def __exit__(self, exc_val, exc_type, traceback): self.release() def is_locked(self): - cuda_locked = self._cuda_lock.is_locked() - if self._numba_lock: - return cuda_locked and self._numba_lock.is_locked() - return cuda_locked + return self._cuda_lock.is_locked() and self._numba_lock.is_locked() def __call__(self, func): @functools.wraps(func) @@ -107,9 +102,3 @@ def _acquire_compile_lock(*args, **kwargs): ) else: global_compiler_lock = _numba_cuda_compiler_lock - - -def require_global_compiler_lock(): - """Sentry that checks the global_compiler_lock is acquired.""" - # Use assert to allow turning off this check - assert global_compiler_lock.is_locked() From 83a8a73a5a4da22d91a20425cd3ddfe0c3117574 Mon Sep 17 00:00:00 2001 From: Atmn Patel Date: Mon, 3 Nov 2025 04:41:38 -0800 Subject: [PATCH 4/4] remove is_locked, is_owned --- numba_cuda/numba/cuda/core/compiler_lock.py | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/numba_cuda/numba/cuda/core/compiler_lock.py b/numba_cuda/numba/cuda/core/compiler_lock.py index 51ba1e1a5..f39cf5cce 100644 --- a/numba_cuda/numba/cuda/core/compiler_lock.py +++ b/numba_cuda/numba/cuda/core/compiler_lock.py @@ -33,12 +33,6 @@ def __enter__(self): def __exit__(self, exc_val, exc_type, traceback): self.release() - def is_locked(self): - is_owned = getattr(self._lock, "_is_owned") - if not callable(is_owned): - is_owned = self._is_owned - return is_owned() - def __call__(self, func): @functools.wraps(func) def _acquire_compile_lock(*args, **kwargs): @@ -47,16 +41,6 @@ def _acquire_compile_lock(*args, **kwargs): return _acquire_compile_lock - def _is_owned(self): - # This method is borrowed from threading.Condition. - # Return True if lock is owned by current_thread. - # This method is called only if _lock doesn't have _is_owned(). - if self._lock.acquire(0): - self._lock.release() - return False - else: - return True - _numba_cuda_compiler_lock = _CompilerLock() @@ -83,9 +67,6 @@ def __enter__(self): def __exit__(self, exc_val, exc_type, traceback): self.release() - def is_locked(self): - return self._cuda_lock.is_locked() and self._numba_lock.is_locked() - def __call__(self, func): @functools.wraps(func) def _acquire_compile_lock(*args, **kwargs):