Step 6 — Extend the MPI Solver to CUDA

Dec 15, 2025·
Hailin Liu
Hailin Liu
· 12 min read
projects

The MPI solver already gives us the distributed structure of the problem. Each process owns one rectangular subdomain, exchanges halo values with its neighbors, and participates in global PCG reductions.

The CUDA version keeps that decomposition and changes where the local numerical work is executed.

The execution hierarchy becomes

global grid
MPI subdomains
one MPI rank per subdomain
CUDA device memory
GPU kernels for local PCG work

MPI continues to handle communication between subdomains. CUDA takes over the stencil, vector updates, preconditioning, and rank-local reductions.

This makes the GPU version a direct extension of the distributed solver from Step 4 instead of a separate numerical implementation.

1. Keep the MPI domain decomposition

The top-level CUDA program reuses the same two-dimensional decomposition used by the MPI solver.

The global grid is still divided into

\[ P=p_xp_y \]

rectangular subdomains. The same DomainDecomposer chooses a valid process grid, the same Cartesian communicator is created with

int dims[2] = {px, py};
int periods[2] = {0, 0};
MPI_Cart_create(MPI_COMM_WORLD, 2, dims, periods, 1, &cart_comm);

and the same neighbor relationships are obtained with

MPI_Cart_shift(cart_comm, 0, 1, &left, &right);
MPI_Cart_shift(cart_comm, 1, 1, &bottom, &top);

Each MPI rank therefore keeps the same responsibilities it had before:

MPI rank
├── owns one local subdomain
├── knows left/right/top/bottom neighbors
├── participates in halo exchange
├── participates in MPI_Allreduce
└── contributes its local solution to the final global result

The CUDA solver is constructed over that local subdomain:

MPICudaPoissonSolver solver(
    x_end - x_start,
    y_end - y_start,
    X_MIN + hx * x_start,
    X_MIN + hx * x_end,
    Y_MIN + hy * y_start,
    Y_MIN + hy * y_end,
    region,
    f_func,
    timer,
    world_rank,
    cart_comm
);

So the global distribution strategy remains an MPI concern. CUDA operates inside the subdomain already assigned to the rank.

2. Bind each MPI rank to a CUDA device

Once several MPI processes may run on the same node, each process needs to select a GPU.

The program first creates a communicator containing processes that share the same physical node:

MPI_Comm local_comm;

MPI_Comm_split_type(
    MPI_COMM_WORLD,
    MPI_COMM_TYPE_SHARED,
    0,
    MPI_INFO_NULL,
    &local_comm
);

It then obtains a node-local rank:

int local_rank = 0;
MPI_Comm_rank(local_comm, &local_rank);

and asks CUDA how many devices are visible:

int dev_count = 0;
cudaGetDeviceCount(&dev_count);

The device is selected with

int dev_id = local_rank % dev_count;
cudaSetDevice(dev_id);

This creates a simple mapping:

rank on node 0 ─► GPU 0
rank on node 1 ─► GPU 1
rank on node 2 ─► GPU 2
...

If there are more local ranks than visible GPUs, the modulo operation causes devices to be reused.

The program also prints

world rank
hostname
local rank
device id
number of visible devices

which makes the rank-to-device mapping visible when the program starts.

3. Allocate a device-side solver state

The CPU solver stores its state in arrays such as

\[ u,\quad r,\quad z,\quad k,\quad p,\quad Ap,\quad M^{-1}. \]

The CUDA implementation allocates corresponding device arrays:

cudaMalloc(&d_u,     size * sizeof(double));
cudaMalloc(&d_r,     size * sizeof(double));
cudaMalloc(&d_z,     size * sizeof(double));
cudaMalloc(&d_k,     size * sizeof(double));
cudaMalloc(&d_p,     size * sizeof(double));
cudaMalloc(&d_A_p,   size * sizeof(double));
cudaMalloc(&d_M_inv, size * sizeof(double));

where

size = (M + 1) * (N + 1);

The GPU representation is one-dimensional. A helper

idx(i, j, N + 1)

maps a two-dimensional grid coordinate to

\[ i(N+1)+j. \]

The solver also allocates

cudaMalloc(&d_partial, NUM_PARTIALS * sizeof(double));

for the GPU reduction results, where NUM_PARTIALS is the number of partial values produced by the reduction kernels (one per launched reduction thread), and creates a CUDA stream:

cudaStreamCreate(&stream);

The resulting local memory layout is

host                              device

u       ───────────────────────►  d_u
r       ───────────────────────►  d_r
z       ───────────────────────►  d_z
k       ───────────────────────►  d_k
p       ───────────────────────►  d_p
A_p     ───────────────────────►  d_A_p
M_inv   ───────────────────────►  d_M_inv
                                  d_partial

