filoferra · Oct 28, 2026 · Oct 28, 2026 · Oct 28, 2026 · Oct 28, 2026 · Oct 28, 2026
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -50,3 +50,28 @@ jobs:
           # anything still on `Base.@deprecate` is silent without this.
           depwarn: yes
 
+
+  # The suite against PETSc built from source and loaded through set_library!,
+  # the path a cluster or GPU installation takes.
+  custom-library:
+    name: Julia 1 - ubuntu-latest - PETSc built from source
+    runs-on: ubuntu-latest
+    timeout-minutes: 120
+    steps:
+      - uses: actions/checkout@v7
+      - uses: julia-actions/setup-julia@v3
+        with:
+          version: '1'
+      - uses: julia-actions/cache@v3
+      - uses: julia-actions/julia-buildpkg@latest
+      - uses: actions/cache@v6
+        with:
+          path: ~/petsc
+          key: petsc-${{ runner.os }}-${{ hashFiles('scripts/build_petsc.sh') }}
+      - name: Build PETSc
+        run: test -f ~/petsc/lib/libpetsc.so || scripts/build_petsc.sh ~/petsc -j 4
+      - name: Select the library
+        run: julia --project -e 'using PETSc; PETSc.set_library!(expanduser("~/petsc/lib/libpetsc.so"))'
+      - uses: julia-actions/julia-runtest@latest
+        with:
+          depwarn: yes
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -12,6 +12,12 @@
 
 - `with_local_arrays!(f!, g_fx, l_x)` calls `f!(fx, lx)` with the local arrays of the two Vecs, on the device when a GPU extension is loaded, and hands them back afterwards, also when `f!` throws.
 - `with_local_coordinate_array(f, dm)` calls `f(coord)` with the local coordinates of a `DMDA` or `DMStag`, addressed with global indices, and hands the array back when `f` returns or throws.
+- AMD GPUs: with AMDGPU.jl loaded and a PETSc built with HIP, `with_local_array!` and `with_local_arrays!` return `ROCArray`s for Vecs in HIP memory (`PETScAMDGPUExt`). `with_local_arrays!` also handles one Vec on the host and the other on the device, as on coarser multigrid levels. `with_field_views!` on a DMStag with device vectors gives views of device arrays for kernels; broadcast over `parent(view)`, since `OffsetArray` broadcasts on the host.
+
+### Fixed
+
+- CUDA: `with_local_arrays!` with one Vec on the host and the other on the device threw a `MethodError`, and arrays went back to PETSc before CUDA.jl's kernels on them had finished unless one was a host copy. Both now work as with AMDGPU: the host Vec is copied to the device and back, and the device is synchronized before every hand-back. Not tested on an NVIDIA GPU.
+- `examples/dmda/ex19.jl` runs on the GPU again, with CUDA or HIP (`const gpu = :cuda` or `:hip`). Its GPU branch passed a raw pointer to `MatSetValuesCOO`, which takes a `Vector`, and its residual got a host and a device array on coarser multigrid levels.
 
 ## v0.5.4
 

diff --git a/Project.toml b/Project.toml
@@ -13,12 +13,15 @@ Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [weakdeps]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 
 [extensions]
+PETScAMDGPUExt = "AMDGPU"
 PETScCUDAExt = "CUDA"
 
 [compat]
+AMDGPU = "1, 2"
 CUDA = "5, 6"
 ForwardDiff = "0.10, 1"
 Libdl = "^1.10"

diff --git a/docs/src/man/dmstag.md b/docs/src/man/dmstag.md
@@ -116,6 +116,14 @@ end
 
 Without `fields`, the block gets each whole array, indexed `[I..., slot]` with `slot` from `dof_slot`. `write = false` checks out read-only.
 
