test: use MemoryPool in testing actx

inducer · Jul 13, 2024 · b28d938 · b28d938
1 parent fb971ab
commit b28d938
Showing 1 changed file with 24 additions and 0 deletions.
diff --git a/grudge/array_context.py b/grudge/array_context.py
@@ -164,6 +164,7 @@ def __init__(self, queue, allocator=None,
             warn("No memory allocator specified, please pass one. "
                  "(Preferably a pyopencl.tools.MemoryPool in order "
                  "to reduce device allocations)", stacklevel=2)
+
         super().__init__(queue, allocator,
                 compile_trace_callback=compile_trace_callback)
 
@@ -509,11 +510,34 @@ class PytestPyOpenCLArrayContextFactory(
         _PytestPyOpenCLArrayContextFactoryWithClass):
     actx_class = PyOpenCLArrayContext
 
+    def __call__(self):
+        actx = super().__call__()
+        if actx.allocator is not None:
+            return actx
+
+        from pyopencl.tools import ImmediateAllocator, MemoryPool
+        alloc = MemoryPool(ImmediateAllocator(actx.queue))
+
+        return self.actx_class(
+                actx.queue,
+                allocator=alloc,
+                force_device_scalars=self.force_device_scalars)
+
 
 class PytestPytatoPyOpenCLArrayContextFactory(
         _PytestPytatoPyOpenCLArrayContextFactory):
     actx_class = PytatoPyOpenCLArrayContext
 
+    def __call__(self):
+        actx = super().__call__()
+        if actx.allocator is not None:
+            return actx
+
+        from pyopencl.tools import ImmediateAllocator, MemoryPool
+        alloc = MemoryPool(ImmediateAllocator(actx.queue))
+
+        return self.actx_class(actx.queue, allocator=alloc)
+
 
 # deprecated
 class PytestPyOpenCLArrayContextFactoryWithHostScalars(