diff --git a/firedrake/preconditioners/offload.py b/firedrake/preconditioners/offload.py index 7810e24a53..da9deb9d2f 100644 --- a/firedrake/preconditioners/offload.py +++ b/firedrake/preconditioners/offload.py @@ -13,6 +13,8 @@ } } +_no_offload_mat_types = ("python", "schurcomplement") + class OffloadPC(PCBase): """Offload PC from CPU to GPU and back. @@ -53,8 +55,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.getType() 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()) @@ -74,7 +83,7 @@ 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: + if A_dev.handle != P_dev.handle and A.getType() not in _no_offload_mat_types: A.copy(A_dev) # Convert vectors to CUDA, solve and get solution on CPU back diff --git a/tests/firedrake/offload/test_poisson_offloading_pc.py b/tests/firedrake/offload/test_poisson_offloading_pc.py index 6d4490580c..9340153a54 100644 --- a/tests/firedrake/offload/test_poisson_offloading_pc.py +++ b/tests/firedrake/offload/test_poisson_offloading_pc.py @@ -55,6 +55,9 @@ 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() + + assert solver.snes.ksp.pc.getOperators()[1].getType() == "seqaij" + return errornorm(u_f, sol)