+On a DM whose vectors live on the GPU (`dm_vec_type = "hip"` or `"cuda"`, with AMDGPU.jl or CUDA.jl loaded), the views wrap device arrays. Pass them to a KernelAbstractions kernel as they are, with the same indices, and write the loop above as the kernel body. A broadcast over a view runs element by element on the host, because `OffsetArray` doesn't pass its parent's broadcast style on, so broadcast over `parent(view)` instead, which drops the offsets and stays on the device:
+
+```julia
+PETSc.with_field_views!(dm, r_local; fields = flow) do (Rx, Ry, Rp)
+    parent(Rp) .= 0
+end
+```
+
 ### Locations and stencils
 
 DMStag names a point by where it sits on an element: `DMSTAG_LEFT`, `DMSTAG_DOWN`, `DMSTAG_BACK_DOWN_LEFT`, and so on. Those names shift meaning between dimensions (`DOWN` is the second axis, whatever it's called in your model), so the location functions count axes instead:

diff --git a/docs/src/man/gpu.md b/docs/src/man/gpu.md
@@ -1,23 +1,28 @@
-# GPU Support (CUDA + KernelAbstractions)
+# GPU Support (CUDA, HIP + KernelAbstractions)
 
 Julia has outstanding support for GPUs as it compiles machine code for the particular devices. Importantly, all modern GPUs are supported, which implies that it is quite straightforward to write GPU kernels in Julia, for example using packages such as [KernelAbstractions](https://github.com/JuliaGPU/KernelAbstractions.jl) with no need to write new code when you use AMD instead of NVIDIA GPU's
 
 PETSc has also added GPU support in recent years, and PETSc vector and matrix objects, along with many of the solvers, can be moved to the GPU.
 
-GPU support in PETSc.jl requires a **locally built PETSc** with CUDA or HIP enabled — the precompiled `PETSc_jll` binaries do not include GPU support. See [Installation](@ref) for instructions on pointing PETSc.jl to a local library.
-The example below are given for CUDA. Doing this on AMD machines (HIP) will likely work the same, but will require a specific extension to be added.
+GPU support in PETSc.jl requires a **locally built PETSc** with CUDA or HIP enabled: the precompiled `PETSc_jll` binaries do not include GPU support. See [Installation](@ref) for instructions on pointing PETSc.jl to a local library.
+
+NVIDIA GPUs go through CUDA.jl (`PETScCUDAExt`), AMD GPUs through AMDGPU.jl (`PETScAMDGPUExt`). The examples below use CUDA; for AMD, load AMDGPU instead, use PETSc's HIP types (`-vec_type hip`, `-mat_type aijhipsparse`) and expect `ROCArray` where the text says `CuArray`.
 
 ## Prerequisites
 
-1. **PETSc built with CUDA** — configure with `--with-cuda=1` (and the matching `--with-cuda-dir`). Confirm with:
+1. **PETSc built with CUDA or HIP**: configure with `--with-cuda=1` (and the matching `--with-cuda-dir`), or `--with-hip=1 --with-hip-dir=$ROCM_DIR --with-hip-arch=gfx1100` (your GPU's architecture). Confirm with:
    ```bash
-   grep -i cuda $PETSC_DIR/$PETSC_ARCH/include/petscconf.h
-   # should show: #define PETSC_HAVE_CUDA 1
+   grep -i -e cuda -e hip $PETSC_DIR/$PETSC_ARCH/include/petscconf.h
+   # should show: #define PETSC_HAVE_CUDA 1  (or PETSC_HAVE_HIP 1)
    ```
+   PETSc and the Julia GPU package must use the same GPU runtime: for HIP, AMDGPU.jl loads the system ROCm, so build PETSc against that ROCm. CUDA.jl downloads its own CUDA runtime by default; point it at the toolkit PETSc was built with by running `CUDA.set_runtime_version!(local_toolkit = true)` once in your environment, then restarting Julia.
+
+   A PETSc built with CUDA stops at startup with "your MPI is not GPU-aware" when MPI.jl uses its default MPI (MPICH_jll), which cannot pass device memory. On one rank that costs nothing; turn the check off with the option `-use_gpu_aware_mpi 0`, for example `PETSC_OPTIONS="-use_gpu_aware_mpi 0"` in the environment. On several ranks, use a GPU-aware system MPI instead (see [Running on HPC Systems](hpc.md)).
 
-2. **CUDA.jl** in your Julia environment:
+2. **CUDA.jl** or **AMDGPU.jl** in your Julia environment:
    ```julia
-   pkg> add CUDA
+   pkg> add CUDA      # NVIDIA
+   pkg> add AMDGPU    # AMD
    ```
 
 3. **KernelAbstractions.jl** if you want to write portable GPU/CPU kernels:
@@ -27,7 +32,7 @@ The example below are given for CUDA. Doing this on AMD machines (HIP) will like
 
 ## How it works
 
-When CUDA.jl is loaded alongside PETSc.jl, the `PETScCUDAExt` extension is activated automatically. It registers CUDA-aware implementations for the functions below via Julia's package extension mechanism — no extra configuration is needed.
+When CUDA.jl or AMDGPU.jl is loaded alongside PETSc.jl, the matching extension (`PETScCUDAExt` or `PETScAMDGPUExt`) is activated automatically and adds device-aware methods for the functions below. No extra configuration is needed.
 
 PETSc manages where vector data lives (host or device). The extension inspects the `PetscMemType` returned by `VecGetArray*AndMemType` calls and either wraps the device pointer as a `CuArray` (zero-copy) or allocates a bounce buffer if the data needs to move between host and device.
 
@@ -55,9 +60,12 @@ PETSc.with_local_array!(g_fx, l_x; read=(true, true), write=(true, false)) do fx
 end
 ```
 
-The array type (`CuArray` or `Vector`) is determined automatically from the `PetscMemType`
-of each Vec, so the same code works on both CPU and GPU.  On coarser multigrid levels where
-Vecs may be host-resident even in a CUDA run, the CPU path is taken transparently.
+The array type (`CuArray` or `Vector`) is determined automatically from the `PetscMemType` of each Vec, so the same code works on both CPU and GPU.  Each Vec keeps its own memory, so when one Vec is on the host and the other on the device the arrays are of different types.
+
+### `with_local_arrays!`
+
+For a residual callback, `with_local_arrays!(g_fx, l_x) do fx, lx … end` checks out the output and the input Vec as a pair and always hands the kernel arrays on the same side.
+With both on the device they wrap device memory; when one is on the host (coarser multigrid levels, PETSc's finite-difference coloring) it is copied to the device and, for the output, copied back on return. Without a GPU extension loaded both are host arrays. The device is synchronized before the arrays go back to PETSc, since the GPU packages launch kernels on their own streams.
 
 ## Writing portable kernels with KernelAbstractions
 
@@ -86,21 +94,21 @@ The same kernel runs on CPU (when Vecs are host-resident) and GPU (when Vecs are
 
 [`examples/dmda/ex19.jl`](https://github.com/JuliaParallel/PETSc.jl/blob/main/examples/dmda/ex19.jl) is a full 2D driven-cavity example (velocity–vorticity–temperature) that demonstrates:
 
-- Switching between CPU and GPU with a single `useCUDA` flag.
+- Switching between CPU, CUDA and HIP with a single `gpu` flag.
 - FD coloring-based Jacobian assembly running entirely on-device.
-- `with_local_array!` in the residual callback for transparent CPU/GPU dispatch.
+- `with_local_arrays!` in the residual callback, which also covers the multigrid levels where one Vec is on the host.
 
 - Using [KernelAbstractions](https://github.com/JuliaGPU/KernelAbstractions.jl) to run kernels on various flavors of GPUs or CPUs.
 
 - Multigrid preconditioning with coarser levels falling back to a CPU Jacobian.
 
 To run it on a GPU:
 
-1. Set `const useCUDA = true` near the top of `ex19.jl`.
-2. Ensure your local PETSc build has CUDA support and is linked via `PETSc.set_library!`.
+1. Set `const gpu = :cuda` (or `:hip`) near the top of `ex19.jl`.
+2. Ensure your local PETSc build has CUDA (or HIP) support and is linked via `PETSc.set_library!`.
 3. Launch:
    ```bash
-   julia --project ex19.jl -da_grid_x 256 -da_grid_y 256 \
+   julia --project ex19.jl -da_grid_x 257 -da_grid_y 257 \
        -pc_type mg -pc_mg_levels 4 \
        -mg_levels_ksp_type chebyshev -mg_levels_pc_type jacobi \
        -snes_monitor -ksp_monitor

diff --git a/docs/src/man/installation.md b/docs/src/man/installation.md
@@ -32,7 +32,16 @@ PETSc.set_petscint!(Int32)   # stored in LocalPreferences.toml; restart Julia
 
 ## Using a custom PETSc build
 
-Sometimes, you may be interested in a PETSc installation that comes with additional external packages, or that you compiled yourself. Ensure the library is compiled as a **dynamic** (not static) library.
+Sometimes, you may be interested in a PETSc installation that comes with additional external packages, or that you compiled yourself. Ensure the library is compiled as a **dynamic** (not static) library, against the same MPI library that MPI.jl uses.
+
+`scripts/build_petsc.sh` in the repository builds such a library: PETSc 3.25 with `Int64` indices, linked against the MPI library MPI.jl loads, with BLAS and LAPACK from Julia's libblastrampoline and with Triangle and CTetGen for simplex meshes. CI uses it to run the test suite on a custom library.
+
+```
+scripts/build_petsc.sh ~/petsc -j 4
+scripts/build_petsc.sh ~/petsc -- --download-hypre   # with extra configure options
+```
+
+It prints the library path on its last line, ready for `set_library!`.
 
 Use `set_library!` to configure the path once — it is stored in `LocalPreferences.toml` and no environment variables are needed afterwards:
 

diff --git a/examples/dmda/ex19.jl b/examples/dmda/ex19.jl
@@ -35,12 +35,12 @@
 
     # MPI parallel (4 ranks)
     mpiexec -n 4 julia --project ex19.jl \\
-        -da_grid_x 256 -da_grid_y 256 \\
+        -da_grid_x 257 -da_grid_y 257 \\
         -pc_type mg -pc_mg_levels 4 \\
         -mg_levels_ksp_type chebyshev -mg_levels_pc_type jacobi
 
-    # GPU (CUDA) — set  useCUDA = true  at the top of the file, then:
-    julia --project ex19.jl -da_grid_x 256 -da_grid_y 256 \\
+    # GPU: set  gpu = :cuda  or  gpu = :hip  at the top of the file, then:
+    julia --project ex19.jl -da_grid_x 257 -da_grid_y 257 \\
         -pc_type mg -pc_mg_levels 4 \\
         -mg_levels_ksp_type chebyshev -mg_levels_pc_type jacobi
 
@@ -54,21 +54,27 @@
   Requires: LocalPreferences.toml in examples/ with PetscInt = "Int32" matching
     the PETSc build (check with: grep sizeof_PetscInt petscconf.h).
 
-  GPU usage: set  useCUDA = true  at the top of this file.
-    Requires PETSc built with --with-cuda, and CUDA.jl in the environment.
+  GPU usage: set  gpu = :cuda  (NVIDIA) or  gpu = :hip  (AMD) at the top of this
+    file. Requires PETSc built with --with-cuda or --with-hip, and CUDA.jl or
+    AMDGPU.jl in the environment.
 =#
 
 # ── GPU switch ────────────────────────────────────────────────────────────────
-const useCUDA = false
+const gpu = :none   # :cuda (CUDA.jl) or :hip (AMDGPU.jl)
 
 using MPI
 using PETSc
 using KernelAbstractions
 
-if useCUDA
+if gpu === :cuda
     using CUDA
+elseif gpu === :hip
+    using AMDGPU
 end
 
+# A copy of `a` on the GPU, or `a` itself on the CPU
+to_device(a) = gpu === :cuda ? CuArray(a) : gpu === :hip ? ROCArray(a) : a
+
 
 # ── Physical parameters ──────────────────────────────────────────────────────
 Base.@kwdef mutable struct AppCtx{T}
@@ -209,8 +215,8 @@ end
 # F(x + h·eₖ), then accumulate (F1 − F0)/h into val_dev at the matching slots.
 # Does NOT call MatSetValuesCOO or MatAssembly; those remain with the caller.
 #
-# Captures from module scope: useCUDA, backend, CuArray, CuPtr,
-#   scatter_perturb_kernel!, fd_accumulate_kernel!, KernelAbstractions.
+# Captures from module scope: scatter_perturb_kernel!, fd_accumulate_kernel!,
+#   KernelAbstractions.
 #
 function fd_coloring_jac!(
     petsclib,
@@ -284,9 +290,12 @@ da = PETSc.DMDA(
 
 # GPU vecs and GPU matrix enable a fully GPU-resident FD coloring path via
 # COO preallocation.  No host↔device bouncing in residual or Jacobian.
-if useCUDA
+if gpu === :cuda
     LibPETSc.DMSetVecType(petsclib, da, "cuda")
     LibPETSc.DMSetMatType(petsclib, da, "aijcusparse")
+elseif gpu === :hip
+    LibPETSc.DMSetVecType(petsclib, da, "hip")
+    LibPETSc.DMSetMatType(petsclib, da, "aijhipsparse")
 end
 
 snes = PETSc.SNES(petsclib, comm; opts...)
@@ -363,9 +372,10 @@ PETSc.set_function!(snes, r) do g_fx, snes, g_x
     hydhx_ = hy_ * dhx_;    hxdhy_ = hx_ * dhy_
     lid_   = _T(1) / dhx_    # lidvelocity = 1/(mx-1)
 
-    # with_local_array! handles CPU/GPU dispatch: returns Vector on host, CuArray
-    # on device.  Both vecs must be on the same device (guaranteed per MG level).
-    PETSc.with_local_array!(g_fx, l_x; read=(true, true), write=(true, false)) do fx, lx
+    # with_local_arrays! hands out both arrays on the same side: on the GPU when
+    # either Vec is there (coarser MG levels and PETSc's FD coloring can pass a
+    # host Vec), on the CPU otherwise.
+    PETSc.with_local_arrays!(g_fx, l_x) do fx, lx
         kern  = KernelAbstractions.get_backend(fx)
         x_par = reshape(lx, 4, nx_g,   ny_g)
         f_par = reshape(fx, 4, nx_own, ny_own)
@@ -426,17 +436,14 @@ perturb_cols_1b = coloring.perturb_cols_1b
 coo_idxs_1b     = coloring.coo_idxs_1b
 local_rows_1b   = coloring.local_rows_1b
 
-if useCUDA
-    perturb_cols_dev = [CuArray(v) for v in perturb_cols_1b]
-    coo_idxs_dev     = [CuArray(v) for v in coo_idxs_1b]
-    local_rows_dev   = [CuArray(v) for v in local_rows_1b]
-    val_dev          = CUDA.zeros(_T, nnz_coo)
-else
-    perturb_cols_dev = perturb_cols_1b
-    coo_idxs_dev     = coo_idxs_1b
-    local_rows_dev   = local_rows_1b
-    val_dev          = zeros(_T, nnz_coo)
-end
+perturb_cols_dev = map(to_device, perturb_cols_1b)
+coo_idxs_dev     = map(to_device, coo_idxs_1b)
+local_rows_dev   = map(to_device, local_rows_1b)
+val_dev          = to_device(zeros(_T, nnz_coo))
+# What MatSetValuesCOO takes: on the GPU an Array over the device memory (the
+# wrapper wants a Vector; PETSc reads device values for a device matrix)
+coo_vals = gpu === :none ? val_dev :
+    unsafe_wrap(Array, Ptr{_T}(UInt(pointer(val_dev))), length(val_dev))
 
 # ── Scratch vectors for the FD loop ────────────────────────────────────────
 x_pert_vec = LibPETSc.VecDuplicate(petsclib, x)
@@ -465,19 +472,14 @@ PETSc.set_snes_jacobian!(snes, J) do Jmat, actual_snes, g_x
     )
 
     # ── Assemble J via COO ─────────────────────────────────────────────────────
-    # On GPU: pass a raw device pointer so cuSPARSE scatters directly on device.
-    # On CPU: pass the Vector{T} directly (uses the vector overload).
-    # Both paths use the COO pattern registered by MatSetPreallocationCOOLocal.
-    if useCUDA
-        LibPETSc.MatSetValuesCOO(petsclib, Jmat, Ptr{_T}(UInt64(pointer(val_dev))), LibPETSc.INSERT_VALUES)
-    else
-        LibPETSc.MatSetValuesCOO(petsclib, Jmat, val_dev, LibPETSc.INSERT_VALUES)
-    end
+    # On GPU the sparse library scatters the device values directly. Both paths
+    # use the COO pattern registered by MatSetPreallocationCOOLocal.
+    LibPETSc.MatSetValuesCOO(petsclib, Jmat, coo_vals, LibPETSc.INSERT_VALUES)
     LibPETSc.MatAssemblyBegin(petsclib, Jmat, LibPETSc.MAT_FINAL_ASSEMBLY)
     LibPETSc.MatAssemblyEnd(petsclib, Jmat, LibPETSc.MAT_FINAL_ASSEMBLY)
-    if useCUDA
+    if gpu !== :none
         # Force GPU→CPU sync so both copies are valid.
-        # MatBindToCPU(PETSC_TRUE) triggers MatSeqAIJCUSPARSECopyFromGPU when
+        # MatBindToCPU(PETSC_TRUE) triggers the device-to-host copy when
         # offloadmask==PETSC_OFFLOAD_GPU, making the CPU CSR correct.
         # MatBindToCPU(PETSC_FALSE) then releases the CPU-only restriction,
         # leaving offloadmask==PETSC_OFFLOAD_BOTH so that:

diff --git a/examples/dmstag/dmstag_stokes_test.jl b/examples/dmstag/dmstag_stokes_test.jl
@@ -19,15 +19,17 @@ function iterative(n; levels, options...)
 end
 
 # LU on A. In a closed box A is singular, since the pressure is defined up to a constant,
-# so the pressure in the first cell is fixed at zero.
+# so the pressure in the first cell is fixed at zero. The zero pressure block needs a
+# pivoting LU, so A is factored dense through LAPACK, which every PETSc build has.
 function direct(n)
     dm, A, P, b = stokes_system(petsclib, comm, n, 1e4)
     PETSc.zero_rows_local!(A, dm, [PETSc.stencil(dm, PC, (1, 1))]; diag = 1.0)
-    ksp = PETSc.KSP(A; ksp_type = "preonly", pc_type = "lu", pc_factor_mat_solver_type = "umfpack")
+    Ad = LibPETSc.MatConvert(petsclib, A, LibPETSc.MATDENSE, LibPETSc.MAT_INITIAL_MATRIX)
+    ksp = PETSc.KSP(Ad; ksp_type = "preonly", pc_type = "lu")
     x = PETSc.global_vec(dm)
     PETSc.solve!(x, ksp, b; check = true)
     sol = owned_fields(dm, x)
-    foreach(PETSc.destroy!, (ksp, x, b, P, A, dm))
+    foreach(PETSc.destroy!, (ksp, x, b, P, Ad, A, dm))
     return sol
 end