diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c68503ffd..311aaaa48 100644 --- 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 index 71207cd9b..f1aa16954 100644 --- 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 index 9bafe039b..c8233b46b 100644 --- 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 index 78a373c34..56729cc19 100644 --- 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 index 140eddc54..113e734e7 100644 --- 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,9 +94,9 @@ 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. @@ -96,11 +104,11 @@ The same kernel runs on CPU (when Vecs are host-resident) and GPU (when Vecs are 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 index 4d0ea4bab..49b50a35e 100644 --- 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 index 42a1a77b6..131dd7231 100644 --- 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 index 31d6ee42c..071921d27 100644 --- 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 diff --git a/ext/PETScAMDGPUExt.jl b/ext/PETScAMDGPUExt.jl new file mode 100644 index 000000000..96e2b0cce --- /dev/null +++ b/ext/PETScAMDGPUExt.jl @@ -0,0 +1,63 @@ +module PETScAMDGPUExt + +using PETSc +using PETSc: LibPETSc, AbstractPetscVec, AbstractPetscMemBackend +using PETSc.LibPETSc: PETSC_MEMTYPE_HIP +using AMDGPU + +# ============================================================================ +# HIP memory backend +# ============================================================================ + +struct HIPMemBackend <: AbstractPetscMemBackend end + +PETSc.memtype_backend(::Val{PETSC_MEMTYPE_HIP}) = HIPMemBackend() +PETSc.array_type(::Val{PETSC_MEMTYPE_HIP}) = ROCArray + +# AMDGPU.jl launches kernels on its own stream, PETSc reads on its own. +PETSc.synchronize_device(::HIPMemBackend) = AMDGPU.synchronize() + +# A ROCArray over the device memory PETSc handed out. PETSc keeps ownership. +device_array(arr) = unsafe_wrap(ROCArray, pointer(arr), length(arr); own = false) + +# ============================================================================ +# with_local_array! +# ============================================================================ + +PETSc.make_local_array(cpu_arr, ::HIPMemBackend) = device_array(cpu_arr) + +# ============================================================================ +# with_local_arrays! +# ============================================================================ + +# Both Vecs on the device: wrap both, nothing to copy. +function PETSc._local_arrays( + petsclib, g_fx, l_x, ::Type, fx_arr, lx_arr, ::HIPMemBackend, ::HIPMemBackend, +) + return device_array(fx_arr), device_array(lx_arr), fx_arr, lx_arr, nothing +end + +# One Vec on the host (coarse multigrid levels, FD coloring): the kernel still +# runs on the device. A host `l_x` is copied over; a host `g_fx` gets a device +# copy (the bounce array) that the restore copies back. +function PETSc._local_arrays( + petsclib, g_fx, l_x, ::Type, fx_arr, lx_arr, + fx_b::Union{Nothing, HIPMemBackend}, lx_b::Union{Nothing, HIPMemBackend}, +) + lx = lx_b === nothing ? ROCArray(lx_arr) : device_array(lx_arr) + fx_bounce = fx_b === nothing ? ROCArray(fx_arr) : nothing + fx = fx_bounce === nothing ? device_array(fx_arr) : fx_bounce + return fx, lx, fx_arr, lx_arr, fx_bounce +end + +function PETSc._restore_local_arrays!( + petsclib, g_fx, l_x, fx::ROCArray, lx, fx_arr::Array, lx_arr::Array, fx_bounce, +) + AMDGPU.synchronize() + fx_bounce === nothing || copyto!(fx_arr, fx_bounce) + LibPETSc.VecRestoreArrayAndMemType(petsclib, PETSc.as_petsc_vec(g_fx), fx_arr) + LibPETSc.VecRestoreArrayReadAndMemType(petsclib, PETSc.as_petsc_vec(l_x), lx_arr) + return nothing +end + +end # module PETScAMDGPUExt diff --git a/ext/PETScCUDAExt.jl b/ext/PETScCUDAExt.jl index b5b6b207c..64a2e4237 100644 --- a/ext/PETScCUDAExt.jl +++ b/ext/PETScCUDAExt.jl @@ -1,53 +1,39 @@ module PETScCUDAExt using PETSc -using PETSc: LibPETSc, AbstractPetscVec +using PETSc: LibPETSc, AbstractPetscVec, AbstractPetscMemBackend using PETSc.LibPETSc: PETSC_MEMTYPE_DEVICE using CUDA -# ── CUDA memory backend ─────────────────────────────────────────────────────── +# ============================================================================ +# CUDA memory backend +# ============================================================================ -struct CUDAMemBackend <: PETSc.AbstractPetscMemBackend end +struct CUDAMemBackend <: AbstractPetscMemBackend end PETSc.memtype_backend(::Val{PETSC_MEMTYPE_DEVICE}) = CUDAMemBackend() PETSc.array_type(::Val{PETSC_MEMTYPE_DEVICE}) = CuArray -# ── No-finalizer acquire/release for with_local_array! ───────────────────────── +# CUDA.jl launches kernels on its own stream, PETSc reads on its own. +PETSc.synchronize_device(::CUDAMemBackend) = CUDA.synchronize() -function PETSc.make_local_array(cpu_arr, ::CUDAMemBackend) - T = eltype(cpu_arr) - n = length(cpu_arr) - ptr = reinterpret(CuPtr{T}, UInt(pointer(cpu_arr))) - return CUDA.unsafe_wrap(CuArray, ptr, n; own = false) -end +# A CuArray over the device memory PETSc handed out. PETSc keeps ownership. +device_array(arr) = CUDA.unsafe_wrap(CuArray, + reinterpret(CuPtr{eltype(arr)}, UInt(pointer(arr))), length(arr); own = false) -function PETSc.release_local_array( - cpu_arr, ::CUDAMemBackend, vec::AbstractPetscVec{PLib}; read::Bool, write::Bool, -) where {PLib} - pv = PETSc.as_petsc_vec(vec) - if write && read - LibPETSc.VecRestoreArrayAndMemType(PLib, pv, cpu_arr) - elseif write - LibPETSc.VecRestoreArrayWriteAndMemType(PLib, pv, cpu_arr) - else - LibPETSc.VecRestoreArrayReadAndMemType(PLib, pv, cpu_arr) - end - return nothing -end +# ============================================================================ +# with_local_array! +# ============================================================================ -# ── wrap_local_array: device branch (legacy, kept for backward compat) ───────── -# -# No longer called by with_local_array! (which uses acquire/release instead). -# Retained in case external code calls unsafe_local_array directly. +PETSc.make_local_array(cpu_arr, ::CUDAMemBackend) = device_array(cpu_arr) +# The finalizer-based checkout, handed back when the CuArray is finalized. +# with_local_array! uses make_local_array and release_local_array instead. function PETSc.wrap_local_array( cpu_arr, ::CUDAMemBackend, vec::AbstractPetscVec{PetscLib}; read::Bool, write::Bool, ) where {PetscLib} - T = eltype(cpu_arr) - n = length(cpu_arr) - ptr = reinterpret(CuPtr{T}, UInt(pointer(cpu_arr))) - dev_arr = CUDA.unsafe_wrap(CuArray, ptr, n; own = false) + dev_arr = device_array(cpu_arr) pv = PETSc.as_petsc_vec(vec) finalizer(dev_arr) do _ if write && read @@ -62,66 +48,38 @@ function PETSc.wrap_local_array( return dev_arr end -# ── _local_arrays: CUDA cases ──────────────────────────────────────── -# -# Two methods cover all GPU sub-cases: -# -# CUDAMemBackend × CUDAMemBackend → both Vecs on device: zero-copy wrap, no bounce -# any other mix → at least one Vec is host-resident: -# lx is wrapped zero-copy if on device, or -# copied H2D if on host; -# fx always gets a fresh scratch CuArray (bounce) -# so the kernel writes there and restore copies D2H. -# -# nothing × nothing (host) is handled entirely in base (vec.jl) and never -# reaches these methods. +# ============================================================================ +# with_local_arrays! +# ============================================================================ -# Both Vecs on the device: zero-copy wrap, no scratch needed. +# Both Vecs on the device: wrap both, nothing to copy. function PETSc._local_arrays( - petsclib, g_fx, l_x, ::Type{T}, fx_arr, lx_arr, - ::CUDAMemBackend, ::CUDAMemBackend, -) where {T} - fx = CUDA.unsafe_wrap(CuArray, - reinterpret(CuPtr{T}, UInt(pointer(fx_arr))), length(fx_arr)) - lx = CUDA.unsafe_wrap(CuArray, - reinterpret(CuPtr{T}, UInt(pointer(lx_arr))), length(lx_arr)) - return fx, lx, fx_arr, lx_arr, nothing + petsclib, g_fx, l_x, ::Type, fx_arr, lx_arr, ::CUDAMemBackend, ::CUDAMemBackend, +) + return device_array(fx_arr), device_array(lx_arr), fx_arr, lx_arr, nothing end -# At least one Vec is host-resident (e.g. MG coarser levels, FD-coloring path). -# Catch-all: less specific than (CUDAMemBackend, CUDAMemBackend), so Julia prefers -# the method above when both are on the device. +# One Vec on the host (coarse multigrid levels, FD coloring): the kernel still +# runs on the device. A host `l_x` is copied over; a host `g_fx` gets a device +# copy (the bounce array) that the restore copies back. function PETSc._local_arrays( - petsclib, g_fx, l_x, ::Type{T}, fx_arr, lx_arr, - fx_b::PETSc.AbstractPetscMemBackend, lx_b::PETSc.AbstractPetscMemBackend, -) where {T} - lx_gpu = if lx_b isa CUDAMemBackend - CUDA.unsafe_wrap(CuArray, - reinterpret(CuPtr{T}, UInt(pointer(lx_arr))), length(lx_arr)) - else - tmp = CuArray{T}(undef, length(lx_arr)) - copyto!(tmp, lx_arr) # H2D: send ghost input to GPU - tmp - end - fx_gpu = CuArray{T}(undef, length(fx_arr)) # scratch buffer for residual - return fx_gpu, lx_gpu, fx_arr, lx_arr, fx_gpu + petsclib, g_fx, l_x, ::Type, fx_arr, lx_arr, + fx_b::Union{Nothing, CUDAMemBackend}, lx_b::Union{Nothing, CUDAMemBackend}, +) + lx = lx_b === nothing ? CuArray(lx_arr) : device_array(lx_arr) + fx_bounce = fx_b === nothing ? CuArray(fx_arr) : nothing + fx = fx_bounce === nothing ? device_array(fx_arr) : fx_bounce + return fx, lx, fx_arr, lx_arr, fx_bounce end -# ── _restore_local_arrays!: CUDA ────────────────────────────────────────── -# -# When fx is a CuArray (returned by the GPU _local_arrays above): -# - if fx_bounce !== nothing, sync the device and copy the scratch D2H -# - call VecRestoreArray*AndMemType on both raw PETSc arrays - function PETSc._restore_local_arrays!( - petsclib, g_fx, l_x, fx::CuArray, lx, fx_arr, lx_arr, fx_bounce, + petsclib, g_fx, l_x, fx::CuArray, lx, fx_arr::Array, lx_arr::Array, fx_bounce, ) - if fx_bounce !== nothing - CUDA.synchronize() - copyto!(fx_arr, fx_bounce) # D2H: copy residual back to host PETSc array - end + CUDA.synchronize() + fx_bounce === nothing || copyto!(fx_arr, fx_bounce) LibPETSc.VecRestoreArrayAndMemType(petsclib, PETSc.as_petsc_vec(g_fx), fx_arr) LibPETSc.VecRestoreArrayReadAndMemType(petsclib, PETSc.as_petsc_vec(l_x), lx_arr) + return nothing end end # module PETScCUDAExt diff --git a/scripts/build_petsc.sh b/scripts/build_petsc.sh new file mode 100755 index 000000000..351653923 --- /dev/null +++ b/scripts/build_petsc.sh @@ -0,0 +1,120 @@ +#!/usr/bin/env bash +# Build PETSc as a shared library for `PETSc.set_library!`, linked against the +# MPI library that MPI.jl uses, so the two never disagree on MPI. +# +# Usage: +# scripts/build_petsc.sh PREFIX [-j N] [-- CONFIGURE_OPTIONS...] +# +# PREFIX install directory; the source is unpacked in PREFIX/src +# -j N parallel make jobs (default 2) +# -- ... extra options passed to PETSc's configure +# +# Environment: +# PETSC_VERSION PETSc release to build (default 3.25.4, the version the +# wrappers are generated for) +# JULIA_PROJECT the Julia environment whose MPI.jl is used (default: this +# repository) +# +# The last line printed is the path of the built library, for set_library!. +set -euo pipefail + +# ============================================================================ +# Arguments +# ============================================================================ + +# Print the usage block above and stop. +usage() { sed -n '5,18p' "$0" | sed 's/^# \{0,1\}//'; exit 1; } + +[ $# -ge 1 ] || usage +prefix=$(realpath -m "$1"); shift +jobs=2 +extra=() +while [ $# -gt 0 ]; do + case "$1" in + -j) jobs="$2"; shift 2 ;; + --) shift; extra=("$@"); break ;; + *) usage ;; + esac +done + +version="${PETSC_VERSION:-3.25.4}" +repo=$(cd "$(dirname "$0")/.." && pwd) +project="${JULIA_PROJECT:-$repo}" + +# Run a Julia snippet in the chosen project and print its output. +jl() { julia --startup-file=no --project="$project" -e "$1"; } + +# ============================================================================ +# Libraries from the Julia side +# ============================================================================ + +# MPI: the library MPI.jl loads. Its headers sit in ../include next to it. +libmpi=$(jl 'using MPI; print(MPI.API.libmpi)') +mpi_dir=$(dirname "$(dirname "$libmpi")") +[ -f "$mpi_dir/include/mpi.h" ] || { echo "no mpi.h in $mpi_dir/include (MPI library $libmpi)" >&2; exit 1; } + +# Directories of the artifacts libmpi depends on (hwloc, …). The linker needs +# them through -rpath-link; at run time Julia has them loaded already. +link_dirs=$(jl 'using MPI, Libdl + print(join(unique(dirname.(filter(contains("artifacts"), Libdl.dllist()))), ":"))') + +# BLAS and LAPACK: Julia's libblastrampoline, through its ILP64 names +# (dgemm_64_, …), as in PETSc_jll. A BLAS of PETSc's own would export dgemm_, +# and inside Julia those calls bind to libblastrampoline's LP64 entry points, +# which have no library behind them. +lbt=$(jl 'using Libdl; print(realpath(only(filter(contains("libblastrampoline"), Libdl.dllist()))))') + +# ============================================================================ +# Source +# ============================================================================ + +# Download and unpack the release once; later runs reuse it. +src="$prefix/src/petsc-$version" +if [ ! -d "$src" ]; then + mkdir -p "$prefix/src" + curl -fsSL "https://web.cels.anl.gov/projects/petsc/download/release-snapshots/petsc-$version.tar.gz" | + tar -xz -C "$prefix/src" +fi + +# ============================================================================ +# Configure options +# ============================================================================ + +opts=( + # Optimized shared library with Int64 indices, as PETSc.jl expects + --prefix="$prefix" + --with-shared-libraries=1 + --with-debugging=0 + --with-64-bit-indices=1 + --with-fc=0 + COPTFLAGS=-O2 CXXOPTFLAGS=-O2 + + # MPI from MPI.jl + --with-mpi-include="$mpi_dir/include" + --with-mpi-lib="$libmpi" + LDFLAGS="-Wl,-rpath-link,$link_dirs" + + # BLAS and LAPACK from libblastrampoline + --with-blaslapack-lib="$lbt" + --with-blaslapack-suffix=_64_ + --known-64-bit-blas-indices=1 + + # Mesh generators for simplex meshes in the DMPlex tests + --download-triangle=1 + --download-ctetgen=1 + + # CPU only: configure enables CUDA or HIP on its own when it finds them + --with-cuda=0 + --with-hip=0 +) + +# ============================================================================ +# Build and install +# ============================================================================ + +cd "$src" +./configure "${opts[@]}" "${extra[@]}" +make -j "$jobs" all +make install + +echo "$prefix/lib/libpetsc.so" diff --git a/scripts/renames.jl b/scripts/renames.jl index a7cba9363..3f3080286 100644 --- a/scripts/renames.jl +++ b/scripts/renames.jl @@ -879,6 +879,7 @@ const INTERNAL = Set{Symbol}([ :stencil_type_enum, :petsclib_of, :make_local_array, + :synchronize_device, :with_broadcast_arrays, :aliases_dest, :check_undistributed, diff --git a/src/dmstag.jl b/src/dmstag.jl index 4bef1d56b..402038806 100644 --- a/src/dmstag.jl +++ b/src/dmstag.jl @@ -743,6 +743,10 @@ vector. `write = false` checks out read-only; the views have the same types eith way. The vectors must be local vectors of `dm` ([`local_vec`](@ref)); anything else throws an `ArgumentError` before any vector is checked out. +Vectors on a GPU, with its extension loaded, give views of device arrays, for use in +a kernel. Broadcast over `parent(view)`: an `OffsetArray` broadcasts element by +element on the host whatever its parent. + ```julia flow = (face_location(dm, 1) => 0, face_location(dm, 2) => 0, element_location(dm) => 0) c = corners(dm) @@ -776,7 +780,7 @@ end ``` # External Links -$(doc_external("DMStag/DMStagVecGetArray")) +$(doc_external("Vec/VecGetArrayAndMemType")) $(doc_external("DMStag/DMStagGetLocationSlot")) """ function with_field_views!( diff --git a/src/vec.jl b/src/vec.jl index 66e92f5f9..7d9252dfa 100644 --- a/src/vec.jl +++ b/src/vec.jl @@ -624,12 +624,14 @@ make_local_array(_, b::AbstractPetscMemBackend) = """ release_local_array(cpu_arr, backend, vec; read, write) -Restore a previously acquired local array. Called in `finally` blocks by -`with_local_array!`. Extensions overload this for GPU backends. +Restore a local array taken with `acquire_local_array`, after waiting for the device to +finish with it. Called in `finally` blocks by `with_local_array!`. """ function release_local_array( - cpu_arr, ::Nothing, vec::AbstractPetscVec{PLib}; read::Bool, write::Bool, + cpu_arr, backend::Union{Nothing, AbstractPetscMemBackend}, vec::AbstractPetscVec{PLib}; + read::Bool, write::Bool, ) where {PLib} + synchronize_device(backend) pv = as_petsc_vec(vec) if write && read LibPETSc.VecRestoreArrayAndMemType(PLib, pv, cpu_arr) @@ -640,9 +642,11 @@ function release_local_array( end return nothing end -release_local_array(cpu_arr, b::AbstractPetscMemBackend, vec; kw...) = - error("release_local_array not implemented for backend $(typeof(b)) — " * - "load the corresponding GPU package (e.g. CUDA.jl)") + +# Wait until the kernels Julia launched on the arrays are done. The GPU packages +# launch on their own streams, not on PETSc's, so PETSc could otherwise read the +# memory before a kernel has written it. Extensions add a method per backend. +synchronize_device(::Union{Nothing, AbstractPetscMemBackend}) = nothing # The auto-generated *AndMemType wrappers are typed `x::PetscVec`, but # `AbstractPetscVec` also includes `VecPtr`. Convert transparently. diff --git a/test/gpu/device_arrays.jl b/test/gpu/device_arrays.jl new file mode 100644 index 000000000..1ce2d2bcc --- /dev/null +++ b/test/gpu/device_arrays.jl @@ -0,0 +1,138 @@ +# test/gpu/device_arrays.jl +# PETScCUDAExt and PETScAMDGPUExt: Vecs in device memory come out of +# with_local_array! and with_local_arrays! as CuArrays or ROCArrays, DMStag +# field views wrap them, and what a kernel writes reaches PETSc. +# +# Not part of runtests.jl: it needs a GPU and a PETSc built with CUDA or HIP. +# Run it in an environment with PETSc.jl (dev'd), MPI.jl, KernelAbstractions.jl +# and CUDA.jl or AMDGPU.jl, whose LocalPreferences.toml points set_library! at +# that PETSc: +# +# julia --project=ENV test/gpu/device_arrays.jl cuda # or hip + +using Test +using PETSc, MPI, KernelAbstractions +using PETSc: LibPETSc + +const gpu = get(ARGS, 1, "") +# The device array type and PETSc's names for a sequential and a DM vector type +if gpu == "cuda" + using CUDA + const DeviceArray, seq_type, dm_vec_type = CuArray, :seqcuda, "cuda" +elseif gpu == "hip" + using AMDGPU + const DeviceArray, seq_type, dm_vec_type = ROCArray, :seqhip, "hip" +else + error("pass the GPU backend as the argument: cuda or hip, got $gpu") +end + +MPI.Initialized() || MPI.Init() + +# Rp = ∂Vx/∂x + ∂Vy/∂y on the owned elements (unit spacing), offset by `lo` to +# the views' ghost-inclusive indices +@kernel function divergence!(Rp, Vx, Vy, lo) + I = @index(Global, Cartesian) + lo + Rp[I] = Vx[I + CartesianIndex(1, 0)] - Vx[I] + Vy[I + CartesianIndex(0, 1)] - Vy[I] +end + +petsclib = PETSc.petsclibs[1] +PETSc.initialize(petsclib) +const T = petsclib.PetscScalar +const n = 100 + +# A sequential Vec of the given type (`seq_type` on the device, :seq on the host). +function vec_of_type(type, values) + v = LibPETSc.VecCreate(petsclib, MPI.COMM_SELF) + LibPETSc.VecSetSizes(petsclib, v, petsclib.PetscInt(n), petsclib.PetscInt(n)) + PETSc.set_type!(v, type) + PETSc.with_local_array!(a -> copyto!(a, values), v) + return v +end + +@testset "device arrays ($gpu)" begin + @testset "with_local_array! on a device Vec" begin + x = rand(T, n) + v = vec_of_type(seq_type, x) + @test PETSc.memtype(v) === DeviceArray + PETSc.with_local_array!(v) do a + @test a isa DeviceArray{T} + a .= 2 .* a + end + @test v[1:n] ≈ 2x + @test PETSc.with_local_array!(sum, DeviceArray, v) ≈ sum(2x) + @test_throws "expected array of type Vector" PETSc.with_local_array!(sum, Vector, v) + # plain VecGetArray: PETSc copies to the host + a = PETSc.unsafe_local_array(v; write = false) + @test a isa Vector{T} + @test a ≈ 2x + finalize(a) + PETSc.destroy!(v) + end + + @testset "with_local_arrays!, $fx_type × $lx_type" for + (fx_type, lx_type) in ((seq_type, seq_type), (:seq, seq_type), (seq_type, :seq)) + x = rand(T, n) + g = vec_of_type(fx_type, ones(T, n)) + l = vec_of_type(lx_type, x) + PETSc.with_local_arrays!(g, l) do fx, lx + @test fx isa DeviceArray{T} + @test lx isa DeviceArray{T} + fx .+= 3 .* lx + end + @test g[1:n] ≈ 1 .+ 3x + @test l[1:n] == x + PETSc.destroy!(g) + PETSc.destroy!(l) + end + + # with_field_views! on a DMStag with device Vecs, against the same steps on the host + @testset "DMStag field views on the device" begin + none = ntuple(_ -> LibPETSc.DM_BOUNDARY_NONE, 2) + function stag_result(vec_type) + dm = PETSc.DMStag(petsclib, MPI.COMM_SELF, none, (8, 6), (0, 1, 1); dm_vec_type = vec_type) + x = PETSc.local_vec(dm) + r = PETSc.local_vec(dm) + m = length(x) + PETSc.with_local_array!(a -> copyto!(a, T.(1:m) .^ 2), x) + PETSc.with_local_array!(a -> fill!(a, 0), r) + flow = (PETSc.face_location(dm, 1) => 0, PETSc.face_location(dm, 2) => 0, + PETSc.element_location(dm) => 0) + c = PETSc.corners(dm) + lo = c.lower - CartesianIndex(1, 1) + # OffsetArray broadcasts element by element whatever its parent, so + # whole-view broadcasts go through `parent`; kernels take the views as they are + PETSc.with_field_views!(dm, x, r; fields = flow, write = (false, true)) do (Vx, Vy, P), (Rx, Ry, Rp) + vec_type == dm_vec_type && @test parent(parent(Rp)) isa DeviceArray{T} + backend = get_backend(Rp) + parent(Rx) .= 3 + divergence!(backend, 64)(Rp, Vx, Vy, lo; ndrange = Tuple(c.upper - c.lower) .+ 1) + KernelAbstractions.synchronize(backend) + end + # the whole array, indexed [I..., slot] + PETSc.with_field_views!(dm, r) do R + parent(R) .*= 2 + end + result = r[1:m] + PETSc.destroy!(x) + PETSc.destroy!(r) + PETSc.destroy!(dm) + return result + end + host = stag_result("standard") + @test stag_result(dm_vec_type) == host + @test count(!iszero, host) > 0 + end + + @testset "host Vecs stay on the host" begin + g = vec_of_type(:seq, ones(T, n)) + l = vec_of_type(:seq, ones(T, n)) + PETSc.with_local_arrays!(g, l) do fx, lx + @test fx isa Vector{T} + @test lx isa Vector{T} + end + PETSc.destroy!(g) + PETSc.destroy!(l) + end +end + +PETSc.finalize(petsclib) diff --git a/test/runtests.jl b/test/runtests.jl index 1c79503e7..d5cbebef2 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -5,9 +5,9 @@ using PETSc, Pkg # Make sure that all dependencies are installed also on a clean system Pkg.instantiate() -# When set_library! has been used, petsc_library is a path string and PETSc_jll -# is not loaded. Only import JLL-specific symbols when using the default binaries. -const USING_CUSTOM_LIB = PETSc.petsclibs[1].petsc_library isa AbstractString +# When set_library! has been used, PETSc_jll is not loaded. Only import +# JLL-specific symbols when using the default binaries. +const USING_CUSTOM_LIB = PETSc.library_info().source === :preferences import MPIPreferences # always import to ensure local MPI API is configured