Skip to content

GPU stack 5/8: Poisson/diffusion solve with fewer passes over the vectors - #786

Open
max-models wants to merge 8 commits into
gpu-stack-4-kronecker-massfrom
gpu-stack-5-diffusion-solve
Open

max-models wants to merge 8 commits into
gpu-stack-4-kronecker-massfrom
gpu-stack-5-diffusion-solve

Conversation

@max-models

@max-models max-models commented Oct 16, 2026 •

Copy link
Copy Markdown
Member

Part 5 of 8 of the split of #772 (cuda-aware-mpi-fixes) into stacked PRs. Requires struphy-hub/feectools#104 to be merged into feectools devel-tiny first (PR 1 bumps the submodule).

  • BoundaryOperator.dot on device Stencil/Block vectors computes out = v * mask, one kernel per block (one for all blocks when they are contiguous float64). The masks are found once from apply_essential_bc_to_array, so they zero exactly the same entries. This replaces a copy and one fill per boundary and component.

  • WeightedMassOperator.without_boundary_ops() / MassWithoutBoundaryOps: the mass operator without its boundary operators. In grad.T @ M1 @ grad, grad already applies the same 0/1 masks, so this is the same operator with two fewer passes per product. It is not used with the multigrid preconditioner.

  • ImplicitDiffusion: the automatic 1e-14 stabilization stays out of the solver operator when a Dirichlet boundary already makes grad.T M1 grad positive definite. There is no stab_mat product for the right-hand side when sigma_2 = 0 (Poisson).

  • struphy.feec.fused_derivatives: fuse_boundary_gradient(derham.grad) and fuse_boundary_gradient(derham.grad.T) apply B_out G B_in with one CUDA kernel each on the GPU (fused_derivatives_cuda.cu), instead of a masked copy, three derivatives and a masked copy. The kernel does the same operations in the same order, so the products are bit for bit those of the composed operators. It reads the input with its ghost regions exchanged, times the input mask with exchanged ghost regions, which equals the exchanged masked input in serial and in parallel. Host data and the fake CuPy use the composed operators. ImplicitDiffusion uses it for grad and grad.T (not with the multigrid preconditioner). The original branch imported this module, but it was never committed, so it is written new here. test_fused_derivatives compares the kernels with the composed operators on the fake CuPy with emulated launches, and guards that the kernels actually run.

Also added: the fused mask kernel is skipped on cunumpy's fake CuPy, which cannot launch raw kernels (test_mass_preconditioners_fake_cupy failed).

Tests

Local, macOS without a GPU, in a fresh environment with feectools 1d8d47c and the pinned cunumpy, after merging devel (with #774 and #775). The PR passes test_cuda_parity, test_cuda_emulation, test_kernel_backends, test_sorting, test_accum_matrix_cupy, test_preconditioner_cupy, test_poisson and test_output on its own. The top of the stack also passes the mass-matrix, preconditioner, Derham-GPU, fused-gradient, Kronecker-block, Poisson, shear Alfvén, post-processing, polar fake-CuPy and all pic tests (1477 passed), the examples test for the two profiling cases, the sorting, Poisson and fused-gradient tests under mpirun -n 2, and the device MPI sort on the fake CuPy on 2 and 4 ranks. The Kronecker products match the stencil mass matrices to 5e-16 for M0, M1, M2, M3, Mv, M1n and M2n on a Cuboid, periodic and Dirichlet, on 1 and 2 ranks.
Not tested on a GPU here. The CUDA kernels and the device paths need a run on the cluster.

Stack (merge in order)

  1. Device ghost-exchange plans, fused block exchanges, single-reduction PCG (CUDA-aware MPI) feectools#104 (feectools, into devel-tiny)
# PR base
1 #782 update feectools (device ghost exchange, single-reduction PCG); pcg_sr solver option devel
2 #783 host/device conversions in post-processing; math.sqrt in Species #782
3 #784 CUDA mass-matrix assembly kernel kernel_3d_mat #783
4 #785 apply separable mass matrices as Kronecker products of 1D banded matrices #784
5 #786 Poisson/diffusion solve with fewer passes over the vectors (this PR) #785
6 #787 marker rows prepared inside the push kernels #786
7 #788 chunked marker loading, row-based weight update, leaner MPI sorting; static field in PushVxB #787
8 #789 profiling cases Poisson cube and Vlasov slab (CPU node vs GPU, GPU scaling) #788

🤖 Generated with Claude Code

max-models and others added 2 commits October 16, 2026 10:26
- BoundaryOperator.dot on device Stencil/Block vectors: out = v * mask, one kernel per block (one
  for all blocks when they are contiguous float64), with masks found once from
  apply_essential_bc_to_array, instead of a copy and one fill per boundary and component.
- WeightedMassOperator.without_boundary_ops / MassWithoutBoundaryOps: the mass operator without its
  boundary operators. In grad.T @ M1 @ grad, grad already applies the same 0/1 masks, so this is
  the same operator with two fewer passes per product (not with the multigrid preconditioner).
- ImplicitDiffusion: the automatic 1e-14 stabilization is left out of the solver operator when a
  Dirichlet boundary already makes grad.T M1 grad positive definite (diffusion_mat M1); no
  stab_mat product for the right-hand side when sigma_2 = 0 (Poisson).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
cunumpy's fake CuPy (tests without a GPU) cannot launch raw kernels; the per-block multiplication
is used there (test_mass_preconditioners_fake_cupy failed).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
max-models and others added 6 commits October 16, 2026 12:51
fuse_boundary_gradient(derham.grad) and fuse_boundary_gradient(derham.grad.T) apply B_out G B_in with
one CUDA kernel on the GPU (fused_derivatives_cuda.cu) instead of the composed operators (a masked
copy, three derivatives, a masked copy). The kernel computes the same operations in the same order,
so the products are bit for bit the same; the input is read with its ghost regions exchanged, times
the input mask with exchanged ghost regions, which equals the exchanged masked input in serial and
in parallel. Host data and cunumpy's fake CuPy use the composed operators. ImplicitDiffusion uses it
for grad and grad.T (not with the multigrid preconditioner), as on cuda-aware-mpi-fixes, where the
module was imported but never committed.

test_fused_derivatives: recognition of the gradient and its transpose, and the kernels against the
composed operators on the fake CuPy with emulated launches (Dirichlet at both ends and at one end).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

This branch has not been deployed

No deployments
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