Skip to content

PERF: Compute the regularization on the GPU - #39

Merged
pastewka merged 2 commits into
mainfrom
perf/regularization-on-device
Oct 31, 2026
Merged

pastewka merged 2 commits into
mainfrom
perf/regularization-on-device

Conversation

@pastewka

Copy link
Copy Markdown
Contributor

Step 1 of the plan to keep the optimization loop on the GPU (follows #38).

Why

The phase-field regularizations computed the Laplacian on the device, then copied it to the host. The consistent double well did the same with its 3 moments and their gradients. Sums and gradients were then done in NumPy. In the 4-rank MI300A run at 512³ that was 2.7 s per objective evaluation and 0.8 s per Hessian product: 26.6 s of 842.

What changes

  • PhaseFieldRegularization, NodalPhaseFieldRegularization and ConsistentDoubleWell compute in the fields' array module (CuPy on a GPU), in double precision. Only the local sums reach the host for the MPI reduction.
  • Results come back like the input: a host array after one copy for a host caller (today's), or a device array for a device caller, without a copy (after the device-resident optimizer loop, step 3).
  • HomogenizationConductivity gains the xp property the regularization uses.

Tests

  • New test/test_gpu.py (marker gpu):

    • gradient and Hessian product of a GPU problem match the CPU, for the element and the nodal design;
    • the regularization returns device arrays for device input and host arrays for host input, matching the CPU.

    These tests are skipped without CuPy and a GPU build of muGrid, so they have not run yet. Please run pytest test/test_gpu.py on MI300A. They also cover PERF: Keep the sensitivity work on the GPU #38's GPU path for the first time.

  • The CPU suite passes (249).

🤖 Generated with Claude Code

The phase-field regularizations copied the Laplacian (and the consistent
double well its three moments and their gradients) to the host and did the
sums and gradients in NumPy: in the 4-rank MI300A run at 512^3, 2.7 s per
objective evaluation and 0.8 s per Hessian product, 26.6 s of 842.

Both regularizations and ConsistentDoubleWell now compute in the fields'
array module (cupy on a GPU) in double precision; only the local sums reach
the host for the MPI reduction. Results come back like the input: a host
array for a host caller, after one copy, or a device array for a device
caller, without one. The conduction homogenization gains the xp property
the regularization uses.

New GPU tests compare gradient, Hessian product and regularization of a GPU
problem with the CPU, and check that device input stays on the device. They
are skipped without cupy and a GPU build of muGrid.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
flake8 E303/E302 in regularization.py.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
@pastewka
pastewka merged commit f49f55d into main Oct 31, 2026
8 checks passed
@pastewka
pastewka deleted the perf/regularization-on-device branch October 31, 2026 07:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant