On the GPU: the same script, different hardware

Nothing in the previous tutorials mentions a device, and that is deliberate: the same script runs on host cores or on GPUs, and what changes is set before hIPPyMFEM is imported.

Two things can move, independently:

what

how

needs

the element kernels (residual, Jacobian, every Hessian block)

HIPPYMFEM_DEVICE=gpu

JAX with a CUDA or ROCm build

MFEM and hypre: the matrices, the AMG hierarchy, the Krylov solves

HIPPYMFEM_HYPRE_DEVICE=1

PyMFEM built for CUDA or HIP

Moving only the kernels gains little, because a Newton-CG step is mostly solves. Moving the solves as well is worth more than an order of magnitude. This notebook runs on whatever the machine has, and says which path it took.

1. Choose the device, then import

import os
import sys
import time

sys.path.insert(0, os.environ.get("HIPPYMFEM_BASE_DIR", os.path.abspath("..")))

# The two variables are read once, when hIPPyMFEM (and through it JAX) is imported, so
# they have to be set *before* the kernel starts.  This notebook reads whatever is set and
# reports the path it took; to take the GPU path, start Jupyter with
#
#     HIPPYMFEM_DEVICE=gpu HIPPYMFEM_HYPRE_DEVICE=1 jupyter notebook
#
# Setting them here, after the kernel has started, would be too late for JAX, and asking
# for a GPU that is not there fails at import rather than falling back.
print("HIPPYMFEM_DEVICE       =", os.environ.get("HIPPYMFEM_DEVICE", "(unset: host)"))
print("HIPPYMFEM_HYPRE_DEVICE =", os.environ.get("HIPPYMFEM_HYPRE_DEVICE", "(unset: host)"))

import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

from mpi4py import MPI
import mfem.par as mfem
import jax.numpy as jnp

import hippymfem as hm
from hippymfem import nb
from hippymfem.modeling.variables import STATE, PARAMETER, ADJOINT

HIPPYMFEM_DEVICE       = gpu
HIPPYMFEM_HYPRE_DEVICE = 1

hippymfem: MFEM on cuda device 0 of 1

2. Check, do not assume

Two questions, two answers. hm.config reports where the element kernels will run, and hm.mfem_config() reports what this PyMFEM was built with, read from its installed configuration header rather than probed (MFEM aborts the job when asked to configure a device it does not have).

from hippymfem.fem import kernel as kernel_mod

on_gpu = "cpu" not in str(kernel_mod.device()).lower()
cfg = hm.mfem_config()
hypre_on_device = hm.config.hypre_device and (cfg["MFEM_USE_CUDA"] or cfg["MFEM_USE_HIP"])

print("element kernels on   :", kernel_mod.device())
print("PyMFEM built with    : CUDA=%s  HIP=%s  (MFEM %s)"
      % (cfg["MFEM_USE_CUDA"], cfg["MFEM_USE_HIP"], cfg["version"]))
print("hypre on the device  :", bool(hypre_on_device))
print()
print("=> this notebook is running the %s path" % ("GPU" if on_gpu else "host"))

element kernels on   : cuda:0
PyMFEM built with    : CUDA=True  HIP=False  (MFEM 4.9.0)
hypre on the device  : True

=> this notebook is running the GPU path

3. A problem, timed on this machine

The same coefficient-inversion problem as the earlier tutorials, sized so it runs in a notebook either way. The interesting number is not the absolute time, which depends on the machine, but that the answer does not change with the hardware: the cost functional and the CG count are the same on a GPU as on a host core, to round-off.

n = 48 if on_gpu else 24
mesh = mfem.ParMesh(MPI.COMM_WORLD,
                    mfem.Mesh.MakeCartesian2D(n, n, mfem.Element.QUADRILATERAL))
Vu = hm.FunctionSpace.H1(mesh, 2)          # the adjoint lives in the state space:
Vh = [Vu, hm.FunctionSpace.H1(mesh, 1), Vu]   # pass the same object for both
bottom, right, top, left = 1, 2, 3, 4
bc = hm.DirichletBC(Vh[STATE], 0.0, bdr_attributes=[left, right])
bc0 = hm.DirichletBC(Vh[STATE], 0.0, bdr_attributes=[left, right])


def pde_varf(u, m, p, x):
    f = 1.0 + 8.0 * jnp.exp(-30.0 * ((x[0] - 0.5) ** 2 + (x[1] - 0.5) ** 2))
    return jnp.exp(m.val) * hm.inner(u.grad, p.grad) - f * p.val


