[Feature][QDP] Add MPI-ready distributed amplitude execution scaffolding

qdp/qdp-core/examples/distributed_multigpu_q33_probe.rs

@@ -0,0 +1,135 @@
+//
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::time::Instant;
+
+use qdp_core::gpu::LocalCollectiveCommunicator;
+use qdp_core::{
+    DistributedExecutionContext, DistributionMode, MahoutError, PlacementRequest, Precision,
+    QdpEngine, ShardPolicy,
+};
+
+fn gib(bytes: usize) -> f64 {
+    bytes as f64 / (1024.0 * 1024.0 * 1024.0)
+}
+
+fn parse_device_ids() -> Result<Vec<usize>, MahoutError> {
+    let raw = std::env::var("GPU_IDS").unwrap_or_else(|_| "0,1,2,3,4,5".to_string());
+    let mut ids = Vec::new();
+    for piece in raw.split(',') {
+        let trimmed = piece.trim();
+        if trimmed.is_empty() {
+            continue;
+        }
+        ids.push(trimmed.parse::<usize>().map_err(|err| {
+            MahoutError::InvalidInput(format!("Invalid GPU ID '{trimmed}': {err}"))
+        })?);
+    }
+
+    if ids.is_empty() {
+        return Err(MahoutError::InvalidInput(
+            "GPU_IDS must contain at least one CUDA device ID".to_string(),
+        ));
+    }
+
+    Ok(ids)
+}
+
+fn main() -> Result<(), MahoutError> {
+    let num_qubits = std::env::var("QUBITS")
+        .ok()
+        .and_then(|value| value.parse::<usize>().ok())
+        .unwrap_or(33);
+    let host_len = std::env::var("HOST_LEN")
+        .ok()
+        .and_then(|value| value.parse::<usize>().ok())
+        .unwrap_or(1);
+    let precision = match std::env::var("PRECISION").ok().as_deref() {
+        Some("f64") | Some("float64") => Precision::Float64,
+        _ => Precision::Float32,
+    };
+    let shard_policy = match std::env::var("SHARD_POLICY").ok().as_deref() {
+        Some("equal") => ShardPolicy::Equal,
+        _ => ShardPolicy::BalancedUneven,
+    };
+    let device_ids = parse_device_ids()?;
+    let request =
+        PlacementRequest::new(num_qubits, DistributionMode::ShardedCapacity, shard_policy);
+    let host_data = vec![1.0f64; host_len];
+
+    println!(
+        "Starting distributed multi-GPU probe: qubits={}, host_len={}, gpus={:?}, precision={:?}, shard_policy={:?}, collectives=in-process",
+        num_qubits, host_len, device_ids, precision, shard_policy
+    );
+
+    let collectives = LocalCollectiveCommunicator;
+    let execution = DistributedExecutionContext::single_process(device_ids.clone(), &collectives)?;
+
+    let prepare_start = Instant::now();
+    let prepared = QdpEngine::prepare_distributed_amplitude_on(
+        &execution,
+        &host_data,
+        num_qubits,
+        precision,
+        Some(request.clone()),
+    )?;
+    let prepare_elapsed = prepare_start.elapsed();
+
+    println!(
+        "Prepared in {:.3}s; global_len={}; shards={}; max_local_len={}; estimated_max_shard_gib={:.2}; gather_device={:?}",
+        prepare_elapsed.as_secs_f64(),
+        prepared.plan.global_len,
+        prepared.layout.num_shards(),
+        prepared.plan.max_local_len(),
+        gib(prepared.plan.estimated_max_shard_bytes(precision)?),
+        prepared.layout.recommended_gather_device_id()
+    );
+
+    for shard in prepared.layout.shards() {
+        let shard_bytes = match precision {
+            Precision::Float32 => shard.local_len * 8,
+            Precision::Float64 => shard.local_len * 16,
+        };
+        println!(
+            "  shard {} -> cuda:{} range=[{}, {}) local_len={} (~{:.2} GiB)",
+            shard.shard_id,
+            shard.device_id,
+            shard.start_idx,
+            shard.end_idx,
+            shard.local_len,
+            gib(shard_bytes)
+        );
+    }
+
+    let encode_start = Instant::now();
+    let state = QdpEngine::encode_distributed_amplitude_to_shards_on(
+        &execution,
+        &host_data,
+        num_qubits,
+        precision,
+        Some(request),
+    )?;
+    let encode_elapsed = encode_start.elapsed();
+
+    println!(
+        "Encoded in {:.3}s; state_shards={}; placement={:?}",
+        encode_elapsed.as_secs_f64(),
+        state.num_shards(),
+        state.recommended_placement_device_ids()
+    );
+
+    Ok(())
+}

qdp/qdp-core/src/gpu/communicator.rs

