Make copy engine initialization optional

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}

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()

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()

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,14 +1724,16 @@ 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.
 
     Args:
         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)

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()

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()