NASP/cuda-samples
84
0
Fork 0
mirror of https://github.com/NVIDIA/cuda-samples.git synced 2026-10-13 00:08:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

CUDA 13.2 samples update (#432)

Browse Source 
- Added Python samples for CUDA Python 1.0 release
- Renamed top-level `Samples` directory to `cpp` to accommodate Python samples.
This commit is contained in:
Dheemanth
2026-05-13 15:13:18 -07:00
committed by GitHub
parent a4526d5229
commit aeab82ff30
2072 changed files with 15077 additions and 1970 deletions
Download Patch File Download Diff File Expand all files Collapse all files

140
python/4_DistributedComputing/ipcMemoryPool/README.md Normal file
View File

@@ -0,0 +1,140 @@
# ipcMemoryPool (Python)
## Description
This sample demonstrates how to share GPU memory between Python
processes using CUDA Inter-Process Communication (IPC) and
`cuda.core`'s IPC-enabled memory pools.
By default each process has its own CUDA virtual address space and
cannot see allocations made by another process. With an IPC-enabled
`DeviceMemoryResource` the parent allocates once, and the child
process maps that same physical GPU memory into its own address space
so both read and write the same bytes. The sample performs a
round-trip test:
1. Parent creates an IPC-enabled `DeviceMemoryResource` and allocates
a `Buffer`.
2. Parent fills the buffer with a known pattern.
3. Parent sends the `Buffer` to a child process through an
`multiprocessing.Queue`. cuda.core's pickle reducers re-create the
memory resource and map the buffer in the child.
4. Child verifies the parent's pattern, writes a new pattern, and
signals completion.
5. Parent verifies the child's writes.
## What You'll Learn
- Enabling IPC on a `DeviceMemoryResource` with `ipc_enabled=True`
- Sending `Buffer` objects across process boundaries via `mp.Queue`
- How cuda.core's pickle reducers rebuild the MR and map the buffer
in the receiving process
- Why `multiprocessing` must use the `"spawn"` start method with CUDA
- Detecting IPC support at runtime (POSIX file-descriptor handle
type, memory-pool support, Linux-only)
## Key Libraries
- [`cuda.core`](https://nvidia.github.io/cuda-python/cuda-core/latest/) - IPC-enabled memory resources and buffer reducers
- `cupy` - zero-copy views over the shared device memory via DLPack
- `multiprocessing` - standard library process management
## Key APIs
### From `cuda.core`
- `DeviceMemoryResource(device, options=DeviceMemoryResourceOptions(ipc_enabled=True))` - create an IPC-enabled memory pool
- `DeviceMemoryResourceOptions(max_size=..., ipc_enabled=True)` - configure the underlying pool
- `mr.allocate(nbytes)` - allocate a `Buffer` from the IPC pool
- `Buffer.is_mapped` - True when the buffer is usable in the current process
- `Device.properties.memory_pools_supported` - runtime feature check
- `Device.properties.handle_type_posix_file_descriptor_supported` - runtime feature check
### From `cuda_samples_utils`
- `print_gpu_info()` - print device name and compute capability
## Requirements
### Hardware
- NVIDIA GPU with Compute Capability 7.0 or higher
- Device that supports CUDA memory pools and POSIX file-descriptor IPC handles (the sample detects and reports this at startup)
- Minimum GPU memory: 512 MB
### Software
- Linux x86_64 (POSIX file-descriptor IPC handles are not available on Windows or macOS)
- CUDA Toolkit 13.0 or newer (matches `cuda-python` 13.x)
- Python 3.10 or newer
- `cuda-python` (>=13.0.0)
- `cuda-core` (>=0.6.0)
- `cupy-cuda13x` (>=13.0.0)
## Installation
Install the required packages from `requirements.txt`:
```bash
cd /path/to/cuda-samples/python/4_DistributedComputing/ipcMemoryPool
pip install -r requirements.txt
```
The `requirements.txt` installs:
- `cuda-python` (>=13.0.0)
- `cuda-core` (>=0.6.0)
- `cupy-cuda13x` (>=13.0.0)
## How to Run
### Basic usage
```bash
cd cuda-samples/python/4_DistributedComputing/ipcMemoryPool
python ipcMemoryPool.py
```
### With custom parameters
```bash
# Larger shared buffer
python ipcMemoryPool.py --elements 65536
# Use a specific GPU
python ipcMemoryPool.py --device 1
```
On platforms or devices that do not support CUDA IPC, the sample
prints a diagnostic and exits cleanly with status 0.
## Expected Output
```
Device: <Your GPU Name>
Compute Capability: <X.Y>
Created IPC-enabled DeviceMemoryResource (is_ipc_enabled=True)
Parent wrote pattern (first 5 values): [100. 101. 102. 103. 104.]
Parent sent buffer to child pid=<pid>; waiting...
[child pid=<pid>] received buffer: is_mapped=True, size=4096
Parent sees child's pattern (first 5 values): [-0. -1. -2. -3. -4.]
IPC round-trip: OK
```
**Note:** Device name, compute capability, and child PID will vary
based on your system.
## Files
- `ipcMemoryPool.py` - Python implementation using `cuda.core` IPC memory pools
- `README.md` - This file
- `requirements.txt` - Sample dependencies
- `../../Utilities/cuda_samples_utils.py` - Common utilities (imported by this sample)
## See Also
- [CUDA Python Documentation](https://nvidia.github.io/cuda-python/)
- [`cuda.core` memory API](https://nvidia.github.io/cuda-python/cuda-core/latest/api.html#memory)
- Upstream `cuda.core` IPC tests: [`test_memory_ipc.py`](https://github.com/NVIDIA/cuda-python/blob/main/cuda_core/tests/memory_ipc/test_memory_ipc.py)
- [CUDA IPC programming guide](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#interprocess-communication)

220
python/4_DistributedComputing/ipcMemoryPool/ipcMemoryPool.py Normal file
View File

@@ -0,0 +1,220 @@
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
# * Redistributions in binary form must reproduce the above copyright
# notice, this list of conditions and the following disclaimer in the
# documentation and/or other materials provided with the distribution.
# * Neither the name of NVIDIA CORPORATION nor the names of its
# contributors may be used to endorse or promote products derived
# from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""
IPC Memory Pool with cuda.core
Share GPU memory between Python processes using CUDA Inter-Process
Communication (IPC) and cuda.core's IPC-enabled memory pools. By default
each worker process has its own CUDA virtual address space and cannot see
allocations made by another process. With an IPC-enabled
``DeviceMemoryResource`` the parent can allocate once, and the child
process can map that same physical GPU memory into its own address space
so both read and write the same bytes.
The sample does a round-trip test:
1. Parent creates an IPC-enabled ``DeviceMemoryResource`` and allocates
a ``Buffer``.
2. Parent fills the buffer with a known pattern.
3. Parent sends the ``Buffer`` to a child process through an
``mp.Queue`` - cuda.core's pickle reducers take care of re-creating
the memory resource and mapping the buffer in the child.
4. Child verifies the parent's pattern, writes a new pattern, and
signals completion.
5. Parent verifies the child's writes.
IPC requires Linux (POSIX file-descriptor handles) and device support for
memory pools. On unsupported platforms the sample prints a diagnostic and
exits cleanly.
"""
import multiprocessing as mp
import platform
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "Utilities"))
try:
import cupy as cp
import numpy as np
from cuda.core import (
Device,
DeviceMemoryResource,
DeviceMemoryResourceOptions,
)
from cuda_samples_utils import print_gpu_info # noqa: E402
except ImportError as e:
print(f"Error: Required package not found: {e}")
print("Please install from requirements.txt:")
print(" pip install -r requirements.txt")
sys.exit(1)
CHILD_TIMEOUT_SEC = 30
def check_ipc_support(device) -> bool:
"""Return True if this device/platform supports CUDA IPC memory pools."""
if platform.system() != "Linux":
print(
f"IPC via POSIX file descriptors is only supported on Linux "
f"(detected {platform.system()})."
)
return False
if not device.properties.memory_pools_supported:
print("Device does not support CUDA memory pools.")
return False
if not device.properties.handle_type_posix_file_descriptor_supported:
print("Device/platform does not support POSIX-fd IPC handles.")
return False
return True
def child_worker(q_in, q_out, n_elements, parent_seed, child_seed):
"""Runs in a separate process. Verifies and modifies the shared buffer."""
device = Device(0)
device.set_current()
pid = mp.current_process().pid
# The Buffer (and its MR) are reconstructed and mapped in this process
# when the queued object is unpickled. Both ``is_mapped`` flags are
# True here.
buffer = q_in.get(timeout=CHILD_TIMEOUT_SEC)
print(
f"[child pid={pid}] received buffer: is_mapped={buffer.is_mapped}, "
f"size={buffer.size}"
)
# Build a zero-copy CuPy view of the shared device memory.
arr = cp.from_dlpack(buffer).view(dtype=cp.float32)
# Verify the parent's pattern.
expected_parent = cp.arange(n_elements, dtype=cp.float32) + float(parent_seed)
if not cp.allclose(arr, expected_parent):
print("[child] ERROR: parent's pattern did not match expectation")
buffer.close()
q_out.put("fail")
return
# Write a new pattern for the parent to verify.
arr[:] = cp.arange(n_elements, dtype=cp.float32) * float(child_seed)
device.sync()
buffer.close()
q_out.put("done")
def main() -> int:
import argparse
parser = argparse.ArgumentParser(
description="Share a GPU buffer between two processes via CUDA IPC"
)
parser.add_argument(
"--elements",
type=int,
default=1024,
help="Number of float32 elements in the shared buffer (default: 1024)",
)
parser.add_argument("--device", type=int, default=0, help="CUDA device id")
args = parser.parse_args()
# CUDA is incompatible with the ``fork`` start method because forked
# children inherit a corrupt CUDA state. Always use ``spawn``.
mp.set_start_method("spawn", force=True)
device = Device(args.device)
device.set_current()
print_gpu_info(device)
if not check_ipc_support(device):
print("\nCUDA IPC is not available on this system; exiting cleanly.")
return 0
N = args.elements
nbytes = N * np.dtype(np.float32).itemsize
parent_seed = 100
child_seed = -1.0
# Create an IPC-enabled memory pool. Buffers allocated from this MR
# are picklable and can be shared across processes.
mr = DeviceMemoryResource(
device,
options=DeviceMemoryResourceOptions(
max_size=max(nbytes * 4, 1 << 20),
ipc_enabled=True,
),
)
print(
"Created IPC-enabled DeviceMemoryResource "
f"(is_ipc_enabled={mr.is_ipc_enabled})"
)
buffer = mr.allocate(nbytes)
try:
# Fill the buffer with a known pattern from the parent side.
arr = cp.from_dlpack(buffer).view(dtype=cp.float32)
arr[:] = cp.arange(N, dtype=cp.float32) + float(parent_seed)
device.sync()
print(f"Parent wrote pattern (first 5 values): {arr[:5].get()}")
# Launch the child process and hand the buffer over.
q_to_child = mp.Queue()
q_from_child = mp.Queue()
child = mp.Process(
target=child_worker,
args=(q_to_child, q_from_child, N, parent_seed, child_seed),
)
child.start()
q_to_child.put(buffer)
print(f"Parent sent buffer to child pid={child.pid}; waiting...")
msg = q_from_child.get(timeout=CHILD_TIMEOUT_SEC)
child.join(timeout=CHILD_TIMEOUT_SEC)
if msg != "done" or child.exitcode != 0:
print(f"Child failed: msg={msg!r}, exitcode={child.exitcode}")
return 1
# Verify the child's writes are visible from the parent.
device.sync()
got = arr[:5].get()
expected = (np.arange(N, dtype=np.float32) * child_seed)[:5]
print(f"Parent sees child's pattern (first 5 values): {got}")
if np.allclose(got, expected):
print("IPC round-trip: OK")
return 0
print(f"IPC round-trip: FAILED (expected {expected})")
return 1
finally:
buffer.close()
mr.close()
if __name__ == "__main__":
sys.exit(main())

4
python/4_DistributedComputing/ipcMemoryPool/requirements.txt Normal file
View File

@@ -0,0 +1,4 @@
cuda-python>=13.0.0
cuda-core>=0.6.0
cupy-cuda13x>=13.0.0
numpy>=1.24.0

110
python/4_DistributedComputing/multiGPUGradientAverage/README.md Normal file
View File

@@ -0,0 +1,110 @@
# Sample: multiGPUGradientAverage (Python)
## Description
This sample demonstrates gradient averaging across multiple GPUs using MPI and cuda.core. Each GPU computes local gradients, which are synchronized (averaged) across all GPUs using MPI Allreduce with host-staging (GPU → CPU → MPI → CPU → GPU) for maximum compatibility.
## What you will learn
- How to initialize MPI for multi-process GPU communication
- How to map MPI ranks to CUDA devices consistently
- How to integrate cuda.core streams with CuPy using `ExternalStream`
- How to compile and launch custom CUDA kernels using cuda.core
- How to use cuda.core Event for GPU timing measurements
- How to use MPI Allreduce with host-staging for universal compatibility
## Prerequisites
- Python 3.10+
- CUDA Toolkit 13.0+
- Standard MPI implementation (OpenMPI, MPICH, or Intel MPI)
- Multiple NVIDIA GPUs (tested with 2+ GPUs)
## Installation
```bash
pip install mpi4py cupy-cuda13x cuda-python cuda-core
```
## Running
**IMPORTANT:** This sample **MUST** be run with `mpirun` with at least 2 processes.
```bash
# Single node (2 GPUs)
mpirun -np 2 python multiGPUGradientAverage.py --size 10000
# Single node (4 GPUs)
mpirun -np 4 python multiGPUGradientAverage.py --size 10000
# With specific GPUs
CUDA_VISIBLE_DEVICES=0,2 mpirun -np 2 python multiGPUGradientAverage.py
```
## Sample Output
```
[Rank 0] World size = 4
======================================================================
Multi-GPU Gradient Average Demo
======================================================================
Number of MPI ranks (GPUs): 4
Gradient vector length per GPU: 10000
Device: NVIDIA GeForce RTX 4090
Computation: gradients computed on GPU via cuda.core.
Communication: gradients averaged via MPI_Allreduce on host (CPU) buffers.
======================================================================
Sample averaged gradient values (rank 0):
avg_grad[0] = 1.500000
avg_grad[5000] = 6.500000
avg_grad[9999] = 11.499000
Expected values:
expected[0] = 1.500000
expected[5000] = 6.500000
expected[9999] = 11.499000
Verifying gradient averaging correctness...
[PASS] Gradient averaging is correct.
[PASS] Gradient averaging is correct on all ranks.
Performance:
Kernel time (GPU only): 0.123 ms
MPI communication time (host-staging, end-to-end): 0.456 ms
Total time: 0.579 ms
======================================================================
Demo complete.
======================================================================
```
## Key Technical Details
The sample uses cuda.core streams and makes CuPy use them via `ExternalStream`:
```python
stream = device.create_stream()
cp.cuda.ExternalStream(int(stream.handle)).use()
```
GPU timing is measured using cuda.core Event:
```python
from cuda.core import EventOptions
timing_options = EventOptions(enable_timing=True)
start_event = stream.record(options=timing_options)
# ... GPU work ...
end_event = stream.record(options=timing_options)
end_event.sync()
kernel_time = end_event - start_event # Returns milliseconds
```
The host-staging pattern transfers data GPU → CPU → MPI → CPU → GPU for universal MPI compatibility without requiring CUDA-aware MPI.
## Troubleshooting
**Error: "This sample requires at least 2 MPI processes!"**
Solution: Run with `mpirun -np 2 python multiGPUGradientAverage.py`

416
python/4_DistributedComputing/multiGPUGradientAverage/multiGPUGradientAverage.py Normal file
View File

@@ -0,0 +1,416 @@
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
# * Redistributions in binary form must reproduce the above copyright
# notice, this list of conditions and the following disclaimer in the
# distribution and/or other materials provided with the distribution.
# * Neither the name of NVIDIA CORPORATION nor the names of its
# contributors may be used to endorse or promote products derived
# from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""
Multi-GPU Gradient Average using MPI and cuda.core (Host-staging Allreduce)
Question: How do I synchronize gradients across GPUs?
Answer:
Each GPU (each MPI rank) computes local gradients on device via CUDA.
Gradients are then averaged across ranks via an MPI Allreduce over host
(CPU) buffers, following the classic data-parallel training pattern.
This sample shows how to:
- Initialize MPI for multi-process GPU workloads
- Map MPI ranks to GPUs
- Use cuda.core for kernel compilation and execution
- Integrate cuda.core with CuPy using the stream protocol
- Perform gradient averaging with MPI Allreduce (using host staging)
- Use cuda.core Event for GPU timing measurements
- Verify correctness of distributed gradient synchronization
Key concepts: Allreduce, NCCL collectives (conceptually), distributed training
Note:
- All gradient computation and validation happen on GPUs.
- MPI Allreduce is executed on CPU (host) buffers via a simple
GPU -> CPU -> MPI -> CPU -> GPU staging pattern so that the sample
works on any MPI stack, without requiring CUDA-aware MPI.
- In production deep learning frameworks (e.g., PyTorch DDP), NCCL
usually implements the GPU Allreduce directly; the communication
pattern and semantics are the same as demonstrated here.
"""
import sys
from pathlib import Path
# Add parent directory to path to import utilities
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "Utilities"))
from cuda_samples_utils import verify_array_result # noqa: E402
try:
import cupy as cp
from cuda.core import (
Device,
EventOptions,
LaunchConfig,
Program,
ProgramOptions,
launch,
system,
)
from mpi4py import MPI
except ImportError as e:
print(f"Error: Required package not found: {e}")
print("Please install: pip install mpi4py cupy-cuda12x cuda-python cuda-core")
sys.exit(1)
# ============================================================================
# CUDA device selection and stream management
# ============================================================================
def init_device(rank: int):
"""
Initialize CUDA device and stream for this MPI rank.
For a simple single-node run, we map rank % num_gpus to a device id.
This covers both the common case (world_size == num_gpus) and the case
where multiple ranks share a GPU.
Returns
-------
tuple[Device, Stream]
CUDA device object and stream object.
"""
num_gpus = system.get_num_devices()
if num_gpus == 0:
raise RuntimeError("No CUDA devices available")
dev_id = rank % num_gpus # simple mapping: rank -> GPU in round-robin
try:
device = Device(dev_id)
except (RuntimeError, ValueError) as e:
if rank == 0:
print(f"Warning: Cannot assign GPU {dev_id}, using GPU 0. Error: {e}")
device = Device(0)
device.set_current()
# Align CuPy with cuda.core's chosen device ID
cp.cuda.Device(device.device_id).use()
# Create cuda.core stream and make CuPy use it
stream = device.create_stream()
cp.cuda.ExternalStream(int(stream.handle)).use()
return device, stream
# ============================================================================
# CUDA kernel definition and compilation
# ============================================================================
# Tiny CUDA kernel to initialize local "gradients"
# Uses grid-stride loop to handle arrays larger than grid size
INIT_KERNEL = r"""
extern "C" __global__
void init_grad_kernel(float* grad, int n, int rank)
{
// Grid-stride loop: each thread processes multiple elements
size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
size_t stride = gridDim.x * blockDim.x;
for (size_t i = tid; i < n; i += stride) {
// Gradient value depends on MPI rank so we can verify reduction:
// grad_i = rank + 0.001 * i
grad[i] = rank + 0.001f * i;
}
}
"""
_kernel_cache = {}
def get_init_kernel(device: Device):
"""Compile (or retrieve cached) init_grad_kernel for this device."""
key = device.pci_bus_id
if key not in _kernel_cache:
opts = ProgramOptions(std="c++17", arch=f"sm_{device.arch}")
prog = Program(INIT_KERNEL, code_type="c++", options=opts)
mod = prog.compile("cubin")
_kernel_cache[key] = mod.get_kernel("init_grad_kernel")
return _kernel_cache[key]
# ============================================================================
# Local gradient computation on each GPU
# ============================================================================
def compute_local_gradients(
num_elements: int, device: Device, stream: object, rank: int
) -> cp.ndarray:
"""
Compute a local "gradient" vector on the current GPU.
For demo purposes, we initialize:
grad[i] = rank + 0.001 * i
Parameters
----------
num_elements : int
Length of gradient vector.
device : Device
CUDA device object.
stream : Stream
CUDA stream object (created at device initialization).
rank : int
MPI rank ID.
Returns
-------
cupy.ndarray
Gradient vector on GPU.
"""
# Create gradient array (CuPy uses the stream set at device initialization)
grad = cp.empty(num_elements, dtype=cp.float32)
# Use a CUDA kernel compiled with cuda.core to fill the array
kernel = get_init_kernel(device)
threads_per_block = 256
blocks_per_grid = (num_elements + threads_per_block - 1) // threads_per_block
config = LaunchConfig(grid=blocks_per_grid, block=threads_per_block)
# Launch kernel using cuda.core stream
launch(stream, config, kernel, grad.data.ptr, num_elements, rank)
return grad
# ============================================================================
# MPI Allreduce to average gradients (host-staging)
# ============================================================================
def average_gradients(
local_grad: cp.ndarray, comm: object, world_size: int
) -> cp.ndarray:
"""
Average gradients across all MPI ranks using host-staging Allreduce.
Steps:
1. Copy local gradients from GPU to CPU (NumPy).
2. Perform MPI_Allreduce on host buffers.
3. Divide by world_size to obtain the average.
4. Copy the averaged gradients back to GPU.
This pattern is environment-agnostic and works on any MPI stack.
"""
assert local_grad.dtype == cp.float32
# GPU -> CPU
local_host = local_grad.get() # NumPy array on host
avg_host = local_host.copy()
# Allreduce on host buffers
comm.Allreduce(local_host, avg_host, op=MPI.SUM)
# Average
avg_host /= world_size
# CPU -> GPU
avg_grad = cp.asarray(avg_host)
return avg_grad
# ============================================================================
# Testing and verification
# ============================================================================
def main():
"""Demo: Multi-GPU gradient averaging with MPI (host-staging Allreduce)."""
import argparse
# Initialize MPI
comm = MPI.COMM_WORLD
world_size = comm.Get_size()
rank = comm.Get_rank()
parser = argparse.ArgumentParser(
description=(
"Multi-GPU Gradient Average with mpi4py + cuda.core "
"(host-staging Allreduce)"
)
)
parser.add_argument(
"--size",
type=int,
default=1024,
help="Number of gradient elements per GPU (default: 1024)",
)
args = parser.parse_args()
num_elements = args.size
# Initialize device and stream
device = None
stream = None
try:
device, stream = init_device(rank)
if rank == 0:
print(f"[Rank 0] World size = {world_size}")
comm.Barrier()
# Validate world size
if world_size < 2:
if rank == 0:
print("=" * 70)
print("ERROR: This sample requires at least 2 MPI processes!")
print("=" * 70)
print("\nPlease run with mpirun:")
print(" mpirun -np 2 python multiGPUGradientAverage.py")
print(" mpirun -np 4 python multiGPUGradientAverage.py --size 10000")
print("\nFor multi-GPU systems:")
print(" mpirun -np N python multiGPUGradientAverage.py")
print(" (where N = number of GPUs)")
print("=" * 70)
sys.exit(1)
# Validate input
if num_elements <= 0:
if rank == 0:
print("Error: --size must be positive")
sys.exit(1)
if rank == 0:
print("\n" + "=" * 70)
print("Multi-GPU Gradient Average Demo")
print("=" * 70)
print(f"Number of MPI ranks (GPUs): {world_size}")
print(f"Gradient vector length per GPU: {num_elements}")
print(f"Device: {device.name}")
print("Computation: gradients computed on GPU via cuda.core.")
print(
"Communication: gradients averaged via MPI_Allreduce on host "
"(CPU) buffers."
)
print("=" * 70)
# Step 1: Compute local gradients on each GPU
# Use cuda.core Event for GPU timing measurements
timing_options = EventOptions(enable_timing=True)
start_event = stream.record(options=timing_options)
local_grad = compute_local_gradients(num_elements, device, stream, rank)
# Record end event and synchronize to ensure timing is complete
end_event = stream.record(options=timing_options)
end_event.sync()
# Calculate elapsed time: Event subtraction returns milliseconds
kernel_time = end_event - start_event
# Step 2: Average gradients across all ranks (host-staging Allreduce)
# Use CPU timing for MPI communication (host-staging includes GPU↔CPU transfers)
import time
comm_start = time.time()
avg_grad = average_gradients(local_grad, comm, world_size)
comm_time = (time.time() - comm_start) * 1000 # Convert to ms
# Step 3: Sanity check on rank 0
# For each element i:
# local_grad_r[i] = r + 0.001 * i, r = 0..world_size-1
# Sum over ranks:
# sum[i] = sum_r r + 0.001 * i * world_size
# Average:
# avg[i] = (0 + ... + (world_size-1))/world_size + 0.001 * i
# = (world_size - 1)/2 + 0.001 * i
#
# We verify this formula.
expected_base = (world_size - 1) / 2.0
i0 = 0
i1 = num_elements // 2
i2 = num_elements - 1
# Copy a few sample elements back to host for printing on rank 0
if rank == 0:
avg_host_samples = avg_grad[[i0, i1, i2]].get()
print("\nSample averaged gradient values (rank 0):")
print(f" avg_grad[{i0}] = {avg_host_samples[0]:.6f}")
print(f" avg_grad[{i1}] = {avg_host_samples[1]:.6f}")
print(f" avg_grad[{i2}] = {avg_host_samples[2]:.6f}")
expected0 = expected_base + 0.001 * i0
expected1 = expected_base + 0.001 * i1
expected2 = expected_base + 0.001 * i2
print("\nExpected values:")
print(f" expected[{i0}] = {expected0:.6f}")
print(f" expected[{i1}] = {expected1:.6f}")
print(f" expected[{i2}] = {expected2:.6f}")
# All ranks perform a full-array correctness check on GPU
expected_full = expected_base + 0.001 * cp.arange(
num_elements, dtype=cp.float32
)
# Use utility function to verify results
if rank == 0:
print("\nVerifying gradient averaging correctness...")
ok = verify_array_result(
avg_grad, expected_full, rtol=1e-5, atol=1e-8, verbose=(rank == 0)
)
# Ensure all ranks agree on correctness
ok_all = comm.allreduce(ok, op=MPI.LAND)
if rank == 0:
if ok_all:
print("[PASS] Gradient averaging is correct on all ranks.")
else:
print(
"[FAIL] Gradient averaging mismatch detected on one or more ranks."
)
print("\nPerformance:")
print(f" Kernel time (GPU only): {kernel_time:.3f} ms")
print(
" MPI communication time (host-staging, end-to-end): "
f"{comm_time:.3f} ms"
)
print(f" Total time: {kernel_time + comm_time:.3f} ms")
print("\n" + "=" * 70)
print("Demo complete.")
print("=" * 70)
return 0 if ok_all else 1
finally:
# Clean up stream resources
if stream is not None:
stream.close()
cp.cuda.Stream.null.use() # Reset CuPy's current stream to the null stream
if __name__ == "__main__":
sys.exit(main())

21
python/4_DistributedComputing/multiGPUGradientAverage/requirements.txt Normal file
View File

@@ -0,0 +1,21 @@
# Multi-GPU Gradient Average Sample Requirements
# MPI Python bindings for distributed communication
mpi4py>=3.1.4
# GPU array library (NumPy-compatible arrays on CUDA)
# Use cupy-cuda11x, cupy-cuda12x, or cupy-cuda13x depending on your CUDA version
cupy-cuda13x>=13.0.0
# CUDA Python bindings (low-level CUDA driver API)
cuda-python>=13.0.0
# cuda.core - Modern Python interface for CUDA
# Provides Program, LaunchConfig, Device, and launch APIs
cuda-core>=0.6.0
# Note: This sample uses host-staging for MPI communication
# Standard MPI installation is sufficient (no CUDA-aware MPI required)
# Install MPI using system package manager:
# Ubuntu/Debian: sudo apt-get install openmpi-bin libopenmpi-dev
# Or build from source: https://www.open-mpi.org/software/ompi/

190
python/4_DistributedComputing/simpleP2P/README.md Normal file
View File

@@ -0,0 +1,190 @@
# Sample: simpleP2P (Python)
## Description
This sample demonstrates peer-to-peer (P2P) memory access between multiple GPUs in CUDA using the cuda.core Python library. P2P allows GPUs to directly access each other's memory without routing data through the host (CPU), enabling efficient multi-GPU applications. This sample detects P2P-capable GPUs, enables peer access, measures bandwidth using CUDA events for accurate GPU-side timing, and launches kernels (using grid-stride loops) that read from one GPU's memory and write to another GPU's memory.
## What you will learn
- How to detect multiple CUDA-capable GPUs using `system.get_num_devices()` and `Device(id)`
- How to check P2P capability between GPU pairs using `device.can_access_peer()`
- How to enable and disable peer access using `DeviceMemoryResource.peer_accessible_by`
- How to allocate device memory on specific GPUs using `DeviceMemoryResource`
- How to perform direct GPU-to-GPU memory transfers with explicit event-based synchronization
- How to measure P2P bandwidth using CUDA events for accurate GPU-side timing
- How to use event-based synchronization between streams for sequential bandwidth measurement
- How to launch kernels on one GPU that access memory from another GPU
- How to compile and launch CUDA kernels using cuda.core's `Program` and `launch` APIs with grid-stride loops
- How to validate multi-GPU computation results
- How to properly clean up resources using try/finally blocks
## Key libraries
- `numpy` - CPU array operations and data initialization
- `cuda-core` - Modern Python interface to CUDA runtime with full P2P support
## Key APIs
**From cuda.core:**
- `system` – Pre-instantiated singleton for system-level CUDA information
- `system.get_num_devices()` – Get number of CUDA-capable devices
- `Device(id)` – Get specific CUDA device handle
- `device.can_access_peer(peer)` – Check if this device can access peer device memory
- `device.set_current()` – Set active device for subsequent operations
- `device.create_stream()` – Create CUDA stream for kernel execution
- `DeviceMemoryResource(device)` – Create memory resource for specific GPU
- `memory_resource.peer_accessible_by` – Get/set which devices can access this memory pool's allocations
- Example: `mr.peer_accessible_by = [1]` grants device 1 access
- Example: `mr.peer_accessible_by = []` revokes all access
- `PinnedMemoryResource()` – Allocate pinned (page-locked) host memory
- `EventOptions(enable_timing=True)` – Create options for CUDA events with timing enabled
- `stream.record(options=event_options)` – Record a CUDA event on a stream
- `event.elapsed_time(start_event)` – Get elapsed time in milliseconds between two events
- `stream.wait_event(event)` – Make a stream wait for an event to complete
- `stream.close()` – Clean up stream resources
- `Program()` – Compile CUDA C++ kernel code
- `LaunchConfig()` – Configure kernel launch parameters (grid, block)
- `launch()` – Launch compiled kernel with arguments
- `buffer.copy_from(src, stream=stream)` – Copy data from source buffer asynchronously
- `buffer.copy_to(dst, stream=stream)` – Copy data to destination buffer asynchronously
**From DLPack:**
- `numpy.from_dlpack()` – Create NumPy array view of memory buffer
**Memory Management:**
- Resources (streams, buffers) should be cleaned up using try/finally blocks to ensure proper cleanup even if errors occur
- Streams should be explicitly closed with `stream.close()` in finally blocks
## Peer-to-Peer (P2P): When to Use
### Benefits
- **Direct GPU-to-GPU transfers**: Bypass host memory for faster communication
- **Higher bandwidth**: PCIe or NVLink bandwidth between GPUs (up to 600 GB/s with NVLink)
- **Lower latency**: No CPU involvement in data transfers
- **Efficient multi-GPU**: Essential for scaling deep learning, HPC, and simulation workloads
- **Simplified programming**: Kernels can directly access remote GPU memory
### Requirements
- **Two or more GPUs**: System must have multiple CUDA-capable GPUs
- **P2P support**: GPUs must be P2P-capable (check with `can_access_peer()`)
- **PCIe topology**: Usually requires GPUs on same PCIe root complex
- **Platform support**: Not available on Mac OSX, limited on ARM platforms
### Best Use Cases
1. Multi-GPU deep learning training (model parallelism, data parallelism)
2. Large-scale scientific simulations across multiple GPUs
3. Real-time rendering with multiple GPUs
4. GPU clusters with direct GPU communication
5. Reducing CPU-GPU traffic in multi-GPU systems
## Requirements
1. **Two or more NVIDIA Graphics Cards** with CUDA support and P2P capability
2. **CUDA Drivers** installed on your system
3. **CUDA Toolkit 13.0+** installed on your system
4. **Python 3.10 or newer**
5. **Proper PCIe topology** (GPUs should be on same PCIe root complex for best performance)
**Note**: This sample will gracefully exit if fewer than 2 GPUs are detected or if P2P is not supported between any GPU pair.
**Install packages:**
```bash
pip install -r requirements.txt
```
Or manually:
```bash
pip install numpy>=2.3.2 cuda-core>=0.6.0 cuda-python>=13.0.0
```
## How to run
Basic usage:
```bash
# Run with default parameters (16M elements = 64MB)
python simpleP2P.py
```
With custom parameters:
```bash
# Use 32M elements (128MB)
python simpleP2P.py --num_elements 33554432
# Show help
python simpleP2P.py --help
```
### Command line arguments
- `--num_elements`: Number of elements in arrays (default: 16777216)
- Each array uses `num_elements * 4 bytes` (float32)
- Default: 64 MB per array
- Sample allocates 2 device buffers + 1 host buffer
## Expected Output
```
======================================================================
simpleP2P - CUDA Python Sample
======================================================================
Starting...
Checking for multiple GPUs...
CUDA-capable device count: 2
Checking GPU(s) for support of peer to peer memory access...
> Peer access from Tesla T10 (GPU0) -> Tesla T10 (GPU1): Yes
> Peer access from Tesla T10 (GPU1) -> Tesla T10 (GPU0): Yes
Using GPU0 (Tesla T10) and GPU1 (Tesla T10)
Allocating buffers (64MB on GPU0, GPU1 and CPU Host)...
Peer access enabled: GPU0 <-> GPU1
Peer access status: MR0 accessible by (1,), MR1 accessible by (0,)
Memory allocated successfully
Measuring P2P bandwidth...
Performing 100 ping-pong copies between GPUs...
P2P bandwidth: 12.37 GB/s
Preparing host buffer and memcpy to GPU0...
Data initialized and copied to GPU
Compiling CUDA kernel...
Kernels compiled successfully
Run kernel on GPU1, taking source data from GPU0 and writing to GPU1...
Kernel execution complete
Run kernel on GPU0, taking source data from GPU1 and writing to GPU0...
Kernel execution complete
Copy data back to host from GPU0 and verify results...
Checking results...
Comparing 16,777,216 elements...
Test PASSED
[PASS] Validation PASSED
Disabling peer access...
Peer access revoked: MR0 accessible by (), MR1 accessible by ()
======================================================================
simpleP2P completed successfully!
======================================================================
Shutting down...
```
**Note**: P2P bandwidth varies based on:
- PCIe generation
- NVLink
- System topology and configuration
## Files
- `simpleP2P.py` – Main Python implementation
- `README.md` – This file
- `requirements.txt` – Python package dependencies

6
python/4_DistributedComputing/simpleP2P/requirements.txt Normal file
View File

@@ -0,0 +1,6 @@
# simpleP2P - Requirements
# Install with: pip install -r requirements.txt
numpy>=2.3.2
cuda-python>=13.0.0
cuda-core>=0.6.0

377
python/4_DistributedComputing/simpleP2P/simpleP2P.py Normal file
View File

@@ -0,0 +1,377 @@
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
# * Redistributions in binary form must reproduce the above copyright
# notice, this list of conditions and the following disclaimer in the
# distribution and/or other materials provided with the distribution.
# * Neither the name of NVIDIA CORPORATION nor the names of its
# contributors may be used to endorse or promote products derived
# from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import argparse
import sys
from pathlib import Path
try:
import numpy as np
from cuda.core import (
Device,
DeviceMemoryResource,
EventOptions,
LaunchConfig,
PinnedMemoryResource,
Program,
ProgramOptions,
launch,
system,
)
except ImportError as e:
print(f"Error: Required package not found: {e}")
print("Please install from requirements.txt:")
print(" pip install -r requirements.txt")
sys.exit(1)
# Add parent directory to path to import utilities
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "Utilities"))
from cuda_samples_utils import verify_array_result # noqa: E402
# CUDA kernel for simple P2P operation
SIMPLE_P2P_KERNEL = """
extern "C" __global__
void SimpleKernel(float *src, float *dst, int N) {
// Grid-stride loop pattern for canonical CUDA kernel
size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
size_t stride = gridDim.x * blockDim.x;
for (size_t i = tid; i < N; i += stride) {
dst[i] = src[i] * 2.0f;
}
}
"""
def run(num_elements=1024 * 1024 * 16):
"""
Demonstrates peer-to-peer (P2P) memory access between multiple GPUs using cuda.core.
This function shows how to:
1. Detect and verify multiple GPUs with P2P capability
2. Enable peer access between GPUs
3. Perform direct GPU-to-GPU memory transfers
4. Launch kernels that access memory from other GPUs
5. Measure P2P bandwidth
6. Validate results
Parameters
----------
num_elements : int
Number of elements in arrays (default: 16M elements = 64MB)
"""
print("\n" + "=" * 70)
print("simpleP2P - CUDA Python Sample")
print("=" * 70)
print("\nStarting...")
# Check for multiple GPUs
print("\nChecking for multiple GPUs...")
num_devices = system.get_num_devices()
print(f"CUDA-capable device count: {num_devices}")
if num_devices < 2:
print(
"Two or more GPUs with Peer-to-Peer access capability are "
"required, waiving this sample."
)
return 2
# Get device properties
devices = [Device(i) for i in range(num_devices)]
# Check for P2P capability
print("\nChecking GPU(s) for support of peer to peer memory access...")
p2p_capable_gpus = [-1, -1]
for i in range(num_devices):
p2p_capable_gpus[0] = i
for j in range(num_devices):
if i == j:
continue
# Check peer access capability using cuda.core
i_access_j = devices[i].can_access_peer(devices[j])
j_access_i = devices[j].can_access_peer(devices[i])
print(
f"> Peer access from {devices[i].name} (GPU{i}) -> "
f"{devices[j].name} (GPU{j}): {'Yes' if i_access_j else 'No'}"
)
print(
f"> Peer access from {devices[j].name} (GPU{j}) -> "
f"{devices[i].name} (GPU{i}): {'Yes' if j_access_i else 'No'}"
)
if i_access_j and j_access_i:
p2p_capable_gpus[1] = j
break
if p2p_capable_gpus[1] != -1:
break
if p2p_capable_gpus[0] == -1 or p2p_capable_gpus[1] == -1:
print("\nTwo or more GPUs with Peer-to-Peer access capability are required.")
print(
"Peer to Peer access is not available amongst GPUs in the system, "
"waiving test."
)
return 2
# Use first pair of P2P capable GPUs detected
gpuid = [p2p_capable_gpus[0], p2p_capable_gpus[1]]
dev0 = devices[gpuid[0]]
dev1 = devices[gpuid[1]]
print(f"\nUsing GPU{gpuid[0]} ({dev0.name}) and GPU{gpuid[1]} ({dev1.name})")
# Allocate buffers with P2P access
buf_size = num_elements * np.dtype(np.float32).itemsize
print(
f"\nAllocating buffers ({int(buf_size / 1024 / 1024)}MB on "
f"GPU{gpuid[0]}, GPU{gpuid[1]} and CPU Host)..."
)
# Allocate on GPU 0 and grant access to GPU 1
dev0.set_current()
mr0 = DeviceMemoryResource(dev0)
mr0.peer_accessible_by = [gpuid[1]] # Grant GPU 1 access to GPU 0's memory
g0 = mr0.allocate(buf_size)
# Allocate on GPU 1 and grant access to GPU 0
dev1.set_current()
mr1 = DeviceMemoryResource(dev1)
mr1.peer_accessible_by = [gpuid[0]] # Grant GPU 0 access to GPU 1's memory
g1 = mr1.allocate(buf_size)
print(f" Peer access enabled: GPU{gpuid[0]} <-> GPU{gpuid[1]}")
print(
f" Peer access status: MR0 accessible by {mr0.peer_accessible_by}, "
f"MR1 accessible by {mr1.peer_accessible_by}"
)
# Allocate pinned host memory
pinned_mr = PinnedMemoryResource()
h0 = pinned_mr.allocate(buf_size)
print(" Memory allocated successfully")
# Create streams
stream0 = dev0.create_stream()
stream1 = dev1.create_stream()
try:
# P2P bandwidth test using CUDA events for accurate GPU-side timing
print("\nMeasuring P2P bandwidth...")
print(" Performing 100 ping-pong copies between GPUs...")
event_options = EventOptions(enable_timing=True)
sync_event0 = None
sync_event1 = None
# Record start event on stream0
start_event = stream0.record(options=event_options)
for i in range(100):
# Ping-pong copy between GPUs with explicit event-based synchronization
if i % 2 == 0:
# Wait for previous stream1 copy to complete (if any)
if sync_event1 is not None:
stream0.wait(sync_event1)
# Copy g0 -> g1 on stream0
g1.copy_from(g0, stream=stream0)
# Record event on stream0 to signal completion of this copy
sync_event0 = stream0.record(options=EventOptions(enable_timing=False))
else:
# Wait for previous stream0 copy to complete
if sync_event0 is not None:
stream1.wait(sync_event0)
# Copy g1 -> g0 on stream1
g0.copy_from(g1, stream=stream1)
# Record event on stream1 to signal completion of this copy
sync_event1 = stream1.record(options=EventOptions(enable_timing=False))
# Wait for last stream1 copy to complete
if sync_event1 is not None:
stream0.wait(sync_event1)
# Record end event on stream0 after all copies have been enqueued
end_event = stream0.record(options=event_options)
end_event.sync()
# Elapsed time in milliseconds (using subtraction operator)
time_memcpy = end_event - start_event
bandwidth = (1.0 / (time_memcpy / 1000.0)) * (100.0 * buf_size) / (1024.0**3)
print(f" P2P bandwidth: {bandwidth:.2f} GB/s")
# Prepare host buffer and initialize data
print(f"\nPreparing host buffer and memcpy to GPU{gpuid[0]}...")
# Create numpy view and initialize
h0_array = np.from_dlpack(h0).view(dtype=np.float32)
h0_array[:] = (np.arange(num_elements, dtype=np.float32) % 4096).astype(
np.float32
)
# Copy to GPU 0
dev0.set_current()
g0.copy_from(h0, stream=stream0)
stream0.sync()
print(" Data initialized and copied to GPU")
# Compile kernel for both GPUs
print("\nCompiling CUDA kernel...")
dev0.set_current()
program_options = ProgramOptions(std="c++17", arch=f"sm_{dev0.arch}")
prog = Program(SIMPLE_P2P_KERNEL, code_type="c++", options=program_options)
mod0 = prog.compile("cubin")
kernel0 = mod0.get_kernel("SimpleKernel")
dev1.set_current()
program_options = ProgramOptions(std="c++17", arch=f"sm_{dev1.arch}")
prog = Program(SIMPLE_P2P_KERNEL, code_type="c++", options=program_options)
mod1 = prog.compile("cubin")
kernel1 = mod1.get_kernel("SimpleKernel")
print(" Kernels compiled successfully")
# Launch configuration
threads = 512
blocks = (num_elements + threads - 1) // threads
config = LaunchConfig(grid=blocks, block=threads)
# Run kernel on GPU 1, reading from GPU 0, writing to GPU 1
print(
f"\nRun kernel on GPU{gpuid[1]}, taking source data from "
f"GPU{gpuid[0]} and writing to GPU{gpuid[1]}..."
)
dev1.set_current()
launch(stream1, config, kernel1, g0, g1, np.int32(num_elements))
stream1.sync()
print(" Kernel execution complete")
# Run kernel on GPU 0, reading from GPU 1, writing to GPU 0
print(
f"\nRun kernel on GPU{gpuid[0]}, taking source data from "
f"GPU{gpuid[1]} and writing to GPU{gpuid[0]}..."
)
dev0.set_current()
launch(stream0, config, kernel0, g1, g0, np.int32(num_elements))
stream0.sync()
print(" Kernel execution complete")
# Copy data back to host and verify
print(f"\nCopy data back to host from GPU{gpuid[0]} and verify results...")
g0.copy_to(h0, stream=stream0)
stream0.sync()
# Verify results
print("\nChecking results...")
print(f" Comparing {num_elements:,} elements...")
# Input data goes through two kernels, each multiplying by 2.0.
expected = (np.arange(num_elements, dtype=np.float32) % 4096) * 4.0
# Use utility function for verification (handles both numpy and cupy arrays)
if verify_array_result(h0_array, expected, rtol=1e-5, atol=1e-6, verbose=True):
print(" [PASS] Validation PASSED")
success = True
else:
print(" [FAIL] Validation FAILED")
# Show first few errors for debugging
errors = np.where(~np.isclose(h0_array, expected, rtol=1e-5, atol=1e-6))[0]
print(f" Number of mismatches: {len(errors)}")
for idx in errors[:10]:
print(
f" Error @ element {idx}: got {h0_array[idx]}, "
f"expected {expected[idx]}"
)
success = False
# Disable peer access
print("\nDisabling peer access...")
mr0.peer_accessible_by = [] # Revoke GPU 1's access to GPU 0's memory
mr1.peer_accessible_by = [] # Revoke GPU 0's access to GPU 1's memory
print(
f" Peer access revoked: MR0 accessible by {mr0.peer_accessible_by}, "
f"MR1 accessible by {mr1.peer_accessible_by}"
)
print("\n" + "=" * 70)
if success:
print("simpleP2P completed successfully!")
else:
print("simpleP2P FAILED!")
print("=" * 70 + "\n")
return 0 if success else 1
finally:
# Cleanup streams and buffers
print("Shutting down...")
stream0.close()
stream1.close()
def main():
"""Main entry point with argument parsing."""
parser = argparse.ArgumentParser(
description=(
"Demonstrate peer-to-peer (P2P) memory access between "
"multiple GPUs with CUDA"
)
)
parser.add_argument(
"--num_elements",
type=int,
default=1024 * 1024 * 16, # 16M elements = 64MB
help="Number of elements in arrays (default: 16777216 = 64MB)",
)
args = parser.parse_args()
# Validate arguments
if args.num_elements <= 0:
print("Error: num_elements must be positive")
return 1
try:
exit_code = run(num_elements=args.num_elements)
sys.exit(exit_code)
except Exception as e:
print(f"\nError: {e}")
import traceback
traceback.print_exc()
sys.exit(1)
if __name__ == "__main__":
main()