@@ -0,0 +1,123 @@
+//
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::ffi::c_void;
+
+use crate::error::{MahoutError, Result};
+
+/// Abstracts cross-shard collective operations.
+///
+/// Implementations expose rank-local scalar semantics: each rank contributes
+/// one local value and receives the globally reduced scalar.
+pub trait CollectiveCommunicator: Send + Sync {
+    /// Rank of the current process in the collective world.
+    fn rank(&self) -> usize;
+
+    /// Number of ranks participating in the collective world.
+    fn world_size(&self) -> usize;
+
+    /// Sum one rank-local contribution into one global scalar.
+    fn all_reduce_sum_f64(&self, local_value: f64) -> Result<f64>;
+}
+
+/// Device collective backend selected for GPU-resident reductions.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum DeviceCollectiveBackend {
+    Local,
+    CudaAwareMpi,
+    Nccl,
+    Unavailable,
+}
+
+/// Abstracts GPU-resident collective operations.
+pub trait DeviceCollectiveCommunicator: Send + Sync {
+    fn backend_kind(&self) -> DeviceCollectiveBackend;
+
+    /// # Safety
+    ///
+    /// Callers must ensure the raw pointers and stream are valid for the
+    /// selected CUDA device and that `recv_ptr` can hold `count` `f32` values.
+    unsafe fn all_reduce_sum_f32_device(
+        &self,
+        send_ptr: *const c_void,
+        recv_ptr: *mut c_void,
+        count: usize,
+        device_id: usize,
+        stream: *mut c_void,
+    ) -> Result<()>;
+}
+
+/// In-process collective implementation for the single-rank path.
+#[derive(Default, Debug, Clone, Copy)]
+pub struct LocalCollectiveCommunicator;
+
+impl CollectiveCommunicator for LocalCollectiveCommunicator {
+    fn rank(&self) -> usize {
+        0
+    }
+
+    fn world_size(&self) -> usize {
+        1
+    }
+
+    fn all_reduce_sum_f64(&self, local_value: f64) -> Result<f64> {
+        Ok(local_value)
+    }
+}
+
+#[derive(Default, Debug, Clone, Copy)]
+pub struct MpiDeviceCollectiveCommunicator;
+
+impl DeviceCollectiveCommunicator for MpiDeviceCollectiveCommunicator {
+    fn backend_kind(&self) -> DeviceCollectiveBackend {
+        DeviceCollectiveBackend::CudaAwareMpi
+    }
+
+    unsafe fn all_reduce_sum_f32_device(
+        &self,
+        _send_ptr: *const c_void,
+        _recv_ptr: *mut c_void,
+        _count: usize,
+        _device_id: usize,
+        _stream: *mut c_void,
+    ) -> Result<()> {
+        Err(MahoutError::NotImplemented(
+            "CUDA-aware MPI device collectives are reserved but not implemented".to_string(),
+        ))
+    }
+}
+
+#[derive(Default, Debug, Clone, Copy)]
+pub struct NcclDeviceCollectiveCommunicator;
+
+impl DeviceCollectiveCommunicator for NcclDeviceCollectiveCommunicator {
+    fn backend_kind(&self) -> DeviceCollectiveBackend {
+        DeviceCollectiveBackend::Nccl
+    }
+
+    unsafe fn all_reduce_sum_f32_device(
+        &self,
+        _send_ptr: *const c_void,
+        _recv_ptr: *mut c_void,
+        _count: usize,
+        _device_id: usize,
+        _stream: *mut c_void,
+    ) -> Result<()> {
+        Err(MahoutError::NotImplemented(
+            "NCCL device collectives are reserved but not implemented".to_string(),
+        ))
+    }
+}

qdp/qdp-core/src/gpu/cuda_ffi.rs

@@ -67,7 +67,14 @@ unsafe extern "C" {
         ptr: *const c_void,
     ) -> i32;
 
+    pub(crate) fn cudaGetDevice(device: *mut i32) -> i32;
+    pub(crate) fn cudaSetDevice(device: i32) -> i32;
     pub(crate) fn cudaMemGetInfo(free: *mut usize, total: *mut usize) -> i32;
+    pub(crate) fn cudaDeviceCanAccessPeer(
+        can_access_peer: *mut i32,
+        device: i32,
+        peer_device: i32,
+    ) -> i32;
 
     pub(crate) fn cudaMemcpyAsync(
         dst: *mut c_void,
@@ -76,7 +83,6 @@ unsafe extern "C" {
         kind: u32,
         stream: *mut c_void,
     ) -> i32;
-
     pub(crate) fn cudaEventCreateWithFlags(event: *mut *mut c_void, flags: u32) -> i32;
     pub(crate) fn cudaEventRecord(event: *mut c_void, stream: *mut c_void) -> i32;
     pub(crate) fn cudaEventDestroy(event: *mut c_void) -> i32;
@@ -149,10 +155,26 @@ mod no_cuda_stubs {
         QDP_CUDA_UNAVAILABLE
     }
 
+    pub(crate) unsafe fn cudaGetDevice(_device: *mut i32) -> i32 {
+        QDP_CUDA_UNAVAILABLE
+    }
+
+    pub(crate) unsafe fn cudaSetDevice(_device: i32) -> i32 {
+        QDP_CUDA_UNAVAILABLE
+    }
+
     pub(crate) unsafe fn cudaMemGetInfo(_free: *mut usize, _total: *mut usize) -> i32 {
         QDP_CUDA_UNAVAILABLE
     }
 
+    pub(crate) unsafe fn cudaDeviceCanAccessPeer(
+        _can_access_peer: *mut i32,
+        _device: i32,
+        _peer_device: i32,
+    ) -> i32 {
+        QDP_CUDA_UNAVAILABLE
+    }
+
     pub(crate) unsafe fn cudaMemcpyAsync(
         _dst: *mut c_void,
         _src: *const c_void,