PERF: Compute the regularization on the GPU

CHANGELOG.md

@@ -33,6 +33,10 @@ v1.5.0 (not yet released)
 - PERF: The element-wise sensitivity work of the gradient and of the Hessian-vector
   products runs on the GPU, in double precision, instead of in NumPy on the host;
   only the gradient and the product are copied to the host
+- PERF: The phase-field regularizations, including the consistent double well, also
+  compute on the GPU; their results come back as host or device arrays like their input
+- TST: GPU tests compare gradient, Hessian product and regularization with the CPU;
+  they are skipped without CuPy and a GPU build of muGrid
 
 v1.4.0 (28Oct26)
 ----------------

muTopOpt/conduction.py

@@ -301,6 +301,12 @@ def nb_pixels(self):
         density array on this rank."""
         return self._nb_pixels
 
+    @property
+    def xp(self):
+        """The array module of the fields: cupy on a device, numpy on the
+        host."""
+        return self._xp
+
     def to_host(self, a):
         return a.get() if hasattr(a, "get") else np.asarray(a)
 

muTopOpt/nodal.py

@@ -207,17 +207,25 @@ def __init__(self, nodal_map: NodalElementMap):
 
     def value_and_gradient(self, rho):
         """Return ``(∫ W(ρ_h) dx, ∂/∂ρ)`` for a nodal density array; the value
