add MPICupyArrayContext

grudge/array_context.py

@@ -3,6 +3,7 @@
 .. autoclass:: PytatoPyOpenCLArrayContext
 .. autoclass:: MPIBasedArrayContext
 .. autoclass:: MPIPyOpenCLArrayContext
+.. autoclass:: MPICupyArrayContext
 .. class:: MPIPytatoArrayContext
 .. autofunction:: get_reasonable_array_context_class
 """
@@ -97,6 +98,8 @@
 from arraycontext.container import ArrayContainer
 from arraycontext.impl.pytato.compile import LazilyPyOpenCLCompilingFunctionCaller
 
+from arraycontext import CupyArrayContext
+
 if TYPE_CHECKING:
     import pytato as pt
     from pytato.partition import PartId
@@ -421,6 +424,26 @@ def clone(self):
 # }}}
 
 
+# {{{
+
+class MPICupyArrayContext(CupyArrayContext, MPIBasedArrayContext):
+    """An array context for using distributed computation with :mod:`cupy`
+    eager evaluation.
+
+    .. autofunction:: __init__
+    """
+
+    def __init__(self, mpi_communicator):
+        super().__init__()
+
+        self.mpi_communicator = mpi_communicator
+
+    def clone(self):
+        return type(self)(self.mpi_communicator)
+
+# }}}
+
+
 # {{{ distributed + pyopencl
 
 class MPIPyOpenCLArrayContext(PyOpenCLArrayContext, MPIBasedArrayContext):