[Runtime] Make set_allocator thread safe (#7685)

This changes the allocator state to a `ContextVar` which is
thread-local, so calling it from different threads won't cause a race.

Note that this does mean you won't inherit allocators in child threads,
and instead would need to set it in each thread independently. For most
use cases this is probably fine though.

Author: peterbell10

python/triton/runtime/_allocation.py

@@ -1,4 +1,5 @@
 from typing import Optional, Protocol
+from contextvars import ContextVar
 
 
 class Buffer(Protocol):
@@ -20,7 +21,7 @@ def __call__(self, size: int, alignment: int, stream: Optional[int]) -> Buffer:
                            "Use triton.set_allocator to specify an allocator.")
 
 
-_allocator: Allocator = NullAllocator()
+_allocator: ContextVar[Allocator] = ContextVar("_allocator", default=NullAllocator())
 
 
 def set_allocator(allocator: Allocator):
@@ -29,4 +30,4 @@ def set_allocator(allocator: Allocator):
     require additional global memory workspace.
     """
     global _allocator
-    _allocator = allocator
+    _allocator.set(allocator)

third_party/nvidia/backend/driver.py

@@ -702,7 +702,8 @@ def __call__(self, gridX, gridY, gridZ, stream, function, *args):
         if self.global_scratch_size > 0:
             grid_size = gridX * gridY * gridZ
             alloc_size = grid_size * self.num_ctas * self.global_scratch_size
-            global_scratch = _allocation._allocator(alloc_size, self.global_scratch_align, stream)
+            alloc_fn = _allocation._allocator.get()
+            global_scratch = alloc_fn(alloc_size, self.global_scratch_align, stream)
         else:
             global_scratch = None
         self.launch(gridX, gridY, gridZ, stream, function, self.launch_cooperative_grid, self.launch_pdl,