From f084463741e1d4b78f721f97155de4ae6c5ca682 Mon Sep 17 00:00:00 2001
From: Dale Roberts <dale.roberts1@anu.edu.au>
Date: Fri, 6 Mar 2026 16:50:51 +1100
Subject: [PATCH 01/12] GPU Offload PC

First pass at firedrake-configure for GPU-enabled PETSc builds, as well
as a github actions test. OffloadPC is more or less unchanged from
olender's work, but with a few extra checks for whether GPU offload is
available. No dedicated GPU tests yet. CUDA only so far.
---
 .github/workflows/core.yml            | 127 ++++++++++++++
 firedrake/preconditioners/__init__.py |   1 +
 firedrake/preconditioners/offload.py  |  92 ++++++++++
 scripts/firedrake-configure           | 239 ++++++++++++++++++++++----
 4 files changed, 422 insertions(+), 37 deletions(-)
 create mode 100644 firedrake/preconditioners/offload.py

diff --git a/.github/workflows/core.yml b/.github/workflows/core.yml
index c0682518a4..949cb8724d 100644
--- a/.github/workflows/core.yml
+++ b/.github/workflows/core.yml
@@ -23,6 +23,10 @@ on:
         description: Whether to test using macOS
         type: boolean
         default: false
+      test_cuda:
+        description: Whether to test using CUDA-enabled PETSc
+        type: boolean
+        default: false
       deploy_website:
         description: Whether to deploy the website
         type: boolean
@@ -54,6 +58,10 @@ on:
         description: Whether to test using macOS
         type: boolean
         default: false
+      test_cuda:
+        description: Whether to test using CUDA-enabled PETSc
+        type: boolean
+        default: false
       deploy_website:
         description: Whether to deploy the website
         type: boolean
@@ -473,6 +481,125 @@ jobs:
         run: |
           find . -delete
 
+  test_cuda:
+    name: Build and test Firedrake (Linux CUDA)
+    runs-on: [self-hosted, Linux, gpu]
+    container:
+      image: ubuntu:latest
+      options: --gpus all
+    if: inputs.test_cuda
+    env:
+      OMPI_ALLOW_RUN_AS_ROOT: 1
+      OMPI_ALLOW_RUN_AS_ROOT_CONFIRM: 1
+      OMP_NUM_THREADS: 1
+      OPENBLAS_NUM_THREADS: 1
+      FIREDRAKE_CI: 1
+      PYOP2_SPMD_STRICT: 1
+      # Disable fast math as it exposes compiler bugs
+      PYOP2_CFLAGS: -fno-fast-math
+      # NOTE: One should occasionally update test_durations.json by running
+      # 'make test_durations' inside a 'firedrake:latest' Docker image.
+      EXTRA_PYTEST_ARGS: --splitting-algorithm least_duration --timeout=600 --timeout-method=thread -o faulthandler_timeout=660 --durations-path=./firedrake-repo/tests/test_durations.json --durations=50
+      PYTEST_MPI_MAX_NPROCS: 8
+    steps:
+      - name: Confirm Nvidia GPUs are enabled
+        run: nvidia-smi
+
+      - name: Fix HOME
+        # For unknown reasons GitHub actions overwrite HOME to /github/home
+        # which will break everything unless fixed
+        # (https://github.com/actions/runner/issues/863)
+        run: echo "HOME=/root" >> "$GITHUB_ENV"
+
+      - name: Pre-run cleanup
+        # Make sure the current directory is empty
+        run: find . -delete
+
+      # Use a different mirror to fetch apt packages from to get around
+      # temporary outage.
+      # (https://askubuntu.com/questions/1549622/problem-with-archive-ubuntu-com-most-of-the-servers-are-not-responding)
+      # The mirror was chosen from https://launchpad.net/ubuntu/+archivemirrors.
+      - name: Configure apt
+        run: |
+          sed -i 's|http://archive.ubuntu.com/ubuntu|http://www.mirrorservice.org/sites/archive.ubuntu.com/ubuntu/|g' /etc/apt/sources.list.d/ubuntu.sources
+          apt-get update
+
+      # Git is needed for actions/checkout and Python for firedrake-configure
+      # curl/gnupg2 and ca-certificates needed for adding new deb repositories to ubunut
+      - name: Install system dependencies (1)
+        run: apt-get -y install git python3 gnupg2 curl ca-certificates
+
+      - name: Add Nvidia CUDA deb repositories
+        run: |
+          curl -fsSL $( python3 ./firedrake-repo/scripts/firedrake-configure --gpu-arch cuda --show-extra-repo-key ) | apt-key add -
+          python3 ./firedrake-repo/scripts/firedrake-configure --gpu-arch cuda --show-extra-repo-file-contents > \
+            /etc/apt/sources.list.d/$( python3 ./firedrake-repo/scripts/firedrake-configure --gpu-arch cuda --show-extra-repo-file-name )
+          apt-get update
+
+      - name: Install PETSc
+        run: |
+          if [ ${{ inputs.target_branch }} = 'release' ]; then
+            git clone --depth 1 \
+              --branch $(python3 ./firedrake-repo/scripts/firedrake-configure --gpu-arch cuda --show-petsc-version) \
+              https://gitlab.com/petsc/petsc.git
+          else
+            git clone --depth 1 https://gitlab.com/petsc/petsc.git
+          fi
+          cd petsc
+          python3 ../firedrake-repo/scripts/firedrake-configure \
+            --arch default --gpu-arch cuda --show-petsc-configure-options | \
+            xargs -L1 ./configure --with-make-np=4
+          make
+          make check
+          {
+            echo "PETSC_DIR=/__w/firedrake/firedrake/petsc"
+            echo "PETSC_ARCH=arch-firedrake-default-cuda"
+            echo "SLEPC_DIR=/__w/firedrake/firedrake/petsc/arch-firedrake-default-cuda"
+          } >> "$GITHUB_ENV"
+
+      - name: Install Firedrake
+        id: install
+        run: |
+          export $(python3 ./firedrake-repo/scripts/firedrake-configure --arch default --gpu-arch cuda --show-env)
+          python3 -m venv venv
+          . venv/bin/activate
+
+          : # Empty the pip cache to ensure that everything is compiled from scratch
+          pip cache purge
+
+          : # Fix for petsc4py+slepc4py build
+          echo 'setuptools<81' > constraints.txt
+          export PIP_CONSTRAINT=constraints.txt
+
+          if [ ${{ inputs.target_branch }} = 'release' ]; then
+            EXTRA_PIP_FLAGS=''
+          else
+            : # Install build dependencies
+            pip install "$PETSC_DIR"/src/binding/petsc4py
+            pip install -r ./firedrake-repo/requirements-build.txt
+
+            : # We have to pass '--no-build-isolation' to use a custom petsc4py
+            EXTRA_PIP_FLAGS='--no-build-isolation'
+          fi
+
+          pip install --verbose $EXTRA_PIP_FLAGS \
+            --no-binary h5py \
+            './firedrake-repo[check]'
+
+          firedrake-clean
+          pip list
+
+      - name: Run smoke tests
+        run: |
+          . venv/bin/activate
+          firedrake-check
+        timeout-minutes: 10
+
+      - name: Post-run cleanup
+        if: always()
+        run: |
+          find . -delete
+
   lint:
     name: Lint codebase
     runs-on: ubuntu-latest
diff --git a/firedrake/preconditioners/__init__.py b/firedrake/preconditioners/__init__.py
index eb37e4f61b..e8bf6e5f63 100644
--- a/firedrake/preconditioners/__init__.py
+++ b/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
diff --git a/firedrake/preconditioners/offload.py b/firedrake/preconditioners/offload.py
new file mode 100644
index 0000000000..eae2c9c8b8
--- /dev/null
+++ b/firedrake/preconditioners/offload.py
@@ -0,0 +1,92 @@
+from firedrake.preconditioners.assembled import AssembledPC
+from firedrake.petsc import PETSc
+
+import petsc4py
+import petsctools
+import firedrake.dmhooks as dmhooks
+
+__all__ = ("OffloadPC",)
+
+
+class OffloadPC(AssembledPC):
+    """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_"
+
+    device_mat_type_map = {"cuda": "aijcusparse"}
+
+    def initialize(self, pc):
+        # Check if our PETSc installation is GPU enabled
+        for device, mat_type in self.device_mat_type_map.items():
+            if device in petsctools.get_external_packages():
+                break
+        else:
+            raise NotImplementedError(
+                "This installation of Firedrake is not GPU-compatible, therefore "
+                "OffloadPC can not be used. PETSc will need to be rebuilt with "
+                "some GPU capability (e.g. '--with-cuda=1') to use this functionality."
+            )
+        # Check if we are we on a machine with a GPU.
+        try:
+            dev = PETSc.Device.Create()
+        except petsc4py.PETSc.Error:
+            raise RuntimeError(
+                "This installation of Firedrake is GPU-Compatible, but no GPU device "
+                "has been detected. OffloadPC can not be used on this host."
+            )
+
+        # Some other reason we could not initialise a device.
+        if dev.getDeviceType() == "HOST":
+            raise RuntimeError(
+                "A GPU-enabled Firedrake build has been detected, but a GPU device was "
+                "unable to be initialised. Cannot use OffloadPC."
+            )
+        dev.destroy()
+
+        super().initialize(pc)
+
+        with PETSc.Log.Event("Event: initialize offload"):
+            A, P = pc.getOperators()
+
+            # Convert matrix to ajicusparse
+            with PETSc.Log.Event("Event: matrix offload"):
+                P_cu = P.convert(mat_type)  # todo
+
+            # Transfer nullspace
+            P_cu.setNullSpace(P.getNullSpace())
+            P_cu.setTransposeNullSpace(P.getTransposeNullSpace())
+            P_cu.setNearNullSpace(P.getNearNullSpace())
+
+            # Update preconditioner with GPU matrix
+            self.pc.setOperators(A, P_cu)
+
+    # Convert vectors to CUDA, solve and get solution on CPU back
+    def apply(self, pc, x, y):
+        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)
+                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)  #
+
+    def applyTranspose(self, pc, X, Y):
+        raise NotImplementedError
+
+    def view(self, pc, viewer=None):
+        super().view(pc, viewer)
+        if hasattr(self, "pc"):
+            viewer.printfASCII("PC to solve on GPU\n")
+            self.pc.view(viewer)
diff --git a/scripts/firedrake-configure b/scripts/firedrake-configure
index f1ac8ca940..c867fb70e4 100755
--- a/scripts/firedrake-configure
+++ b/scripts/firedrake-configure
@@ -30,6 +30,15 @@ LINUX_APT_AARCH64 = PackageManager.LINUX_APT_AARCH64
 MACOS_HOMEBREW_ARM64 = PackageManager.MACOS_HOMEBREW_ARM64
 
 
+class GPUArch(enum.Enum):
+    NO_GPU = "none"
+    CUDA = "cuda"
+
+
+NO_GPU = GPUArch.NO_GPU
+CUDA = GPUArch.CUDA
+
+
 class FiredrakeArch(enum.Enum):
     DEFAULT = "default"
     COMPLEX = "complex"
@@ -40,6 +49,24 @@ ARCH_COMPLEX = FiredrakeArch.COMPLEX
 
 
 SUPPORTED_PETSC_VERSION = "v3.24.0"
+# SuperLU_DIST built via PETSc does not support CUDA 13
+SUPPORTED_CUDA_VERSION = "12.9"
+
+
+CUDA_ARCH_MAP = {
+    "aarch64": "sbsa"
+}
+# Structure is ( deb_repo_filename, file_contents, GPG_key_URL )
+# A file named /etc/apt/sources.list.d/deb_repo_filename containing file_contents will be created
+# The output of curl -fsSL GPG_Key_URL will be passed to 'apt-key add'
+EXTRA_LINUX_APT_REPO = {
+    NO_GPU: ("", "", ""),
+    CUDA: (
+        "cuda.list",
+        f"deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/{CUDA_ARCH_MAP.get(platform.machine(), platform.machine())} /",
+        f"https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/{CUDA_ARCH_MAP.get(platform.machine(), platform.machine())}/3bf863cc.pub",
+    ),
+}
 
 
 def main():
@@ -77,6 +104,12 @@ Please see https://firedrakeproject.org/install for more information."""
         default=ARCH_DEFAULT,
         help="The target configuration to install.",
     )
