JDBetteridge/cuda runner

.github/actionlint.yaml

@@ -0,0 +1,4 @@
+self-hosted-runner:
+  labels:
+    - gpu
+    - nvidia

.github/workflows/build.yml

@@ -6,6 +6,7 @@ on:
     branches:
       - master
   pull_request:
+    types: [ labeled ]
   schedule:
     - cron: '0 0 * * 0'
     - cron: '0 0 1 * *'  # Monthly release
@@ -21,6 +22,45 @@ env:
   RELEASE_TAG: latest
 
 jobs:
+  build-gpu:
+    name: "Build Firedrake for GPU"
+    runs-on: [self-hosted, gpu]
+    #if: ${{ github.event.label.name == 'gpu' }}
+    if: contains(github.event.pull_request.labels.*.name, 'gpu')
+    container:
+      image: firedrakeproject/firedrake-env:latest
+    env:
+      # PETSC_DIR and MPICH_DIR are set inside the docker image
+      FIREDRAKE_CI_TESTS: 1
+      PYOP2_CI_TESTS: 1
+      PETSC_ARCH: default
+      OMP_NUM_THREADS: 1
+      OPENBLAS_NUM_THREADS: 1
+      COMPLEX: ""
+      RDMAV_FORK_SAFE: 1
+    steps:
+      - uses: actions/checkout@v3
+      - name: Cleanup
+        if: ${{ always() }}
+        run: |
+          cd ..
+          rm -rf firedrake_venv
+      - name: Build Firedrake
+        run: |
+          cd ..
+          # Linting should ignore unquoted shell variable $COMPLEX
+          # shellcheck disable=SC2086
+          ./firedrake/scripts/firedrake-install \
+            $COMPLEX \
+            --honour-petsc-dir \
+            --mpicc="$MPICH_DIR"/mpicc \
+            --mpicxx="$MPICH_DIR"/mpicxx \
+            --mpif90="$MPICH_DIR"/mpif90 \
+            --mpiexec="$MPICH_DIR"/mpiexec \
+            --venv-name firedrake_venv \
+            --no-package-manager \
+            --disable-ssh \
+            || (cat firedrake-install.log && /bin/false)
   build:
     name: "Build Firedrake"
     # The type of runner that the job will run on

firedrake/linear_solver.py

@@ -55,6 +55,8 @@
         solver_parameters = solving_utils.set_defaults(solver_parameters,
                                                        A.arguments(),
                                                        ksp_defaults=self.DEFAULT_KSP_PARAMETERS)
+        # todo: add offload to solver parameters - how? prefix?
+
         self.A = A
         self.comm = A.comm
         self._comm = internal_comm(self.comm, self)
@@ -163,6 +165,18 @@
         else:
             acc = x.dat.vec_wo
 
+        # if "cu" in self.A.petscmat.type:  # todo: cuda or cu?
+        #     with self.inserted_options(), b.dat.vec_ro as rhs, acc as solution, dmhooks.add_hooks(self.ksp.dm, self):
+        #         b_cu = PETSc.Vec()
+        #         b_cu.createCUDAWithArrays(rhs)
+        #         u = PETSc.Vec()
+        #         u.createCUDAWithArrays(solution)
+        #         self.ksp.solve(b_cu, u)
+        #         u.getArray()
+
+        # else:
+        # instead: preconditioner
+
         with self.inserted_options(), b.dat.vec_ro as rhs, acc as solution, dmhooks.add_hooks(self.ksp.dm, self):
             self.ksp.solve(rhs, solution)
 

firedrake/preconditioners/__init__.py

@@ -12,3 +12,4 @@
 from firedrake.preconditioners.fdm import *          # noqa: F401
 from firedrake.preconditioners.hiptmair import *     # noqa: F401
 from firedrake.preconditioners.facet_split import *  # noqa: F401
+from firedrake.preconditioners.offload import *      # noqa: F401

firedrake/preconditioners/offload.py

