OffloadPC subclasses PCBase. Add Offloading test to firedrake-check

Author: dsroberts

firedrake/__init__.py

@@ -89,7 +89,7 @@ def init_petsc():
     MassInvPC, PCDPC, PatchPC, PlaneSmoother, PatchSNES, P1PC, P1SNES,
     LORPC, GTMGPC, PMGPC, PMGSNES, HypreAMS, HypreADS, FDMPC,
     PoissonFDMPC, TwoLevelPC, HiptmairPC, FacetSplitPC, BDDCPC,
-    CovariancePC
+    CovariancePC, OffloadPC
 )
 from firedrake.mesh import (  # noqa: F401
     Mesh, ExtrudedMesh, VertexOnlyMesh, RelabeledMesh,

firedrake/preconditioners/__init__.py

@@ -9,6 +9,7 @@
     AssembledPC, AuxiliaryOperatorPC
 )
 from firedrake.preconditioners.massinv import MassInvPC  # noqa: F401
+from firedrake.preconditioners.offload import OffloadPC  # noqa: F401
 from firedrake.preconditioners.pcd import PCDPC  # noqa: F401
 from firedrake.preconditioners.patch import (  # noqa: F401
     PatchPC, PlaneSmoother, PatchSNES

firedrake/preconditioners/offload.py

@@ -1,45 +1,20 @@
-from firedrake.preconditioners.assembled import AssembledPC
+from firedrake.preconditioners.assembled import PCBase
 from firedrake.petsc import PETSc
-from firedrake.utils import device_matrix_type
-from firedrake.logging import logger
-from functools import cache
-import warnings
+from firedrake.utils import device_matrix_type, get_device_type
 
 import firedrake.dmhooks as dmhooks
 
 __all__ = ("OffloadPC",)
 
 
-@cache
-def offload_mat_type(pc_comm_rank) -> str | None:
-    mat_type = device_matrix_type()
-    if mat_type is None:
-        if pc_comm_rank == 0:
-            warnings.warn(
-                "This installation of Firedrake is not GPU-enabled, therefore OffloadPC"
-                "will do nothing. For this preconditioner to function correctly PETSc"
-                "will need to be rebuilt with some GPU capability (e.g. '--with-cuda=1')."
-            )
-        return None
-    try:
-        dev = PETSc.Device.create()
-    except PETSc.Error:
-        if pc_comm_rank == 0:
-            logger.warning(
-                "This installation of Firedrake is GPU-enabled, but no GPU device has"
-                "been detected. OffloadPC will do nothing on this host"
-            )
-        return None
-    if dev.getDeviceType() == "HOST":
-        raise RuntimeError(
-            "A GPU-enabled Firedrake build has been detected, and GPU hardware has been"
-            "detected but a GPU device was unable to be initialised."
-        )
-    dev.destroy()
-    return mat_type
+_device_vector_impls = {
+    "CUDA": {
+        "createWithArrays": "createCUDAWithArrays",
+    }
+}
 
 
-class OffloadPC(AssembledPC):
+class OffloadPC(PCBase):
     """Offload PC from CPU to GPU and back.
 
     Internally this makes a PETSc PC object that can be controlled by
@@ -50,44 +25,84 @@ class OffloadPC(AssembledPC):
 
     def initialize(self, pc):
         # Check if our PETSc installation is GPU enabled
-        super().initialize(pc)
-        self.offload_mat_type = offload_mat_type(pc.comm.rank)
-        if self.offload_mat_type is not None:
+
+        A, P = pc.getOperators()
+
+        if pc.type != "python":
+            raise ValueError("Expecting PC type python")
+        opc = pc
+        if P.type == "python":
+            context = P.getPythonContext()
+            # It only makes sense to preconditioner/invert a diagonal
+            # block in general.  That's all we're going to allow.
+            if not context.on_diag:
+                raise ValueError("Only makes sense to invert diagonal block")
+
+        prefix = pc.getOptionsPrefix() or ""
+        options_prefix = prefix + self._prefix
+
+        self.device_mat = device_matrix_type(pc.comm.rank == 0)
+        dm = opc.getDM()
+
+        pc = PETSc.PC().create(comm=opc.comm)
+        pc.setDM(dm)
+        pc.setOptionsPrefix(options_prefix)
+        if self.device_mat is not None:
             with PETSc.Log.Event("Event: initialize offload"):
-                A, P = pc.getOperators()
+                P_dev = P.convert(mat_type=self.device_mat)
+                A_dev = A.convert(mat_type=self.device_mat)
+            P_dev.setNullSpace(P.getNullSpace())
+            P_dev.setTransposeNullSpace(P.getTransposeNullSpace())
+            P_dev.setNearNullSpace(P.getNearNullSpace())
+            self.vector_impls = _device_vector_impls[get_device_type()]
+            pc.setOperators(A_dev, P_dev)
+        else:
+            pc.setOperators(A, P)
 
-                # Convert matrix to ajicusparse
-                with PETSc.Log.Event("Event: matrix offload"):
-                    P_cu = P.convert(self.offload_mat_type)  # todo
+        # Simplest reconstruction we can manage
+        octx = dmhooks.get_appctx(dm)
+        self._ctx_ref = octx.reconstruct(
+            problem=None, mat_type=self.device_mat, pmat_type=self.device_mat
+        )
+        self.pc = pc
 
-            # Transfer nullspace
-            P_cu.setNullSpace(P.getNullSpace())
-            P_cu.setTransposeNullSpace(P.getTransposeNullSpace())
-            P_cu.setNearNullSpace(P.getNearNullSpace())
+        with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref, save=False):
+            self.pc.setFromOptions()
 
-            # Update preconditioner with GPU matrix
-            self.pc.setOperators(A, P_cu)
+    def update(self, pc):
+        A, P = pc.getOperators()
+        A_dev, P_dev = self.pc.getOperators()
+        P.copy(P_dev)
+        A.copy(A_dev)
 
     # Convert vectors to CUDA, solve and get solution on CPU back
     def apply(self, pc, x, y):
-        if self.offload_mat_type is None:
+        if self.device_mat is None:
             self.pc.apply(x, y)
         else:
             with PETSc.Log.Event("Event: apply offload"):  #
                 dm = pc.getDM()
                 with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref):
                     with PETSc.Log.Event("Event: vectors offload"):
-                        y_cu = PETSc.Vec()  # begin
-                        y_cu.createCUDAWithArrays(y)
-                        x_cu = PETSc.Vec()
-                        # Passing a vec into another vec doesnt work because original is locked
-                        x_cu.createCUDAWithArrays(x.array_r)
+                        # Create the to-be-offloaded vector
+                        y_dev = PETSc.Vec()
+                        # Use device implementation of 'createWithArrays' function
+                        getattr(y_dev, self.vector_impls["createWithArrays"])(
+                            y.array_r, None
+                        )
+                        # Create the to-be-offloaded vector
+                        x_dev = PETSc.Vec()
+                        # Use device implementation of 'createWithArrays' function
+                        getattr(x_dev, self.vector_impls["createWithArrays"])(
+                            x.array_r, None
+                        )
                     with PETSc.Log.Event("Event: solve"):
-                        self.pc.apply(x_cu, y_cu)
-                        # Calling data to synchronize vector
-                        tmp = y_cu.array_r  # noqa: F841
-                with PETSc.Log.Event("Event: vectors copy back"):
-                    y.copy(y_cu)  #
+                        self.pc.apply(x_dev, y_dev)
+                    with PETSc.Log.Event("Event: vectors copy back"):
+                        # y is already designated as host storage for y_dev, so calling
+                        # getArray is sufficient to synchronise the vector on the device
+                        # with y on the host
+                        y_dev.getArray(True)
 
     def applyTranspose(self, pc, X, Y):
         raise NotImplementedError

firedrake/solving_utils.py

@@ -279,7 +279,7 @@ def __init__(self, problem,
 
     def reconstruct(self, problem=None, mat_type=None, pmat_type=None, **kwargs):
         """Reconstruct this _SNESContext instance with new arguments."""
-        problem = problem or self.problem
+        problem = problem or self._problem
         mat_type = mat_type or self.mat_type
         pmat_type = pmat_type or self.pmat_type
         kwargs.setdefault("sub_mat_type", self.sub_mat_type)

firedrake/utils.py

@@ -26,6 +26,29 @@
 SLATE_SUPPORTS_COMPLEX = False
 
 
+@cache
+def get_device_type() -> str | None:
+    r"""Get PETSc device type
+
+    Attempt to initialise a GPU and return the type of GPU
+    identified by PETSc
+
+    Returns
+    -------
+    str | None
+        The PETSc device type
+    """
+    try:
+        dev = PETSc.Device.create()
+    except PETSc.Error:
+        # Could not initialise device - not a failure condition as this could
+        # be a GPU-enabled PETSc installation running on a CPU-only host.
+        return None
+    dev_type = dev.getDeviceType()
+    dev.destroy()
+    return dev_type
+
+
 @cache
 def device_matrix_type(warn: bool = True) -> str | None:
     r"""Get device matrix type
@@ -39,7 +62,7 @@ def device_matrix_type(warn: bool = True) -> str | None:
     ----------
     warn
         Emit a warning containing the reason a device mat_type
-        has not been returned. Defaults to False.
+        has not been returned. Defaults to True.
 
     Raises
     ------
@@ -55,18 +78,13 @@ def device_matrix_type(warn: bool = True) -> str | None:
 
     """
     _device_mat_type_map = {"HOST": None, "CUDA": "aijcusparse"}
-    try:
-        dev = PETSc.Device.create()
-    except PETSc.Error:
-        # Could not initialise device - not a failure condition as this could
-        # be a GPU-enabled PETSc installation running on a CPU-only host.
+    dev_type = get_device_type()
+    if dev_type is None:
         if warn:
             warnings.warn(
                 "This installation of Firedrake is GPU-enabled, but no GPU device has been detected"
             )
         return None
-    dev_type = dev.getDeviceType()
-    dev.destroy()
     if dev_type not in _device_mat_type_map:
         raise UnrecognisedDeviceError(
             f"Unknown device type: {dev_type} initialised by PETSc. Firedrake "

scripts/firedrake-check

@@ -23,6 +23,8 @@ TESTS = {
         "tests/firedrake/regression/test_matrix_free.py::test_fieldsplitting[parameters3-cofunc_rhs-variational]",
         # near nullspace
         "tests/firedrake/regression/test_nullspace.py::test_near_nullspace",
+        # GPU offload
+        "tests/firedrake/offload/test_poisson_offloading_pc.py::test_poisson_offload",
     ),
     2: (
         # HDF5/checkpointing

setup.py

@@ -240,6 +240,7 @@ def extensions():
     "tests/firedrake/regression/test_dg_advection.py",
     "tests/firedrake/regression/test_interpolate_cross_mesh.py",
     "tests/firedrake/output/test_io_function.py",
+    "tests/firedrake/offload/test_poisson_offloading_pc.py",
 )