The important design goal is to keep the iterative PCG vectors on the GPU during the solve.

4. Initialize on the CPU, then move persistent data to the GPU

The current implementation does not move every initialization step to CUDA.

It first reuses the existing CPU routines:

initialize_f();

initialize_k();
exchange_halo(k);

initialize_M_inv();

exchange_halo(u);
initialize_r();

Then the persistent fields needed by the GPU solver are copied to device memory:

to_device(k,     d_k);
to_device(M_inv, d_M_inv);
to_device(u,     d_u);
to_device(r,     d_r);

to_device() first flattens the two-dimensional host field into a contiguous host buffer and then calls

cudaMemcpy(
    d_field,
    h_buf.data(),
    size * sizeof(double),
    cudaMemcpyHostToDevice
);

After the initial residual has reached the GPU, the preconditioner and initial search direction are created on the device:

apply_preconditioner();
initialize_p();

initialize_p() uses a device-to-device copy:

cudaMemcpy(
    d_p,
    d_z,
    size * sizeof(double),
    cudaMemcpyDeviceToDevice
);

This leaves the main PCG state ready for GPU execution before the iteration begins.

5. Move the finite-difference stencil to a CUDA kernel

The most important local operation is still

\[ Ap=-\nabla\cdot(k\nabla p). \]

The CUDA implementation preserves the same five-point stencil used by the CPU solver.

Each GPU thread computes one interior grid point:

int j = blockIdx.x * blockDim.x + threadIdx.x;
int i = blockIdx.y * blockDim.y + threadIdx.y;

if (i < 1 || i >= M) return;
if (j < 1 || j >= N) return;

The thread maps the center and four neighbors to linear addresses:

int id_c = idx(i,   j,   N + 1);
int id_l = idx(i-1, j,   N + 1);
int id_r = idx(i+1, j,   N + 1);
int id_d = idx(i,   j-1, N + 1);
int id_u = idx(i,   j+1, N + 1);

It then computes the same arithmetic edge averages:

double kx_plus  = 0.5 * (k_c + k_r);
double kx_minus = 0.5 * (k_c + k_l);
double ky_plus  = 0.5 * (k_c + k_u);
double ky_minus = 0.5 * (k_c + k_d);

and evaluates the stencil:

out[id_c] = -1.0 * (
    (kx_plus  * (v_r - v_c)
   - kx_minus * (v_c - v_l)) / hx2
  +
    (ky_plus  * (v_u - v_c)
   - ky_minus * (v_c - v_d)) / hy2
);

The wrapper launches this kernel on a two-dimensional CUDA grid:

dim3 block(CUDA_BLOCK_X, CUDA_BLOCK_Y);

dim3 grid(
    (N + block.x - 1) / block.x,
    (M + block.y - 1) / block.y
);

apply_A_kernel<<<grid, block, 0, stream>>>(
    d_v,
    d_k,
    d_out,
    M,
    N,
    hx2,
    hy2
);

In the current configuration,

CUDA_BLOCK_X = 256
CUDA_BLOCK_Y = 1

unless these values are overridden at compile time.

The mathematical operator is therefore unchanged. The local grid traversal is now mapped to CUDA blocks and threads.

6. Move the PCG vector updates to CUDA

The pointwise operations from the OpenMP version map naturally to GPU kernels.

For the solution update,

\[ u_{ij} \leftarrow u_{ij}+\alpha p_{ij}, \]

the kernel is

u[id] += alpha * p[id];

For the residual,

\[ r_{ij} \leftarrow r_{ij}-\alpha(Ap)_{ij}, \]

the GPU performs

r[id] -= alpha * A_p[id];

The search direction update,

\[ p_{ij} \leftarrow z_{ij}+\beta p_{ij}, \]

becomes

p[id] = z[id] + beta * p[id];

and the diagonal preconditioner is

z[id] = M_inv[id] * r[id];

The host-side wrappers expose these operations as

cuda_update_u()
cuda_update_r()
cuda_update_p()
cuda_apply_preconditioner()

so the PCG control flow remains close to the CPU implementation.

The mapping from the earlier solver to CUDA is therefore direct:

CPU operationCUDA operation
apply_A()cuda_apply_A()
update_u()cuda_update_u()
update_r()cuda_update_r()
update_p()cuda_update_p()
apply_preconditioner()cuda_apply_preconditioner()

This separation keeps the numerical algorithm readable while placing the regular local loops on the GPU.

7. Perform rank-local reductions on the GPU

PCG also needs

\[ (r,r),\qquad (r,z),\qquad (p,Ap). \]

The CUDA implementation handles these in two device-side stages.

First, a reduction kernel computes partial sums.

For example, the \(r^2\) kernel assigns each CUDA thread a grid-stride loop:

