Fix the published API reference and document the missing device APIs

.gitignore

@@ -69,3 +69,7 @@ docker/rocm-systems/
 .intellikit
 .github/agents/skills/
 docs/benchmark-results/*.png
+
+# Local Sphinx build
+docs/.venv/
+docs/_build/

docs/conf.py

@@ -11,6 +11,7 @@
 # documentation root, use os.path.abspath to make it absolute, like shown here.
 #
 import os
+import shutil
 import sys
 from unittest.mock import MagicMock
 
@@ -32,6 +33,11 @@
 # ones.
 extensions = [
     "rocm_docs",
+    # Sphinx removed its bundled jQuery in 6.0, but four rocm-docs-core theme
+    # scripts (code_word_breaks, renameVersionLinks, rdcMisc,
+    # theme_mode_captions) still call $(). Without this every page throws
+    # "$ is not defined" and those scripts silently do nothing.
+    "sphinxcontrib.jquery",
     "sphinx.ext.autodoc",
     "sphinx.ext.napoleon",
     "sphinx.ext.viewcode",
@@ -96,6 +102,11 @@
     "iris.host.distributed.helpers",
     "iris.host.platform.hip",
     "tritonblas",
+    # rocm-xio. iris/host/iris.py imports it unconditionally at module load, so
+    # without this every autodoc directive fails with ModuleNotFoundError and the
+    # whole API reference renders empty. The docs CI installs only
+    # docs/sphinx/requirements.txt, so xio is never present there either.
+    "xio",
 ]
 
 # Custom mocks that preserve docstrings for Triton Gluon
@@ -115,6 +126,26 @@ def __call__(self, func=None, **kwargs):
 
 # Mock triton.language first
 triton_language_mock = MagicMock()
+
+
+# These tl.* names are used as parameter ANNOTATIONS throughout the device APIs,
+# so autodoc resolves them and prints their repr into the rendered signature. A
+# bare MagicMock attribute reprs as
+#     <MagicMock name='mock.constexpr' id='139779281861264'>
+# which leaks into the parameter list, and the id changes on every build so the
+# HTML churns too. Giving each one a real class makes it repr as its own name.
+# The list is every `: tl.<name>` annotation appearing in iris/ -- if a new one
+# is introduced it will show up as a MagicMock on the rendered page.
+def _name_mock_types(mock, module, names):
+    """Give a mocked module real classes for the names used as annotations."""
+    for n in names:
+        cls = type(n, (), {"__doc__": "Triton %s." % n})
+        cls.__module__ = module
+        setattr(mock, n, cls)
+
+
+_name_mock_types(triton_language_mock, "triton.language", ("constexpr", "tensor"))
+
 sys.modules["triton.language"] = triton_language_mock
 sys.modules["triton.language.core"] = MagicMock()
 sys.modules["triton.language.core"]._aggregate = lambda cls: cls  # Preserve class
@@ -142,7 +173,13 @@ class GluonMock:
 sys.modules["triton.experimental"] = MagicMock()
 sys.modules["triton.experimental"].gluon = GluonMock()
 sys.modules["triton.experimental.gluon"] = GluonMock()
-sys.modules["triton.experimental.gluon"].language = MagicMock()
+_gluon_language_mock = MagicMock()
+# Gluon annotates with gl.constexpr, which resolves through a *different* mock
+# module than tl.constexpr -- so it needs the same treatment or it reprs as a
+# MagicMock in the rendered signature.
+_name_mock_types(_gluon_language_mock, "triton.experimental.gluon.language", ("constexpr", "tensor"))
+sys.modules["triton.experimental.gluon"].language = _gluon_language_mock
+sys.modules["triton.experimental.gluon.language"] = _gluon_language_mock
 
 # Napoleon settings for Google/NumPy docstring parsing
 napoleon_google_docstring = True
@@ -164,6 +201,7 @@ class GluonMock:
 # Table of contents
 external_toc_path = "./sphinx/_toc.yml"
 
+
 # Theme options for AMD ROCm theme
 html_theme_options = {
     "flavor": "instinct",
@@ -185,3 +223,29 @@ class GluonMock:
     "copybutton_hide": copybutton_hide,
     "copybutton_remove_prompts": copybutton_remove_prompts,
 }
+
+# ---------------------------------------------------------------------------
+# Copy docs/images/ verbatim into the built site.
+#
+# Several conceptual pages reference images from raw HTML blocks, e.g.
+#   <img class="dark-theme-img only-dark" src="../images/iris-model.png">
+# which is how the light/dark theme variants are selected. Sphinx only tracks
+# and copies images referenced through markdown or image directives, so every
+# raw-HTML reference 404s on the published site -- 18 of them, across
+# conceptual/programming-model, conceptual/taxonomy and index.
+#
+# Copying the directory wholesale fixes all of them without rewriting the
+# theme-switching markup.
+# ---------------------------------------------------------------------------
+
+
+def _copy_images(app, exception):
+    if exception is not None or app.builder.name != "html":
+        return
+    src = os.path.join(os.path.dirname(__file__), "images")
+    if os.path.isdir(src):
+        shutil.copytree(src, os.path.join(app.outdir, "images"), dirs_exist_ok=True)
+
+
+def setup(app):
+    app.connect("build-finished", _copy_images)

docs/reference/api-reference.md

@@ -7,12 +7,13 @@ Explore Iris APIs. The host-side API (class, tensor creation, CCL) is identical
 - [Iris Class](host/class.md)
 - [Tensor Creation](host/tensor-creation.md)
 - [Collective Communication (CCL)](host/ccl.md)
+- [Copy Engine (SDMA)](host/copy-engine.md)
 
 ## Triton Backend
 
 - [Overview](triton/overview.md)
 - [Device Functions](triton/device-functions.md)
-- [Fused GEMM + CCL Operations](triton/ops.md)
+- [Fused Communication and Compute Operations](triton/ops.md)
 
 ## Gluon Backend (Experimental)
 

docs/reference/host/class.md

@@ -68,3 +68,27 @@ Broadcast data from a source rank to all ranks. This method automatically detect
 ```{eval-rst}
 .. automethod:: iris.host.iris.Iris.broadcast
 ```
+
+## allocate_symmetric
+
+```{eval-rst}
+.. automethod:: iris.host.iris.Iris.allocate_symmetric
+```
+
+## device_barrier
+
+```{eval-rst}
+.. automethod:: iris.host.iris.Iris.device_barrier
+```
+
+## get_device_id
+
+```{eval-rst}
+.. automethod:: iris.host.iris.Iris.get_device_id
+```
+
+## ops
+
+```{eval-rst}
+.. autoproperty:: iris.host.iris.Iris.ops
+```

docs/reference/host/copy-engine.md

@@ -0,0 +1,47 @@
+# Copy Engine (SDMA)
+
+Host-initiated transfers that run on the GPU's SDMA copy engines rather than on
+compute units. The host enqueues a descriptor and the DMA engine moves the
+bytes, so no kernel launch is involved and no CUs are consumed.
+
+```{note}
+These operations require the copy engine to be initialised. See
+[`iris.iris()`](class.md) for how to enable it, and the
+[Installation guide](../../getting-started/installation.md) for the optional
+`rocm-xio` dependency they rely on.
+```
+
+Unlike the device-side operations, there is no engine choice to make here: a CU
+store is an instruction a wave executes, so the host cannot issue one. Anything
+host-initiated goes through the copy engine.
+
+## put
+
+```{eval-rst}
+.. automethod:: iris.host.iris.Iris.put
+```
+
+## put_tile
+
+```{eval-rst}
+.. automethod:: iris.host.iris.Iris.put_tile
+```
+
+## put_tiles
+
+```{eval-rst}
+.. automethod:: iris.host.iris.Iris.put_tiles
+```
+
+## quiet
+
+```{eval-rst}
+.. automethod:: iris.host.iris.Iris.quiet
+```
+
+## get_copy_engine_ctx
+
+```{eval-rst}
+.. automethod:: iris.host.iris.Iris.get_copy_engine_ctx
+   :no-index:
+```

docs/reference/triton/device-context.md

@@ -0,0 +1,165 @@
+# Device Context
+
+Device-side methods provided by Iris for remote memory operations, atomics and
+collectives. These methods are part of the `DeviceContext` aggregate used inside
+`@triton.jit` kernels, and are reached through the context tensor passed in from
+the host:
+
+```python
+from iris import DeviceContext
+
+@triton.jit
+def producer(context_tensor, target_buffer, ..., rank: tl.constexpr, world_size: tl.constexpr):
+    ctx = DeviceContext.initialize(context_tensor, rank, world_size)
+    ...
+    ctx.store(target_buffer + offsets, values, to_rank=consumer_rank, mask=mask)
+    ctx.atomic_cas(flag + pid, 0, 1, to_rank=consumer_rank, sem="release", scope="sys")
+```
+
+Taken from `examples/06_message_passing/message_passing_device_context.py`.
+
+```{note}
+The collectives below are **tile-level** -- they run inside your kernel, on the
+tile a program is already holding, and fuse with whatever compute surrounds
+them. They are a different surface from the host collectives in
+[Collective Communication Operations](../host/ccl.md), which launch their own
+kernels and are called from Python.
+```
+
+Every operation here is also available as a module-level function -- see
+[Device Functions](device-functions.md). The module-level form takes the heap
+bases and both ranks explicitly; the context already knows the current rank and
+holds the heap bases, so its methods take fewer arguments:
+
+```python
+iris.load(pointer, to_rank, from_rank, heap_bases)   # module-level function
+ctx.load(pointer, from_rank)                         # method on the context
+```
+
+## Initialization
+
+### initialize
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.initialize
+```
+
+## Memory transfer operations
+
+### load
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.load
+```
+
+### store
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.store
+```
+
+### get
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.get
+```
+
+### put
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.put
+```
+
+### copy
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.copy
+```
+
+## Atomic operations
+
+### atomic_add
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.atomic_add
+```
+
+### atomic_sub
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.atomic_sub
+```
+
+### atomic_cas
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.atomic_cas
+```
+
+### atomic_xchg
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.atomic_xchg
+```
+
+### atomic_xor
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.atomic_xor
+```
+
+### atomic_and
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.atomic_and
+```
+
+### atomic_or
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.atomic_or
+```
+
+### atomic_min
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.atomic_min
+```
+
+### atomic_max
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.atomic_max
+```
+
+## Collective operations
+
+### all_reduce_atomic
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.all_reduce_atomic
+```
+
+### all_reduce_spinlock
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.all_reduce_spinlock
+```
+
+### all_reduce_one_shot
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.all_reduce_one_shot
+```
+
+### all_reduce_ring
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.all_reduce_ring
+```
+
+### all_reduce_two_shot
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.all_reduce_two_shot
+```
+
+### all_gather
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.all_gather
+```
+
+### gather
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.gather
+```
+
+### all_to_all
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.all_to_all
+```
+
+### reduce_scatter
+```{eval-rst}
+.. automethod:: iris.mem.triton.context.Context.reduce_scatter
+```

docs/reference/triton/device-functions.md

@@ -37,6 +37,11 @@ Device-side functions provided by Iris for remote memory operations and atomics.
 .. autofunction:: iris.mem.triton.ops.atomic_add
 ```
 
+### atomic_sub
+```{eval-rst}
+.. autofunction:: iris.mem.triton.ops.atomic_sub
+```
+
 ### atomic_cas
 ```{eval-rst}
 .. autofunction:: iris.mem.triton.ops.atomic_cas
@@ -72,4 +77,9 @@ Device-side functions provided by Iris for remote memory operations and atomics.
 .. autofunction:: iris.mem.triton.ops.atomic_max
 ```
 
+## Synchronization
 
+### quiet
+```{eval-rst}
+.. autofunction:: iris.mem.triton.ops.quiet
+```