@@ -0,0 +1,111 @@
+from firedrake.preconditioners.base import PCBase
+from firedrake.functionspace import FunctionSpace, MixedFunctionSpace
+from firedrake.petsc import PETSc
+from firedrake.ufl_expr import TestFunction, TrialFunction
+import firedrake.dmhooks as dmhooks
+from firedrake.dmhooks import get_function_space
+
+__all__ = ("OffloadPC",)
+
+
+class OffloadPC(PCBase):
+    """Offload PC from CPU to GPU and back.
+
+    Internally this makes a PETSc PC object that can be controlled by
+    options using the extra options prefix ``offload_``.
+    """
+
+    _prefix = "offload_"
+
+    def initialize(self, pc):
+        A, P = pc.getOperators()  # P preconditioner
+
+        outer_pc = pc
+        appctx = self.get_appctx(pc)
+        fcp = appctx.get("form_compiler_parameters")
+
+        V = get_function_space(pc.getDM())
+        if len(V) == 1:
+            V = FunctionSpace(V.mesh(), V.ufl_element())
+        else:
+            V = MixedFunctionSpace([V_ for V_ in V])
+        test = TestFunction(V)
+        trial = TrialFunction(V)
+
+        (a, bcs) = self.form(pc, test, trial)
+
+        if P.type == "assembled":
+            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()
+        options_prefix = prefix + self._prefix
+
+        mat_type = PETSc.Options().getString(options_prefix + "mat_type", "cusparse")
+
+        # Convert matrix to ajicusparse
+        P_cu = P.convert(mat_type='aijcusparse')
+
+        # Transfer nullspace
+        P_cu.setNullSpace(P.getNullSpace())
+        tnullsp = P.getTransposeNullSpace()
+        if tnullsp.handle != 0:
+            P_cu.setTransposeNullSpace(tnullsp)
+        P_cu.setNearNullSpace(P.getNearNullSpace())
+
+        # PC object set-up
+        pc = PETSc.PC().create(comm=outer_pc.comm)
+        pc.incrementTabLevel(1, parent=outer_pc)
+
+        # We set a DM and an appropriate SNESContext on the constructed PC
+        # so one can do e.g. multigrid or patch solves.
+        dm = outer_pc.getDM()
+        self._ctx_ref = self.new_snes_ctx(
+            outer_pc, a, bcs, mat_type,
+            fcp=fcp, options_prefix=options_prefix
+        )
+
+        pc.setDM(dm)
+        pc.setOptionsPrefix(options_prefix)
+        pc.setOperators(A, P_cu)
+        self.pc = pc
+        with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref, save=False):
+            pc.setFromOptions()
+
+    def update(self, pc):
+        _, P = pc.getOperators()
+        _, P_cu = self.pc.getOperators()
+        P.copy(P_cu)
+
+    def form(self, pc, test, trial):
+        _, P = pc.getOperators()
+        if P.getType() == "python":
+            context = P.getPythonContext()
+            return (context.a, context.row_bcs)
+        else:
+            context = dmhooks.get_appctx(pc.getDM())
+            return (context.Jp or context.J, context._problem.bcs)
+
+    # Convert vectors to CUDA, solve and get solution on CPU back
+    def apply(self, pc, x, y):
+        dm = pc.getDM()
+        with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref):
+            y_cu = PETSc.Vec()
+            y_cu.createCUDAWithArrays(y)
+            x_cu = PETSc.Vec()
+            x_cu.createCUDAWithArrays(x)
+            self.pc.apply(x_cu, y_cu)
+        y.copy(y_cu)
+
+    def applyTranspose(self, pc, X, Y):
+        raise NotImplementedError
+
+    def view(self, pc, viewer=None):
+        super().view(pc, viewer)
+        print("viewing PC")
+        if hasattr(self, "pc"):
+            viewer.printfASCII("PC to solve on GPU\n")
+            self.pc.view(viewer)

firedrake/solving.py

@@ -235,19 +235,22 @@
                              options_prefix=options_prefix)
     if isinstance(x, firedrake.Vector):
         x = x.function
-    # linear MG doesn't need RHS, supply zero.
-    lvp = vs.LinearVariationalProblem(a=A.a, L=0, u=x, bcs=A.bcs)
-    mat_type = A.mat_type
-    appctx = solver_parameters.get("appctx", {})
-    ctx = solving_utils._SNESContext(lvp,
-                                     mat_type=mat_type,
-                                     pmat_type=mat_type,
-                                     appctx=appctx,
-                                     options_prefix=options_prefix)
+    if not isinstance(A, firedrake.matrix.AssembledMatrix):
+        # linear MG doesn't need RHS, supply zero.
+        lvp = vs.LinearVariationalProblem(a=A.a, L=0, u=x, bcs=A.bcs)
+        mat_type = A.mat_type
+        appctx = solver_parameters.get("appctx", {})
+        ctx = solving_utils._SNESContext(lvp,
+                                        mat_type=mat_type,
+                                        pmat_type=mat_type,
+                                        appctx=appctx,
+                                        options_prefix=options_prefix)
+    else:
+        ctx = None
     dm = solver.ksp.dm
 
     with dmhooks.add_hooks(dm, solver, appctx=ctx):
         solver.solve(x, b)
 
 
 def _extract_linear_solver_args(*args, **kwargs):