diff --git a/.github/actionlint.yaml b/.github/actionlint.yaml new file mode 100644 index 0000000000..deafdc92ad --- /dev/null +++ b/.github/actionlint.yaml @@ -0,0 +1,4 @@ +self-hosted-runner: + labels: + - gpu + - nvidia diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 69768243e0..2560021550 100644 --- a/.github/workflows/build.yml +++ b/.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 diff --git a/firedrake/linear_solver.py b/firedrake/linear_solver.py index d423b36bbb..d0e9dc81a4 100644 --- a/firedrake/linear_solver.py +++ b/firedrake/linear_solver.py @@ -55,6 +55,8 @@ def __init__(self, A, *, P=None, solver_parameters=None, 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 @@ def solve(self, x, b): 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) diff --git a/firedrake/preconditioners/__init__.py b/firedrake/preconditioners/__init__.py index cd75ae7380..ca04bd9cbd 100644 --- a/firedrake/preconditioners/__init__.py +++ b/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 diff --git a/firedrake/preconditioners/offload.py b/firedrake/preconditioners/offload.py new file mode 100644 index 0000000000..c62515853d --- /dev/null +++ b/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) diff --git a/firedrake/solving.py b/firedrake/solving.py index b5ece29988..369ada111c 100644 --- a/firedrake/solving.py +++ b/firedrake/solving.py @@ -235,15 +235,18 @@ def _la_solve(A, x, b, **kwargs): 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):