diff --git a/CHANGELOG.md b/CHANGELOG.md index c3364a6..2f2c8a4 100644 --- a/CHANGELOG.md +++ b/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) ---------------- diff --git a/muTopOpt/conduction.py b/muTopOpt/conduction.py index 5910ddb..b362958 100644 --- a/muTopOpt/conduction.py +++ b/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) diff --git a/muTopOpt/nodal.py b/muTopOpt/nodal.py index a98e306..cd563cd 100644 --- a/muTopOpt/nodal.py +++ b/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)) diff --git a/muTopOpt/regularization.py b/muTopOpt/regularization.py index 9d7ad2d..73153d1 100644 --- a/muTopOpt/regularization.py +++ b/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,15 +370,17 @@ 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 @@ -373,8 +388,8 @@ def value_and_gradient(self, rho): 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) diff --git a/test/conftest.py b/test/conftest.py index 02b844c..7ccb7e6 100644 --- a/test/conftest.py +++ b/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") diff --git a/test/test_gpu.py b/test/test_gpu.py new file mode 100644 index 0000000..4154c17 --- /dev/null +++ b/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)