Fix segmentation fault in distributed batched state evolution with store_intermediate_results

[huaweil-nv]

Description

This PR fixes a segmentation fault that occurs when running batched quantum state evolution with store_intermediate_results=IntermediateResultSave.ALL in a multi-GPU MPI environment.

Problem Description

When users run batched cudaq.evolve() with multiple GPUs and request intermediate results, the program crashes with a segmentation fault:

# This crashes on multi-GPU systems
results = cudaq.evolve(
    hamiltonian,
    dimensions,
    schedule,
    [state1, state2, state3, state4],  # Batched states
    store_intermediate_results=cudaq.IntermediateResultSave.ALL,  # Triggers the bug
    ...
)

Error output:

===================================================================================
=   BAD TERMINATION OF ONE OF YOUR APPLICATION PROCESSES
=   PROCESS ID 12345 RUNNING AT hostname
=   EXIT CODE: 11 (Segmentation fault)
===================================================================================

Reproduction

Prerequisites

Requirement	Description
Hardware	Multi-GPU system (2+ NVIDIA GPUs, e.g., 2x A100)
Software	CUDA-Q with dynamics target (cudaq.has_target('dynamics'))
MPI	MPI runtime with GPU support (e.g., OpenMPI)
cuDensityMat plugin	Installed (typically via CUDAQ_DYNAMICS_MPI_COMM_LIB)

Conditions to Trigger the Bug

All three conditions must be met simultaneously:

1. Multi-GPU MPI execution: mpirun -np N where N > 1
2. Batched initial states: Pass a list of states to cudaq.evolve() (e.g., [state1, state2, state3, state4])
3. Store intermediate results: store_intermediate_results=cudaq.IntermediateResultSave.ALL

Note: If any one condition is missing, the bug does NOT occur:

• Single GPU → works fine (no distribution needed)
• Single initial state → works fine (no batching)
• store_intermediate_results=NONE or LAST_ONLY → works fine (doesn't call splitBatchedState)

Minimal Reproducer

# reproducer.py
import cudaq
import cupy as cp
import numpy as np
from cudaq import spin, Schedule, RungeKuttaIntegrator

cudaq.mpi.initialize()
cudaq.set_target('dynamics')

# Simple single-qubit Hamiltonian
hamiltonian = 2 * np.pi * 0.1 * spin.x(0)
dimensions = {0: 2}
schedule = Schedule(np.linspace(0, 1, 11), ['time'])

# Create 4 distinct initial states (batched)
initial_states = []
for i in range(4):
    theta = i * np.pi / 8
    state_data = cp.array([np.cos(theta), np.sin(theta)], dtype=cp.complex128)
    initial_states.append(cudaq.State.from_data(state_data))

# This crashes before the fix (Segmentation fault or incorrect results)
results = cudaq.evolve(
    hamiltonian,
    dimensions,
    schedule,
    initial_states,  # batched states
    store_intermediate_results=cudaq.IntermediateResultSave.ALL,  # triggers splitBatchedState
    integrator=RungeKuttaIntegrator())

# Verify results (won't reach here before fix)
for i, result in enumerate(results):
    states = result.intermediate_states()
    print(f"Rank {cudaq.mpi.rank()}: State {i} has {len(states)} intermediate states")

cudaq.mpi.finalize()

Run Command

# Set MPI comm library if needed
export CUDAQ_DYNAMICS_MPI_COMM_LIB=/path/to/libcudaq_distributed_interface_mpi.so

# Run with 2 GPUs
mpirun -np 2 python reproducer.py

Expected Behavior (Before Fix)

===================================================================================
=   BAD TERMINATION OF ONE OF YOUR APPLICATION PROCESSES
=   EXIT CODE: 11 (Segmentation fault)
===================================================================================

Expected Behavior (After Fix)

Rank 0: State 0 has 11 intermediate states
Rank 0: State 1 has 11 intermediate states
Rank 1: State 2 has 11 intermediate states
Rank 1: State 3 has 11 intermediate states

Root Cause Analysis

The bug is in the existing distributed batched state implementation. The code path for distributing batched states across GPUs exists (distributeBatchedStateData), but the reverse operation (splitBatchedState) was never updated to handle distributed states correctly.

Bug 1: CuDensityMatState::splitBatchedState (C++ layer)

Original buggy code (line 785-793):

const int64_t stateSize = batchedState.dimension / batchedState.batchSize;
// ...
for (int i = 0; i < batchedState.batchSize; ++i) {
    // Read from ptr + i * stateSize
}

Problem: In distributed mode:

• batchedState.dimension = local buffer size (e.g., 4 for 2 states on this rank)
• batchedState.batchSize = total batch size (e.g., 4 states across all ranks)
• stateSize = 4 / 4 = 1 ❌ (should be 2)
• Loop tries to read 4 states from a buffer that only holds 2 → segmentation fault

Bug 2: cudm_solver.py (Python layer)

Original buggy code:

split_states = bindings.splitBatchedState(state)
for i in range(batch_size):  # Assumes len(split_states) == batch_size
    intermediate_states[i].append(split_states[i])  # IndexError in distributed mode

Problem: In distributed mode, splitBatchedState only returns the local subset of states, not all batch_size states. Indexing with batch_size causes an IndexError.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[schweitzpgi]

This looks like it should be merged into the python rewrite branch to me.

Good to go!

[sacpis]

LGTM. Thanks @huaweil-nv. May be @1tnguyen can take another look at it.

[1tnguyen]

/ok to test 1c11fa1

Command Bot: Processing...

[1tnguyen]

LGTM 👍

I've run the newly added tests and confirmed the fix. Thanks @huaweil-nv!

[github-actions]

CUDA Quantum Docs Bot: A preview of the documentation can be found here.

[1tnguyen]

/ok to test 2cc2d70

Command Bot: Processing...

[github-actions]

CUDA Quantum Docs Bot: A preview of the documentation can be found here.