Minor updates to OffloadPC

firedrake/preconditioners/offload.py

@@ -16,6 +16,11 @@
     },
 }
 
+# These matrix types require an expensive implicit -> dense -> sparse conversion when
+# offloaded to a GPU. The A matrix does not need to be offloaded, therefore if the A
+# matrix is any of these types, do not offload it.
+_no_offload_mat_types = ("python", "schurcomplement")
+
 
 class OffloadPC(PCBase):
     """Offload PC from CPU to GPU and back.
@@ -56,8 +61,15 @@ def initialize(self, pc):
         pc.setOptionsPrefix(options_prefix)
         if self.device_mat is not None:
             with PETSc.Log.Event("Event: initialize offload"):
-                P_dev = P.convert(mat_type=self.device_mat)
-                A_dev = P_dev if A.handle == P.handle else A.convert(mat_type=self.device_mat)
+                P_dev = PETSc.Mat()
+                P_dev = P.convert(mat_type=self.device_mat, out=P_dev)
+                if A.handle == P.handle:
+                    A_dev = P_dev
+                elif A.type in _no_offload_mat_types:
+                    A_dev = A
+                else:
+                    A_dev = PETSc.Mat()
+                    A_dev = A.convert(mat_type=self.device_mat, out=A_dev)
             P_dev.setNullSpace(P.getNullSpace())
             P_dev.setTransposeNullSpace(P.getTransposeNullSpace())
             P_dev.setNearNullSpace(P.getNearNullSpace())
@@ -76,9 +88,11 @@ def initialize(self, pc):
     def update(self, pc):
         A, P = pc.getOperators()
         A_dev, P_dev = self.pc.getOperators()
-        P.copy(P_dev)
-        if A_dev.handle != P_dev.handle:
-            A.copy(A_dev)
+        # Perform a value-only copy
+        P.copy(P_dev, structure=PETSc.Mat.Structure.SAME_NONZERO_PATTERN)
+        if A_dev.handle != P_dev.handle and A.type not in _no_offload_mat_types:
+            # Perform a value-only copy
+            A.copy(A_dev, structure=PETSc.Mat.Structure.SAME_NONZERO_PATTERN)
 
     # Convert vectors to CUDA, solve and get solution on CPU back
     def apply(self, pc, x, y, transpose=False):

tests/firedrake/offload/test_advection_offload.py

@@ -0,0 +1,64 @@
+
+from firedrake import *
+import pytest
+import numpy as np
+
+
+@pytest.mark.skipnogpu
+@pytest.mark.parametrize(
+    "ksp_type, pc_type",
+    [("cg", "sor"), ("cg", "gamg"), ("preonly", "lu")]
+)
+@pytest.mark.parallel([1, 2])
+def test_advection_offload(ksp_type, pc_type):
+    mesh = PeriodicIntervalMesh(100, length=2)
+    x = SpatialCoordinate(mesh)[0]
+    u_init = sin(2*pi*x)
+
+    nested_parameters = {
+        "pc_type": "ksp",
+        "ksp": {
+            "ksp_type": ksp_type,
+            "ksp_max_it": 50,
+            "ksp_rtol": 1e-5,
+            "ksp_monitor": None,
+            "pc_type": pc_type,
+            'ksp_converged_reason': None,
+        },
+    }
+    parameters = {
+        "snes_type": "ksponly",
+        "ksp_type": "preonly",
+        "pc_type": "python",
+        "pc_python_type": "firedrake.OffloadPC",
+        "offload": nested_parameters,
+    }
+
+    nu = Constant(1e-2)
+
+    V = FunctionSpace(mesh, "Lagrange", 2)
+
+    u_n1 = Function(V, name="u^{n+1}")
+    u_n = Function(V, name="u^{n}")
+    v = TestFunction(V)
+
+    u_n.interpolate(u_init)
+    dt = 0.01
+
+    F = (((u_n1 - u_n)/dt) * v + u_n1 * u_n1.dx(0) * v + nu*u_n1.dx(0)*v.dx(0))*dx
+
+    problem = NonlinearVariationalProblem(F, u_n1)
+    solver = NonlinearVariationalSolver(problem, solver_parameters=parameters)
+
+    t = 0
+    steps = 600
+    t_end = steps * dt
+
+    while t <= t_end:
+        solver.solve()
+        maxchange = sqrt(assemble((u_n - u_n1)**2 * dx))
+        u_n.assign(u_n1)
+        t += dt
+        if maxchange < 1e-4 or np.isnan(maxchange):
+            break
+    assert maxchange < 1e-4

tests/firedrake/offload/test_no_offload_A.py

@@ -0,0 +1,81 @@
+from firedrake import *
+import pytest
+
+
+@pytest.mark.skipnogpu
+def test_matfree_A_does_not_offload():
+    length = 10
+    n = 3
+    mesh = RectangleMesh(2**n, 2**n, length, 1)
+
+    P1 = FiniteElement("CG", cell=mesh.ufl_cell(), degree=1)
+    B = FiniteElement("B", cell=mesh.ufl_cell(), degree=3)
+    mini = P1 + B
+    V = VectorFunctionSpace(mesh, mini)
+    P = FunctionSpace(mesh, 'CG', 1)
+
+    W = V*P
+
+    u, p = TrialFunctions(W)
+    v, q = TestFunctions(W)
+
+    a = inner(grad(u), grad(v)) * dx - inner(p, div(v)) * dx + inner(div(u), q) * dx
+
+    f = Constant((0, 0))
+    L = inner(f, v) * dx
+
+    # No-slip velocity boundary condition on top and bottom,
+    # y == 0 and y == 1
+    noslip = Constant((0, 0))
+    bc0 = DirichletBC(W[0], noslip, (3, 4))
+
+    # Parabolic inflow y(1-y) at x = 0 in positive x direction
+    x = SpatialCoordinate(W.mesh())
+    inflow = as_vector((x[1]*(1 - x[1]), 0.0))
+    bc1 = DirichletBC(W[0], inflow, 1)
+
+    # Zero pressure at outflow at x = 1
+    bc2 = DirichletBC(W[1], 0.0, 2)
+
+    bcs = [bc0, bc1, bc2]
+
+    w = Function(W)
+
+    u, p = w.subfunctions
+
+    iterative_stokes_solver_parameters = {
+        "snes_type": "ksponly",
+        "mat_type": "matfree",
+        "ksp_type": "preonly",
+        "pc_type": "fieldsplit",
+        "pc_fieldsplit_type": "schur",
+        "pc_fieldsplit_schur_type": "full",
+        "fieldsplit_0": {
+            "ksp_type": "preonly",
+            "pc_type": "python",
+            "pc_python_type": "firedrake.AssembledPC",
+            "assembled": {
+                "pc_type": "python",
+                "pc_python_type": "firedrake.OffloadPC",
+                "offload": {
+                    "pc_type": "none",
+                    "ksp_type": "preonly",
+                }
+            }
+        },
+        "fieldsplit_1": {
+            "ksp_type": "preonly",
+            "pc_type": "none",
+        },
+    }
+
+    problem = LinearVariationalProblem(a, L, w, bcs=bcs)
+    solver = LinearVariationalSolver(problem, solver_parameters=iterative_stokes_solver_parameters)
+    solver.solve()
+
+    ksp0 = solver.snes.ksp.pc.getFieldSplitSchurGetSubKSP()[0]
+    assembled_ctx = ksp0.pc.getPythonContext()
+    offload_ctx = assembled_ctx.pc.getPythonContext()
+    A, P = offload_ctx.pc.getOperators()
+    assert P.type == "seqaijcusparse"
+    assert A.type == "python"

tests/firedrake/offload/test_poisson_offloading_pc.py

@@ -55,6 +55,10 @@ def run_test_poisson_offload(ksp_type, pc_type, homogeneous_bcs):
     problem = LinearVariationalProblem(L, R, u_f, bcs=bcs)
     solver = LinearVariationalSolver(problem, solver_parameters=parameters)
     solver.solve()
+
+    # Ensure the offload has not been done in-place
+    assert solver.snes.ksp.pc.getOperators()[1].type == "seqaij"
+
     return errornorm(u_f, sol)