test: use MemoryPool in testing actx

inducer · Jul 14, 2024 · fe4d0b1 · fe4d0b1
1 parent fb971ab
commit fe4d0b1
Showing 1 changed file with 21 additions and 0 deletions.
diff --git a/grudge/array_context.py b/grudge/array_context.py
@@ -131,6 +131,7 @@ def __init__(self, queue: "pyopencl.CommandQueue",
             force_device_scalars: bool = False) -> None:
 
         if allocator is None:
+            breakpoint()
             warn("No memory allocator specified, please pass one. "
                  "(Preferably a pyopencl.tools.MemoryPool in order "
                  "to reduce device allocations)", stacklevel=2)
@@ -164,6 +165,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 +511,30 @@ class PytestPyOpenCLArrayContextFactory(
         _PytestPyOpenCLArrayContextFactoryWithClass):
     actx_class = PyOpenCLArrayContext
 
+    def __call__(self):
+        from pyopencl.tools import ImmediateAllocator, MemoryPool
+
+        _ctx, queue = self.get_command_queue()
+        alloc = MemoryPool(ImmediateAllocator(queue))
+
+        return self.actx_class(
+                queue,
+                allocator=alloc,
+                force_device_scalars=self.force_device_scalars)
+
 
 class PytestPytatoPyOpenCLArrayContextFactory(
         _PytestPytatoPyOpenCLArrayContextFactory):
     actx_class = PytatoPyOpenCLArrayContext
 
+    def __call__(self):
+        _ctx, queue = self.get_command_queue()
+
+        from pyopencl.tools import ImmediateAllocator, MemoryPool
+        alloc = MemoryPool(ImmediateAllocator(queue))
+
+        return self.actx_class(queue, allocator=alloc)
+
 
 # deprecated
 class PytestPyOpenCLArrayContextFactoryWithHostScalars(