diff --git a/docs/reference/host/class.md b/docs/reference/host/class.md index 87a0187bb..a0106ca9c 100644 --- a/docs/reference/host/class.md +++ b/docs/reference/host/class.md @@ -8,6 +8,28 @@ Prefer using the convenience factory over calling the constructor directly: .. autofunction:: iris.host.iris.iris ``` +## Optional SDMA Copy Engine + +The `copy_engine` parameter defaults to `None`. In this mode, Iris skips SDMA +initialization and queue creation, and `get_copy_engine_ctx()` returns `None`. +Regular Triton and Gluon memory operations and collectives work without SDMA. + +Pass `copy_engine="sdma"` to initialize both host- and device-initiated SDMA queues: + +```python +ctx = iris.iris(heap_size=1 << 30, copy_engine="sdma") +copy_engine_ctx = ctx.get_copy_engine_ctx() +``` + +This is required for the host methods `put`, `put_tile`, `put_tiles`, and `quiet`, +and for device operations using `use_copy_engine=True`. Host SDMA methods raise +`RuntimeError` if the copy engine is disabled. Pass the returned context to +device operations through their `copy_engine_ctx` argument. + +```{eval-rst} +.. automethod:: iris.host.iris.Iris.get_copy_engine_ctx +``` + ## Core Methods ```{eval-rst} diff --git a/examples/06_message_passing/message_passing_host_initiated.py b/examples/06_message_passing/message_passing_host_initiated.py index 51b59fd5b..29be2afd9 100644 --- a/examples/06_message_passing/message_passing_host_initiated.py +++ b/examples/06_message_passing/message_passing_host_initiated.py @@ -174,7 +174,7 @@ def _worker(local_rank: int, world_size: int, init_url: str, args: dict): ) # Main benchmark logic - shmem = iris.iris(args["heap_size"]) + shmem = iris.iris(args["heap_size"], copy_engine="sdma") dtype = torch_dtype_from_str(args["datatype"]) cur_rank = shmem.get_rank() world_size = shmem.get_num_ranks() diff --git a/examples/06_message_passing/message_passing_put.py b/examples/06_message_passing/message_passing_put.py index 6cad6928f..8c4855db7 100755 --- a/examples/06_message_passing/message_passing_put.py +++ b/examples/06_message_passing/message_passing_put.py @@ -157,7 +157,7 @@ def _worker(local_rank: int, world_size: int, init_url: str, args: dict): ) # Main benchmark logic - shmem = iris.iris(args["heap_size"]) + shmem = iris.iris(args["heap_size"], copy_engine="sdma" if args["use_copy_engine"] else None) dtype = torch_dtype_from_str(args["datatype"]) cur_rank = shmem.get_rank() world_size = shmem.get_num_ranks() diff --git a/iris/host/iris.py b/iris/host/iris.py index aa1674334..3c07fd89c 100644 --- a/iris/host/iris.py +++ b/iris/host/iris.py @@ -83,6 +83,8 @@ class Iris: Args: heap_size (int): Size of the symmetric heap in bytes. Default: 1GB (2^30) allocator_type (str): Type of allocator to use. Options: "torch" (default), "vmem" + copy_engine (str, optional): Copy engine to initialize. Use "sdma" for host- and + device-initiated SDMA operations. Defaults to None (disabled). Example: >>> ctx = iris.iris(heap_size=2**31) # 2GB heap with torch allocator @@ -93,7 +95,10 @@ class Iris: >>> ctx = iris.iris(heap_size=2**31, allocator_type="vmem") """ - def __init__(self, heap_size=1 << 30, allocator_type="torch"): + def __init__(self, heap_size=1 << 30, allocator_type="torch", copy_engine=None): + if copy_engine not in (None, "sdma"): + raise ValueError(f"Unsupported copy engine: {copy_engine!r}. Expected None or 'sdma'.") + # Initialize distributed environment comm, cur_rank, num_ranks = init_distributed() num_gpus = count_devices() @@ -137,36 +142,40 @@ 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 + self.copy_engines_device_ctx = None + if copy_engine == "sdma": + 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) @@ -950,9 +959,14 @@ def get_heap_bases(self): """ return self.heap_bases - def get_copy_engine_ctx(self): + def get_copy_engine_ctx(self) -> torch.Tensor | None: + """Return the device copy-engine context, or None when copy engines are disabled.""" return self.copy_engines_device_ctx + def _require_copy_engine(self): + if self.copy_engines_device_ctx is None: + raise RuntimeError('SDMA is disabled; initialize Iris with copy_engine="sdma".') + @staticmethod def _dtype_to_flag_bits(dtype: torch.dtype) -> int: if dtype in (torch.int32, torch.int): @@ -1003,6 +1017,8 @@ def put( """ One-sided put operation with optional wait (POLL) and signal (ATOMIC). + Requires initialization with ``copy_engine="sdma"``. + Supports: - Simple copy: put(src, to_rank) - Copy + signal: put(src, to_rank, signal_flag=flag) @@ -1035,6 +1051,7 @@ def put( ... wait_flag=batch_ready, wait_value=256, ... signal_flag=transfer_done, signal_value=1) """ + self._require_copy_engine() if to_tensor is None: to_tensor = from_tensor @@ -1094,6 +1111,8 @@ def put_tile( """ 2D tile transfer with optional wait/signal (sub-window copy). + Requires initialization with ``copy_engine="sdma"``. + Low-level API - caller provides pre-translated pointers for performance. Args: @@ -1125,6 +1144,7 @@ def put_tile( >>> shmem.put_tile(tile, to_rank=1, to_ptr=dst_ptr, to_stride=dst_stride, ... wait_flag=wait_ptr, wait_value=256, signal_flag=signal_ptr) """ + self._require_copy_engine() from_rank = self.get_rank() wait_ptr, wait_bits = self._flag_pointer_and_bits(wait_flag, default_bits=32) @@ -1176,6 +1196,8 @@ def put_tiles( """ Batched 2D tile transfer with optional shared wait/signal. + Requires initialization with ``copy_engine="sdma"``. + Args: tiles: Sequence of pre-configured sdma_ep.Tile objects to_rank: Destination rank @@ -1188,6 +1210,7 @@ def put_tiles( async_op: If True, don't wait for completion channel: SDMA channel to use """ + self._require_copy_engine() from_rank = self.get_rank() if len(tiles) != len(to_ptrs) or len(tiles) != len(to_strides): @@ -1232,6 +1255,8 @@ def quiet(self, to_rank: int = None, channel: int = 0): """ Wait for all outstanding SDMA operations to complete. + Requires initialization with ``copy_engine="sdma"``. + Args: to_rank: If specified, wait only for ops to this rank. If None, wait for ops to all ranks. @@ -1242,6 +1267,7 @@ def quiet(self, to_rank: int = None, channel: int = 0): >>> shmem.quiet(to_rank=1) # Wait for completion >>> shmem.quiet() # Wait for all ranks """ + self._require_copy_engine() src_rank = self.get_rank() if to_rank is not None: sdma_ep.quiet(src_rank, to_rank, channel) @@ -1698,7 +1724,7 @@ def reduce_scatter(self, output_tensor, input_tensor, op=None, group=None, async ) -def iris(heap_size=1 << 30, allocator_type="torch"): +def iris(heap_size=1 << 30, allocator_type="torch", copy_engine=None): """ Create and return an Iris instance with the specified heap size. @@ -1706,6 +1732,8 @@ def iris(heap_size=1 << 30, allocator_type="torch"): heap_size (int): Size of the heap in bytes. Defaults to 1GB. allocator_type (str): Type of allocator to use. Options: "torch" (default), "vmem". Can be overridden with IRIS_ALLOCATOR environment variable. + copy_engine (str, optional): Copy engine to initialize. Use "sdma" for host- and + device-initiated SDMA operations. Defaults to None (disabled). Returns: Iris: An initialized Iris instance. @@ -1719,4 +1747,4 @@ def iris(heap_size=1 << 30, allocator_type="torch"): >>> iris_ctx = iris.iris(2**30, allocator_type="vmem") >>> tensor = iris_ctx.zeros(1024, 1024) """ - return Iris(heap_size, allocator_type) + return Iris(heap_size, allocator_type, copy_engine=copy_engine) diff --git a/tests/examples/test_message_passing.py b/tests/examples/test_message_passing.py index 6df7d6886..6c7e38c48 100644 --- a/tests/examples/test_message_passing.py +++ b/tests/examples/test_message_passing.py @@ -40,7 +40,7 @@ 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: - shmem = iris.iris(args["heap_size"]) + shmem = iris.iris(args["heap_size"], copy_engine="sdma" if use_copy_engine else None) dtype = module.torch_dtype_from_str(args["datatype"]) cur_rank = shmem.get_rank() world_size = shmem.get_num_ranks() @@ -203,7 +203,7 @@ 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"]) + shmem = iris.iris(args["heap_size"], copy_engine="sdma") dtype = module.torch_dtype_from_str(args["datatype"]) cur_rank = shmem.get_rank() world_size = shmem.get_num_ranks() diff --git a/tests/unittests/test_copy_engine_ops.py b/tests/unittests/test_copy_engine_ops.py index c2c4220e7..83e65dfbc 100644 --- a/tests/unittests/test_copy_engine_ops.py +++ b/tests/unittests/test_copy_engine_ops.py @@ -90,7 +90,7 @@ def _make_grid(n, block): @pytest.mark.parametrize("num_elements", [256, 1024]) def test_copy_engine_device_linear_put(num_elements): - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -130,7 +130,7 @@ def test_copy_engine_device_linear_put(num_elements): def test_copy_engine_device_linear_put_no_mask(): num_elements = 512 block_size = 128 - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -166,7 +166,7 @@ def test_copy_engine_device_linear_put_no_mask(): @pytest.mark.parametrize("num_elements", [512, 2048]) def test_copy_engine_host_put(num_elements): - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -192,7 +192,7 @@ def test_copy_engine_host_put(num_elements): def test_copy_engine_host_put_loopback(): """Test a host-initiated copy-engine put when source and destination are the same rank.""" - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, copy_engine="sdma") num_elements = 512 src = _allocate_symmetric_range(shmem, num_elements, torch.float32) @@ -218,7 +218,7 @@ def test_copy_engine_host_put_loopback(): def test_copy_engine_device_linear_put_loopback(): """Test a device-initiated copy-engine put when source and destination share one rank.""" - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, copy_engine="sdma") rank = shmem.get_rank() num_elements = 256 @@ -271,7 +271,7 @@ def _copy_engine_atomic_kernel( # For int64, adding to 0xFFFFFFFF carries into the upper dword only if the SDMA add is 64-bit. @pytest.mark.parametrize("dtype, initial", [(torch.int32, 0), (torch.int64, 0xFFFFFFFF)]) def test_copy_engine_atomic_add(dtype, initial): - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -327,7 +327,7 @@ def _copy_engine_atomic_cas_kernel( def test_copy_engine_atomic_cas(): - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -449,7 +449,7 @@ def _copy_engine_2d_no_mask_kernel( @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) + shmem = iris.iris(1 << 20, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -492,7 +492,7 @@ def test_copy_engine_2d_tiled(M, N): def test_copy_engine_2d_tiled_no_mask(): """Test full-tile 2D copy using shape-inferred SDMA packet size.""" - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -579,7 +579,7 @@ def _copy_engine_put_signal_kernel( def test_copy_engine_put_with_signal(): """Test copy followed by atomic signal.""" - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -667,7 +667,7 @@ def _copy_engine_multi_block_kernel( @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) + shmem = iris.iris(1 << 20, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -712,7 +712,7 @@ def test_copy_engine_multi_block_concurrent(num_blocks): def test_copy_engine_zero_size(): """Test copy engine with zero-size transfer (should be no-op).""" - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -738,7 +738,7 @@ def test_copy_engine_zero_size(): def test_copy_engine_device_zero_size(): """Test device-side copy engine puts with fully-masked tiles (should be no-op).""" - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -792,7 +792,7 @@ def test_copy_engine_device_zero_size(): @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) + shmem = iris.iris(1 << 20, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() @@ -831,7 +831,7 @@ def test_copy_engine_different_dtypes(dtype): def test_copy_engine_bidirectional(): """Test both ranks doing copy engine operations simultaneously.""" - shmem = iris.iris(1 << 20) + shmem = iris.iris(1 << 20, copy_engine="sdma") _require_two_ranks(shmem) rank = shmem.get_rank() diff --git a/tests/unittests/test_device_context.py b/tests/unittests/test_device_context.py index 874afa102..5c628da31 100644 --- a/tests/unittests/test_device_context.py +++ b/tests/unittests/test_device_context.py @@ -478,6 +478,7 @@ def test_device_context_tracing_1d_address(): def test_device_context_initialize(): """Test DeviceContext.initialize() creates valid context.""" ctx = iris.iris(1 << 20) + assert ctx.get_copy_engine_ctx() is None cur_rank = ctx.get_rank() num_ranks = ctx.get_num_ranks()