int tid = blockIdx.x * blockDim.x + threadIdx.x;
int stride = blockDim.x * gridDim.x;

double sum = 0.0;

for (int i = tid; i < total; i += stride) {
    sum += r[i] * r[i];
}

d_partial[tid] = sum;

The same structure is used for

r[i] * z[i]
p[i] * Ap[i]

The current configuration launches

DOT_BLOCKS  = 128
DOT_THREADS = 256

so d_partial contains one partial value per launched reduction thread.

The second stage uses Thrust:

thrust::reduce(
    thrust::cuda::par.on(stream),
    begin,
    end,
    0.0,
    thrust::plus<double>()
);

to combine the partial array into one rank-local scalar.

The reduction path is therefore

local device vectors
CUDA grid-stride reduction kernel
d_partial
Thrust reduction
one scalar for this MPI rank

The custom partial kernels shown here do not use CUDA shared memory; each thread accumulates into a register-local scalar and writes one value to d_partial.

8. Keep MPI_Allreduce for the global PCG scalars

CUDA only replaces the local part of the reduction.

The distributed PCG solver still needs one global scalar shared by all MPI ranks.

For the residual norm, the code performs

cuda_reduce_r2_partials(d_partial, d_r, M, N, stream);
double local_sum =
    cuda_reduce_sum(d_partial, NUM_PARTIALS, stream);

double global_sum = 0.0;

MPI_Allreduce(
    &local_sum,
    &global_sum,
    1,
    MPI_DOUBLE,
    MPI_SUM,
    cart_comm
);

and returns

\[ \sqrt{\text{global\_sum}}. \]

The same structure is used for \((r,z)\) and \((p,Ap)\).

The reduction hierarchy is now

GPU threads
rank-local GPU reduction
local scalar on the host
MPI_Allreduce
global PCG scalar

So MPI and CUDA have clearly separated roles:

CUDA:
    reduce values inside one subdomain

MPI:
    combine one scalar from every subdomain

The values \(\alpha\) and \(\beta\) are then computed from those global scalars and used by the CUDA update kernels.

9. Exchange GPU halo values through host buffers

The stencil still needs halo values from neighboring MPI ranks.

The current implementation does not pass device pointers directly to MPI. Instead, cuda_exchange_halo() stages the boundary data through host memory.

The transfer sizes follow the local array layout: Ny = N + 1 grid lines per column, row_bytes = Ny * sizeof(double) bytes per row, and pitch denotes that same row size in bytes, used as the stride of the 2D copies.

For a left or right boundary, the device row is copied to a host send buffer:

cudaMemcpy(
    send_left.data(),
    d_field + idx(1, 0, Ny),
    row_bytes,
    cudaMemcpyDeviceToHost
);

For the bottom and top boundaries, the data is strided in device memory, so the implementation uses cudaMemcpy2D() to copy one column:

cudaMemcpy2D(
    send_bottom.data(),
    sizeof(double),
    d_field + idx(0, 1, Ny),
    pitch,
    sizeof(double),
    M + 1,
    cudaMemcpyDeviceToHost
);

After the device-to-host transfers, MPI performs the same nonblocking neighbor exchange used by the CPU solver:

MPI_Irecv(recv_left.data(), Ny, MPI_DOUBLE, left_rank, 0, cart_comm, &reqs[rq++]);
MPI_Isend(send_left.data(), Ny, MPI_DOUBLE, left_rank, 1, cart_comm, &reqs[rq++]);

with the other three directions following the same pattern, and finally

MPI_Waitall(rq, reqs, MPI_STATUSES_IGNORE);

The received host buffers are then copied into the GPU halo cells with host-to-device cudaMemcpy or cudaMemcpy2D.

The complete path is

neighbor boundary on GPU
Device → Host copy
host send buffer
MPI_Isend / MPI_Irecv
host receive buffer
Host → Device copy
GPU halo cells

This is an important performance difference from the MPI+OpenMP version. Communication now includes both MPI transfer time and GPU/CPU data movement.

The solver records these separately as

mem_to_host
mpi_exchange_halo
mem_to_device

which will be useful when we analyze performance later.

10. Follow the complete MPI+CUDA PCG iteration

After initialization, the main iteration is

for (iter = 0; r_norm > tolerance && iter < max_iter; ++iter) {
    cuda_exchange_halo(d_p);

    cuda_apply_A(d_p, d_k, d_A_p, M, N, hx2, hy2, stream);

    double p_Ap = compute_p_Ap();
    double alpha = rz_prev / p_Ap;

    update_u(alpha);
    update_r(alpha);

    r_norm = compute_l2_norm();

    apply_preconditioner();

    double rz = compute_rz();
    double beta = rz / rz_prev;

    update_p(beta);
    rz_prev = rz;
}