pde = hm.PDEVariationalProblem(Vh, pde_varf, bc, bc0, is_fwd_linear=True)
mtrue = Vh[PARAMETER].project(
    lambda x: 0.8 * np.exp(-12.0 * ((x[0] - 0.35) ** 2 + (x[1] - 0.6) ** 2)))

utrue = pde.generate_state()
pde.solveFwd(utrue, [utrue, mtrue, None])

rng = np.random.default_rng(1)
targets = np.column_stack((rng.uniform(0.05, 0.95, 120), rng.uniform(0.05, 0.95, 120)))
B = hm.assemblePointwiseObservation(Vh[STATE], targets)
data = B.createVecLeft(); B.mult(utrue, data)
noise_std = 0.005 * data.norm("linf")
hm.parRandom.set_seed(2); B.perturb(data, noise_std)

model = hm.Model(pde, hm.BiLaplacianPrior(Vh[PARAMETER], 1.0, 8.0, robin_bc=True),
                 hm.DiscreteStateObservation(B, data, noise_std ** 2))

params = hm.ReducedSpaceNewtonCG_ParameterList()
params["rel_tolerance"] = 1e-8
params["max_iter"] = 8
params["print_level"] = -1
solver = hm.ReducedSpaceNewtonCG(model, params)

t0 = time.perf_counter()
x = solver.solve([None, model.prior.mean.copy(), None])
dt = time.perf_counter() - t0

print("%s path, %d state dofs" % ("GPU" if on_gpu else "host", Vh[STATE].GlobalTrueVSize()))
print("%d Newton iterations, %d CG, J = %.8e, %.1f s"
      % (solver.it, solver.total_cg_iter, solver.final_cost, dt))

GPU path, 9409 state dofs
8 Newton iterations, 55 CG, J = 1.07884372e+02, 4.7 s

4. What the device path actually costs

One assembly is one host-to-device copy of the local dof vectors and one device-to-host copy of the assembled values; the (ne, nd, nd) element-matrix array is never formed on the host. The measured picture, from docs/source/guide/gpu.rst, for two Newton-CG steps at \(128^3\) (36 M unknowns):

where

time

hIPPYlibx on 4 CPU cores

2669 s

hIPPyMFEM on 4 CPU cores

1409 s

hIPPyMFEM on 4 L40S GPUs

23.4 s

The card matters as well: the element kernels, which are double precision throughout, run 2.3 to 6.3 times faster on an H100 than on an L40S. Measure, do not assume.

5. Memory: three knobs worth knowing

A GPU run is bounded by memory far more often than by arithmetic, and three settings decide how the card is shared between JAX and hypre.

  • HIPPYMFEM_GPU_MEM_FRACTION is JAX’s share of the card. The default follows the ranks per device, and halves when hypre shares the card, because JAX’s allocator never gives back what it took. A share of 0.20 saves about 3 GiB per card at \(128^3\) on four L40S for 14 % more time.

  • HIPPYMFEM_ELEMENT_CHUNK pins how many elements go to one kernel launch. The planner sizes this from the free budget, which is usually right; pin it when a run has to be predictable, since a plan made when the card was empty is reused later when hypre has filled it.

  • XLA_PYTHON_CLIENT_PREALLOCATE=true takes JAX’s arena as one region. Worth setting when a single array is a large share of the arena: the arena otherwise grows in pieces and counts every piece against its cap, so a large contiguous request can fail while the cap reads almost empty. The assembly warns and names this setting when it sees the situation coming.

for knob in ("device", "hypre_device", "element_chunk", "geometry_stream",
             "fused_keep", "fused_keep_share", "geometry_slice"):
    print("%-18s %s" % (knob, getattr(hm.config, knob)))

device             gpu
hypre_device       True
element_chunk      0
geometry_stream    0.25
fused_keep         True
fused_keep_share   0.25
geometry_slice     0

6. More than one card

One MPI rank per card. The rank must choose its card before CUDA starts, which tools/mpirun_pinned.sh does from the launcher’s node-local rank:

HIPPYMFEM_DEVICE=gpu HIPPYMFEM_HYPRE_DEVICE=1 \
    mpirun -n 4 tools/mpirun_pinned.sh python my_inversion.py

Without the pinning every rank opens a context on every card, which shows up as memory disappearing rather than as an error. The library also pins itself when it is imported before mpi4py and MFEM, and says so on rank 0 when it came too late.

Where to look next

  • docs/source/guide/gpu.rst has the per-stage times, the memory per rank, the AMD/ROCm build, and the limits that appear past a few million elements per rank.

  • 08_ResidualsAsPrograms is the density side of the same story; the kernels it generates are exactly what moves to the card here.