From b28d93826f6deea222dea36ac8882efc15331e28 Mon Sep 17 00:00:00 2001
From: Alexandru Fikl <alexfikl@gmail.com>
Date: Sat, 13 Jul 2024 11:09:38 +0300
Subject: [PATCH] test: use MemoryPool in testing actx

---
 grudge/array_context.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/grudge/array_context.py b/grudge/array_context.py
index 6fc44bf66..d303a4818 100644
--- 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(