+    parser.add_argument(
+        "--gpu-arch",
+        choices=[arch.value for arch in GPUArch],
+        default="none",
+        help="Target GPU architecture"
+    )
     cmd_group = parser.add_mutually_exclusive_group(required=True)
     cmd_group.add_argument(
         "--show-system-packages",
@@ -113,6 +146,24 @@ Please see https://firedrakeproject.org/install for more information."""
         action="store_true",
         help="Print out the environment variables that need to be exported to install Firedrake.",
     )
+    cmd_group.add_argument(
+        "--show-extra-repo-key",
+        "--repokey",
+        action="store_true",
+        help="Print out the URL of the GPG key of any non-OS package repository required for this build"
+    )
+    cmd_group.add_argument(
+        "--show-extra-repo-file-contents",
+        "--repofile",
+        action="store_true",
+        help="Print out the contents of any non-OS package repository source file required for this build"
+    )
+    cmd_group.add_argument(
+        "--show-extra-repo-file-name",
+        "--repofilename",
+        action="store_true",
+        help="Print out the name of the any non-OS package repository source file required for this build"
+    )
 
     args = parser.parse_args()
     if args.package_manager is not None:
@@ -124,13 +175,20 @@ Please see https://firedrakeproject.org/install for more information."""
         package_manager = sniff_package_manager()
     arch = FiredrakeArch(args.arch)
 
+    gpu_arch = GPUArch(args.gpu_arch)
+    if gpu_arch != NO_GPU and package_manager == MACOS_HOMEBREW_ARM64:
+        raise RuntimeError(
+            "GPU-compatible PETSc builds are currently only supported"
+            "on Linux"
+        )
+
     if args.show_system_packages:
         if package_manager is None:
             raise RuntimeError(
                 "Cannot install Firedrake dependencies without a package manager, "
                 "please install them manually"
             )
-        print(" ".join(SYSTEM_PACKAGES[package_manager, arch]), end="")
+        print(" ".join(SYSTEM_PACKAGES[package_manager, arch, gpu_arch]), end="")
     elif args.show_minimal_system_packages:
         if package_manager is None:
             raise RuntimeError(
@@ -139,12 +197,18 @@ Please see https://firedrakeproject.org/install for more information."""
             )
         print(" ".join(MINIMAL_SYSTEM_PACKAGES[package_manager]), end="")
     elif args.show_petsc_configure_options:
-        print(" ".join(PETSC_CONFIGURE_OPTIONS[package_manager, arch]), end="")
+        print(" ".join(PETSC_CONFIGURE_OPTIONS[package_manager, arch, gpu_arch]), end="")
     elif args.show_petsc_version:
         print(SUPPORTED_PETSC_VERSION, end="")
+    elif args.show_extra_repo_key:
+        print(EXTRA_LINUX_APT_REPO[gpu_arch][2], end="")
+    elif args.show_extra_repo_file_contents:
+        print(EXTRA_LINUX_APT_REPO[gpu_arch][1], end="")
+    elif args.show_extra_repo_file_name:
+        print(EXTRA_LINUX_APT_REPO[gpu_arch][0], end="")
     else:
         assert args.show_env
-        print(" ".join(ENVIRONMENT_VARS[package_manager, arch]), end="")
+        print(" ".join(ENVIRONMENT_VARS[package_manager, arch, gpu_arch]), end="")
 
 
 def sniff_package_manager() -> Optional[PackageManager]:
@@ -215,8 +279,38 @@ PETSC_EXTRAS_LINUX_APT_PACKAGES = (
     "libsuperlu-dist-dev",
 )
 
+cuda_ver_str = SUPPORTED_CUDA_VERSION.replace(".", "-")
+
+CUDA_EXTRAS_LINUX_APT_PACKAGES = (
+    f"cuda-compat-{cuda_ver_str}",
+    f"cuda-nvtx-{cuda_ver_str}",
+    f"cuda-cudart-dev-{cuda_ver_str}",
+    f"cuda-command-line-tools-{cuda_ver_str}",
+    f"cuda-minimal-build-{cuda_ver_str}",
+    f"cuda-libraries-dev-{cuda_ver_str}",
+    f"cuda-nvml-dev-{cuda_ver_str}",
+    f"libnpp-dev-{cuda_ver_str}",
+    f"libcusparse-dev-{cuda_ver_str}",
+    f"libcublas-dev-{cuda_ver_str}",
+)
+
+CUDA_PETSC_LINUX_APT_SKIP_PACKAGES = (
+    "libsuitesparse-dev",
+    "libsuperlu-dev",
+    "libsuperlu-dist-dev",
+)
+
 LINUX_APT_PACKAGES = BASE_LINUX_APT_PACKAGES + PETSC_EXTRAS_LINUX_APT_PACKAGES
 