The loop calls the solver methods (update_u, update_r, update_p, apply_preconditioner, compute_p_Ap, …). In MPICudaPoissonSolver these override the CPU versions and dispatch to the CUDA wrappers from the table above — for example, update_u(alpha) internally launches cuda_update_u(d_u, d_p, alpha, M, N, stream).

Its execution structure can be read as

GPU p
D2H boundary copies
MPI halo exchange
H2D halo copies
CUDA stencil: Ap
CUDA local reduction
MPI_Allreduce
alpha
CUDA update u
CUDA update r
CUDA local residual reduction
MPI_Allreduce
global residual norm
CUDA preconditioner
CUDA local (r,z) reduction
MPI_Allreduce
beta
CUDA update p

The PCG recurrence remains the same as in the sequential, OpenMP, MPI, and hybrid CPU versions. The implementation changes the execution location of the local numerical kernels.

11. Copy the final solution back to the host

The solution remains in d_u during the GPU iteration.

After convergence, the solver calls

to_host(u, d_u);

which performs

cudaMemcpy(
    h_buf.data(),
    d_u,
    size * sizeof(double),
    cudaMemcpyDeviceToHost
);

and reconstructs the two-dimensional host array.

The top-level MPI program then uses the same global reconstruction strategy as the CPU MPI version:

local owned solution
flatten each rank's block
MPI_Gather domain metadata
MPI_Gatherv local values
rank 0 reconstructs global U
write CSV

This gives the MPI+CUDA implementation the same output format as the earlier versions, which is useful for numerical validation.

12. Build the MPI+CUDA executable

The repository builds the GPU version with its Makefile.

The CUDA architecture can be selected through

ARCH

and defaults to

ARCH ?= sm_60

The Makefile uses

NVCCFLAGS := -std=c++11 -O3 -arch=$(ARCH) -Xcompiler "-fPIC"

for CUDA sources and compiles the C++ side with mpicxx.

The source set includes

task_mpi_cuda.cpp
src/conjugate_gradient.cpp
src/mpi_conjugate_gradient.cpp
src/cuda_kernels.cu
src/cuda_operators.cu
src/mpi_cuda_conjugate_gradient.cu

The resulting executable is

task_mpi_cuda

A build can therefore be requested with a target architecture such as

make ARCH=sm_60

or another architecture supported by the environment in which the code is being compiled.

The repository’s local experiment script currently demonstrates the same interface with

make ARCH=sm_90

before running the generated task_mpi_cuda binary.

The exact architecture used for the course cluster should be chosen according to that machine’s required CUDA target; the build interface itself is already parameterized through ARCH.

13. Measure the new GPU-specific costs

The CUDA version introduces several timing categories in addition to the MPI ones:

CPU Laplace operator
CPU update
CPU local reduction

CUDA Laplace operator
CUDA update kernels
CUDA reduction kernels

Host → Device memcpy
Device → Host memcpy

MPI halo exchange
MPI Allreduce

Finalize
Total runtime

This makes it possible to separate three major costs:

\[ T_{\mathrm{compute}}, \qquad T_{\mathrm{memory\ transfer}}, \qquad T_{\mathrm{communication}}. \]

For the GPU solver, kernel execution time alone is not enough to explain total runtime.

The halo path, for example, contains

\[ T_{\mathrm{D2H}} + T_{\mathrm{MPI}} + T_{\mathrm{H2D}}. \]

The global reductions contain a GPU-local reduction followed by MPI_Allreduce.

These categories should be preserved when collecting performance results so that later comparisons can explain where time is actually spent.

14. What changed when we moved the local solver to CUDA?

The distributed ownership model remains the same:

global grid
MPI process grid
local subdomain

The local execution model changes:

MPI version
    local CPU loops

MPI + OpenMP
    local CPU thread teams

MPI + CUDA
    local GPU kernels

The mapping of responsibilities is now

MPI
├── domain decomposition
├── neighbor communication
├── global reductions
└── final reconstruction

CUDA
├── local stencil
├── local vector updates
├── local preconditioner
└── local reduction kernels

The current implementation also exposes an important practical boundary between the two models: halo values are staged through host buffers before MPI communication.

That gives the CUDA solver a clear set of components to evaluate later:

GPU computation
GPU ↔ CPU transfers
MPI halo exchange
MPI global reductions

The next step moves from the implementation itself to the target execution environment: building, submitting, and running the solver on IBM Polus.

Hailin Liu
Authors
PhD Researcher in Agentic AI and Multi-Agent Systems
Hailin Liu is a PhD researcher in Artificial Intelligence and Machine Learning, focusing on Agentic AI, Multi-Agent Systems, and AI Security. His research explores runtime governance mechanisms for autonomous intelligent systems, including agent safety, long-horizon reasoning, and adaptive control.