diff --git a/examples/06_message_passing/message_passing_host_initiated.py b/examples/06_message_passing/message_passing_host_initiated.py new file mode 100644 index 000000000..d2c1b88c0 --- /dev/null +++ b/examples/06_message_passing/message_passing_host_initiated.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. +""" +Host-Initiated Message Passing Example + +This example demonstrates message passing where the producer (GPU 0) is +controlled by the HOST (Python/CPU) instead of a device kernel, while +the consumer (GPU 1) remains a device kernel. + +Key difference from message_passing_put.py: +- Producer: Host uses sdma_ep (rocm-xio) to initiate SDMA transfers from Python +- Consumer: Same device kernel waiting for data + +This shows how to orchestrate GPU-to-GPU transfers from Python without +requiring kernel launches on the source GPU. +""" + +import argparse + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import triton +import triton.language as tl +import random + +import iris + + +@triton.jit +def consumer_kernel( + buffer, # tl.tensor: pointer to shared buffer (read from target_rank) + flag, # tl.tensor: sync flag per block + buffer_size, # int32: total number of elements + consumer_rank: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + heap_bases_ptr: tl.tensor, # tl.tensor: pointer to heap bases pointers +): + pid = tl.program_id(0) + + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < buffer_size + + # Spin-wait until writer sets flag[pid] = 1 + done = 0 + while done == 0: + done = iris.atomic_cas( + flag + pid, 1, 0, consumer_rank, consumer_rank, heap_bases_ptr, sem="acquire", scope="sys" + ) + + # Read from the target buffer (written by producer) + values = tl.load(buffer + offsets, mask=mask) + + # Do something with values... + # (Here you might write to output, do computation, etc.) + values = values * 2 + + # Store chunk to target buffer + tl.store( + buffer + offsets, + values, + mask=mask, + ) + + # Optionally reset the flag for next iteration + tl.store(flag + pid, 0) + + +torch.manual_seed(123) +random.seed(123) + + +def torch_dtype_from_str(datatype: str) -> torch.dtype: + dtype_map = { + "fp16": torch.float16, + "fp32": torch.float32, + "int8": torch.int8, + "bf16": torch.bfloat16, + } + try: + return dtype_map[datatype] + except KeyError: + print(f"Unknown datatype: {datatype}") + exit(1) + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Host-Initiated SDMA Message Passing Example", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "-t", + "--datatype", + type=str, + default="fp32", + choices=["fp16", "fp32", "int8", "bf16"], + help="Datatype of computation", + ) + parser.add_argument("-s", "--buffer_size", type=int, default=4096, help="Buffer Size") + parser.add_argument("-b", "--block_size", type=int, default=512, help="Block Size") + parser.add_argument("-p", "--heap_size", type=int, default=1 << 33, help="Iris heap size") + parser.add_argument("-r", "--num_ranks", type=int, default=2, help="Number of ranks/processes") + + return vars(parser.parse_args()) + + +def host_initiated_producer(shmem, source_buffer, destination_buffer, flags, consumer_rank, block_size, verbose=True): + """ + Producer rank logic for host-initiated SDMA transfers. + + Args: + shmem: Iris instance + source_buffer: Source buffer (symmetric) + destination_buffer: Destination buffer (symmetric) + flags: Flag buffer for synchronization (symmetric) + consumer_rank: Destination rank + block_size: Block size for chunking + verbose: Whether to print timing information + """ + n_elements = source_buffer.numel() + num_blocks = triton.cdiv(n_elements, block_size) + + if verbose: + shmem.info(f"Rank {shmem.get_rank()} (HOST) is sending data to rank {consumer_rank}.") + + # Initialize CUDA context even though we're doing host-side operations + # This is needed for the barrier to work + torch.cuda.current_device() + + if verbose: + import time + + start_time = time.time() + + for block_id in range(num_blocks): + block_start = block_id * block_size + block_end = min(block_start + block_size, n_elements) + block_slice = slice(block_start, block_end) + + # Views remain symmetric, so Iris can translate remote pointers automatically + src_chunk = source_buffer[block_slice] + dst_chunk = destination_buffer[block_slice] + flag_view = flags[block_id : block_id + 1] + + shmem.put( + src_chunk, + dst_rank=consumer_rank, + dst_tensor=dst_chunk, + signal_flag=flag_view, + async_op=True, + ) + + shmem.quiet(dst_rank=consumer_rank) + + if verbose: + end_time = time.time() + elapsed_ms = (end_time - start_time) * 1000 + shmem.info( + f"Host SDMA loop took {elapsed_ms:.2f} ms for {num_blocks} blocks ({elapsed_ms / num_blocks:.2f} ms/block)" + ) + + +def _worker(local_rank: int, world_size: int, init_url: str, args: dict): + """Worker function for PyTorch distributed execution.""" + backend = "nccl" if torch.cuda.is_available() else "gloo" + dist.init_process_group( + backend=backend, + init_method=init_url, + world_size=world_size, + rank=local_rank, + device_id=torch.device(f"cuda:{local_rank}"), + ) + + # Main benchmark logic + shmem = iris.iris(args["heap_size"]) + dtype = torch_dtype_from_str(args["datatype"]) + cur_rank = shmem.get_rank() + world_size = shmem.get_num_ranks() + + # Allocate source and destination buffers on the symmetric heap + destination_buffer = shmem.zeros(args["buffer_size"], device="cuda", dtype=dtype) + if dtype.is_floating_point: + source_buffer = shmem.randn(args["buffer_size"], device="cuda", dtype=dtype) + else: + ii = torch.iinfo(dtype) + source_buffer = shmem.randint(ii.min, ii.max, (args["buffer_size"],), device="cuda", dtype=dtype) + + if world_size != 2: + raise ValueError("This example requires exactly two processes.") + + producer_rank = 0 + consumer_rank = 1 + + n_elements = source_buffer.numel() + BLOCK_SIZE = args["block_size"] + num_blocks = triton.cdiv(n_elements, BLOCK_SIZE) + grid = (num_blocks,) + + # Allocate flags on the symmetric heap + flags = shmem.zeros((num_blocks,), device="cuda", dtype=torch.int32) + + if cur_rank == producer_rank: + host_initiated_producer( + shmem, source_buffer, destination_buffer, flags, consumer_rank, BLOCK_SIZE, verbose=True + ) + else: + shmem.info(f"Rank {cur_rank} is receiving data from rank {producer_rank}.") + kk = consumer_kernel[grid]( + destination_buffer, flags, n_elements, consumer_rank, BLOCK_SIZE, shmem.get_heap_bases() + ) + + shmem.barrier() + shmem.info(f"Rank {cur_rank} has finished sending/receiving data.") + shmem.info("Validating output...") + + success = True + if cur_rank == consumer_rank: + expected = source_buffer * 2 + diff_mask = ~torch.isclose(destination_buffer, expected, atol=1) + breaking_indices = torch.nonzero(diff_mask, as_tuple=False) + + if not torch.allclose(destination_buffer, expected, atol=1): + max_diff = (destination_buffer - expected).abs().max().item() + shmem.info(f"Max absolute difference: {max_diff}") + for idx in breaking_indices: + idx = tuple(idx.tolist()) + computed_val = destination_buffer[idx] + expected_val = expected[idx] + shmem.info(f"Mismatch at index {idx}: C={computed_val}, expected={expected_val}") + success = False + break + + if success: + shmem.info("Validation successful.") + else: + shmem.info(f"Validation failed with {len(breaking_indices)} errors / {destination_buffer.numel()}") + + shmem.barrier() + + dist.barrier() + dist.destroy_process_group() + + +def main(): + args = parse_args() + + num_ranks = args["num_ranks"] + + init_url = "tcp://127.0.0.1:29500" + mp.spawn( + fn=_worker, + args=(num_ranks, init_url, args), + nprocs=num_ranks, + join=True, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/06_message_passing/message_passing_put.py b/examples/06_message_passing/message_passing_put.py index 4f7269695..5f440f288 100755 --- a/examples/06_message_passing/message_passing_put.py +++ b/examples/06_message_passing/message_passing_put.py @@ -23,6 +23,8 @@ def producer_kernel( consumer_rank: tl.constexpr, BLOCK_SIZE: tl.constexpr, heap_bases_ptr: tl.tensor, # tl.tensor: pointer to heap bases pointers + copy_engine_handle_ptr, + USE_COPY_ENGINE: tl.constexpr, ): pid = tl.program_id(0) @@ -34,10 +36,31 @@ def producer_kernel( mask = offsets < buffer_size # Put chunk into remote buffer - iris.put(source_buffer + offsets, target_buffer + offsets, producer_rank, consumer_rank, heap_bases_ptr, mask=mask) + iris.put( + source_buffer + offsets, + target_buffer + offsets, + producer_rank, + consumer_rank, + heap_bases_ptr, + mask=mask, + copy_engine_ctx=copy_engine_handle_ptr, + USE_COPY_ENGINE=USE_COPY_ENGINE, + CONTIGUOUS_COPY=True, + ) # Set flag to signal completion - iris.atomic_cas(flag + pid, 0, 1, producer_rank, consumer_rank, heap_bases_ptr, sem="release", scope="sys") + iris.atomic_cas( + flag + pid, + 0, + 1, + producer_rank, + consumer_rank, + heap_bases_ptr, + sem="release", + scope="sys", + USE_COPY_ENGINE=USE_COPY_ENGINE, + copy_engine_ctx=copy_engine_handle_ptr, + ) @triton.jit @@ -113,9 +136,11 @@ def parse_args(): ) parser.add_argument("-s", "--buffer_size", type=int, default=4096, help="Buffer Size") parser.add_argument("-b", "--block_size", type=int, default=512, help="Block Size") - parser.add_argument("-p", "--heap_size", type=int, default=1 << 33, help="Iris heap size") parser.add_argument("-r", "--num_ranks", type=int, default=2, help="Number of ranks/processes") + parser.add_argument( + "-c", "--use_copy_engine", action="store_true", help="Use copy engine for device-to-device copies" + ) return vars(parser.parse_args()) @@ -138,12 +163,12 @@ def _worker(local_rank: int, world_size: int, init_url: str, args: dict): world_size = shmem.get_num_ranks() # Allocate source and destination buffers on the symmetric heap - source_buffer = shmem.zeros(args["buffer_size"], device="cuda", dtype=dtype) + destination_buffer = shmem.zeros(args["buffer_size"], device="cuda", dtype=dtype) if dtype.is_floating_point: - destination_buffer = shmem.randn(args["buffer_size"], device="cuda", dtype=dtype) + source_buffer = shmem.randn(args["buffer_size"], device="cuda", dtype=dtype) else: ii = torch.iinfo(dtype) - destination_buffer = shmem.randint(ii.min, ii.max, (args["buffer_size"],), device="cuda", dtype=dtype) + source_buffer = shmem.randint(ii.min, ii.max, (args["buffer_size"],), device="cuda", dtype=dtype) if world_size != 2: raise ValueError("This example requires exactly two processes.") @@ -158,6 +183,9 @@ def _worker(local_rank: int, world_size: int, init_url: str, args: dict): # Allocate flags on the symmetric heap flags = shmem.zeros((num_blocks,), device="cuda", dtype=torch.int32) + # Get copy engine context + copy_engine_ctx = shmem.get_copy_engine_ctx() + if cur_rank == producer_rank: shmem.info(f"Rank {cur_rank} is sending data to rank {consumer_rank}.") kk = producer_kernel[grid]( @@ -169,6 +197,8 @@ def _worker(local_rank: int, world_size: int, init_url: str, args: dict): consumer_rank, args["block_size"], shmem.get_heap_bases(), + copy_engine_ctx, + USE_COPY_ENGINE=args["use_copy_engine"], ) else: shmem.info(f"Rank {cur_rank} is receiving data from rank {producer_rank}.") @@ -199,7 +229,7 @@ def _worker(local_rank: int, world_size: int, init_url: str, args: dict): if success: shmem.info("Validation successful.") else: - shmem.info("Validation failed.") + shmem.info(f"Validation failed with {len(breaking_indices)} errors / {destination_buffer.numel()}") shmem.barrier() diff --git a/iris/__init__.py b/iris/__init__.py index 7e4047ec4..38125e1ed 100644 --- a/iris/__init__.py +++ b/iris/__init__.py @@ -58,6 +58,7 @@ copy, get, put, + quiet, atomic_add, atomic_cas, atomic_xchg, @@ -111,6 +112,7 @@ "copy", "get", "put", + "quiet", "atomic_add", "atomic_cas", "atomic_xchg", diff --git a/iris/device/__init__.py b/iris/device/__init__.py index 588d446b3..4d1e1cd56 100644 --- a/iris/device/__init__.py +++ b/iris/device/__init__.py @@ -3,3 +3,6 @@ """Backward compat shim — canonical code moved to iris.mem.""" from iris.mem import * # noqa: F401,F403 +from . import sdma_utils + +__all__ = ["sdma_utils"] diff --git a/iris/device/sdma_utils.py b/iris/device/sdma_utils.py new file mode 100644 index 000000000..dba55e0bc --- /dev/null +++ b/iris/device/sdma_utils.py @@ -0,0 +1,413 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2025-2026 Advanced Micro Devices, Inc. All rights reserved. + +""" +SDMA (System DMA) device-side utilities for Triton kernels. + +This module provides low-level Triton device functions for directly managing +SDMA queues from GPU kernels, including packet construction, queue reservation, +and submission operations. +""" + +import triton +import triton.language as tl +from xio import sdma_ep + + +@triton.jit +def wait_cnt(): + tl.inline_asm_elementwise("s_waitcnt vmcnt(0)", "=r", [], dtype=tl.int32, is_pure=False, pack=1) + + +@triton.jit +def wrap_into_ring(index: tl.uint64): + queue_size_u32 = sdma_ep.SDMA_QUEUE_SIZE + queue_size = queue_size_u32.to(tl.uint64) + return index.to(tl.uint64) % queue_size + + +@triton.jit +def can_write_up_to(rptr, up_to_index: tl.uint64): + """Check if there's space to write up to the given index in the ring buffer.""" + hw_read_ptr = tl.load(rptr, cache_modifier=".cv", volatile=True) + return (up_to_index - hw_read_ptr) < sdma_ep.SDMA_QUEUE_SIZE + + +@triton.jit +def acquire( + queue_ptr_u32, + read_ptr, + write_ptr, + doorbell_ptr, + cached_write_ptr: tl.pointer_type(tl.uint64), + committed_write_ptr, + command_in_bytes: tl.uint64, +): + """ + Reserve space in the SDMA queue. + Returns (base_index, offset) where: + - base_index: the index where the packet should be written (cur_index initially) + - offset: padding bytes added for wraparound (0 if no wraparound) + + Based on ReserveQueueSpace from anvil_device.hpp. + """ + queue_size_u32 = sdma_ep.SDMA_QUEUE_SIZE + queue_size_in_bytes = queue_size_u32.to(tl.uint64) + + base_u32 = 0 + base = (base_u32).to(tl.uint64) + offset_u32 = 0 + offset = (offset_u32).to(tl.uint64) + + stop_loop = False + while not stop_loop: + cur_index = tl.load(cached_write_ptr, volatile=True) + offset = (offset_u32).to(tl.uint64) + + # Calculate current position in ring buffer + cur_ring_pos = wrap_into_ring(cur_index) + + # Check if we need to wrap around + if (cur_ring_pos + command_in_bytes) > queue_size_in_bytes: + # Need to pad to end of ring before wrap around + offset = queue_size_in_bytes - cur_ring_pos + + # Calculate new index including any wraparound padding + new_index = cur_index + command_in_bytes + offset + base = cur_index + + # Check if queue has space + if can_write_up_to(read_ptr, new_index): + # Try to atomically claim this space + if tl.atomic_cas(cached_write_ptr, cur_index, new_index, sem="relaxed", scope="gpu") == cur_index: + stop_loop = True + return base, offset + + +# acquire function using atomic_add instead of atomic_cas +@triton.jit +def acquire_fadd( + queue_ptr_u32, + read_ptr, + write_ptr, + doorbell_ptr, + cached_write_ptr: tl.pointer_type(tl.uint64), + committed_write_ptr, + command_in_bytes: tl.uint64, +): + """ + Reserve space in the SDMA queue using atomic_add. + Returns (base_index, 0) where base_index is where the packet should be written. + + Uses atomic_add instead of CAS. Immediately acquires space, and if wraparound + is detected, places padding NOP packet, submits it, and tries again. + Always returns a non-wrapping allocation. + """ + queue_size_u32 = sdma_ep.SDMA_QUEUE_SIZE + queue_size_in_bytes = queue_size_u32.to(tl.uint64) + + stop_loop = False + base_u32 = 0 + base = (base_u32).to(tl.uint64) + offset_u32 = 0 + offset = (offset_u32).to(tl.uint64) + + while not stop_loop: + # Atomically acquire space for the command + base = tl.atomic_add(cached_write_ptr, command_in_bytes, sem="relaxed", scope="gpu") + end_index = base + command_in_bytes + # Calculate current position in ring buffer + cur_ring_pos = wrap_into_ring(base) + + # Block until there is space in the queue to write the command + while not can_write_up_to(read_ptr, end_index): + pass + + # Check if we need to wrap around + if (cur_ring_pos + command_in_bytes) > queue_size_in_bytes: + # Wrap detected - need to pad to end of ring + padding_bytes = queue_size_in_bytes - cur_ring_pos + + # Place NOP packet at end of ring + place_nop_packet(queue_ptr_u32, base, padding_bytes) + + # Place remaining NOP padding at the beginning of the ring. + remaining_padding_bytes = command_in_bytes - padding_bytes + place_nop_packet(queue_ptr_u32, base + padding_bytes, remaining_padding_bytes) + + # Submit the padding - update committed write pointer + # This allows other threads to proceed past this padding + submit(write_ptr, doorbell_ptr, committed_write_ptr, base, end_index) + + # Continue loop to acquire space for the actual command (will be at ring start) + else: + # No wrap - this allocation is good + stop_loop = True + + return base, offset + + +@triton.jit +def submit(write_ptr, doorbell_ptr, committed_write_ptr, base, pending_wptr): + """ + Submit SDMA commands to the hardware by updating write pointer and ringing the doorbell. + + Waits for previous threads to commit, then updates write pointer, rings doorbell, + and updates committed pointer to allow subsequent threads to proceed. + """ + while tl.load(committed_write_ptr, cache_modifier=".cv", volatile=True) != base: + pass + + wait_cnt() + tl.debug_barrier() + + tl.store(write_ptr, pending_wptr, cache_modifier=".wt") + wait_cnt() + tl.debug_barrier() + + # Ring doorbell + tl.store(doorbell_ptr, pending_wptr, cache_modifier=".wt") + wait_cnt() + tl.debug_barrier() + + tl.store(committed_write_ptr, pending_wptr, cache_modifier=".wt") + + +@triton.jit +def quiet(read_ptr, cached_write_ptr: tl.pointer_type(tl.uint64)): + """ + Wait for all submitted SDMA operations to complete. + + Polls the hardware read pointer until it catches up to the cached write pointer, + ensuring all previously submitted SDMA packets have been processed. + + Args: + read_ptr: Pointer to hardware read pointer + cached_write_ptr: Pointer to cached write pointer + """ + target_wptr = tl.load(cached_write_ptr, cache_modifier=".cv", volatile=True) + while tl.load(read_ptr, cache_modifier=".cv", volatile=True) != target_wptr: + pass + + +@triton.jit +def place_nop_packet(queue_ptr_u32, offset_bytes: tl.uint64, padding_bytes): + """Place a NOP (no operation) packet for ring buffer padding.""" + num_padding_dwords = (padding_bytes // 4).to(tl.int32) + offset_ring_pos = wrap_into_ring(offset_bytes) + offset_in_dwords = (offset_ring_pos // 4).to(tl.int32) + for i in range(num_padding_dwords): + if i == 0: + tl.store(queue_ptr_u32 + offset_in_dwords, ((num_padding_dwords - 1) & 0xFFFF) << 16, cache_modifier=".wt") + else: + tl.store(queue_ptr_u32 + offset_in_dwords + i, 0, cache_modifier=".wt") + + +@triton.jit +def place_copy_packet(queue_ptr_u32, offset_bytes: tl.uint64, size_bytes: tl.uint32, src_ptr_val, dst_ptr_val): + """Place a SDMA_PKT_COPY_LINEAR packet for 1D linear memory copy.""" + slot_ptr_u32 = queue_ptr_u32 + (wrap_into_ring(offset_bytes) // 4) + # offset 0: op + sub_op + tl.store(slot_ptr_u32 + 0, 1, cache_modifier=".wt") + # offset 1: count + tl.store(slot_ptr_u32 + 1, size_bytes - 1, cache_modifier=".wt") + # offset 2: parameters + tl.store(slot_ptr_u32 + 2, 0, cache_modifier=".wt") + # offset 3: src address 31:0 + tl.store(slot_ptr_u32 + 3, src_ptr_val.to(tl.uint32), cache_modifier=".wt") + # offset 4: src address 63:32 + tl.store(slot_ptr_u32 + 4, (src_ptr_val >> 32).to(tl.uint32), cache_modifier=".wt") + # offset 5: dst address 31:0 + tl.store(slot_ptr_u32 + 5, dst_ptr_val.to(tl.uint32), cache_modifier=".wt") + # offset 6: dst address 63:32 + tl.store(slot_ptr_u32 + 6, (dst_ptr_val >> 32).to(tl.uint32), cache_modifier=".wt") + + +# atomic op codes and operation +# atomic add 32bit w/rtn: op 10, operation 15 +# atomic add 64bit w/rtn: op 10, operation: 47 -> 32 + 15 +# atomic add 32bit w/o rtn: op 10, operation: 31 -> 64 + 15 +# atomic add 64bit w/o rtn: op 10, operation: 63 -> 96 + 15 +# atomic cmp&swap 32bit w/rtn: op 10, operation: 8 +# atomic cmp&swap 64bit w/rtn: op 10, operation: -> 32 + 8 +# atomic cmp&swap 32bit w/o rtn: op 10, operation -> 64 + 8 +# atomic cmp&swap 64bit w/o rtn: op 10, operation 56 -> 06 + 8 +@triton.jit +def place_atomic_packet( + queue_ptr_u32, + offset_bytes: tl.uint64, + dst_ptr_val, + src_data, + comp_data, + OP: tl.constexpr, + RETURN: tl.constexpr = False, + IS_64_BIT: tl.constexpr = False, +): + """ + Place a SDMA_PKT_ATOMIC packet for atomic memory operations. + + OP codes: + 15: atomic add (32/64-bit with/without return) + 8: atomic compare-and-swap (32/64-bit with/without return) + Flags are encoded via IS_64_BIT (bit 4) and RETURN (bit 5). + """ + slot_ptr_u32 = queue_ptr_u32 + (wrap_into_ring(offset_bytes) // 4) + if IS_64_BIT: + OP = OP | (0x1 << 4) + if not RETURN: + OP = OP | (0x1 << 5) + tl.store(slot_ptr_u32 + 0, ((OP & 0x7F) << 25) | (0xA & 0xFF), cache_modifier=".wt") + # offset 1: dst address 31:0 + tl.store(slot_ptr_u32 + 1, dst_ptr_val.to(tl.uint32), cache_modifier=".wt") + # offset 2: dst address 63:32 + tl.store(slot_ptr_u32 + 2, (dst_ptr_val >> 32).to(tl.uint32), cache_modifier=".wt") + # offset 3: src data 31:0 + tl.store(slot_ptr_u32 + 3, src_data, cache_modifier=".wt") + # offset 4: src data 63:32 + if IS_64_BIT: + tl.store(slot_ptr_u32 + 4, (src_data >> 32).to(tl.uint32), cache_modifier=".wt") + else: + tl.store(slot_ptr_u32 + 4, 0, cache_modifier=".wt") + # offset 5: compare data 31:0 + tl.store(slot_ptr_u32 + 5, comp_data, cache_modifier=".wt") + # offset 6: compare data 63:32 + if IS_64_BIT: + tl.store(slot_ptr_u32 + 6, (comp_data >> 32), cache_modifier=".wt") + else: + tl.store(slot_ptr_u32 + 6, 0, cache_modifier=".wt") + # offset 7: loop timer + loop interval + tl.store(slot_ptr_u32 + 7, 0, cache_modifier=".wt") + + +@triton.jit +def place_atomic_add_packet(queue_ptr_u32, offset_bytes: tl.uint64, dst_ptr_val, val): + """Place an atomic add packet (OP=15, with return).""" + place_atomic_packet(queue_ptr_u32, offset_bytes, dst_ptr_val, val, 0, 15, True) + + +@triton.jit +def place_atomic_cas_packet( + queue_ptr_u32, + offset_bytes: tl.uint64, + dst_ptr_val, + compare_val, + swap_val, +): + """Place an atomic compare-and-swap packet (OP=8, with return).""" + place_atomic_packet(queue_ptr_u32, offset_bytes, dst_ptr_val, swap_val, compare_val, 8, True) + + +@triton.jit +def place_poll_regmem_packet( + queue_ptr_u32, + offset_bytes: tl.uint64, + flag_ptr_val, + expected_value, + interval: tl.constexpr = 10, + retry_count: tl.constexpr = 0xFFF, +): + """ + Place a SDMA_PKT_POLL_REGMEM packet for memory polling. + + Polls memory location until (value >= expected_value). + """ + slot_ptr_u32 = queue_ptr_u32 + (wrap_into_ring(offset_bytes) // 4) + header = ((1 & 0x1) << 31) | ((5 & 0x7) << 28) | (8 & 0xFF) + dw5 = ((retry_count & 0xFFF) << 16) | (interval & 0xFFFF) + + tl.store(slot_ptr_u32 + 0, header, cache_modifier=".wt") + tl.store(slot_ptr_u32 + 1, flag_ptr_val.to(tl.uint32), cache_modifier=".wt") + tl.store(slot_ptr_u32 + 2, (flag_ptr_val.to(tl.uint64) >> 32).to(tl.uint32), cache_modifier=".wt") + tl.store(slot_ptr_u32 + 3, expected_value.to(tl.uint32), cache_modifier=".wt") + tl.store(slot_ptr_u32 + 4, 0xFFFFFFFF, cache_modifier=".wt") + tl.store(slot_ptr_u32 + 5, dw5, cache_modifier=".wt") + + +@triton.jit +def place_sub_window_copy_packet( + queue_ptr_u32, + offset_bytes: tl.uint64, + src_ptr_val, + dst_ptr_val, + tile_width: tl.uint32, + tile_height: tl.uint32, + src_buffer_pitch: tl.uint32, + dst_buffer_pitch: tl.uint32, + src_x: tl.uint32, + src_y: tl.uint32, + dst_x: tl.uint32, + dst_y: tl.uint32, +): + """ + Place a SDMA_PKT_LINEAR_LARGE_SUB_WINDOW_COPY packet for 2D tile transfer. + + Copies a rectangular tile with arbitrary source/destination offsets. + Note: pitch, slice_pitch and rect fields are 1-based (subtract 1 before writing). + Args: + queue_ptr_u32: Pointer to the SDMA queue buffer (as uint32 array) + offset_bytes: Byte offset in the queue where to place the packet + src_ptr_val: Source buffer base address + dst_ptr_val: Destination buffer base address + tile_width: Width of the tile to copy in bytes + tile_height: Height of the tile to copy in rows + src_buffer_pitch: Row stride of the source buffer in bytes + dst_buffer_pitch: Row stride of the destination buffer in bytes + src_x: Source X offset in bytes + src_y: Source Y offset in rows + dst_x: Destination X offset in bytes + dst_y: Destination Y offset in rows + """ + slot_ptr_u32 = queue_ptr_u32 + (wrap_into_ring(offset_bytes) // 4) + + # DW 0: Header (op=1, sub_op=0x24) + # op[7:0] = 1 (SDMA_OP_COPY), sub_op[15:8] = 0x24 (SDMA_SUBOP_COPY_LINEAR_SUB_WINDOW) + tl.store(slot_ptr_u32 + 0, ((0x24 & 0xFF) << 8) | (0x1 & 0xFF), cache_modifier=".wt") + + # DW 1-2: Source base address + tl.store(slot_ptr_u32 + 1, src_ptr_val.to(tl.uint32), cache_modifier=".wt") + tl.store(slot_ptr_u32 + 2, (src_ptr_val >> 32).to(tl.uint32), cache_modifier=".wt") + + # DW 3: Source X offset (bytes) + tl.store(slot_ptr_u32 + 3, src_x, cache_modifier=".wt") + + # DW 4: Source Y offset (rows) + tl.store(slot_ptr_u32 + 4, src_y, cache_modifier=".wt") + + # DW 5: Source Z offset (0 for 2D) + tl.store(slot_ptr_u32 + 5, 0, cache_modifier=".wt") + + # DW 6: Source pitch (1-based, so subtract 1) + tl.store(slot_ptr_u32 + 6, src_buffer_pitch - 1, cache_modifier=".wt") + + # DW 7-8: Source slice pitch (1-based, 0 means slice_pitch of 1, for 2D) + tl.store(slot_ptr_u32 + 7, 0, cache_modifier=".wt") + tl.store(slot_ptr_u32 + 8, 0, cache_modifier=".wt") + + # DW 9-10: Destination base address + tl.store(slot_ptr_u32 + 9, dst_ptr_val.to(tl.uint32), cache_modifier=".wt") + tl.store(slot_ptr_u32 + 10, (dst_ptr_val >> 32).to(tl.uint32), cache_modifier=".wt") + + # DW 11: Destination X offset (bytes) + tl.store(slot_ptr_u32 + 11, dst_x, cache_modifier=".wt") + + # DW 12: Destination Y offset (rows) + tl.store(slot_ptr_u32 + 12, dst_y, cache_modifier=".wt") + + # DW 13: Destination Z offset (0 for 2D) + tl.store(slot_ptr_u32 + 13, 0, cache_modifier=".wt") + + # DW 14: Destination pitch (1-based, so subtract 1) + tl.store(slot_ptr_u32 + 14, dst_buffer_pitch - 1, cache_modifier=".wt") + + # DW 15-16: Destination slice pitch (1-based, 0 means slice_pitch of 1, for 2D) + tl.store(slot_ptr_u32 + 15, 0, cache_modifier=".wt") + tl.store(slot_ptr_u32 + 16, 0, cache_modifier=".wt") + + # DW 17: Rectangle X (width in bytes, 1-based) + tl.store(slot_ptr_u32 + 17, tile_width - 1, cache_modifier=".wt") + + # DW 18: Rectangle Y (height in rows, 1-based) + tl.store(slot_ptr_u32 + 18, tile_height - 1, cache_modifier=".wt") + + # DW 19: Rectangle Z (depth, 1-based, 0 for 2D means depth of 1) + tl.store(slot_ptr_u32 + 19, 0, cache_modifier=".wt") diff --git a/iris/host/iris.py b/iris/host/iris.py index 1c9843fbe..783c5bc2f 100644 --- a/iris/host/iris.py +++ b/iris/host/iris.py @@ -52,6 +52,8 @@ get_cu_count, count_devices, ) + +from xio import sdma_ep from iris.host.memory.symmetric_heap import SymmetricHeap import numpy as np from typing import Any @@ -135,6 +137,36 @@ def __init__(self, heap_size=1 << 30, allocator_type="torch"): distributed_barrier() + # initialize copy engines + sdma_ep.init() + + context_size = sdma_ep.QUEUE_DEVICE_CTX_SIZE + self.copy_engines_device_ctx = torch.zeros((num_ranks, context_size), dtype=torch.uint64, device=self.device) + + num_local_ranks = min(num_gpus, num_ranks) + cur_local_rank = cur_rank % num_local_ranks + + for local_rank in range(num_local_ranks): + # Device-initiated queues + sdma_ep.create_queue(cur_local_rank, local_rank) + # Host-initiated queues + sdma_ep.create_host_queue(cur_local_rank, local_rank) + + handle = sdma_ep.get_queue_device_ctx(cur_local_rank, local_rank) + self.debug(f"---- Queue {local_rank} ------------") + self.debug(f"queue_buf {handle.queue_buf:#x} at {id(handle.queue_buf):#x}") + self.debug(f"rptr {handle.rptr:#x} at {id(handle.rptr):#x}") + self.debug(f"wptr {handle.wptr:#x} at {id(handle.wptr):#x}") + self.debug(f"doorbell {handle.doorbell:#x} at {id(handle.doorbell):#x}") + self.debug(f"cached_write_ptr {handle.cached_wptr:#x} at {id(handle.cached_wptr):#x}") + self.debug(f"committed_write_ptr {handle.committed_wptr:#x} at {id(handle.committed_wptr):#x}") + + self.copy_engines_device_ctx[local_rank][0] = handle.queue_buf + self.copy_engines_device_ctx[local_rank][1] = handle.rptr + self.copy_engines_device_ctx[local_rank][2] = handle.wptr + self.copy_engines_device_ctx[local_rank][3] = handle.doorbell + self.copy_engines_device_ctx[local_rank][4] = handle.cached_wptr + self.copy_engines_device_ctx[local_rank][5] = handle.committed_wptr # Initialize CCL interface self.ccl = self.CCL(self) @@ -918,6 +950,306 @@ def get_heap_bases(self): """ return self.heap_bases + def get_copy_engine_ctx(self): + return self.copy_engines_device_ctx + + @staticmethod + def _dtype_to_flag_bits(dtype: torch.dtype) -> int: + if dtype in (torch.int32, torch.int): + return 32 + if dtype in (torch.int64, torch.long): + return 64 + raise ValueError(f"Unsupported flag tensor dtype: {dtype}") + + def _flag_pointer_and_bits( + self, + flag, + *, + translate: bool = False, + dst_rank: int | None = None, + default_bits: int = 32, + ) -> tuple[int, int]: + if flag is None: + return 0, 0 + + if isinstance(flag, torch.Tensor): + bits = self._dtype_to_flag_bits(flag.dtype) + ptr = flag.data_ptr() + else: + ptr = int(flag) + bits = default_bits + + if translate: + if dst_rank is None: + raise ValueError("dst_rank must be provided when translate=True") + ptr = self.heap.translate(ptr, self.get_rank(), dst_rank) + + if ptr == 0: + return 0, 0 + return ptr, bits + + def put( + self, + src_tensor: torch.Tensor, + dst_rank: int, + dst_tensor: torch.Tensor = None, + wait_flag: torch.Tensor = None, + wait_value: int = None, + signal_flag: torch.Tensor = None, + signal_value: int = 1, + async_op: bool = False, + channel: int = 0, + ): + """ + One-sided put operation with optional wait (POLL) and signal (ATOMIC). + + Supports: + - Simple copy: put(src, dst_rank) + - Copy + signal: put(src, dst_rank, signal_flag=flag) + - Wait + copy: put(src, dst_rank, wait_flag=flag, wait_value=N) + - Wait + copy + signal: put(src, dst_rank, wait_flag=..., signal_flag=...) + + Args: + src_tensor: Source tensor (local, must be symmetric) + dst_rank: Destination rank + dst_tensor: Destination tensor (symmetric). If None, uses src_tensor. + wait_flag: Optional LOCAL flag tensor to poll before transfer (POLL packet) + wait_value: Expected value for wait_flag + signal_flag: Optional flag tensor to atomic-add on REMOTE rank after transfer (will be translated) + signal_value: Value to add to signal_flag (default 1) + async_op: If True, don't wait for completion + channel: SDMA channel to use + + Examples: + >>> # Simple copy + >>> shmem.put(data, dst_rank=1) + + >>> # Copy with completion signal + >>> shmem.put(data, dst_rank=1, signal_flag=completion_flag) + + >>> # Wait for ready signal, then copy + >>> shmem.put(data, dst_rank=1, wait_flag=ready_flag, wait_value=1) + + >>> # Full pipeline: wait, copy, signal + >>> shmem.put(data, dst_rank=1, + ... wait_flag=batch_ready, wait_value=256, + ... signal_flag=transfer_done, signal_value=1) + """ + if dst_tensor is None: + dst_tensor = src_tensor + + src_rank = self.get_rank() + src_ptr = src_tensor.data_ptr() + dst_ptr = self.heap.translate(dst_tensor.data_ptr(), src_rank, dst_rank) + size = src_tensor.numel() * src_tensor.element_size() + + # Early return for zero-size transfers (no-op) + if size == 0: + return + + wait_ptr, wait_bits = self._flag_pointer_and_bits(wait_flag) + signal_ptr, signal_bits = self._flag_pointer_and_bits(signal_flag, translate=True, dst_rank=dst_rank) + + has_wait = wait_ptr != 0 + has_signal = signal_ptr != 0 + + if has_wait and has_signal: + # Wait + copy + signal (two calls) + wait_val = int(wait_value if wait_value is not None else 0) + signal_val = int(signal_value) + sdma_ep.wait_flag_then_put( + src_rank, dst_rank, channel, wait_ptr, wait_val, src_ptr, dst_ptr, size, wait_bits + ) + sdma_ep.signal(src_rank, dst_rank, channel, signal_ptr, signal_val, signal_bits) + elif has_wait: + # Wait + copy + wait_val = int(wait_value if wait_value is not None else 0) + sdma_ep.wait_flag_then_put( + src_rank, dst_rank, channel, wait_ptr, wait_val, src_ptr, dst_ptr, size, wait_bits + ) + elif has_signal: + # Copy + signal + signal_val = int(signal_value) + sdma_ep.put_signal(src_rank, dst_rank, channel, src_ptr, dst_ptr, size, signal_ptr, signal_val, signal_bits) + else: + # Simple copy + sdma_ep.put(src_rank, dst_rank, channel, src_ptr, dst_ptr, size) + + if not async_op: + sdma_ep.quiet(src_rank, dst_rank, channel) + + def put_tile( + self, + tile, + dst_rank: int, + dst_ptr: int, + dst_stride: int, + wait_flag: int = None, + wait_value: int = None, + signal_flag: int = None, + signal_value: int = 1, + async_op: bool = False, + channel: int = 0, + ): + """ + 2D tile transfer with optional wait/signal (sub-window copy). + + Low-level API - caller provides pre-translated pointers for performance. + + Args: + tile: Pre-configured sdma_ep.Tile object with data pointer and dimensions set + dst_rank: Destination rank + dst_ptr: Destination pointer (already translated to remote address space) + dst_stride: Destination row stride in bytes + wait_flag: Optional LOCAL flag pointer to poll before transfer + wait_value: Expected value for wait_flag + signal_flag: Optional REMOTE flag pointer to atomic-add after transfer (already translated) + signal_value: Value to add to signal_flag + async_op: If True, don't wait for completion + channel: SDMA channel to use + + Examples: + >>> from xio import sdma_ep + >>> tile = sdma_ep.Tile() + >>> tile.pid_m = 0 + >>> tile.pid_n = 0 + >>> tile.block_m = 256 + >>> tile.block_n = 256 + >>> tile.elem_size = A.element_size() + >>> tile.src_stride = A.stride(0) * tile.elem_size + >>> tile.data = A.data_ptr() + >>> dst_ptr = shmem.translate(A.data_ptr(), src_rank, dst_rank) + >>> dst_stride = A.stride(0) * tile.elem_size + >>> wait_ptr = flag.data_ptr() + >>> signal_ptr = shmem.translate(flag.data_ptr(), src_rank, dst_rank) + >>> shmem.put_tile(tile, dst_rank=1, dst_ptr=dst_ptr, dst_stride=dst_stride, + ... wait_flag=wait_ptr, wait_value=256, signal_flag=signal_ptr) + """ + src_rank = self.get_rank() + + wait_ptr, wait_bits = self._flag_pointer_and_bits(wait_flag, default_bits=32) + signal_ptr, signal_bits = self._flag_pointer_and_bits(signal_flag, default_bits=32) + + has_wait = wait_ptr != 0 + has_signal = signal_ptr != 0 + + if has_wait and has_signal: + # Wait + tile copy + signal (two calls) + wait_val = int(wait_value if wait_value is not None else 0) + signal_val = int(signal_value) + sdma_ep.wait_flag_then_put_tile( + src_rank, dst_rank, channel, wait_ptr, wait_val, tile, int(dst_ptr), int(dst_stride), wait_bits + ) + sdma_ep.signal(src_rank, dst_rank, channel, signal_ptr, signal_val, signal_bits) + elif has_wait: + # Wait + tile copy + wait_val = int(wait_value if wait_value is not None else 0) + sdma_ep.wait_flag_then_put_tile( + src_rank, dst_rank, channel, wait_ptr, wait_val, tile, int(dst_ptr), int(dst_stride), wait_bits + ) + elif has_signal: + # Tile copy + signal + signal_val = int(signal_value) + sdma_ep.put_tile_signal( + src_rank, dst_rank, channel, tile, int(dst_ptr), int(dst_stride), signal_ptr, signal_val, signal_bits + ) + else: + # Simple tile copy + sdma_ep.put_tile(src_rank, dst_rank, channel, tile, int(dst_ptr), int(dst_stride)) + + if not async_op: + sdma_ep.quiet(src_rank, dst_rank, channel) + + def put_tiles( + self, + tiles, + dst_rank: int, + dst_ptrs, + dst_strides, + wait_flag: int = None, + wait_value: int = None, + signal_flag: int = None, + signal_value: int = 1, + async_op: bool = False, + channel: int = 0, + ): + """ + Batched 2D tile transfer with optional shared wait/signal. + + Args: + tiles: Sequence of pre-configured sdma_ep.Tile objects + dst_rank: Destination rank + dst_ptrs: Sequence of translated destination pointers + dst_strides: Sequence of destination row strides in bytes + wait_flag: Optional LOCAL flag pointer to poll before all transfers + wait_value: Expected value for wait_flag + signal_flag: Optional REMOTE flag pointer to atomic-add after all transfers + signal_value: Value to add to signal_flag + async_op: If True, don't wait for completion + channel: SDMA channel to use + """ + src_rank = self.get_rank() + + if len(tiles) != len(dst_ptrs) or len(tiles) != len(dst_strides): + raise ValueError("tiles, dst_ptrs, and dst_strides must have the same length") + + wait_ptr, wait_bits = self._flag_pointer_and_bits(wait_flag, default_bits=32) + signal_ptr, signal_bits = self._flag_pointer_and_bits(signal_flag, default_bits=32) + + has_wait = wait_ptr != 0 + has_signal = signal_ptr != 0 + + dst_ptr_list = [int(p) for p in dst_ptrs] + dst_stride_list = [int(s) for s in dst_strides] + + if has_wait and has_signal: + # Wait + tiles copy + signal (two calls) + wait_val = int(wait_value if wait_value is not None else 0) + signal_val = int(signal_value) + sdma_ep.wait_flag_then_put_tiles( + src_rank, dst_rank, channel, wait_ptr, wait_val, list(tiles), dst_ptr_list, dst_stride_list, wait_bits + ) + sdma_ep.signal(src_rank, dst_rank, channel, signal_ptr, signal_val, signal_bits) + elif has_wait: + # Wait + tiles copy + wait_val = int(wait_value if wait_value is not None else 0) + sdma_ep.wait_flag_then_put_tiles( + src_rank, dst_rank, channel, wait_ptr, wait_val, list(tiles), dst_ptr_list, dst_stride_list, wait_bits + ) + elif has_signal: + # Tiles copy + signal (loop + signal) + signal_val = int(signal_value) + sdma_ep.put_tiles(src_rank, dst_rank, channel, list(tiles), dst_ptr_list, dst_stride_list) + sdma_ep.signal(src_rank, dst_rank, channel, signal_ptr, signal_val, signal_bits) + else: + # Simple tiles copy + sdma_ep.put_tiles(src_rank, dst_rank, channel, list(tiles), dst_ptr_list, dst_stride_list) + + if not async_op: + sdma_ep.quiet(src_rank, dst_rank, channel) + + def quiet(self, dst_rank: int = None, channel: int = 0): + """ + Wait for all outstanding SDMA operations to complete. + + Args: + dst_rank: If specified, wait only for ops to this rank. + If None, wait for ops to all ranks. + channel: SDMA channel + + Example: + >>> shmem.put(tensor, dst_rank=1, async_op=True) + >>> shmem.quiet(dst_rank=1) # Wait for completion + >>> shmem.quiet() # Wait for all ranks + """ + src_rank = self.get_rank() + if dst_rank is not None: + sdma_ep.quiet(src_rank, dst_rank, channel) + else: + # Quiet to all ranks + for rank in range(self.get_num_ranks()): + sdma_ep.quiet(src_rank, rank, channel) + def _build_device_context(self): """ Build and cache the device context tensor. diff --git a/iris/host/memory/symmetric_heap.py b/iris/host/memory/symmetric_heap.py index 28165288b..262a66af1 100644 --- a/iris/host/memory/symmetric_heap.py +++ b/iris/host/memory/symmetric_heap.py @@ -179,6 +179,10 @@ def __init__( else: self.heap_bases = torch.tensor(heap_bases_array, device=device, dtype=torch.int64) + # Pre-fetch heap_bases to CPU for host-side address translation + # This avoids GPU->CPU transfer on every translate() call + self.heap_bases_cpu = self.heap_bases.cpu().numpy() + self._peer_refresh_failed = False self.refresh_peer_access() @@ -327,6 +331,36 @@ def get_heap_bases(self) -> torch.Tensor: self._ensure_peer_refresh_healthy() return self.heap_bases + def translate(self, ptr: int, from_rank: int, to_rank: int) -> int: + """ + Translate a pointer address from one rank's address space to another. + + This is useful for host-side SDMA operations where you need to convert + peer-mapped addresses to the target GPU's local address space. + + Args: + ptr (int): The pointer address in from_rank's address space + from_rank (int): Source rank (address space of ptr) + to_rank (int): Target rank (desired address space) + + Returns: + int: Translated pointer address in to_rank's address space + + Example: + >>> ctx = iris.iris() + >>> buffer = ctx.zeros(1024, dtype=torch.float32) + >>> # Translate buffer address from rank 0 to rank 1's address space + >>> remote_addr = ctx.heap.translate(buffer.data_ptr(), 0, 1) + """ + # Ensure heap is in healthy state before translating + self._ensure_peer_refresh_healthy() + + # Use pre-cached CPU copy to avoid GPU->CPU transfer on every call + from_base = int(self.heap_bases_cpu[from_rank]) + to_base = int(self.heap_bases_cpu[to_rank]) + offset = ptr - from_base + return to_base + offset + def refresh_peer_access(self): """ Refresh peer imports for the active allocator backend. @@ -392,6 +426,9 @@ def _refresh_peer_access_impl(self): else: return + # Update CPU cache after heap_bases modifications + self.heap_bases_cpu = self.heap_bases.cpu().numpy() + if dist.is_initialized(): dist.barrier() diff --git a/iris/mem/triton/ops.py b/iris/mem/triton/ops.py index 12233d2e8..847e5c029 100644 --- a/iris/mem/triton/ops.py +++ b/iris/mem/triton/ops.py @@ -10,6 +10,8 @@ import triton import triton.language as tl from iris.mem.triton.context import __translate +from xio import sdma_ep +from iris.device import sdma_utils @triton.jit @@ -287,6 +289,13 @@ def put( load_cache_modifier=None, store_cache_modifier=None, hint: tl.constexpr = None, + copy_engine_ctx: tl.tensor = None, + src_row_stride: tl.constexpr = 0, + dst_row_stride: tl.constexpr = 0, + USE_COPY_ENGINE: tl.constexpr = False, + CONTIGUOUS_COPY: tl.constexpr = False, + from_base_ptr=None, + to_base_ptr=None, ): """ Copies data from the current rank's local memory to the specified rank's memory. @@ -294,16 +303,22 @@ def put( rank's `from_ptr`, translating the `to_ptr` from the current rank's address space to the `to_rank`'s address space, and storing the data to the `to_rank` memory location. + Supports both 1D (flat/linear) and 2D (tiled) copies: + - 1D copies: Used for 1D pointer blocks, uses linear SDMA packets + - 2D copies: Used for 2D pointer blocks, uses sub-window SDMA packets for better performance The load is **always local** (reading from the current rank's own ``from_ptr``), while the store is **remote** when ``from_rank != to_rank`` (writing to a peer GPU). Args: from_ptr (triton.PointerType, or block of dtype=triton.PointerType): Pointer in the current rank's local memory from which to read data. - to_ptr (triton.PointerType, or block of dtype=triton.PointerType): Pointer in the current rank's address space that will be translated to the `to_rank`'s address space. Must be the current rank where the pointer is local. + to_ptr (triton.PointerType, or block of dtype=triton.PointerType): Pointer in the current rank's address space that will be translated to the `to_rank`'s address space. from_rank (int): The current rank ID from which to read the data. - to_rank (int): The `to_rank` ID to which the data will be written. + to_rank (int): The rank ID to which the data will be written. heap_bases (triton.PointerType): Array containing the heap base addresses for all ranks. - mask (Block of triton.int1, optional): If mask[idx] is false, do not load the data at address from_ptr[idx] and do not store to to_ptr[idx]. Defaults to None. + copy_engine_ctx (tl.tensor): Copy engine context for SDMA operations. + mask (Block of triton.int1, optional): If mask[idx] is false, do not load/copy data at that index. + When ``USE_COPY_ENGINE`` and ``CONTIGUOUS_COPY`` are true, ``mask=None`` copies the full pointer block. + Defaults to None. other (Block, optional): Value to return for masked-out elements during the load operation. If not provided, the result for masked-out elements is undefined. Defaults to None. load_cache_modifier (str, optional): Controls cache behavior of the load (always local). Supported values are: @@ -319,27 +334,186 @@ def put( - ".cs": Cache Streaming. Bypasses L1, streamed through L2, not retained in LLC. - ".wt": Write-Through. Bypasses L1 and L2 (coherent cache bypass), may hit in LLC with LRU. hint (int or tuple, optional): Vectorization hint passed to tl.multiple_of / tl.max_contiguous on the translated pointer. Use a scalar for 1-D (e.g. 16) or a tuple for N-D (e.g. (1, 16)). Defaults to None (no hint). + copy_engine_ctx (tl.tensor, optional): Copy engine context for SDMA operations. Required for SDMA bulk copies. + src_row_stride (int, optional): Source row stride in elements for 2D SDMA copies. Defaults to 0. + dst_row_stride (int, optional): Destination row stride in elements for 2D SDMA copies. Defaults to 0. + USE_COPY_ENGINE (tl.constexpr, optional): Whether to use SDMA copy engine. Defaults to False (uses regular load/store). + CONTIGUOUS_COPY (tl.constexpr, optional): Opt-in assertion that the masked pointer block represents one contiguous + 1D span or one rectangular 2D tile. SDMA bulk copies are only used when this is True; otherwise the function + falls back to regular load/store semantics. + from_base_ptr (triton.PointerType, optional): Base pointer of the source buffer. Required for 2D copies when USE_COPY_ENGINE is True. + to_base_ptr (triton.PointerType, optional): Base pointer of the destination buffer. Required for 2D copies when USE_COPY_ENGINE is True. Returns: None - Example: + Examples: + 1D (flat) copy: >>> @triton.jit - >>> def kernel(local_ptr, remote_ptr, heap_bases): + >>> def kernel(local_ptr, remote_ptr, heap_bases, copy_engine_ctx): >>> from_rank = 0 >>> to_rank = 1 - >>> iris.put(local_ptr, remote_ptr, from_rank, to_rank, heap_bases) + >>> offsets = tl.arange(0, 256) + >>> iris.put(local_ptr + offsets, remote_ptr + offsets, + >>> from_rank, to_rank, heap_bases, + >>> mask=offsets < 256, copy_engine_ctx=copy_engine_ctx, + >>> USE_COPY_ENGINE=True, CONTIGUOUS_COPY=True) + + 2D (tiled) copy: + >>> @triton.jit + >>> def kernel(local_ptr, remote_ptr, heap_bases, copy_engine_ctx, base_ptr): + >>> from_rank = 0 + >>> to_rank = 1 + >>> iris.put(local_ptr, remote_ptr, from_rank, to_rank, heap_bases, + >>> dst_row_stride=1024, src_row_stride=1024, + >>> mask=mask, copy_engine_ctx=copy_engine_ctx, + >>> USE_COPY_ENGINE=True, CONTIGUOUS_COPY=True, + >>> from_base_ptr=base_ptr, to_base_ptr=base_ptr) """ translated_to_ptr = __translate(to_ptr, from_rank, to_rank, heap_bases, hint) - data = tl.load(from_ptr, mask=mask, other=other, cache_modifier=load_cache_modifier) - - tl.store(translated_to_ptr, data, mask=mask, cache_modifier=store_cache_modifier) + if not USE_COPY_ENGINE or not CONTIGUOUS_COPY: + data = tl.load(from_ptr, mask=mask, other=other, cache_modifier=load_cache_modifier) + + tl.store(translated_to_ptr, data, mask=mask, cache_modifier=store_cache_modifier) + else: + ctx = copy_engine_ctx + (sdma_ep.QUEUE_DEVICE_CTX_SIZE * to_rank) + queue_ptr_u32 = tl.load(ctx + 0).to(tl.pointer_type(tl.uint32)) + read_ptr = tl.load(ctx + 1).to(tl.pointer_type(tl.uint64)) + write_ptr = tl.load(ctx + 2).to(tl.pointer_type(tl.uint64)) + doorbell_ptr = tl.load(ctx + 3).to(tl.pointer_type(tl.uint64)) + cached_write_ptr = tl.load(ctx + 4).to(tl.pointer_type(tl.uint64)) + committed_write_ptr = tl.load(ctx + 5).to(tl.pointer_type(tl.uint64)) + + dst_ptr_val = tl.min(translated_to_ptr.to(tl.uint64)) + # Extract source address (min of pointer block where data is stored) + src_ptr_u64 = from_ptr.to(tl.uint64) + src_ptr_val = tl.min(src_ptr_u64) + + # Infer element size from pointer type + # src_ptr is a block of pointers with a specific element type (e.g., pointer) + # The pointer dtype tells us the element type, which has a known size + # Map Triton dtypes to their byte sizes + ptr_dtype = from_ptr.dtype.element_ty # Get the element type that the pointer points to + + # Get element size in bytes from the dtype + # tl.float16 -> 2, tl.float32 -> 4, tl.float64 -> 8, etc. + if ptr_dtype == tl.float16 or ptr_dtype == tl.bfloat16: + element_size_bytes = 2 + elif ptr_dtype == tl.float32 or ptr_dtype == tl.int32 or ptr_dtype == tl.uint32: + element_size_bytes = 4 + elif ptr_dtype == tl.float64 or ptr_dtype == tl.int64 or ptr_dtype == tl.uint64: + element_size_bytes = 8 + elif ptr_dtype == tl.int8 or ptr_dtype == tl.uint8: + element_size_bytes = 1 + elif ptr_dtype == tl.int16 or ptr_dtype == tl.uint16: + element_size_bytes = 2 + else: + # Default to 4 bytes for unknown types + element_size_bytes = 4 + + is_2d_copy: tl.constexpr = len(from_ptr.shape) == 2 + + # Determine packet size based on copy type + # Linear copy packet: 32 bytes for 1D, Sub-window copy packet: 80 bytes for 2D + command_in_bytes = ( + sdma_ep.COPY_LINEAR_SUB_WINDOW_COMMAND_BYTES if is_2d_copy else sdma_ep.COPY_LINEAR_COMMAND_BYTES + ) + + # Acquire space in the queue + base, offset = sdma_utils.acquire_fadd( + queue_ptr_u32, + read_ptr, + write_ptr, + doorbell_ptr, + cached_write_ptr, + committed_write_ptr, + command_in_bytes, + ) + + # Write padding NOPs if we wrapped around + sdma_utils.place_nop_packet(queue_ptr_u32, base, offset) + + # Place the appropriate packet type + packet_offset_bytes = base + offset + + if not is_2d_copy: + if mask is None: + size_bytes = tl.full((), from_ptr.numel * element_size_bytes, dtype=tl.uint32) + else: + # For 1D copies, mask is 1D, so just sum all elements + mask_int = mask.to(tl.int32) + num_elements = tl.sum(mask_int, axis=0) + size_bytes = (num_elements * element_size_bytes).to(tl.uint32) + + # Place linear copy packet for 1D/flat copies + sdma_utils.place_copy_packet( + queue_ptr_u32, + packet_offset_bytes, + size_bytes, + src_ptr_val, + dst_ptr_val, + ) + else: + if mask is None: + num_elements_per_stride = tl.full((), from_ptr.shape[1], dtype=tl.uint32) + num_strides = tl.full((), from_ptr.shape[0], dtype=tl.uint32) + else: + # For 2D copies, mask is 2D [M, N], use axis operations + mask_int = mask.to(tl.int32) + num_elements_per_stride = tl.max(tl.sum(mask_int, axis=-1)) + num_strides = tl.max(tl.sum(mask_int, axis=0)) + size_bytes = (num_elements_per_stride * element_size_bytes).to(tl.uint32) + src_stride = (src_row_stride * element_size_bytes).to(tl.uint32) + dst_stride = (dst_row_stride * element_size_bytes).to(tl.uint32) + + # Place sub-window copy packet for 2D tiled copies + # Calculate base addresses and offsets for sub-window copy + src_base = from_base_ptr.to(tl.uint64) + dst_base = __translate(to_base_ptr, from_rank, to_rank, heap_bases).to(tl.uint64) + + # Calculate tile offset from base + tile_offset_bytes = src_ptr_val - src_base + src_y_val = (tile_offset_bytes // src_stride).to(tl.uint32) + src_x_val = (tile_offset_bytes % src_stride).to(tl.uint32) + + tile_offset_bytes_dst = dst_ptr_val - dst_base + dst_y_val = (tile_offset_bytes_dst // dst_stride).to(tl.uint32) + dst_x_val = (tile_offset_bytes_dst % dst_stride).to(tl.uint32) + + sdma_utils.place_sub_window_copy_packet( + queue_ptr_u32, + packet_offset_bytes, + src_base, + dst_base, + tile_width=size_bytes, + tile_height=num_strides, + src_buffer_pitch=src_stride, + dst_buffer_pitch=dst_stride, + src_x=src_x_val, + src_y=src_y_val, + dst_x=dst_x_val, + dst_y=dst_y_val, + ) + + # Submit the command to the queue + pending_wptr = base + offset + command_in_bytes + sdma_utils.submit(write_ptr, doorbell_ptr, committed_write_ptr, base, pending_wptr) @triton.jit def atomic_add( - pointer, val, from_rank, to_rank, heap_bases, mask=None, sem=None, scope=None, hint: tl.constexpr = None + pointer, + val, + from_rank, + to_rank, + heap_bases, + mask=None, + sem=None, + scope=None, + hint: tl.constexpr = None, + copy_engine_ctx=None, + USE_COPY_ENGINE: tl.constexpr = False, ): """ Performs an atomic add at the specified rank's memory location. @@ -359,6 +533,8 @@ def atomic_add( sem (str, optional): Specifies the memory semantics for the operation. Acceptable values are "acquire", "release", "acq_rel" (stands for "ACQUIRE_RELEASE"), and "relaxed". If not provided, the function defaults to using "acq_rel" semantics. scope (str, optional): Defines the scope of threads that observe the synchronizing effect of the atomic operation. Acceptable values are "gpu" (default), "cta" (cooperative thread array, thread block), or "sys" (stands for "SYSTEM"). The default value is "gpu". hint (int or tuple, optional): Vectorization hint passed to tl.multiple_of / tl.max_contiguous on the translated pointer. Defaults to None (no hint). + copy_engine_ctx (tl.tensor, optional): Copy engine context used when issuing SDMA-backed atomics. Required when ``USE_COPY_ENGINE`` is True. + USE_COPY_ENGINE (tl.constexpr, optional): Whether to route the atomic through the SDMA copy engine. Defaults to False. Returns: Block: The data stored at pointer before the atomic operation. @@ -373,7 +549,43 @@ def atomic_add( >>> old_val = iris.atomic_add(ptr, increment, cur_rank, remote_rank, heap_bases) """ translated_ptr = __translate(pointer, from_rank, to_rank, heap_bases, hint) - return tl.atomic_add(translated_ptr, val, mask=mask, sem=sem, scope=scope) + if not USE_COPY_ENGINE: + return tl.atomic_add(translated_ptr, val, mask=mask, sem=sem, scope=scope) + else: + handle = copy_engine_ctx + (sdma_ep.QUEUE_DEVICE_CTX_SIZE * to_rank) + queue_ptr_u32 = tl.load(handle + 0).to(tl.pointer_type(tl.uint32)) + read_ptr = tl.load(handle + 1).to(tl.pointer_type(tl.uint64)) + write_ptr = tl.load(handle + 2).to(tl.pointer_type(tl.uint64)) + doorbell_ptr = tl.load(handle + 3).to(tl.pointer_type(tl.uint64)) + cached_write_ptr = tl.load(handle + 4).to(tl.pointer_type(tl.uint64)) + committed_write_ptr = tl.load(handle + 5).to(tl.pointer_type(tl.uint64)) + + dst_ptr_val = translated_ptr.to(tl.uint64) + + command_in_bytes = sdma_ep.ATOMIC_COMMAND_BYTES + # Acquire space (returns base index and wraparound offset) + base, offset = sdma_utils.acquire_fadd( + # base = sdma_utils.acquire( + queue_ptr_u32, + read_ptr, + write_ptr, + doorbell_ptr, + cached_write_ptr, + committed_write_ptr, + command_in_bytes, + ) + # Write padding NOPs if we wrapped around + sdma_utils.place_nop_packet(queue_ptr_u32, base, offset) + + # Calculate packet position (base + offset for wraparound) + packet_offset_bytes = base + offset + + # Place command packet + sdma_utils.place_atomic_add_packet(queue_ptr_u32, packet_offset_bytes, dst_ptr_val, val) + + # Submit command + pending_wptr = base + offset + command_in_bytes + sdma_utils.submit(write_ptr, doorbell_ptr, committed_write_ptr, base, pending_wptr) @triton.jit @@ -416,7 +628,19 @@ def atomic_sub( @triton.jit -def atomic_cas(pointer, cmp, val, from_rank, to_rank, heap_bases, sem=None, scope=None, hint: tl.constexpr = None): +def atomic_cas( + pointer, + cmp, + val, + from_rank, + to_rank, + heap_bases, + sem=None, + scope=None, + hint: tl.constexpr = None, + copy_engine_ctx=None, + USE_COPY_ENGINE: tl.constexpr = False, +): """ Atomically compares and exchanges the specified rank's memory location. @@ -435,10 +659,17 @@ def atomic_cas(pointer, cmp, val, from_rank, to_rank, heap_bases, sem=None, scop sem (str, optional): Specifies the memory semantics for the operation. Acceptable values are "acquire", "release", "acq_rel" (stands for "ACQUIRE_RELEASE"), and "relaxed". Defaults to "acq_rel". scope (str, optional): Defines the scope of threads that observe the synchronizing effect of the atomic operation. Acceptable values are "gpu" (default), "cta" (cooperative thread array, thread block), or "sys" (stands for "SYSTEM"). Defaults to "gpu". hint (int or tuple, optional): Vectorization hint passed to tl.multiple_of / tl.max_contiguous on the translated pointer. Defaults to None (no hint). + copy_engine_ctx (tl.tensor, optional): Copy engine context used when issuing SDMA-backed atomics. Required when ``USE_COPY_ENGINE`` is True. + USE_COPY_ENGINE (tl.constexpr, optional): Whether to route the CAS through the SDMA copy engine. Defaults to False. Returns: Block: The value contained at the memory location before the atomic operation attempt. + Note: + The SDMA implementation used when ``USE_COPY_ENGINE`` is True does not + provide the previous value stored at ``pointer``. In that case this + function returns the compare operand ``cmp``. + Example: >>> @triton.jit >>> def kernel(ptr, heap_bases): @@ -450,7 +681,44 @@ def atomic_cas(pointer, cmp, val, from_rank, to_rank, heap_bases, sem=None, scop >>> old_val = iris.atomic_cas(ptr, expected, new_val, cur_rank, remote_rank, heap_bases) """ translated_ptr = __translate(pointer, from_rank, to_rank, heap_bases, hint) - return tl.atomic_cas(translated_ptr, cmp, val, sem=sem, scope=scope) + if not USE_COPY_ENGINE: + return tl.atomic_cas(translated_ptr, cmp, val, sem=sem, scope=scope) + else: + handle = copy_engine_ctx + (sdma_ep.QUEUE_DEVICE_CTX_SIZE * to_rank) + queue_ptr_u32 = tl.load(handle + 0).to(tl.pointer_type(tl.uint32)) + read_ptr = tl.load(handle + 1).to(tl.pointer_type(tl.uint64)) + write_ptr = tl.load(handle + 2).to(tl.pointer_type(tl.uint64)) + doorbell_ptr = tl.load(handle + 3).to(tl.pointer_type(tl.uint64)) + cached_write_ptr = tl.load(handle + 4).to(tl.pointer_type(tl.uint64)) + committed_write_ptr = tl.load(handle + 5).to(tl.pointer_type(tl.uint64)) + + dst_ptr_val = translated_ptr.to(tl.uint64) + + command_in_bytes = sdma_ep.ATOMIC_COMMAND_BYTES + # Acquire space (returns base index and wraparound offset) + base, offset = sdma_utils.acquire_fadd( + queue_ptr_u32, + read_ptr, + write_ptr, + doorbell_ptr, + cached_write_ptr, + committed_write_ptr, + command_in_bytes, + ) + # Write padding NOPs if we wrapped around + sdma_utils.place_nop_packet(queue_ptr_u32, base, offset) + + # Calculate packet position (base + offset for wraparound) + packet_offset_bytes = base + offset + + # Place command packet + sdma_utils.place_atomic_cas_packet(queue_ptr_u32, packet_offset_bytes, dst_ptr_val, cmp, val) + + # Submit command + pending_wptr = base + offset + command_in_bytes + sdma_utils.submit(write_ptr, doorbell_ptr, committed_write_ptr, base, pending_wptr) + + return cmp @triton.jit @@ -683,3 +951,40 @@ def atomic_max( """ translated_ptr = __translate(pointer, from_rank, to_rank, heap_bases, hint) return tl.atomic_max(translated_ptr, val, mask=mask, sem=sem, scope=scope) + + +@triton.jit +def quiet(copy_engine_ctx: tl.tensor, to_rank): + """ + Wait for all submitted SDMA operations to complete for the specified destination rank. + + Polls the hardware read pointer until it catches up to the committed write pointer, + ensuring all previously submitted SDMA packets to the destination rank have been processed. + + Args: + copy_engine_ctx: Copy engine context tensor containing queue metadata + to_rank: The destination rank whose SDMA queue should be drained + + Example: + >>> @triton.jit + >>> def kernel(src, dst, heap_bases, copy_engine_ctx): + >>> # Submit SDMA operations + >>> iris.put( + >>> src, dst, 0, 1, heap_bases, mask=mask, + >>> copy_engine_ctx=copy_engine_ctx, USE_COPY_ENGINE=True, CONTIGUOUS_COPY=True + >>> ) + >>> # Wait for completion + >>> iris.quiet(copy_engine_ctx, 1) + """ + # Extract queue pointers from context + # Context layout: [queue_buf, rptr, wptr, doorbell, cached_wptr, committed_wptr] + handle = copy_engine_ctx + (sdma_ep.QUEUE_DEVICE_CTX_SIZE * to_rank) + read_ptr = tl.load(handle + 1).to(tl.pointer_type(tl.uint64)) + committed_wptr = tl.load(handle + 5).to(tl.pointer_type(tl.uint64)) + + # Read current committed write pointer (all submitted packets from all blocks) + target = tl.load(committed_wptr, cache_modifier=".cv", volatile=True) + + # Poll until hardware read pointer catches up + while tl.load(read_ptr, cache_modifier=".cv", volatile=True) < target: + pass diff --git a/pyproject.toml b/pyproject.toml index b457fe022..8a4f04bb4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,7 +22,7 @@ dependencies = [ "requests", "ruff", "tritonblas @ git+https://github.com/ROCm/tritonBLAS.git@muhaawad/iris", - + "rocm-xio @ git+https://github.com/ROCm/rocm-xio.git", ] [project.urls] diff --git a/tests/examples/test_message_passing.py b/tests/examples/test_message_passing.py index aa173dead..cc8ce9392 100644 --- a/tests/examples/test_message_passing.py +++ b/tests/examples/test_message_passing.py @@ -12,19 +12,23 @@ current_dir = Path(__file__).parent -# Import message_passing_load_store module -load_store_file_path = (current_dir / "../../examples/06_message_passing/message_passing_load_store.py").resolve() -load_store_module_name = "message_passing_load_store" -load_store_spec = importlib.util.spec_from_file_location(load_store_module_name, load_store_file_path) -load_store_module = importlib.util.module_from_spec(load_store_spec) -load_store_spec.loader.exec_module(load_store_module) -# Import message_passing_put module -put_file_path = (current_dir / "../../examples/06_message_passing/message_passing_put.py").resolve() -put_module_name = "message_passing_put" -put_spec = importlib.util.spec_from_file_location(put_module_name, put_file_path) -put_module = importlib.util.module_from_spec(put_spec) -put_spec.loader.exec_module(put_module) +def load_example_module(relative_path: str, module_name: str): + file_path = (current_dir / relative_path).resolve() + spec = importlib.util.spec_from_file_location(module_name, file_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +# Import message passing example modules +load_store_module = load_example_module( + "../../examples/06_message_passing/message_passing_load_store.py", "message_passing_load_store" +) +put_module = load_example_module("../../examples/06_message_passing/message_passing_put.py", "message_passing_put") +host_initiated_module = load_example_module( + "../../examples/06_message_passing/message_passing_host_initiated.py", "message_passing_host_initiated" +) def create_test_args(dtype_str, buffer_size, heap_size, block_size): @@ -32,7 +36,7 @@ def create_test_args(dtype_str, buffer_size, heap_size, block_size): return {"datatype": dtype_str, "buffer_size": buffer_size, "heap_size": heap_size, "block_size": block_size} -def run_message_passing_kernels(module, args): +def run_message_passing_kernels(module, args, *, use_copy_engine: bool = False): """Run the core message passing logic without command line argument parsing.""" shmem = None try: @@ -63,9 +67,18 @@ def run_message_passing_kernels(module, args): # Allocate flags on the symmetric heap flags = shmem.zeros((num_blocks,), device="cuda", dtype=torch.int32) + copy_engine_ctx = shmem.get_copy_engine_ctx() + + producer_fn = getattr(module.producer_kernel, "fn", None) + producer_params = ( + producer_fn.__code__.co_varnames if producer_fn and hasattr(producer_fn, "__code__") else tuple() + ) + needs_copy_engine_arg = any(param in producer_params for param in ("copy_engine_handle_ptr", "copy_engine_ctx")) + has_use_copy_engine = "USE_COPY_ENGINE" in producer_params + if cur_rank == producer_rank: # Run producer kernel - module.producer_kernel[grid]( + kernel_args = [ source_buffer, destination_buffer, flags, @@ -74,7 +87,12 @@ def run_message_passing_kernels(module, args): consumer_rank, args["block_size"], shmem.get_heap_bases(), - ) + ] + if needs_copy_engine_arg: + kernel_args.append(copy_engine_ctx) + + launch_kwargs = {"USE_COPY_ENGINE": use_copy_engine} if has_use_copy_engine else {} + module.producer_kernel[grid](*kernel_args, **launch_kwargs) else: # Run consumer kernel module.consumer_kernel[grid]( @@ -107,6 +125,8 @@ def run_message_passing_kernels(module, args): import gc gc.collect() + # Clear CUDA cache to free GPU memory between tests + torch.cuda.empty_cache() @pytest.mark.parametrize( @@ -167,3 +187,93 @@ def test_message_passing_put(dtype_str, buffer_size, heap_size, block_size): args = create_test_args(dtype_str, buffer_size, heap_size, block_size) success = run_message_passing_kernels(put_module, args) assert success, "Message passing put validation failed" + + +@pytest.mark.parametrize("dtype_str", ["fp16", "fp32"]) +@pytest.mark.parametrize("buffer_size, heap_size", [(4096, 1 << 20)]) +@pytest.mark.parametrize("block_size", [512]) +def test_message_passing_copy_engine(dtype_str, buffer_size, heap_size, block_size): + """Test message passing with device-initiated copy engine.""" + args = create_test_args(dtype_str, buffer_size, heap_size, block_size) + success = run_message_passing_kernels(put_module, args, use_copy_engine=True) + assert success, "Message passing copy-engine validation failed" + + +def run_host_initiated_copy_engine(module, args): + """Execute the host-initiated message passing example logic.""" + shmem = None + try: + shmem = iris.iris(args["heap_size"]) + dtype = module.torch_dtype_from_str(args["datatype"]) + cur_rank = shmem.get_rank() + world_size = shmem.get_num_ranks() + + if world_size != args.get("num_ranks", 2): + pytest.skip("Host-initiated message passing example requires two ranks.") + + # Allocate buffers + destination_buffer = shmem.zeros(args["buffer_size"], device="cuda", dtype=dtype) + if dtype.is_floating_point: + source_buffer = shmem.randn(args["buffer_size"], device="cuda", dtype=dtype) + else: + ii = torch.iinfo(dtype) + source_buffer = shmem.randint(ii.min, ii.max, (args["buffer_size"],), device="cuda", dtype=dtype) + + producer_rank = 0 + consumer_rank = 1 + + n_elements = source_buffer.numel() + block_size = args["block_size"] + num_blocks = triton.cdiv(n_elements, block_size) + grid = (num_blocks,) + + flags = shmem.zeros((num_blocks,), device="cuda", dtype=torch.int32) + + # Use the example's producer function for the producer rank + if cur_rank == producer_rank: + module.host_initiated_producer( + shmem, source_buffer, destination_buffer, flags, consumer_rank, block_size, verbose=False + ) + else: + # Consumer uses the kernel (same as other tests) + module.consumer_kernel[grid]( + destination_buffer, flags, n_elements, consumer_rank, block_size, shmem.get_heap_bases() + ) + + shmem.barrier() + + # Validation + success = True + if cur_rank == consumer_rank: + expected = source_buffer * 2 + if not torch.allclose(destination_buffer, expected, atol=1): + success = False + + shmem.barrier() + return success + finally: + if shmem is not None: + try: + shmem.barrier() + except Exception: + pass + import gc + + del shmem + gc.collect() + + +@pytest.mark.parametrize("dtype_str", ["fp16", "fp32"]) +@pytest.mark.parametrize("buffer_size, heap_size", [(4096, 1 << 20)]) +@pytest.mark.parametrize("block_size", [512]) +def test_message_passing_host_initiated(dtype_str, buffer_size, heap_size, block_size): + """Test host-initiated copy engine example.""" + args = { + "datatype": dtype_str, + "buffer_size": buffer_size, + "heap_size": heap_size, + "block_size": block_size, + "num_ranks": 2, + } + success = run_host_initiated_copy_engine(host_initiated_module, args) + assert success, "Host-initiated message passing validation failed" diff --git a/tests/unittests/test_copy_engine_ops.py b/tests/unittests/test_copy_engine_ops.py new file mode 100644 index 000000000..bf3013caa --- /dev/null +++ b/tests/unittests/test_copy_engine_ops.py @@ -0,0 +1,757 @@ +# SPDX-License-Identifier: MIT + +import pytest +import torch +import triton +import triton.language as tl + +import iris + + +@triton.jit +def _copy_engine_linear_kernel( + src, + dst, + flag, + num_elements, + from_rank: tl.constexpr, + to_rank: tl.constexpr, + heap_bases, + copy_engine_ctx, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < num_elements + iris.put( + src + offsets, + dst + offsets, + from_rank, + to_rank, + heap_bases, + mask=mask, + copy_engine_ctx=copy_engine_ctx, + USE_COPY_ENGINE=True, + CONTIGUOUS_COPY=True, + ) + # Each block signals completion to force cache coherency on destination GPU + iris.atomic_add(flag, 1, from_rank, to_rank, heap_bases, copy_engine_ctx=copy_engine_ctx, USE_COPY_ENGINE=True) + # Wait for this block's SDMA operations to complete + iris.quiet(copy_engine_ctx, to_rank) + + +@triton.jit +def _copy_engine_linear_no_mask_kernel( + src, + dst, + flag, + from_rank: tl.constexpr, + to_rank: tl.constexpr, + heap_bases, + copy_engine_ctx, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + iris.put( + src + offsets, + dst + offsets, + from_rank, + to_rank, + heap_bases, + copy_engine_ctx=copy_engine_ctx, + USE_COPY_ENGINE=True, + CONTIGUOUS_COPY=True, + ) + iris.atomic_add(flag, 1, from_rank, to_rank, heap_bases, copy_engine_ctx=copy_engine_ctx, USE_COPY_ENGINE=True) + iris.quiet(copy_engine_ctx, to_rank) + + +def _require_two_ranks(shmem): + if shmem.get_num_ranks() != 2: + pytest.skip("Copy engine tests require exactly two ranks.") + + +def _make_expected(size, dtype, device): + return torch.arange(size, dtype=dtype, device=device) + + +def _allocate_symmetric_range(shmem, size, dtype): + # Each rank gets an identical symmetric tensor with deterministic values. + values = torch.arange(size, dtype=dtype, device=shmem.get_device()) + tensor = shmem.zeros(size, device="cuda", dtype=dtype) + tensor.copy_(values) + return tensor + + +def _make_grid(n, block): + return lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),) + + +@pytest.mark.parametrize("num_elements", [256, 1024]) +def test_copy_engine_device_linear_put(num_elements): + shmem = iris.iris(1 << 20) + _require_two_ranks(shmem) + + rank = shmem.get_rank() + remote_rank = 1 - rank + + src = _allocate_symmetric_range(shmem, num_elements, torch.float32) + dst = shmem.zeros(num_elements, device="cuda", dtype=torch.float32) + completion_flag = shmem.zeros(1, device="cuda", dtype=torch.int32) + + grid = _make_grid(num_elements, 128) + if rank == 0: + _copy_engine_linear_kernel[grid]( + src, + dst, + completion_flag, + num_elements, + rank, + remote_rank, + shmem.get_heap_bases(), + shmem.get_copy_engine_ctx(), + BLOCK_SIZE=128, + ) + + shmem.barrier() + + if rank == 1: + # Verify all blocks completed (one signal per block) + num_blocks = triton.cdiv(num_elements, 128) + assert completion_flag.item() == num_blocks, f"Expected {num_blocks} signals, got {completion_flag.item()}" + expected = _make_expected(num_elements, torch.float32, dst.device) + assert torch.allclose(dst, expected) + + shmem.barrier() + del shmem + + +def test_copy_engine_device_linear_put_no_mask(): + num_elements = 512 + block_size = 128 + shmem = iris.iris(1 << 20) + _require_two_ranks(shmem) + + rank = shmem.get_rank() + remote_rank = 1 - rank + + src = _allocate_symmetric_range(shmem, num_elements, torch.float32) + dst = shmem.zeros(num_elements, device="cuda", dtype=torch.float32) + completion_flag = shmem.zeros(1, device="cuda", dtype=torch.int32) + + if rank == 0: + _copy_engine_linear_no_mask_kernel[(num_elements // block_size,)]( + src, + dst, + completion_flag, + rank, + remote_rank, + shmem.get_heap_bases(), + shmem.get_copy_engine_ctx(), + BLOCK_SIZE=block_size, + ) + + shmem.barrier() + + if rank == 1: + num_blocks = num_elements // block_size + assert completion_flag.item() == num_blocks, f"Expected {num_blocks} signals, got {completion_flag.item()}" + expected = _make_expected(num_elements, torch.float32, dst.device) + assert torch.allclose(dst, expected) + + shmem.barrier() + del shmem + + +@pytest.mark.parametrize("num_elements", [512, 2048]) +def test_copy_engine_host_put(num_elements): + shmem = iris.iris(1 << 20) + _require_two_ranks(shmem) + + rank = shmem.get_rank() + remote_rank = 1 - rank + + src = _allocate_symmetric_range(shmem, num_elements, torch.float32) + dst = shmem.zeros(num_elements, device="cuda", dtype=torch.float32) + completion_flag = shmem.zeros(1, device="cuda", dtype=torch.int32) + + if rank == 0: + shmem.put(src, dst_rank=remote_rank, dst_tensor=dst, signal_flag=completion_flag, signal_value=1, async_op=True) + shmem.quiet(dst_rank=remote_rank) + + shmem.barrier() + + if rank == 1: + expected = _make_expected(num_elements, torch.float32, dst.device) + assert torch.allclose(dst, expected) + + shmem.barrier() + del shmem + + +@triton.jit +def _copy_engine_atomic_kernel( + flag, + from_rank: tl.constexpr, + to_rank: tl.constexpr, + heap_bases, + copy_engine_ctx, + increment: tl.constexpr, +): + iris.atomic_add( + flag, + increment, + from_rank, + to_rank, + heap_bases, + copy_engine_ctx=copy_engine_ctx, + USE_COPY_ENGINE=True, + ) + + +def test_copy_engine_atomic_add(): + shmem = iris.iris(1 << 20) + _require_two_ranks(shmem) + + rank = shmem.get_rank() + remote_rank = 1 - rank + + flag = shmem.zeros((1,), device="cuda", dtype=torch.int32) + + if rank == 0: + _copy_engine_atomic_kernel[(1,)]( + flag, + rank, + remote_rank, + shmem.get_heap_bases(), + shmem.get_copy_engine_ctx(), + increment=5, + ) + + shmem.barrier() + + if rank == 1: + assert flag.item() == 5 + + shmem.barrier() + del shmem + + +# ============================================================================ +# Copy Engine Atomic CAS Tests +# ============================================================================ + + +@triton.jit +def _copy_engine_atomic_cas_kernel( + flag, + from_rank: tl.constexpr, + to_rank: tl.constexpr, + heap_bases, + copy_engine_ctx, + compare: tl.constexpr, + value: tl.constexpr, +): + iris.atomic_cas( + flag, + compare, + value, + from_rank, + to_rank, + heap_bases, + copy_engine_ctx=copy_engine_ctx, + USE_COPY_ENGINE=True, + ) + + +def test_copy_engine_atomic_cas(): + shmem = iris.iris(1 << 20) + _require_two_ranks(shmem) + + rank = shmem.get_rank() + remote_rank = 1 - rank + + flag = shmem.zeros((1,), device="cuda", dtype=torch.int32) + + if rank == 0: + _copy_engine_atomic_cas_kernel[(1,)]( + flag, + rank, + remote_rank, + shmem.get_heap_bases(), + shmem.get_copy_engine_ctx(), + compare=0, + value=1, + ) + + shmem.barrier() + + if rank == 1: + assert flag.item() == 1 + + shmem.barrier() + del shmem + + +# ============================================================================ +# 2D/Tiled Copy Tests +# ============================================================================ + + +@triton.jit +def _copy_engine_2d_kernel( + src_base, + dst_base, + num_rows: tl.constexpr, + num_cols: tl.constexpr, + src_stride: tl.constexpr, + dst_stride: tl.constexpr, + from_rank: tl.constexpr, + to_rank: tl.constexpr, + heap_bases, + copy_engine_ctx, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """2D tiled copy using strided parameters.""" + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + # Calculate offsets for this tile + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + + # Create 2D pointer blocks + src_ptrs = src_base + offs_m[:, None] * src_stride + offs_n[None, :] + dst_ptrs = dst_base + offs_m[:, None] * dst_stride + offs_n[None, :] + + # Create mask + mask = (offs_m[:, None] < num_rows) & (offs_n[None, :] < num_cols) + + # 2D copy with strides + iris.put( + src_ptrs, + dst_ptrs, + from_rank, + to_rank, + heap_bases, + src_row_stride=src_stride, + dst_row_stride=dst_stride, + mask=mask, + copy_engine_ctx=copy_engine_ctx, + USE_COPY_ENGINE=True, + CONTIGUOUS_COPY=True, + from_base_ptr=src_base, + to_base_ptr=dst_base, + ) + + +@triton.jit +def _copy_engine_2d_no_mask_kernel( + src_base, + dst_base, + src_stride: tl.constexpr, + dst_stride: tl.constexpr, + from_rank: tl.constexpr, + to_rank: tl.constexpr, + heap_bases, + copy_engine_ctx, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + + src_ptrs = src_base + offs_m[:, None] * src_stride + offs_n[None, :] + dst_ptrs = dst_base + offs_m[:, None] * dst_stride + offs_n[None, :] + + iris.put( + src_ptrs, + dst_ptrs, + from_rank, + to_rank, + heap_bases, + src_row_stride=src_stride, + dst_row_stride=dst_stride, + copy_engine_ctx=copy_engine_ctx, + USE_COPY_ENGINE=True, + CONTIGUOUS_COPY=True, + from_base_ptr=src_base, + to_base_ptr=dst_base, + ) + + +@pytest.mark.parametrize("M,N", [(16, 16), (32, 64)]) +def test_copy_engine_2d_tiled(M, N): + """Test 2D tiled copy with strides.""" + shmem = iris.iris(1 << 20) + _require_two_ranks(shmem) + + rank = shmem.get_rank() + remote_rank = 1 - rank + + # Allocate 2D matrices with row-major layout + stride = N # Elements per row + src = _allocate_symmetric_range(shmem, M * N, torch.float32).view(M, N) + dst = shmem.zeros(M * N, device="cuda", dtype=torch.float32).view(M, N) + + BLOCK_M, BLOCK_N = 8, 16 + grid_m = triton.cdiv(M, BLOCK_M) + grid_n = triton.cdiv(N, BLOCK_N) + + if rank == 0: + _copy_engine_2d_kernel[(grid_m, grid_n)]( + src, + dst, + M, + N, + stride, + stride, + rank, + remote_rank, + shmem.get_heap_bases(), + shmem.get_copy_engine_ctx(), + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + ) + + shmem.barrier() + + if rank == 1: + expected = _make_expected(M * N, torch.float32, dst.device).view(M, N) + assert torch.allclose(dst, expected) + + shmem.barrier() + del shmem + + +def test_copy_engine_2d_tiled_no_mask(): + """Test full-tile 2D copy using shape-inferred SDMA packet size.""" + shmem = iris.iris(1 << 20) + _require_two_ranks(shmem) + + rank = shmem.get_rank() + remote_rank = 1 - rank + + M, N = 16, 32 + BLOCK_M, BLOCK_N = 8, 16 + stride = N + + src = _allocate_symmetric_range(shmem, M * N, torch.float32).view(M, N) + dst = shmem.zeros(M * N, device="cuda", dtype=torch.float32).view(M, N) + + if rank == 0: + _copy_engine_2d_no_mask_kernel[(M // BLOCK_M, N // BLOCK_N)]( + src, + dst, + stride, + stride, + rank, + remote_rank, + shmem.get_heap_bases(), + shmem.get_copy_engine_ctx(), + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + ) + + shmem.barrier() + + if rank == 1: + expected = _make_expected(M * N, torch.float32, dst.device).view(M, N) + assert torch.allclose(dst, expected) + + shmem.barrier() + del shmem + + +# ============================================================================ +# Combined Operations Tests (put + signal, wait + put) +# ============================================================================ + + +@triton.jit +def _copy_engine_put_signal_kernel( + src, + dst, + flag, + num_elements, + from_rank: tl.constexpr, + to_rank: tl.constexpr, + heap_bases, + copy_engine_ctx, + BLOCK_SIZE: tl.constexpr, +): + """Copy data and signal completion with atomic add.""" + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < num_elements + + # Copy data + iris.put( + src + offsets, + dst + offsets, + from_rank, + to_rank, + heap_bases, + mask=mask, + copy_engine_ctx=copy_engine_ctx, + USE_COPY_ENGINE=True, + CONTIGUOUS_COPY=True, + ) + + # Signal completion (last thread in block) + if pid == 0: + iris.atomic_add( + flag, + 1, + from_rank, + to_rank, + heap_bases, + copy_engine_ctx=copy_engine_ctx, + USE_COPY_ENGINE=True, + ) + + +def test_copy_engine_put_with_signal(): + """Test copy followed by atomic signal.""" + shmem = iris.iris(1 << 20) + _require_two_ranks(shmem) + + rank = shmem.get_rank() + remote_rank = 1 - rank + + num_elements = 512 + src = _allocate_symmetric_range(shmem, num_elements, torch.float32) + dst = shmem.zeros(num_elements, device="cuda", dtype=torch.float32) + flag = shmem.zeros((1,), device="cuda", dtype=torch.int32) + + grid = _make_grid(num_elements, 128) + if rank == 0: + _copy_engine_put_signal_kernel[grid]( + src, + dst, + flag, + num_elements, + rank, + remote_rank, + shmem.get_heap_bases(), + shmem.get_copy_engine_ctx(), + BLOCK_SIZE=128, + ) + + shmem.barrier() + + if rank == 1: + # Check data transferred + expected = _make_expected(num_elements, torch.float32, dst.device) + assert torch.allclose(dst, expected) + # Check signal received + assert flag.item() == 1 + + shmem.barrier() + del shmem + + +# ============================================================================ +# Multi-Block Concurrent Operations +# ============================================================================ + + +@triton.jit +def _copy_engine_multi_block_kernel( + src, + dst, + counters, + num_elements, + from_rank: tl.constexpr, + to_rank: tl.constexpr, + heap_bases, + copy_engine_ctx, + BLOCK_SIZE: tl.constexpr, +): + """Multiple blocks concurrently using copy engine.""" + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < num_elements + + # Each block copies its chunk + iris.put( + src + offsets, + dst + offsets, + from_rank, + to_rank, + heap_bases, + mask=mask, + copy_engine_ctx=copy_engine_ctx, + USE_COPY_ENGINE=True, + CONTIGUOUS_COPY=True, + ) + + # Each block atomically increments its counter + iris.atomic_add( + counters + pid, + 1, + from_rank, + to_rank, + heap_bases, + copy_engine_ctx=copy_engine_ctx, + USE_COPY_ENGINE=True, + ) + + +@pytest.mark.parametrize("num_blocks", [4, 8]) +def test_copy_engine_multi_block_concurrent(num_blocks): + """Test multiple workgroups using copy engine concurrently.""" + shmem = iris.iris(1 << 20) + _require_two_ranks(shmem) + + rank = shmem.get_rank() + remote_rank = 1 - rank + + BLOCK_SIZE = 128 + num_elements = num_blocks * BLOCK_SIZE + src = _allocate_symmetric_range(shmem, num_elements, torch.float32) + dst = shmem.zeros(num_elements, device="cuda", dtype=torch.float32) + counters = shmem.zeros(num_blocks, device="cuda", dtype=torch.int32) + + if rank == 0: + _copy_engine_multi_block_kernel[(num_blocks,)]( + src, + dst, + counters, + num_elements, + rank, + remote_rank, + shmem.get_heap_bases(), + shmem.get_copy_engine_ctx(), + BLOCK_SIZE=BLOCK_SIZE, + ) + + shmem.barrier() + + if rank == 1: + # Check all data transferred + expected = _make_expected(num_elements, torch.float32, dst.device) + assert torch.allclose(dst, expected) + # Check all blocks signaled + assert torch.all(counters == 1).item() + + shmem.barrier() + del shmem + + +# ============================================================================ +# Edge Cases and Error Conditions +# ============================================================================ + + +def test_copy_engine_zero_size(): + """Test copy engine with zero-size transfer (should be no-op).""" + shmem = iris.iris(1 << 20) + _require_two_ranks(shmem) + + rank = shmem.get_rank() + remote_rank = 1 - rank + + src = shmem.zeros(128, device="cuda", dtype=torch.float32) + dst = shmem.zeros(128, device="cuda", dtype=torch.float32) + + if rank == 0: + # Empty slice should be a no-op + shmem.put(src[:0], dst_rank=remote_rank, dst_tensor=dst[:0], async_op=True) + shmem.quiet(dst_rank=remote_rank) + + shmem.barrier() + + # Destination should still be zeros + if rank == 1: + assert torch.all(dst == 0).item() + + shmem.barrier() + del shmem + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.float32, torch.int32]) +def test_copy_engine_different_dtypes(dtype): + """Test copy engine with different data types.""" + shmem = iris.iris(1 << 20) + _require_two_ranks(shmem) + + rank = shmem.get_rank() + remote_rank = 1 - rank + + num_elements = 256 + src = _allocate_symmetric_range(shmem, num_elements, dtype) + dst = shmem.zeros(num_elements, device="cuda", dtype=dtype) + completion_flag = shmem.zeros(1, device="cuda", dtype=torch.int32) + + grid = _make_grid(num_elements, 128) + if rank == 0: + _copy_engine_linear_kernel[grid]( + src, + dst, + completion_flag, + num_elements, + rank, + remote_rank, + shmem.get_heap_bases(), + shmem.get_copy_engine_ctx(), + BLOCK_SIZE=128, + ) + + shmem.barrier() + + if rank == 1: + num_blocks = triton.cdiv(num_elements, 128) + assert completion_flag.item() == num_blocks, f"Expected {num_blocks} signals, got {completion_flag.item()}" + expected = _make_expected(num_elements, dtype, dst.device) + assert torch.allclose(dst, expected) + + shmem.barrier() + del shmem + + +def test_copy_engine_bidirectional(): + """Test both ranks doing copy engine operations simultaneously.""" + shmem = iris.iris(1 << 20) + _require_two_ranks(shmem) + + rank = shmem.get_rank() + remote_rank = 1 - rank + + num_elements = 256 + # Each rank has its own data to send + src = _allocate_symmetric_range(shmem, num_elements, torch.float32) + # Scale by rank to make data different + src.mul_(rank + 1) + dst = shmem.zeros(num_elements, device="cuda", dtype=torch.float32) + completion_flag = shmem.zeros(1, device="cuda", dtype=torch.int32) + + grid = _make_grid(num_elements, 128) + # Both ranks send their data + _copy_engine_linear_kernel[grid]( + src, + dst, + completion_flag, + num_elements, + rank, + remote_rank, + shmem.get_heap_bases(), + shmem.get_copy_engine_ctx(), + BLOCK_SIZE=128, + ) + + shmem.barrier() + + # Each rank should have received the other's data + num_blocks = triton.cdiv(num_elements, 128) + assert completion_flag.item() == num_blocks, f"Expected {num_blocks} signals, got {completion_flag.item()}" + expected = _make_expected(num_elements, torch.float32, dst.device) * (remote_rank + 1) + assert torch.allclose(dst, expected) + + shmem.barrier() + del shmem