+LINUX_APT_PACKAGES_CUDA = (
+    tuple(
+        pkg
+        for pkg in LINUX_APT_PACKAGES
+        if pkg not in CUDA_PETSC_LINUX_APT_SKIP_PACKAGES
+    )
+    + CUDA_EXTRAS_LINUX_APT_PACKAGES
+)
+
 MINIMAL_MACOS_HOMEBREW_PACKAGES = (
     "autoconf",
     "automake",
@@ -255,12 +349,14 @@ MINIMAL_SYSTEM_PACKAGES = {
 }
 
 SYSTEM_PACKAGES = {
-    (LINUX_APT_X86_64, ARCH_DEFAULT): LINUX_APT_PACKAGES,
-    (LINUX_APT_X86_64, ARCH_COMPLEX): LINUX_APT_PACKAGES,
-    (LINUX_APT_AARCH64, ARCH_DEFAULT): LINUX_APT_PACKAGES,
-    (LINUX_APT_AARCH64, ARCH_COMPLEX): LINUX_APT_PACKAGES,
-    (MACOS_HOMEBREW_ARM64, ARCH_DEFAULT): MACOS_HOMEBREW_PACKAGES,
-    (MACOS_HOMEBREW_ARM64, ARCH_COMPLEX): MACOS_HOMEBREW_PACKAGES,
+    (LINUX_APT_X86_64, ARCH_DEFAULT, NO_GPU): LINUX_APT_PACKAGES,
+    (LINUX_APT_X86_64, ARCH_COMPLEX, NO_GPU): LINUX_APT_PACKAGES,
+    (LINUX_APT_AARCH64, ARCH_DEFAULT, NO_GPU): LINUX_APT_PACKAGES,
+    (LINUX_APT_AARCH64, ARCH_COMPLEX, NO_GPU): LINUX_APT_PACKAGES,
+    (MACOS_HOMEBREW_ARM64, ARCH_DEFAULT, NO_GPU): MACOS_HOMEBREW_PACKAGES,
+    (MACOS_HOMEBREW_ARM64, ARCH_COMPLEX, NO_GPU): MACOS_HOMEBREW_PACKAGES,
+    (LINUX_APT_X86_64, ARCH_DEFAULT, CUDA): LINUX_APT_PACKAGES_CUDA,
+    (LINUX_APT_AARCH64, ARCH_DEFAULT, CUDA): LINUX_APT_PACKAGES_CUDA,
 }
 
 COMMON_PETSC_CONFIGURE_OPTIONS = (
@@ -271,11 +367,18 @@ COMMON_PETSC_CONFIGURE_OPTIONS = (
     "--with-strict-petscerrorcode",
 )
 
+
+class PetscPackageAction(enum.IntEnum):
+    PETSC_AUTODETECT = enum.auto()
+    PETSC_DOWNLOAD = enum.auto()
+
+
 # Placeholder value to use when we want PETSc to autodetect the package
-PETSC_AUTODETECT = 333
+PETSC_AUTODETECT = PetscPackageAction.PETSC_AUTODETECT
 
 # Placeholder value to use when we want PETSc to download the package
-PETSC_DOWNLOAD = 666
+PETSC_DOWNLOAD = PetscPackageAction.PETSC_DOWNLOAD
+
 
 # For each package and architecture there are a number of different types of input:
 # 1. PETSC_AUTODETECT - PETSc will be able to find the package itself
@@ -285,7 +388,10 @@ PETSC_DOWNLOAD = 666
 #      'lib' subdirectories)
 # 4. tuple[str, tuple[str, ...]] - a 2-tuple consisting of the includes directory
 #      (location of the header files) and a collection of library files that PETSc needs.
-PETSC_EXTERNAL_PACKAGE_SPECS = {
+PetscSpecValueType = PetscPackageAction | str | tuple[str | None, tuple[str, ...]]
+PetscSpecsDictType = dict[str, dict[PackageManager, PetscSpecValueType]]
+
+PETSC_EXTERNAL_PACKAGE_SPECS: PetscSpecsDictType = {
     "bison": {
         LINUX_APT_X86_64: PETSC_AUTODETECT,
         LINUX_APT_AARCH64: PETSC_AUTODETECT,
@@ -351,6 +457,11 @@ PETSC_EXTERNAL_PACKAGE_SPECS = {
         LINUX_APT_AARCH64: PETSC_AUTODETECT,
         MACOS_HOMEBREW_ARM64: PETSC_DOWNLOAD,
     },
+    "umpire": {
+        LINUX_APT_X86_64: PETSC_DOWNLOAD,
+        LINUX_APT_AARCH64: PETSC_DOWNLOAD,
+        MACOS_HOMEBREW_ARM64: PETSC_DOWNLOAD,
+    },
     "zlib": {
         LINUX_APT_X86_64: PETSC_AUTODETECT,
         LINUX_APT_AARCH64: PETSC_AUTODETECT,
@@ -358,6 +469,22 @@ PETSC_EXTERNAL_PACKAGE_SPECS = {
     },
 }
 
+PetscSpecsDeltaDictType = dict[GPUArch, PetscSpecsDictType]
+
+# Suitesparse and SuperLU-DIST must be built from source to enable GPU features
+PETSC_EXTERNAL_PACKAGE_SPECS_DELTA_GPU: PetscSpecsDeltaDictType = {
+    CUDA: {
+        "suitesparse": {
+            LINUX_APT_X86_64: PETSC_DOWNLOAD,
+            LINUX_APT_AARCH64: PETSC_DOWNLOAD,
+        },
+        "superlu_dist": {
+            LINUX_APT_X86_64: PETSC_DOWNLOAD,
+            LINUX_APT_AARCH64: PETSC_DOWNLOAD,
+        },
+    }
+}
+
 COMMON_PETSC_EXTERNAL_PACKAGES = (
     "bison",
     "fftw",
@@ -375,17 +502,33 @@ COMMON_PETSC_EXTERNAL_PACKAGES = (
 )
 
 
+def blend_package_spec(
+    package_manager: PackageManager, gpu_arch: GPUArch
+) -> dict[str, PetscSpecValueType]:
+    if gpu_arch == NO_GPU:
+        return {k: v[package_manager] for k, v in PETSC_EXTERNAL_PACKAGE_SPECS.items()}
+    else:
+        out = {k: v[package_manager] for k, v in PETSC_EXTERNAL_PACKAGE_SPECS.items()}
+        for k, v in PETSC_EXTERNAL_PACKAGE_SPECS_DELTA_GPU[gpu_arch].items():
+            out[k] = v[package_manager]
+        return out
+
+
 def prepare_external_package_configure_options(
     external_packages: Sequence[str],
-    package_manager: Optional[PackageManager],
+    package_manager: PackageManager | None = None,
+    gpu_arch: GPUArch = NO_GPU,
 ) -> tuple[str, ...]:
+    if package_manager is not None:
+        petsc_package_spec = blend_package_spec(package_manager, gpu_arch)
+
     configure_options = []
     for external_package in external_packages:
         if package_manager is None:
             # Don't know anything about the system, download everything
             package_spec = PETSC_DOWNLOAD
         else:
-            package_spec = PETSC_EXTERNAL_PACKAGE_SPECS[external_package][package_manager]
+            package_spec = petsc_package_spec[external_package]
 
         if package_spec == PETSC_AUTODETECT:
             # PETSc will find the package for us
@@ -408,12 +551,20 @@ def prepare_external_package_configure_options(
     return tuple(configure_options)
 
 
+def get_petsc_arch(arch: FiredrakeArch, gpu_arch: GPUArch) -> str:
+    arr = ["arch", "firedrake", arch.value]
+    if gpu_arch != NO_GPU:
+        arr.append(gpu_arch.value)
+    return "-".join(arr)
+
+
 def prepare_configure_options(
     package_manager: Optional[PackageManager],
     arch: FiredrakeArch,
+    gpu_arch: GPUArch,
 ) -> tuple[str, ...]:
     configure_options = list(COMMON_PETSC_CONFIGURE_OPTIONS)
-    configure_options.append(f"PETSC_ARCH=arch-firedrake-{arch.value}")
+    configure_options.append(f"PETSC_ARCH={get_petsc_arch(arch, gpu_arch)}")
 
     # include/link flags
     if package_manager in (LINUX_APT_X86_64, LINUX_APT_AARCH64):
@@ -458,39 +609,60 @@ def prepare_configure_options(
     if arch == ARCH_COMPLEX:
         configure_options.append("--with-scalar-type=complex")
 
+    if gpu_arch == CUDA:
+        # SuperLU_DIST does not support cuda-arch<70, comma separated multiple arch CUDA builds are not supported by PETSc
+        configure_options.extend(
+            ["--with-cuda=1", "--with-openmp=1", "--with-cuda-arch=70"]
+        )
+
     external_packages = list(COMMON_PETSC_EXTERNAL_PACKAGES)
     if arch != ARCH_COMPLEX:
         external_packages.append("hypre")
+    if gpu_arch == CUDA:
+        external_packages.append("umpire")
     configure_options.extend(
-        prepare_external_package_configure_options(external_packages, package_manager)
+        prepare_external_package_configure_options(
+            external_packages, package_manager, gpu_arch
+        )
     )
     return tuple(configure_options)
 
 
+PETSC_VALID_BUILD_COMBINATIONS = (
+    (LINUX_APT_X86_64, ARCH_DEFAULT, NO_GPU),
+    (LINUX_APT_X86_64, ARCH_COMPLEX, NO_GPU),
+    (LINUX_APT_AARCH64, ARCH_DEFAULT, NO_GPU),
+    (LINUX_APT_AARCH64, ARCH_COMPLEX, NO_GPU),
+    (MACOS_HOMEBREW_ARM64, ARCH_DEFAULT, NO_GPU),
+    (MACOS_HOMEBREW_ARM64, ARCH_COMPLEX, NO_GPU),
+    (None, ARCH_DEFAULT, NO_GPU),
+    (None, ARCH_COMPLEX, NO_GPU),
+    (LINUX_APT_X86_64, ARCH_DEFAULT, CUDA),
+    (LINUX_APT_AARCH64, ARCH_DEFAULT, CUDA),
+    (None, ARCH_DEFAULT, CUDA),
+)
+
+
 PETSC_CONFIGURE_OPTIONS = {
-    (package_manager, arch): prepare_configure_options(package_manager, arch)
-    for (package_manager, arch) in (
-        (LINUX_APT_X86_64, ARCH_DEFAULT),
-        (LINUX_APT_X86_64, ARCH_COMPLEX),
-        (LINUX_APT_AARCH64, ARCH_DEFAULT),
-        (LINUX_APT_AARCH64, ARCH_COMPLEX),
-        (MACOS_HOMEBREW_ARM64, ARCH_DEFAULT),
-        (MACOS_HOMEBREW_ARM64, ARCH_COMPLEX),
-        (None, ARCH_DEFAULT),
-        (None, ARCH_COMPLEX),
+    (package_manager, arch, gpu_arch): prepare_configure_options(
+        package_manager, arch, gpu_arch
     )
+    for (package_manager, arch, gpu_arch) in PETSC_VALID_BUILD_COMBINATIONS
 }
 
 
 def prepare_environment_vars(
     package_manager: Optional[PackageManager],
     arch: FiredrakeArch,
+    gpu_arch: GPUArch,
 ) -> tuple[str, ...]:
     vars = {
         "PETSC_DIR": f"{os.getcwd()}/petsc",
-        "PETSC_ARCH": f"arch-firedrake-{arch.value}",
+        "PETSC_ARCH": get_petsc_arch(arch, gpu_arch),
         "HDF5_MPI": "ON",
     }
+    if gpu_arch == CUDA:
+        vars["PATH"] = f"/usr/local/cuda/bin:{os.environ.get('PATH', '')}"
 
     if package_manager == MACOS_HOMEBREW_ARM64:
         # On macOS h5py cannot find the HDF5 library without help
@@ -503,17 +675,10 @@ def prepare_environment_vars(
 
 
 ENVIRONMENT_VARS = {
-    (package_manager, arch): prepare_environment_vars(package_manager, arch)
-    for (package_manager, arch) in (
-        (LINUX_APT_X86_64, ARCH_DEFAULT),
-        (LINUX_APT_X86_64, ARCH_COMPLEX),
-        (LINUX_APT_AARCH64, ARCH_DEFAULT),
-        (LINUX_APT_AARCH64, ARCH_COMPLEX),
-        (MACOS_HOMEBREW_ARM64, ARCH_DEFAULT),
-        (MACOS_HOMEBREW_ARM64, ARCH_COMPLEX),
-        (None, ARCH_DEFAULT),
-        (None, ARCH_COMPLEX),
+    (package_manager, arch, gpu_arch): prepare_environment_vars(
+        package_manager, arch, gpu_arch
     )
+    for (package_manager, arch, gpu_arch) in PETSC_VALID_BUILD_COMBINATIONS
 }
 
 

From 7a4ef2edf74e5ae868f4e5a7cd77c630a34c63b1 Mon Sep 17 00:00:00 2001
From: Dale Roberts <dale.roberts1@anu.edu.au>
Date: Tue, 10 Mar 2026 11:57:42 +1100
Subject: [PATCH 02/12] Re-add test_poisson_offloading_pc

Have offload preconditioner warn instead of crash if a GPU could not be
initialised for whatever reason. Remove --with-cuda-arch flag from PETSc
config in firedrake-configure.
---
 .github/workflows/core.yml                    |   2 +-
 firedrake/preconditioners/offload.py          | 108 ++++++++++--------
 scripts/firedrake-configure                   |   3 +-
 .../offload/test_poisson_offloading_pc.py     |  67 +++++++++++
 4 files changed, 127 insertions(+), 53 deletions(-)
 create mode 100644 tests/firedrake/offload/test_poisson_offloading_pc.py

diff --git a/.github/workflows/core.yml b/.github/workflows/core.yml
index 949cb8724d..580b00a63b 100644
--- a/.github/workflows/core.yml
+++ b/.github/workflows/core.yml
@@ -525,7 +525,7 @@ jobs:
           apt-get update
 
       # Git is needed for actions/checkout and Python for firedrake-configure
-      # curl/gnupg2 and ca-certificates needed for adding new deb repositories to ubunut
+      # curl/gnupg2 and ca-certificates needed for adding new deb repositories to ubuntu
       - name: Install system dependencies (1)
         run: apt-get -y install git python3 gnupg2 curl ca-certificates
 
diff --git a/firedrake/preconditioners/offload.py b/firedrake/preconditioners/offload.py
index eae2c9c8b8..55fa391983 100644
--- a/firedrake/preconditioners/offload.py
+++ b/firedrake/preconditioners/offload.py
@@ -1,13 +1,45 @@
 from firedrake.preconditioners.assembled import AssembledPC
 from firedrake.petsc import PETSc
+from functools import cache
 
-import petsc4py
 import petsctools
 import firedrake.dmhooks as dmhooks
 
 __all__ = ("OffloadPC",)
 
 
+device_mat_type_map = {"cuda": "aijcusparse"}
+
+
+@cache
+def offload_mat_type() -> str | None:
+    for device, mat_type in device_mat_type_map.items():
+        if device in petsctools.get_external_packages():
+            break
+    else:
+        PETSc.Sys.Print(
+            "This installation of Firedrake is not GPU-compatible, 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:
+        PETSc.Sys.Print(
+            "This installation of Firedrake is GPU-Compatible, but no GPU device "
+            "has been detected. OffloadPC will do nothing on this host"
+        )
+        return None
+    if dev.getDeviceType() == "HOST":
+        PETSc.Sys.Print(
+            "A GPU-enabled Firedrake build has been detected, but a GPU device was "
+            "unable to be initialised. OffloadPC will do nothing."
+        )
+        return None
+    dev.destroy()
+    return mat_type
+
 class OffloadPC(AssembledPC):
     """Offload PC from CPU to GPU and back.
 
@@ -17,44 +49,17 @@ class OffloadPC(AssembledPC):
 
     _prefix = "offload_"
 
-    device_mat_type_map = {"cuda": "aijcusparse"}
-
     def initialize(self, pc):
         # Check if our PETSc installation is GPU enabled
-        for device, mat_type in self.device_mat_type_map.items():
-            if device in petsctools.get_external_packages():
-                break
-        else:
-            raise NotImplementedError(
-                "This installation of Firedrake is not GPU-compatible, therefore "
-                "OffloadPC can not be used. PETSc will need to be rebuilt with "
-                "some GPU capability (e.g. '--with-cuda=1') to use this functionality."
-            )
-        # Check if we are we on a machine with a GPU.
-        try:
-            dev = PETSc.Device.Create()
-        except petsc4py.PETSc.Error:
-            raise RuntimeError(
-                "This installation of Firedrake is GPU-Compatible, but no GPU device "
-                "has been detected. OffloadPC can not be used on this host."
-            )
-
-        # Some other reason we could not initialise a device.
-        if dev.getDeviceType() == "HOST":
-            raise RuntimeError(
-                "A GPU-enabled Firedrake build has been detected, but a GPU device was "
-                "unable to be initialised. Cannot use OffloadPC."
-            )
-        dev.destroy()
-
         super().initialize(pc)
+        self.offload_mat_type = offload_mat_type()
+        if self.offload_mat_type is not None:
+            with PETSc.Log.Event("Event: initialize offload"):
+                A, P = pc.getOperators()
 
-        with PETSc.Log.Event("Event: initialize offload"):
-            A, P = pc.getOperators()
-
-            # Convert matrix to ajicusparse
-            with PETSc.Log.Event("Event: matrix offload"):
-                P_cu = P.convert(mat_type)  # todo
+                # Convert matrix to ajicusparse
+                with PETSc.Log.Event("Event: matrix offload"):
+                    P_cu = P.convert(self.offload_mat_type)  # todo
 
             # Transfer nullspace
             P_cu.setNullSpace(P.getNullSpace())
@@ -66,21 +71,24 @@ def initialize(self, pc):
 
     # Convert vectors to CUDA, solve and get solution on CPU back
     def apply(self, pc, x, y):
-        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)
-                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)  #
+        if self.offload_mat_type 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)
+                    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)  #
 
     def applyTranspose(self, pc, X, Y):
         raise NotImplementedError
diff --git a/scripts/firedrake-configure b/scripts/firedrake-configure
index c867fb70e4..64e1336bb5 100755
--- a/scripts/firedrake-configure
+++ b/scripts/firedrake-configure
@@ -610,9 +610,8 @@ def prepare_configure_options(
         configure_options.append("--with-scalar-type=complex")
 
     if gpu_arch == CUDA:
-        # SuperLU_DIST does not support cuda-arch<70, comma separated multiple arch CUDA builds are not supported by PETSc
         configure_options.extend(
-            ["--with-cuda=1", "--with-openmp=1", "--with-cuda-arch=70"]
+            ["--with-cuda=1", "--with-openmp=1"]
         )
 
     external_packages = list(COMMON_PETSC_EXTERNAL_PACKAGES)
diff --git a/tests/firedrake/offload/test_poisson_offloading_pc.py b/tests/firedrake/offload/test_poisson_offloading_pc.py
new file mode 100644
index 0000000000..75b0ef8b7a
--- /dev/null
+++ b/tests/firedrake/offload/test_poisson_offloading_pc.py
@@ -0,0 +1,67 @@
+from firedrake import *
+import numpy as np
+import pytest
+
+
+# TODO: add marker for cuda pytests and something to check if cuda memory was really used
+@pytest.mark.skipcuda
+@pytest.mark.parametrize(
+    "ksp_type, pc_type", [("cg", "sor"), ("cg", "gamg"), ("preonly", "lu")]
+)
+def test_poisson_offload(ksp_type, pc_type):
+
+    # Different tests for poisoon: cg and pctype sor, --ksp_type=cg --pc_type=gamg
+    print(f"Using ksp_type = {ksp_type}, and pc_type = {pc_type}.", flush=True)
+
+    nested_parameters = {
+        "pc_type": "ksp",
+        "ksp": {
+            "ksp_type": ksp_type,
+            "ksp_view": None,
+            "ksp_rtol": "1e-10",
+            "ksp_monitor": None,
+            "pc_type": pc_type,
+        },
+    }
+    parameters = {
+        "ksp_type": "preonly",
+        "pc_type": "python",
+        "pc_python_type": "firedrake.OffloadPC",
+        "offload": nested_parameters,
+    }
+
+    mesh = UnitSquareMesh(10, 10)
+    V = FunctionSpace(mesh, "CG", 1)
+    u = TrialFunction(V)
+    v = TestFunction(V)
+
+    f = Function(V)
+    x, y = SpatialCoordinate(mesh)
+    f.interpolate(2 * pi**2 * sin(pi * x) * sin(pi * y))
+
+    # Equations
+    L = inner(grad(u), grad(v)) * dx
+    R = inner(v, f) * dx
+
+    # Dirichlet boundary on all sides to 0
+    bcs = DirichletBC(V, 0, "on_boundary")
+
+    # Exact solution
+    sol = Function(V)
+    sol.interpolate(sin(pi * x) * sin(pi * y))
+
+    # Solution function
+    u_f = Function(V)
+
+    problem = LinearVariationalProblem(L, R, u_f, bcs=bcs)
+    solver = LinearVariationalSolver(problem, solver_parameters=parameters)
+    solver.solve()
+    npsol = sol.dat.data[:]
+    npu_f = u_f.dat.data[:]
+    error = errornorm(problem.u, sol)
+    print(f"Error norm = {error}", flush=True)
+    assert error < 1.2e-2
+
+
+if __name__ == "__main__":
+    test_poisson_offload("cg", "gamg")

From 34e7b2198128035ffcf22f415c15e4130511bf4d Mon Sep 17 00:00:00 2001
From: Dale Roberts <dale.roberts1@anu.edu.au>
Date: Tue, 10 Mar 2026 12:49:39 +1100
Subject: [PATCH 03/12] Issue warning on Rank 0 of PC communicator, not
 COMM_WORLD

---
 firedrake/__init__.py                         |  2 +-
 firedrake/preconditioners/offload.py          | 35 +++++++++++--------
 .../offload/test_poisson_offloading_pc.py     |  3 --
 3 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/firedrake/__init__.py b/firedrake/__init__.py
index e0009c2b0e..1b0a78640a 100644
--- a/firedrake/__init__.py
+++ b/firedrake/__init__.py
@@ -88,7 +88,7 @@ def init_petsc():
     ASMLinesmoothPC, ASMExtrudedStarPC, AssembledPC, AuxiliaryOperatorPC,
     MassInvPC, PCDPC, PatchPC, PlaneSmoother, PatchSNES, P1PC, P1SNES,
     LORPC, GTMGPC, PMGPC, PMGSNES, HypreAMS, HypreADS, FDMPC,
-    PoissonFDMPC, TwoLevelPC, HiptmairPC, FacetSplitPC, BDDCPC
+    PoissonFDMPC, TwoLevelPC, HiptmairPC, FacetSplitPC, BDDCPC, OffloadPC
 )
 from firedrake.mesh import (  # noqa: F401
     Mesh, ExtrudedMesh, VertexOnlyMesh, RelabeledMesh,
diff --git a/firedrake/preconditioners/offload.py b/firedrake/preconditioners/offload.py
index 55fa391983..6a11c242ca 100644
--- a/firedrake/preconditioners/offload.py
+++ b/firedrake/preconditioners/offload.py
@@ -1,6 +1,7 @@
 from firedrake.preconditioners.assembled import AssembledPC
 from firedrake.petsc import PETSc
 from functools import cache
+import warnings
 
 import petsctools
 import firedrake.dmhooks as dmhooks
@@ -12,34 +13,38 @@
 
 
 @cache
-def offload_mat_type() -> str | None:
+def offload_mat_type(pc_comm_rank) -> str | None:
     for device, mat_type in device_mat_type_map.items():
         if device in petsctools.get_external_packages():
             break
     else:
-        PETSc.Sys.Print(
-            "This installation of Firedrake is not GPU-compatible, 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')."
-        )
+        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:
-        PETSc.Sys.Print(
-            "This installation of Firedrake is GPU-Compatible, but no GPU device "
-            "has been detected. OffloadPC will do nothing on this host"
-        )
+        if pc_comm_rank == 0:
+            warnings.warn(
+                "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":
-        PETSc.Sys.Print(
-            "A GPU-enabled Firedrake build has been detected, but a GPU device was "
-            "unable to be initialised. OffloadPC will do nothing."
-        )
+        if pc_comm_rank == 0:
+            warnings.warn(
+                "A GPU-enabled Firedrake build has been detected, but a GPU device was"
+                "unable to be initialised. OffloadPC will do nothing."
+            )
         return None
     dev.destroy()
     return mat_type
 
+
 class OffloadPC(AssembledPC):
     """Offload PC from CPU to GPU and back.
 
@@ -52,7 +57,7 @@ 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()
+        self.offload_mat_type = offload_mat_type(pc.comm.rank)
         if self.offload_mat_type is not None:
             with PETSc.Log.Event("Event: initialize offload"):
                 A, P = pc.getOperators()
diff --git a/tests/firedrake/offload/test_poisson_offloading_pc.py b/tests/firedrake/offload/test_poisson_offloading_pc.py
index 75b0ef8b7a..e6e9b488fe 100644
--- a/tests/firedrake/offload/test_poisson_offloading_pc.py
+++ b/tests/firedrake/offload/test_poisson_offloading_pc.py
@@ -1,5 +1,4 @@
 from firedrake import *
-import numpy as np
 import pytest
 
 
@@ -56,8 +55,6 @@ def test_poisson_offload(ksp_type, pc_type):
     problem = LinearVariationalProblem(L, R, u_f, bcs=bcs)
     solver = LinearVariationalSolver(problem, solver_parameters=parameters)
     solver.solve()
-    npsol = sol.dat.data[:]
-    npu_f = u_f.dat.data[:]
     error = errornorm(problem.u, sol)
     print(f"Error norm = {error}", flush=True)
     assert error < 1.2e-2

From f88702161620876789cad946186d0916c1bca52a Mon Sep 17 00:00:00 2001
From: Dale Roberts <dale.roberts1@anu.edu.au>
Date: Wed, 11 Mar 2026 14:39:07 +1100
Subject: [PATCH 04/12] Partially address review comments

Addresses review comments on firedrake-configure, core.yml and
test_poisson_offloading_pc.py. Updates to firedrake-check and skipnogpu
marker still to come.
---
 .github/workflows/core.yml                    |  40 +++--
 .github/workflows/pr.yml                      |   2 +
 scripts/firedrake-configure                   | 138 +++++++-----------
 .../offload/test_poisson_offloading_pc.py     |  12 +-
 4 files changed, 80 insertions(+), 112 deletions(-)

diff --git a/.github/workflows/core.yml b/.github/workflows/core.yml
index 580b00a63b..bbf3843695 100644
--- a/.github/workflows/core.yml
+++ b/.github/workflows/core.yml
@@ -23,7 +23,7 @@ on:
         description: Whether to test using macOS
         type: boolean
         default: false
-      test_cuda:
+      test_gpu:
         description: Whether to test using CUDA-enabled PETSc
         type: boolean
         default: false
@@ -58,7 +58,7 @@ on:
         description: Whether to test using macOS
         type: boolean
         default: false
-      test_cuda:
+      test_gpu:
         description: Whether to test using CUDA-enabled PETSc
         type: boolean
         default: false
@@ -481,13 +481,13 @@ jobs:
         run: |
           find . -delete
 
-  test_cuda:
+  test_gpu:
     name: Build and test Firedrake (Linux CUDA)
     runs-on: [self-hosted, Linux, gpu]
     container:
       image: ubuntu:latest
       options: --gpus all
-    if: inputs.test_cuda
+    if: inputs.test_gpu
     env:
       OMPI_ALLOW_RUN_AS_ROOT: 1
       OMPI_ALLOW_RUN_AS_ROOT_CONFIRM: 1
@@ -515,27 +515,27 @@ jobs:
         # Make sure the current directory is empty
         run: find . -delete
 
-      # Use a different mirror to fetch apt packages from to get around
-      # temporary outage.
-      # (https://askubuntu.com/questions/1549622/problem-with-archive-ubuntu-com-most-of-the-servers-are-not-responding)
-      # The mirror was chosen from https://launchpad.net/ubuntu/+archivemirrors.
-      - name: Configure apt
-        run: |
-          sed -i 's|http://archive.ubuntu.com/ubuntu|http://www.mirrorservice.org/sites/archive.ubuntu.com/ubuntu/|g' /etc/apt/sources.list.d/ubuntu.sources
-          apt-get update
-
       # Git is needed for actions/checkout and Python for firedrake-configure
-      # curl/gnupg2 and ca-certificates needed for adding new deb repositories to ubuntu
+      # curl needed for adding new deb repositories to ubuntu
       - name: Install system dependencies (1)
-        run: apt-get -y install git python3 gnupg2 curl ca-certificates
+        run: apt-get -y install git python3 curl
 
       - name: Add Nvidia CUDA deb repositories
         run: |
-          curl -fsSL $( python3 ./firedrake-repo/scripts/firedrake-configure --gpu-arch cuda --show-extra-repo-key ) | apt-key add -
-          python3 ./firedrake-repo/scripts/firedrake-configure --gpu-arch cuda --show-extra-repo-file-contents > \
-            /etc/apt/sources.list.d/$( python3 ./firedrake-repo/scripts/firedrake-configure --gpu-arch cuda --show-extra-repo-file-name )
+          deburl=$( python3 ./firedrake-repo/scripts/firedrake-configure --show-extra-repo-pkg-url --gpu-arch cuda )
+          debfile=$( basename "${deburl}" )
+          curl -fsSLO "${deburl}"
+          dpkg -i "${debfile}"
           apt-get update
 
+      - name: Install system dependencies (2)
+        run: |
+          apt-get -y install \
+            $(python3 ./firedrake-repo/scripts/firedrake-configure --arch default --gpu-arch cuda --show-system-packages)
+          apt-get -y install python3-venv
+          : # Dependencies needed to run the test suite
+          apt-get -y install fonts-dejavu graphviz graphviz-dev parallel poppler-utils
+
       - name: Install PETSc
         run: |
           if [ ${{ inputs.target_branch }} = 'release' ]; then
@@ -567,10 +567,6 @@ jobs:
           : # Empty the pip cache to ensure that everything is compiled from scratch
           pip cache purge
 
-          : # Fix for petsc4py+slepc4py build
-          echo 'setuptools<81' > constraints.txt
-          export PIP_CONSTRAINT=constraints.txt
-
           if [ ${{ inputs.target_branch }} = 'release' ]; then
             EXTRA_PIP_FLAGS=''
           else
diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml
index a1430b57e7..6b63f97cae 100644
--- a/.github/workflows/pr.yml
+++ b/.github/workflows/pr.yml
@@ -12,4 +12,6 @@ jobs:
       target_branch: ${{ github.base_ref }}
       # Only run macOS tests if the PR is labelled 'macOS'
       test_macos: ${{ contains(github.event.pull_request.labels.*.name, 'macOS') }}
+      # Only run GPU tests if the PR is labelled 'gpu'
+      test_gpu: ${{ contains(github.event.pull_request.labels.*.name, 'gpu') }}
     secrets: inherit
diff --git a/scripts/firedrake-configure b/scripts/firedrake-configure
index 64e1336bb5..bebf77b507 100755
--- a/scripts/firedrake-configure
+++ b/scripts/firedrake-configure
@@ -59,13 +59,9 @@ CUDA_ARCH_MAP = {
 # Structure is ( deb_repo_filename, file_contents, GPG_key_URL )
 # A file named /etc/apt/sources.list.d/deb_repo_filename containing file_contents will be created
 # The output of curl -fsSL GPG_Key_URL will be passed to 'apt-key add'
-EXTRA_LINUX_APT_REPO = {
-    NO_GPU: ("", "", ""),
-    CUDA: (
-        "cuda.list",
-        f"deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/{CUDA_ARCH_MAP.get(platform.machine(), platform.machine())} /",
-        f"https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/{CUDA_ARCH_MAP.get(platform.machine(), platform.machine())}/3bf863cc.pub",
-    ),
+EXTRA_LINUX_APT_PKG_URL = {
+    NO_GPU: "",
+    CUDA: f"https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/{CUDA_ARCH_MAP.get(platform.machine(), platform.machine())}/cuda-keyring_1.1-1_all.deb",
 }
 
 
@@ -147,22 +143,10 @@ Please see https://firedrakeproject.org/install for more information."""
         help="Print out the environment variables that need to be exported to install Firedrake.",
     )
     cmd_group.add_argument(
-        "--show-extra-repo-key",
-        "--repokey",
+        "--show-extra-repo-pkg-url",
+        "--repopkgurl",
         action="store_true",
-        help="Print out the URL of the GPG key of any non-OS package repository required for this build"
-    )
-    cmd_group.add_argument(
-        "--show-extra-repo-file-contents",
-        "--repofile",
-        action="store_true",
-        help="Print out the contents of any non-OS package repository source file required for this build"
-    )
-    cmd_group.add_argument(
-        "--show-extra-repo-file-name",
-        "--repofilename",
-        action="store_true",
-        help="Print out the name of the any non-OS package repository source file required for this build"
+        help="Print out the URL of any package required to enable non-OS repo access for this build",
     )
 
     args = parser.parse_args()
@@ -200,12 +184,8 @@ Please see https://firedrakeproject.org/install for more information."""
         print(" ".join(PETSC_CONFIGURE_OPTIONS[package_manager, arch, gpu_arch]), end="")
     elif args.show_petsc_version:
         print(SUPPORTED_PETSC_VERSION, end="")
-    elif args.show_extra_repo_key:
-        print(EXTRA_LINUX_APT_REPO[gpu_arch][2], end="")
-    elif args.show_extra_repo_file_contents:
-        print(EXTRA_LINUX_APT_REPO[gpu_arch][1], end="")
-    elif args.show_extra_repo_file_name:
-        print(EXTRA_LINUX_APT_REPO[gpu_arch][0], end="")
+    elif args.show_extra_repo_pkg_url:
+        print(EXTRA_LINUX_APT_PKG_URL[gpu_arch], end="")
     else:
         assert args.show_env
         print(" ".join(ENVIRONMENT_VARS[package_manager, arch, gpu_arch]), end="")
@@ -263,7 +243,7 @@ BASE_LINUX_APT_PACKAGES = (
     MINIMAL_LINUX_APT_PACKAGES + ("bison", "cmake", "libopenblas-dev", "libopenmpi-dev")
 )
 
-PETSC_EXTRAS_LINUX_APT_PACKAGES = (
+PETSC_EXTRAS_COMMON_APT_PACKAGES = (
     "libfftw3-dev",
     "libfftw3-mpi-dev",
     "libhwloc-dev",
@@ -274,6 +254,9 @@ PETSC_EXTRAS_LINUX_APT_PACKAGES = (
     "libpnetcdf-dev",
     "libptscotch-dev",
     "libscalapack-openmpi-dev",
+)
+
+PETSC_EXTRAS_LINUX_APT_PACKAGES = PETSC_EXTRAS_COMMON_APT_PACKAGES + (
     "libsuitesparse-dev",
     "libsuperlu-dev",
     "libsuperlu-dist-dev",
@@ -281,7 +264,7 @@ PETSC_EXTRAS_LINUX_APT_PACKAGES = (
 
 cuda_ver_str = SUPPORTED_CUDA_VERSION.replace(".", "-")
 
-CUDA_EXTRAS_LINUX_APT_PACKAGES = (
+PETSC_EXTRAS_LINUX_APT_CUDA_PACKAGES = PETSC_EXTRAS_COMMON_APT_PACKAGES + (
     f"cuda-compat-{cuda_ver_str}",
     f"cuda-nvtx-{cuda_ver_str}",
     f"cuda-cudart-dev-{cuda_ver_str}",
@@ -294,22 +277,9 @@ CUDA_EXTRAS_LINUX_APT_PACKAGES = (
     f"libcublas-dev-{cuda_ver_str}",
 )
 
-CUDA_PETSC_LINUX_APT_SKIP_PACKAGES = (
-    "libsuitesparse-dev",
-    "libsuperlu-dev",
-    "libsuperlu-dist-dev",
-)
-
 LINUX_APT_PACKAGES = BASE_LINUX_APT_PACKAGES + PETSC_EXTRAS_LINUX_APT_PACKAGES
 
-LINUX_APT_PACKAGES_CUDA = (
-    tuple(
-        pkg
-        for pkg in LINUX_APT_PACKAGES
-        if pkg not in CUDA_PETSC_LINUX_APT_SKIP_PACKAGES
-    )
-    + CUDA_EXTRAS_LINUX_APT_PACKAGES
-)
+LINUX_APT_PACKAGES_CUDA = BASE_LINUX_APT_PACKAGES + PETSC_EXTRAS_LINUX_APT_CUDA_PACKAGES
 
 MINIMAL_MACOS_HOMEBREW_PACKAGES = (
     "autoconf",
@@ -391,7 +361,7 @@ PETSC_DOWNLOAD = PetscPackageAction.PETSC_DOWNLOAD
 PetscSpecValueType = PetscPackageAction | str | tuple[str | None, tuple[str, ...]]
 PetscSpecsDictType = dict[str, dict[PackageManager, PetscSpecValueType]]
 
-PETSC_EXTERNAL_PACKAGE_SPECS: PetscSpecsDictType = {
+PETSC_EXTERNAL_PACKAGE_SPECS_COMMON: PetscSpecsDictType = {
     "bison": {
         LINUX_APT_X86_64: PETSC_AUTODETECT,
         LINUX_APT_AARCH64: PETSC_AUTODETECT,
@@ -447,21 +417,6 @@ PETSC_EXTERNAL_PACKAGE_SPECS: PetscSpecsDictType = {
         LINUX_APT_AARCH64: (None, ("-lscalapack-openmpi",)),
         MACOS_HOMEBREW_ARM64: "/opt/homebrew",
     },
-    "suitesparse": {
-        LINUX_APT_X86_64: PETSC_AUTODETECT,
-        LINUX_APT_AARCH64: PETSC_AUTODETECT,
-        MACOS_HOMEBREW_ARM64: "/opt/homebrew",
-    },
-    "superlu_dist": {
-        LINUX_APT_X86_64: PETSC_AUTODETECT,
-        LINUX_APT_AARCH64: PETSC_AUTODETECT,
-        MACOS_HOMEBREW_ARM64: PETSC_DOWNLOAD,
-    },
-    "umpire": {
-        LINUX_APT_X86_64: PETSC_DOWNLOAD,
-        LINUX_APT_AARCH64: PETSC_DOWNLOAD,
-        MACOS_HOMEBREW_ARM64: PETSC_DOWNLOAD,
-    },
     "zlib": {
         LINUX_APT_X86_64: PETSC_AUTODETECT,
         LINUX_APT_AARCH64: PETSC_AUTODETECT,
@@ -469,21 +424,42 @@ PETSC_EXTERNAL_PACKAGE_SPECS: PetscSpecsDictType = {
     },
 }
 
-PetscSpecsDeltaDictType = dict[GPUArch, PetscSpecsDictType]
+PETSC_EXTERNAL_PACKAGE_SPECS: PetscSpecsDictType = (
+    PETSC_EXTERNAL_PACKAGE_SPECS_COMMON
+    | {
+        "suitesparse": {
+            LINUX_APT_X86_64: PETSC_AUTODETECT,
+            LINUX_APT_AARCH64: PETSC_AUTODETECT,
+            MACOS_HOMEBREW_ARM64: "/opt/homebrew",
+        },
+        "superlu_dist": {
+            LINUX_APT_X86_64: PETSC_AUTODETECT,
+            LINUX_APT_AARCH64: PETSC_AUTODETECT,
+            MACOS_HOMEBREW_ARM64: PETSC_DOWNLOAD,
+        },
+    }
+)
 
-# Suitesparse and SuperLU-DIST must be built from source to enable GPU features
-PETSC_EXTERNAL_PACKAGE_SPECS_DELTA_GPU: PetscSpecsDeltaDictType = {
-    CUDA: {
+PETSC_EXTERNAL_PACKAGE_SPECS_CUDA: PetscSpecsDictType = (
+    PETSC_EXTERNAL_PACKAGE_SPECS_COMMON
+    | {
         "suitesparse": {
             LINUX_APT_X86_64: PETSC_DOWNLOAD,
             LINUX_APT_AARCH64: PETSC_DOWNLOAD,
+            MACOS_HOMEBREW_ARM64: "/opt/homebrew",
         },
         "superlu_dist": {
             LINUX_APT_X86_64: PETSC_DOWNLOAD,
             LINUX_APT_AARCH64: PETSC_DOWNLOAD,
+            MACOS_HOMEBREW_ARM64: PETSC_DOWNLOAD,
+        },
+        "umpire": {
+            LINUX_APT_X86_64: PETSC_DOWNLOAD,
+            LINUX_APT_AARCH64: PETSC_DOWNLOAD,
+            MACOS_HOMEBREW_ARM64: PETSC_DOWNLOAD,
         },
     }
-}
+)
 
 COMMON_PETSC_EXTERNAL_PACKAGES = (
     "bison",
@@ -501,17 +477,7 @@ COMMON_PETSC_EXTERNAL_PACKAGES = (
     "zlib",
 )
 
-
-def blend_package_spec(
-    package_manager: PackageManager, gpu_arch: GPUArch
-) -> dict[str, PetscSpecValueType]:
-    if gpu_arch == NO_GPU:
-        return {k: v[package_manager] for k, v in PETSC_EXTERNAL_PACKAGE_SPECS.items()}
-    else:
-        out = {k: v[package_manager] for k, v in PETSC_EXTERNAL_PACKAGE_SPECS.items()}
-        for k, v in PETSC_EXTERNAL_PACKAGE_SPECS_DELTA_GPU[gpu_arch].items():
-            out[k] = v[package_manager]
-        return out
+PETSC_EXTRA_EXTERNAL_PACKAGES_CUDA = ("umpire",)
 
 
 def prepare_external_package_configure_options(
@@ -519,16 +485,16 @@ def prepare_external_package_configure_options(
     package_manager: PackageManager | None = None,
     gpu_arch: GPUArch = NO_GPU,
 ) -> tuple[str, ...]:
-    if package_manager is not None:
-        petsc_package_spec = blend_package_spec(package_manager, gpu_arch)
-
     configure_options = []
     for external_package in external_packages:
         if package_manager is None:
             # Don't know anything about the system, download everything
             package_spec = PETSC_DOWNLOAD
         else:
-            package_spec = petsc_package_spec[external_package]
+            if gpu_arch == NO_GPU:
+                package_spec = PETSC_EXTERNAL_PACKAGE_SPECS[external_package][package_manager]
+            elif gpu_arch == CUDA:
+                package_spec = PETSC_EXTERNAL_PACKAGE_SPECS_CUDA[external_package][package_manager]
 
         if package_spec == PETSC_AUTODETECT:
             # PETSc will find the package for us
@@ -577,10 +543,14 @@ def prepare_configure_options(
         includes = (
             f"{incdir}/hdf5/openmpi",
             f"{incdir}/scotch",
-            f"{incdir}/superlu",
-            f"{incdir}/superlu-dist",
         )
 
+        if gpu_arch == NO_GPU:
+            includes = includes + (
+                f"{incdir}/superlu",
+                f"{incdir}/superlu-dist",
+            )
+
         libraries = (
             f"{libdir}/hdf5/openmpi",
         )
@@ -611,14 +581,14 @@ def prepare_configure_options(
 
     if gpu_arch == CUDA:
         configure_options.extend(
-            ["--with-cuda=1", "--with-openmp=1"]
+            ["--with-cuda=1", "--with-openmp=1", "--with-cxx-dialect=c++17"]
         )
 
     external_packages = list(COMMON_PETSC_EXTERNAL_PACKAGES)
     if arch != ARCH_COMPLEX:
         external_packages.append("hypre")
     if gpu_arch == CUDA:
-        external_packages.append("umpire")
+        external_packages.extend(PETSC_EXTRA_EXTERNAL_PACKAGES_CUDA)
     configure_options.extend(
         prepare_external_package_configure_options(
             external_packages, package_manager, gpu_arch
diff --git a/tests/firedrake/offload/test_poisson_offloading_pc.py b/tests/firedrake/offload/test_poisson_offloading_pc.py
index e6e9b488fe..0c9da03e87 100644
--- a/tests/firedrake/offload/test_poisson_offloading_pc.py
+++ b/tests/firedrake/offload/test_poisson_offloading_pc.py
@@ -3,19 +3,20 @@
 
 
 # TODO: add marker for cuda pytests and something to check if cuda memory was really used
-@pytest.mark.skipcuda
+@pytest.mark.skipnogpu
 @pytest.mark.parametrize(
     "ksp_type, pc_type", [("cg", "sor"), ("cg", "gamg"), ("preonly", "lu")]
 )
 def test_poisson_offload(ksp_type, pc_type):
 
-    # Different tests for poisoon: cg and pctype sor, --ksp_type=cg --pc_type=gamg
+    # Different tests for poisson: cg and pctype sor, --ksp_type=cg --pc_type=gamg
     print(f"Using ksp_type = {ksp_type}, and pc_type = {pc_type}.", flush=True)
 
     nested_parameters = {
         "pc_type": "ksp",
         "ksp": {
             "ksp_type": ksp_type,
+            "ksp_max_it": 50,
             "ksp_view": None,
             "ksp_rtol": "1e-10",
             "ksp_monitor": None,
@@ -40,14 +41,13 @@ def test_poisson_offload(ksp_type, pc_type):
 
     # Equations
     L = inner(grad(u), grad(v)) * dx
-    R = inner(v, f) * dx
 
     # Dirichlet boundary on all sides to 0
     bcs = DirichletBC(V, 0, "on_boundary")
 
     # Exact solution
     sol = Function(V)
-    sol.interpolate(sin(pi * x) * sin(pi * y))
+    R = action(L, sol)
 
     # Solution function
     u_f = Function(V)
@@ -55,9 +55,9 @@ def test_poisson_offload(ksp_type, pc_type):
     problem = LinearVariationalProblem(L, R, u_f, bcs=bcs)
     solver = LinearVariationalSolver(problem, solver_parameters=parameters)
     solver.solve()
-    error = errornorm(problem.u, sol)
+    error = errornorm(u_f, sol)
     print(f"Error norm = {error}", flush=True)
-    assert error < 1.2e-2
+    assert error < 1.0e-9
 
 
 if __name__ == "__main__":

From 5c191d790b2100c7ef6ba2de84b3442fd61b56bf Mon Sep 17 00:00:00 2001
From: Dale Roberts <dale.roberts1@anu.edu.au>
Date: Wed, 11 Mar 2026 17:30:03 +1100
Subject: [PATCH 05/12] More review response

Move device matrix type selection to utils.py. Introduce skipnogpu test
marker.
---
 firedrake/preconditioners/offload.py | 20 +++++++-------------
 firedrake/utils.py                   |  8 ++++++++
 tests/firedrake/conftest.py          | 10 +++++++++-
 3 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/firedrake/preconditioners/offload.py b/firedrake/preconditioners/offload.py
index 6a11c242ca..e9ffb8eb7d 100644
--- a/firedrake/preconditioners/offload.py
+++ b/firedrake/preconditioners/offload.py
@@ -1,23 +1,19 @@
 from firedrake.preconditioners.assembled import AssembledPC
 from firedrake.petsc import PETSc
+from firedrake.utils import device_matrix_type
 from functools import cache
 import warnings
 
-import petsctools
 import firedrake.dmhooks as dmhooks
 
 __all__ = ("OffloadPC",)
 
 
-device_mat_type_map = {"cuda": "aijcusparse"}
-
 
 @cache
 def offload_mat_type(pc_comm_rank) -> str | None:
-    for device, mat_type in device_mat_type_map.items():
-        if device in petsctools.get_external_packages():
-            break
-    else:
+    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"
@@ -35,12 +31,10 @@ def offload_mat_type(pc_comm_rank) -> str | None:
             )
         return None
     if dev.getDeviceType() == "HOST":
-        if pc_comm_rank == 0:
-            warnings.warn(
-                "A GPU-enabled Firedrake build has been detected, but a GPU device was"
-                "unable to be initialised. OffloadPC will do nothing."
-            )
-        return None
+        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
 
diff --git a/firedrake/utils.py b/firedrake/utils.py
index f39f84d478..d1f2bcb260 100644
--- a/firedrake/utils.py
+++ b/firedrake/utils.py
@@ -23,6 +23,14 @@
 SLATE_SUPPORTS_COMPLEX = False
 
 
+def device_matrix_type() -> str | None:
+    _device_mat_type_map = {"cuda": "aijcusparse"}
+    for device, mat_type in _device_mat_type_map.items():
+        if device in petsctools.get_external_packages():
+            return mat_type
+    return None
+
+
 def _new_uid(comm):
     uid = comm.Get_attr(FIREDRAKE_UID)
     if uid is None:
diff --git a/tests/firedrake/conftest.py b/tests/firedrake/conftest.py
index 695239b685..5e8dbfbecd 100644
--- a/tests/firedrake/conftest.py
+++ b/tests/firedrake/conftest.py
@@ -165,10 +165,14 @@ def pytest_configure(config):
         "markers",
         "skipnetgen: mark as skipped if netgen and ngsPETSc is not installed"
     )
+    config.addinivalue_line(
+        "markers",
+        "skipnogpu: mark as skipped when GPU hardware is unavailable"
+    )
 
 
 def pytest_collection_modifyitems(session, config, items):
-    from firedrake.utils import complex_mode, SLATE_SUPPORTS_COMPLEX
+    from firedrake.utils import complex_mode, device_matrix_type, SLATE_SUPPORTS_COMPLEX
 
     for item in items:
         if complex_mode:
@@ -180,6 +184,10 @@ def pytest_collection_modifyitems(session, config, items):
             if item.get_closest_marker("skipreal") is not None:
                 item.add_marker(pytest.mark.skip(reason="Test makes no sense unless in complex mode"))
 
+        if device_matrix_type() is None:
+            if item.get_closest_marker("skipnogpu") is not None:
+                item.add_marker(pytest.mark.skip(reason="Test requires GPU hardware to run."))
+
         for dep, marker, reason in dependency_skip_markers_and_reasons:
             if item.get_closest_marker(marker) is not None and _skip_test_dependency(dep):
                 item.add_marker(pytest.mark.skip(reason))

From f2af668e055affca2febc408a42f568efd1dcd40 Mon Sep 17 00:00:00 2001
From: Dale Roberts <dale.roberts1@anu.edu.au>
Date: Mon, 16 Mar 2026 10:26:57 +1100
Subject: [PATCH 06/12] Add missing apt-get update step. Add actionlint.yaml to
 allow custom runner label to pass linting

---
 .github/actionlint.yaml              |  3 +++
 .github/workflows/core.yml           | 20 ++++++++++++++++----
 firedrake/preconditioners/offload.py |  1 -
 3 files changed, 19 insertions(+), 5 deletions(-)
 create mode 100644 .github/actionlint.yaml

diff --git a/.github/actionlint.yaml b/.github/actionlint.yaml
new file mode 100644
index 0000000000..1f1ba208b3
--- /dev/null
+++ b/.github/actionlint.yaml
@@ -0,0 +1,3 @@
+self-hosted-runner:
+  labels:
+    - gpu
\ No newline at end of file
diff --git a/.github/workflows/core.yml b/.github/workflows/core.yml
index 5cbf1642d2..3031f28e2b 100644
--- a/.github/workflows/core.yml
+++ b/.github/workflows/core.yml
@@ -495,6 +495,9 @@ jobs:
       PYTEST_MPI_MAX_NPROCS: 8
     steps:
       - name: Confirm Nvidia GPUs are enabled
+        # The presence of the nvidia-smi command indicates that the Nvidia drivers have
+        # successfully been imported inot the container, there is no point continuing
+        # if nvidia-smi is not present
         run: nvidia-smi
 
       - name: Fix HOME
@@ -503,14 +506,23 @@ jobs:
         # (https://github.com/actions/runner/issues/863)
         run: echo "HOME=/root" >> "$GITHUB_ENV"
 
-      - name: Pre-run cleanup
-        # Make sure the current directory is empty
-        run: find . -delete
 
       # Git is needed for actions/checkout and Python for firedrake-configure
       # curl needed for adding new deb repositories to ubuntu
       - name: Install system dependencies (1)
-        run: apt-get -y install git python3 curl
+        run: |
+          apt-get update
+          apt-get -y install git python3 curl
+
+
+      - name: Pre-run cleanup
+        # Make sure the current directory is empty
+        run: find . -delete
+
+      - uses: actions/checkout@v5
+        with:
+          path: firedrake-repo
+          ref: ${{ inputs.source_ref }}
 
       - name: Add Nvidia CUDA deb repositories
         run: |
diff --git a/firedrake/preconditioners/offload.py b/firedrake/preconditioners/offload.py
index e9ffb8eb7d..ed073439ed 100644
--- a/firedrake/preconditioners/offload.py
+++ b/firedrake/preconditioners/offload.py
@@ -9,7 +9,6 @@
 __all__ = ("OffloadPC",)
 
 
-
 @cache
 def offload_mat_type(pc_comm_rank) -> str | None:
     mat_type = device_matrix_type()

From b59f1d94247225c37eeed81c210f2b891a854912 Mon Sep 17 00:00:00 2001
From: Dale Roberts <dale.roberts1@anu.edu.au>
Date: Mon, 16 Mar 2026 10:46:48 +1100
Subject: [PATCH 07/12] Set PETSC_OPTIONS='-use_gpu_aware_mpi 0'

---
 .github/workflows/core.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/core.yml b/.github/workflows/core.yml
index 3031f28e2b..87192ae5e8 100644
--- a/.github/workflows/core.yml
+++ b/.github/workflows/core.yml
@@ -493,6 +493,7 @@ jobs:
       # 'make test_durations' inside a 'firedrake:latest' Docker image.
       EXTRA_PYTEST_ARGS: --splitting-algorithm least_duration --timeout=600 --timeout-method=thread -o faulthandler_timeout=660 --durations-path=./firedrake-repo/tests/test_durations.json --durations=50
       PYTEST_MPI_MAX_NPROCS: 8
+      PETSC_OPTIONS: -use_gpu_aware_mpi 0
     steps:
       - name: Confirm Nvidia GPUs are enabled
         # The presence of the nvidia-smi command indicates that the Nvidia drivers have

From 04d12d9fcdeb2f2adfdc99edb5028b35287583ff Mon Sep 17 00:00:00 2001
From: Dale Roberts <dale.roberts1@anu.edu.au>
Date: Mon, 16 Mar 2026 12:28:13 +1100
Subject: [PATCH 08/12] Set EXTRA_OPTIONS='-use_gpu_aware_mpi 0'

---
 .github/workflows/core.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/core.yml b/.github/workflows/core.yml
index 87192ae5e8..d3e98c07c2 100644
--- a/.github/workflows/core.yml
+++ b/.github/workflows/core.yml
@@ -494,6 +494,7 @@ jobs:
       EXTRA_PYTEST_ARGS: --splitting-algorithm least_duration --timeout=600 --timeout-method=thread -o faulthandler_timeout=660 --durations-path=./firedrake-repo/tests/test_durations.json --durations=50
       PYTEST_MPI_MAX_NPROCS: 8
       PETSC_OPTIONS: -use_gpu_aware_mpi 0
+      EXTRA_OPTIONS: -use_gpu_aware_mpi 0
     steps:
       - name: Confirm Nvidia GPUs are enabled
         # The presence of the nvidia-smi command indicates that the Nvidia drivers have

From 336f1a768529bdfd3b0bfa04308bb13ee2ed06b9 Mon Sep 17 00:00:00 2001
From: Dale Roberts <dale.roberts1@anu.edu.au>
Date: Mon, 16 Mar 2026 13:55:03 +1100
Subject: [PATCH 09/12] Remove accidentally committed files. Fix typo in
 core.yml

---
 .github/workflows/core.yml |  2 +-
 do.sh                      | 31 -------------------------------
 thing.txt                  |  1 -
 3 files changed, 1 insertion(+), 33 deletions(-)
 delete mode 100644 do.sh
 delete mode 100644 thing.txt

diff --git a/.github/workflows/core.yml b/.github/workflows/core.yml
index d3e98c07c2..662701a828 100644
--- a/.github/workflows/core.yml
+++ b/.github/workflows/core.yml
@@ -498,7 +498,7 @@ jobs:
     steps:
       - name: Confirm Nvidia GPUs are enabled
         # The presence of the nvidia-smi command indicates that the Nvidia drivers have
-        # successfully been imported inot the container, there is no point continuing
+        # successfully been imported into the container, there is no point continuing
         # if nvidia-smi is not present
         run: nvidia-smi
 
diff --git a/do.sh b/do.sh
deleted file mode 100644
index 0202a74204..0000000000
--- a/do.sh
+++ /dev/null
@@ -1,31 +0,0 @@
-sed -i 's|http://archive.ubuntu.com/ubuntu|http://www.mirrorservice.org/sites/archive.ubuntu.com/ubuntu/|g' /etc/apt/sources.list.d/ubuntu.sources
-apt-get update
-apt-get -y install git python3 gnupg2 curl ca-certificates
-curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/3bf863cc.pub | apt-key add -
-echo "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64 /" > /etc/apt/sources.list.d/cuda.list
-apt-get update
-apt-get -y install $(python3 ./firedrake-configure --arch default --gpu-arch cuda --show-system-packages)
-apt-get -y install fonts-dejavu graphviz graphviz-dev parallel poppler-utils python3-venv
-
-cd /opt
-git clone --depth 1 https://gitlab.com/petsc/petsc.git
-cd petsc
-export LIBRARY_PATH=/usr/local/cuda/lib64/stubs
-python3 ../firedrake-repo/scripts/firedrake-configure --arch default --gpu-arch cuda --show-petsc-configure-options | xargs -L1 ./configure --with-make-np=8 --download-slepc
-make PETSC_DIR=/opt/petsc PETSC_ARCH=arch-firedrake-default-cuda
-make PETSC_DIR=/opt/petsc PETSC_ARCH=arch-firedrake-default-cuda check
-export PETSC_DIR="/opt/petsc"
-export PETSC_ARCH="arch-firedrake-default-cuda"
-export SLEPC_DIR="/opt/petsc/arch-firedrake-default-cuda"
-
-cd /opt
-export $(python3 ./firedrake-configure --arch default --gpu-arch cuda --show-env)
-python3 -m venv venv
-. venv/bin/activate
-pip cache purge
-pip install "$PETSC_DIR"/src/binding/petsc4py
-pip install -r ./firedrake/requirements-build.txt
-pip install --no-build-isolation --no-deps "$PETSC_DIR"/"$PETSC_ARCH"/externalpackages/git.slepc/src/binding/slepc4py
-pip install --no-deps git+https://github.com/NGSolve/ngsPETSc.git netgen-mesher netgen-occt
-pip install --verbose --no-build-isolation --no-binary h5py './firedrake-repo[check]'
-
diff --git a/thing.txt b/thing.txt
deleted file mode 100644
index 55e7b36a3a..0000000000
--- a/thing.txt
+++ /dev/null
@@ -1 +0,0 @@
-deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64 /
\ No newline at end of file

From 4375ed535c55225900fc599f445f494fc0da2b98 Mon Sep 17 00:00:00 2001
From: Dale Roberts <dale.roberts1@anu.edu.au>
Date: Mon, 16 Mar 2026 17:28:12 +1100
Subject: [PATCH 10/12] Add GPU offload test to firedrake-check

---
 scripts/firedrake-check | 2 ++
 setup.py                | 1 +
 2 files changed, 3 insertions(+)

diff --git a/scripts/firedrake-check b/scripts/firedrake-check
index deeb73ca6e..f76bd0788b 100644
--- a/scripts/firedrake-check
+++ b/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
diff --git a/setup.py b/setup.py
index 1d5b588355..692b3bafa4 100644
--- a/setup.py
+++ b/setup.py
@@ -239,6 +239,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"
 )
 
 

From 9ea7a8aa367f9db9243363fcb0506c0bb8f7ae5c Mon Sep 17 00:00:00 2001
From: Dale Roberts <dale.roberts1@anu.edu.au>
Date: Mon, 16 Mar 2026 17:41:56 +1100
Subject: [PATCH 11/12] DROP BEFORE MERGE Verify offload test uses GPU

---
 .github/workflows/core.yml | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/.github/workflows/core.yml b/.github/workflows/core.yml
index 662701a828..6f783eb5b2 100644
--- a/.github/workflows/core.yml
+++ b/.github/workflows/core.yml
@@ -597,6 +597,12 @@ jobs:
           firedrake-check
         timeout-minutes: 10
 
+      - name: Verify GPU usage
+        run: |
+          . venv/bin/activate
+          export PETSC_OPTIONS="${PETSC_OPTIONS} -log_view_gpu_time -log_view"
+          python3 ./firedrake-repo/tests/firedrake/offload/test_poisson_offloading_pc.py
+
       - name: Post-run cleanup
         if: always()
         run: |

From 529ba796013be72d71ddd187f67c83f66e92e8ea Mon Sep 17 00:00:00 2001
From: Dale Roberts <dale.roberts1@anu.edu.au>
Date: Tue, 17 Mar 2026 10:14:53 +1100
Subject: [PATCH 12/12] Add comment to actionlint.yaml, change no-GPU-detected
 warning to logger.warning

---
 .github/actionlint.yaml              | 1 +
 firedrake/preconditioners/offload.py | 3 ++-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/.github/actionlint.yaml b/.github/actionlint.yaml
index 1f1ba208b3..7196b121ca 100644
--- a/.github/actionlint.yaml
+++ b/.github/actionlint.yaml
@@ -1,3 +1,4 @@
 self-hosted-runner:
   labels:
+    # Custom label for GPU-enabled self-hosted runners
     - gpu
\ No newline at end of file
diff --git a/firedrake/preconditioners/offload.py b/firedrake/preconditioners/offload.py
index ed073439ed..57b898de79 100644
--- a/firedrake/preconditioners/offload.py
+++ b/firedrake/preconditioners/offload.py
@@ -1,6 +1,7 @@
 from firedrake.preconditioners.assembled import AssembledPC
 from firedrake.petsc import PETSc
 from firedrake.utils import device_matrix_type
+from firedrake.logging import logger
 from functools import cache
 import warnings
 
@@ -24,7 +25,7 @@ def offload_mat_type(pc_comm_rank) -> str | None:
         dev = PETSc.Device.create()
     except PETSc.Error:
         if pc_comm_rank == 0:
-            warnings.warn(
+            logger.warning(
                 "This installation of Firedrake is GPU-enabled, but no GPU device has"
                 "been detected. OffloadPC will do nothing on this host"
             )