[Feature]: Producer–Consumer Example Across Ranks Using Streams

Suggestion Description

We want an example that shows the following pattern:

Each GPU rank:

• Launches a producer kernel that sends data to its peer GPU
• Launches a consumer kernel that waits for incoming data
• Uses separate HIP streams for the producer and consumer to allow concurrency
• Synchronizes using atomic signaling

This example will showcase how to use Iris to:

• Launch multiple kernels per GPU using separate HIP streams
• Perform direct device-to-device communication between GPUs using iris.store and iris.load
• Coordinate remote execution using atomic compare-and-swap (iris.atomic_cas)
• Demonstrate per-block synchronization using flags in symmetric heap memory

Pseudocode

import torch
import iris
import triton
import triton.language as tl

# === Rank setup ===
rank = iris.get_rank()
peer = (rank + 1) % 2  # only two ranks

# Define number of tiles (blocks)
num_tiles = 32
BLOCK_SIZE = 128
N = num_tiles * BLOCK_SIZE

# === Allocate symmetric heap buffers ===
data = iris.empty((N,), dtype=torch.int32, device='npu', heap=True)
signal = iris.zeros((num_tiles,), dtype=torch.int32, device='npu', heap=True)

# === Create two HIP streams ===
producer_stream = torch.cuda.Stream()
consumer_stream = torch.cuda.Stream()

# === Launch producer kernel in stream 0 ===
with torch.cuda.stream(producer_stream):
    producer_kernel[grid](
        data, signal,
        from_rank=rank,
        to_rank=peer
    )

# === Launch consumer kernel in stream 1 ===
with torch.cuda.stream(consumer_stream):
    consumer_kernel[grid](
        data, signal,
        from_rank=peer,
        to_rank=rank
    )

# === Synchronize both streams ===
producer_stream.synchronize()
consumer_stream.synchronize()

Device Kernels

@triton.jit
def producer_kernel(data_ptr, signal_ptr, from_rank, to_rank):
    pid = tl.program_id(0)
    val = pid + 100 * from_rank  # tag value with sender rank

    # Store data remotely
    iris.store(data_ptr + pid, val, from_rank, to_rank)

    # Signal completion
    cmp = 0
    new = 1
    iris.atomic_cas(signal_ptr + pid, cmp, new, from_rank, to_rank, sem="release")


@triton.jit
def consumer_kernel(data_ptr, signal_ptr, from_rank, to_rank):
    pid = tl.program_id(0)

    # Wait until signaled
    flag = 0
    while flag == 0:
        flag = iris.atomic_cas(signal_ptr + pid, 1, 1, from_rank, to_rank, sem="acquire")

    # Load data from remote
    val = iris.load(data_ptr + pid, from_rank, to_rank)
    # TODO: use val (e.g., print, check, store)

Operating System

No response

GPU

No response

ROCm Component

No response