-        is MPI-reduced, the gradient is the local nodal slice.
+        is MPI-reduced, the gradient is the local nodal slice, an array of
+        the kind `rho` is (host, or in the fields' array module).
 
         The kernel gives each thread its own node, so it writes the nodal
         gradient directly -- no scatter and no ghost reduction here.
         """
         h = self.h
-        self._rho.p[...] = h.to_device(np.asarray(rho, dtype=h.dtype))
+        xp = h.xp
+        shape = rho.shape
+        self._rho.p[...] = h.to_device(rho)
         h.engine.communicate_ghosts(self._rho)
         self.op.compute(self._rho, self._moments, self._grads)
-        c = self.W_COEFFS
-        m = h.to_host(self._moments.p).reshape((self.NB_MOMENTS, -1))
-        g = h.to_host(self._grads.p).reshape((self.NB_MOMENTS, -1))
+        # Combined in the fields' array module, in double precision; only
+        # the NB_MOMENTS local moment sums reach the host for the reduction.
+        c = xp.asarray(self.W_COEFFS)
+        m = xp.asarray(self._moments.p, dtype=float).reshape(
+            (self.NB_MOMENTS, -1))
+        g = xp.asarray(self._grads.p, dtype=float).reshape(
+            (self.NB_MOMENTS, -1))
         f = h.comm.sum(float(c @ m.sum(axis=1)))
-        return f, (c @ g).reshape(np.shape(rho))
+        grad = (c @ g).reshape(shape)
+        return f, (grad if isinstance(rho, xp.ndarray) else h.to_host(grad))

muTopOpt/regularization.py

@@ -68,6 +68,15 @@
 C_W = 1.0 / 3.0
 
 
+def _like(h, a, given):
+    """``a`` (in the fields' array module) as an array of the kind `given`
+    came in as: kept on the device for a device array, copied to the host
+    otherwise. The regularization computes in the fields' module either
+    way, so that a host caller pays one copy of the result and a device
+    caller none."""
+    return a if isinstance(given, h.xp.ndarray) else h.to_host(a)
+
+
 def perimeter_prefactor(homogenization):
     """``1 / (c_W V^((d−1)/d))``: the factor that turns the raw Modica-Mortola
     functional into the interfacial area in units of ``L^(d−1)``, with ``V``
@@ -164,15 +173,17 @@ def value_and_gradient(self, rho):
         ``w/(c_W V^((d−1)/d)) · h^d [2η (−Δ_h ρ) + W'(ρ)/η]`` with
         ``W'(ρ) = 2ρ(1 − ρ)(1 − 2ρ)``.
         """
-        rho = np.asarray(rho)
-        self._rho.p[...] = self.h.to_device(rho)
-        self.h.engine.communicate_ghosts(self._rho)
+        h = self.h
+        given = rho
+        rho = h.xp.asarray(rho, dtype=float)
+        self._rho.p[...] = h.to_device(rho)
+        h.engine.communicate_ghosts(self._rho)
         self.laplace.apply(self._rho, self._lap)  # (-Δ rho)
-        lap = self.h.to_host(self._lap.p)
+        lap = h.xp.asarray(self._lap.p, dtype=float)
 
-        grad_pen = self.h.comm.sum(float(np.sum(rho * lap))) * self.vol_pixel
-        dwell = self.h.comm.sum(
-            float(np.sum(rho**2 * (1.0 - rho) ** 2))
+        grad_pen = h.comm.sum(float((rho * lap).sum())) * self.vol_pixel
+        dwell = h.comm.sum(
+            float((rho**2 * (1.0 - rho) ** 2).sum())
         ) * self.vol_pixel
 
         # Raw Modica-Mortola bracket eta * ∫|∇rho|^2 + (1/eta) * ∫ W; the
@@ -186,7 +197,7 @@ def value_and_gradient(self, rho):
             + dwell_drho * self.vol_pixel / self.eta
         )
         scale = self.scale
-        return scale * f, scale * g
+        return scale * f, _like(h, scale * g, given)
 
     def hessian_vector_product(self, rho, v):
         """Return ``H_reg v`` for a density direction ``v`` (exact).
@@ -196,19 +207,21 @@ def hessian_vector_product(self, rho, v):
         ``h^d W''(ρ)/η`` with ``W''(ρ) = 2 (1 − 6ρ + 6ρ²)`` (both times the
         normalization :attr:`scale`). Since ``W''(1/2) < 0``, the
         regularization Hessian is indefinite in the interfacial region."""
-        rho = np.asarray(rho)
-        v = np.asarray(v)
-        self._rho.p[...] = self.h.to_device(v)
-        self.h.engine.communicate_ghosts(self._rho)
+        h = self.h
+        given = v
+        rho = h.xp.asarray(rho, dtype=float)
+        v = h.xp.asarray(v, dtype=float)
+        self._rho.p[...] = h.to_device(v)
+        h.engine.communicate_ghosts(self._rho)
         self.laplace.apply(self._rho, self._lap)  # (-Delta v)
-        lap_v = self.h.to_host(self._lap.p)
+        lap_v = h.xp.asarray(self._lap.p, dtype=float)
 
         d2well = 2.0 * (1.0 - 6.0 * rho + 6.0 * rho**2)
         hv = (
             self.eta * 2.0 * lap_v * self.vol_pixel
             + d2well * v * self.vol_pixel / self.eta
         )
-        return self.scale * hv
+        return _like(h, self.scale * hv, given)
 
 
 def fe_laplacian_stencil(dim, grid_spacing, element):
@@ -357,24 +370,26 @@ def scale(self):
     def value_and_gradient(self, rho):
         """Return ``(f_reg, ∂f_reg/∂ρ)`` for a nodal density array (value
         MPI-reduced, gradient the local nodal slice)."""
-        rho = np.asarray(rho)
-        self._rho.p[...] = self.h.to_device(rho)
-        self.h.engine.communicate_ghosts(self._rho)
+        h = self.h
+        given = rho
+        rho = h.xp.asarray(rho, dtype=float)
+        self._rho.p[...] = h.to_device(rho)
+        h.engine.communicate_ghosts(self._rho)
         self.laplace.apply(self._rho, self._Lrho)  # L rho
-        Lrho = self.h.to_host(self._Lrho.p)
+        Lrho = h.xp.asarray(self._Lrho.p, dtype=float)
 
         # rhoᵀ L rho = ∫|∇rho|² (L already includes the physical volume
         # weighting); multiplied by eta below.
-        grad_pen = self.h.comm.sum(float(np.sum(rho * Lrho)))
+        grad_pen = h.comm.sum(float((rho * Lrho).sum()))
 
         if self._dwell is not None:
             # Fully consistent Galerkin: exact ∫ W(rho(x)) dx of the nodal
             # interpolant, with its exact nodal gradient.
             dwell, dwell_grad = self._dwell.value_and_gradient(rho)
         else:
             # Lumped nodal quadrature (pointwise).
-            dwell = self.h.comm.sum(
-                float(np.sum(rho**2 * (1.0 - rho) ** 2))
+            dwell = h.comm.sum(
+                float((rho**2 * (1.0 - rho) ** 2).sum())
             ) * self.vol_pixel
             dwell_grad = (
                 2.0 * rho * (1.0 - rho) * (1.0 - 2.0 * rho) * self.vol_pixel
@@ -385,7 +400,7 @@ def value_and_gradient(self, rho):
         f = self.eta * grad_pen + dwell / self.eta
         g = self.eta * 2.0 * Lrho + dwell_grad / self.eta
         scale = self.scale
-        return scale * f, scale * g
+        return scale * f, _like(h, scale * g, given)
 
     def hessian_vector_product(self, rho, v):
         """Return ``H_reg v`` for a nodal density direction ``v``.
@@ -398,16 +413,18 @@ def hessian_vector_product(self, rho, v):
         trust-region method only requires a bounded model Hessian, and the
         lumped diagonal retains the indefiniteness in the interface
         (``W''(1/2) < 0``)."""
-        rho = np.asarray(rho)
-        v = np.asarray(v)
-        self._rho.p[...] = self.h.to_device(v)
-        self.h.engine.communicate_ghosts(self._rho)
+        h = self.h
+        given = v
+        rho = h.xp.asarray(rho, dtype=float)
+        v = h.xp.asarray(v, dtype=float)
+        self._rho.p[...] = h.to_device(v)
+        h.engine.communicate_ghosts(self._rho)
         self.laplace.apply(self._rho, self._Lrho)  # L v
-        Lv = self.h.to_host(self._Lrho.p)
+        Lv = h.xp.asarray(self._Lrho.p, dtype=float)
 
         d2well = 2.0 * (1.0 - 6.0 * rho + 6.0 * rho**2)
         hv = (
             self.eta * 2.0 * Lv
             + d2well * v * self.vol_pixel / self.eta
         )
-        return self.scale * hv
+        return _like(h, self.scale * hv, given)

test/conftest.py

@@ -8,3 +8,9 @@ def comm():
     launcher, serial otherwise. muGrid starts MPI itself, so mpi4py is not
     needed (muGrid.Communicator() alone would always be serial)."""
     return muGrid.Communicator.world()
+
+
+def pytest_configure(config):
+    config.addinivalue_line(
+        "markers", "gpu: compares a GPU run with the CPU; skipped without "
+        "cupy and a GPU build of muGrid")

test/test_gpu.py

@@ -0,0 +1,106 @@
+#
+# Copyright 2026 Lars Pastewka
+#
+# MIT License (see LICENSE)
+#
+"""The same problems on the GPU as on the CPU.
+
+The element-wise work of the gradient, the Hessian-vector product and the
+regularization runs in the fields' array module -- cupy on a GPU -- and only
+the gradient and the product come back to the host. On the CPU that module
+is numpy, so the other tests cannot catch a cupy-specific mistake (such as an
+implicit device-to-host conversion, which cupy refuses); these do. They are
+skipped without cupy and a GPU build of muGrid.
+"""
+
+import numpy as np
+import pytest
+
+from muTopOpt import (
+    Homogenization,
+    NodalPhaseFieldRegularization,
+    PhaseFieldRegularization,
+    SimpMaterial,
+    StressTargetProblem,
+)
+from muTopOpt.loadcases import isotropic_stiffness_tensor, target_load_cases
+
+from mpi_helpers import assert_allclose_on_all_ranks
+
+pytestmark = pytest.mark.gpu
+
+DESIGNS = ["element", "nodal"]
+
+
+def _make_problem(comm, device, design, n=16):
+    material = SimpMaterial(E_solid=1.0, nu=0.3, penalty=3.0, void_ratio=1e-2)
+    homog = Homogenization((n, n), material, comm=comm, cg_tol=1e-12,
+                           cg_maxiter=5000, device=device)
+    cases = target_load_cases(
+        2, isotropic_stiffness_tensor(2, K=0.15, G=0.08), magnitude=0.01)
+    reg = (NodalPhaseFieldRegularization(homog, dwell="consistent")
+           if design == "nodal" else PhaseFieldRegularization(homog))
+    return StressTargetProblem(homog, cases, regularization=reg,
+                               design=design, hessian=True)
+
+
+@pytest.fixture
+def gpu_problem(comm, request):
+    """The GPU problem of the test's design, or a skip without a GPU."""
+    try:
+        import cupy
+
+        cupy.zeros(1)
+        problem = _make_problem(comm, "gpu", request.param)
+    except Exception as err:  # no cupy, no GPU, or a CPU-only muGrid
+        pytest.skip(f"needs cupy and a GPU build of muGrid ({err})")
+    if not problem.h.device.is_device:
+        pytest.skip("muGrid did not place the fields on a GPU")
+    return problem
+
+
+def _inputs(problem, seed=0):
+    rng = np.random.default_rng(seed + problem.h.comm.rank)
+    shape = problem.h.nb_pixels
+    return rng.uniform(0.2, 0.8, shape), rng.standard_normal(shape)
+
+
+@pytest.mark.parametrize("gpu_problem", DESIGNS, indirect=True)
+def test_gradient_and_hessian_match_the_cpu(comm, gpu_problem):
+    cpu_problem = _make_problem(comm, "cpu", gpu_problem.design)
+    rho, v = _inputs(gpu_problem)
+    f_gpu, g_gpu = gpu_problem.objective_and_gradient(rho)
+    f_cpu, g_cpu = cpu_problem.objective_and_gradient(rho)
+    # The gradient comes back to the host for the optimizer
+    assert isinstance(g_gpu, np.ndarray)
+    assert f_gpu == pytest.approx(f_cpu, rel=1e-8)
+    assert_allclose_on_all_ranks(comm, g_gpu, g_cpu, rtol=1e-7,
+                                 atol=1e-7 * np.abs(g_cpu).max())
+    hv_gpu = gpu_problem.hessian_vector_product(v)
+    hv_cpu = cpu_problem.hessian_vector_product(v)
+    assert isinstance(hv_gpu, np.ndarray)
+    assert_allclose_on_all_ranks(comm, hv_gpu, hv_cpu, rtol=1e-6,
+                                 atol=1e-6 * np.abs(hv_cpu).max())
+
+
+@pytest.mark.parametrize("gpu_problem", DESIGNS, indirect=True)
+def test_regularization_keeps_device_arrays_on_the_device(comm, gpu_problem):
+    import cupy
+
+    reg = gpu_problem.regularization
+    cpu_reg = _make_problem(comm, "cpu", gpu_problem.design).regularization
+    rho, v = _inputs(gpu_problem)
+    f_cpu, g_cpu = cpu_reg.value_and_gradient(rho)
+    hv_cpu = cpu_reg.hessian_vector_product(rho, v)
+    # Device in, device out
+    f, g = reg.value_and_gradient(cupy.asarray(rho))
+    hv = reg.hessian_vector_product(cupy.asarray(rho), cupy.asarray(v))
+    assert isinstance(g, cupy.ndarray) and isinstance(hv, cupy.ndarray)
+    assert f == pytest.approx(f_cpu, rel=1e-10)
+    assert_allclose_on_all_ranks(comm, g.get(), g_cpu, rtol=1e-10, atol=0)
+    assert_allclose_on_all_ranks(comm, hv.get(), hv_cpu, rtol=1e-10,
+                                 atol=1e-12 * np.abs(hv_cpu).max())
+    # Host in, host out
+    _, g_host = reg.value_and_gradient(rho)
+    assert isinstance(g_host, np.ndarray)
+    assert_allclose_on_all_ranks(comm, g_host, g_cpu, rtol=1e-10, atol=0)