[Runtime] Make AsyncCompileMode thread safe (#7701)

In the same vein as #7685

We rely on the global variable `active_mode`, so if you were to use
`AsyncCompileMode` in two different threads they may race with each
other.

Author: peterbell10

python/triton/runtime/_allocation.py

@@ -29,5 +29,4 @@ def set_allocator(allocator: Allocator):
     The allocator function is called during kernel launch for kernels that
     require additional global memory workspace.
     """
-    global _allocator
     _allocator.set(allocator)

python/triton/runtime/_async_compile.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 from typing import Callable, Optional
 from concurrent.futures import Executor, as_completed, Future
+from contextvars import ContextVar
 
-active_mode: Optional[AsyncCompileMode] = None
+active_mode: ContextVar[Optional[AsyncCompileMode]] = ContextVar("async_compile_active_mode", default=None)
 
 
 class FutureKernel:
@@ -42,14 +43,13 @@ def submit(self, key, compile_fn, finalize_fn):
         return future_kernel
 
     def __enter__(self):
-        global active_mode
-        if active_mode is not None:
+        if active_mode.get() is not None:
             raise RuntimeError("Another AsyncCompileMode is already active")
-        active_mode = self
+        active_mode.set(self)
+        return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        global active_mode
         # Finalize any outstanding compiles
         for future in as_completed(self.raw_futures):
             self.future_kernels[future._key].result()
-        active_mode = None
+        active_mode.set(None)

python/triton/runtime/jit.py

@@ -757,7 +757,7 @@ def _do_compile(self, key, signature, device, constexprs, options, attrs, warmup
             return None
         src = self.ASTSource(self, signature, constexprs, attrs)
 
-        async_mode = _async_compile.active_mode
+        async_mode = _async_compile.active_mode.get()
         if async_mode is not None:
 
             env_vars = get_cache_invalidating_env_vars()