diff --git a/.gitignore b/.gitignore index 67ebd97bd..4b7306ca7 100644 --- a/.gitignore +++ b/.gitignore @@ -69,3 +69,7 @@ docker/rocm-systems/ .intellikit .github/agents/skills/ docs/benchmark-results/*.png + +# Local Sphinx build +docs/.venv/ +docs/_build/ diff --git a/docs/conf.py b/docs/conf.py index ae01bbb61..0a0660de7 100644 --- a/docs/conf.py +++ b/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 +# +# 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.` 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. +# +# 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) diff --git a/docs/reference/api-reference.md b/docs/reference/api-reference.md index 38ab0ba40..54d6d4d76 100644 --- a/docs/reference/api-reference.md +++ b/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) diff --git a/docs/reference/host/class.md b/docs/reference/host/class.md index 46845d8b8..2122a0714 100644 --- a/docs/reference/host/class.md +++ b/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 +``` diff --git a/docs/reference/host/copy-engine.md b/docs/reference/host/copy-engine.md new file mode 100644 index 000000000..44387072d --- /dev/null +++ b/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: +``` diff --git a/docs/reference/triton/device-context.md b/docs/reference/triton/device-context.md new file mode 100644 index 000000000..eff6ab1b6 --- /dev/null +++ b/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 +``` diff --git a/docs/reference/triton/device-functions.md b/docs/reference/triton/device-functions.md index c9af17dd3..92ef5d07a 100644 --- a/docs/reference/triton/device-functions.md +++ b/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 +``` diff --git a/docs/reference/triton/ops.md b/docs/reference/triton/ops.md index 93be2b31d..08b4e9e12 100644 --- a/docs/reference/triton/ops.md +++ b/docs/reference/triton/ops.md @@ -1,6 +1,10 @@ -# Fused GEMM + CCL Operations +# Fused Communication and Compute Operations -Fused matrix multiplication and collective communication operations accessible via the `ops` property on the `Iris` instance (e.g. `ctx.ops.matmul_all_reduce(...)`). +Operations that fuse communication with compute in a single kernel, accessible +via the `ops` property on the `Iris` instance (e.g. `ctx.ops.matmul_all_reduce(...)`). +The compute side is currently GEMM, supplied by tritonBLAS; the communication side +is an Iris collective. Fusing them means the collective overlaps the matmul instead +of running as a separate kernel after it. ## matmul_all_reduce ```{eval-rst} diff --git a/docs/reference/triton/overview.md b/docs/reference/triton/overview.md index 46f12f986..d82bcb510 100644 --- a/docs/reference/triton/overview.md +++ b/docs/reference/triton/overview.md @@ -36,5 +36,5 @@ Explore the API by section: - [Tensor Creation](../host/tensor-creation.md) - [Device Functions](device-functions.md) - [Collective Communication (CCL)](../host/ccl.md) -- [Fused GEMM + CCL Operations](ops.md) +- [Fused Communication and Compute Operations](ops.md) diff --git a/docs/reference/triton/types.md b/docs/reference/triton/types.md new file mode 100644 index 000000000..3042a01d8 --- /dev/null +++ b/docs/reference/triton/types.md @@ -0,0 +1,48 @@ +# Tile Types + +The tile abstractions the device APIs operate on. All of these are exported from +the top level, so `iris.Tile`, `iris.TensorView` and so on are the import paths. + +A `TensorView` wraps a pointer plus the tensor's shape and strides, and hands out +tile pointers and masks. A `TileView` carries a tile's coordinates; a `Tile` +carries its coordinates *and* the data. `AllReduceConfig` is the device-side +configuration consumed by the all-reduce methods on +[Device Context](device-context.md). + +## TensorView + +```{eval-rst} +.. autoclass:: iris.mem.triton.types.TensorView + :no-undoc-members: + :no-special-members: +``` + +## TileView + +```{eval-rst} +.. autoclass:: iris.mem.triton.types.TileView + :no-undoc-members: + :no-special-members: +``` + +## Tile + +```{eval-rst} +.. autoclass:: iris.mem.triton.types.Tile + :no-undoc-members: + :no-special-members: +``` + +## AllReduceConfig + +```{eval-rst} +.. autoclass:: iris.mem.triton.types.AllReduceConfig + :no-undoc-members: + :no-special-members: +``` + +## make_tensor_view + +```{eval-rst} +.. autofunction:: iris.mem.triton.types.make_tensor_view +``` diff --git a/docs/sphinx/_toc.yml b/docs/sphinx/_toc.yml index afe6c23cf..49ce51d06 100644 --- a/docs/sphinx/_toc.yml +++ b/docs/sphinx/_toc.yml @@ -3,6 +3,7 @@ subtrees: - caption: Getting Started entries: - file: getting-started/installation.md + - file: getting-started/slurm.md - file: reference/examples.md - file: CONTRIBUTING.md - caption: Conceptual @@ -18,9 +19,12 @@ subtrees: - file: reference/host/class.md - file: reference/host/tensor-creation.md - file: reference/host/ccl.md + - file: reference/host/copy-engine.md - file: reference/triton/overview.md entries: + - file: reference/triton/device-context.md - file: reference/triton/device-functions.md + - file: reference/triton/types.md - file: reference/triton/ops.md - file: reference/gluon/overview.md entries: diff --git a/docs/sphinx/_toc.yml.in b/docs/sphinx/_toc.yml.in index afe6c23cf..49ce51d06 100644 --- a/docs/sphinx/_toc.yml.in +++ b/docs/sphinx/_toc.yml.in @@ -3,6 +3,7 @@ subtrees: - caption: Getting Started entries: - file: getting-started/installation.md + - file: getting-started/slurm.md - file: reference/examples.md - file: CONTRIBUTING.md - caption: Conceptual @@ -18,9 +19,12 @@ subtrees: - file: reference/host/class.md - file: reference/host/tensor-creation.md - file: reference/host/ccl.md + - file: reference/host/copy-engine.md - file: reference/triton/overview.md entries: + - file: reference/triton/device-context.md - file: reference/triton/device-functions.md + - file: reference/triton/types.md - file: reference/triton/ops.md - file: reference/gluon/overview.md entries: diff --git a/docs/sphinx/requirements.txt b/docs/sphinx/requirements.txt index cdd7c0c73..4633e2640 100644 --- a/docs/sphinx/requirements.txt +++ b/docs/sphinx/requirements.txt @@ -4,3 +4,5 @@ sphinx-copybutton>=0.5.0 sphinx-autobuild>=2021.3.14 sphinx_external_toc>=1.0.0 rocm-docs-core +# rocm-docs-core's theme JS uses jQuery, which Sphinx stopped bundling in 6.0. +sphinxcontrib-jquery>=4.1 diff --git a/iris/host/iris.py b/iris/host/iris.py index 49d4a9913..081ab37c5 100644 --- a/iris/host/iris.py +++ b/iris/host/iris.py @@ -337,6 +337,7 @@ def broadcast(self, value, src=0): This method automatically detects the type of value and uses the appropriate broadcast mechanism: + - For tensors and arrays: uses efficient PyTorch distributed tensor collectives - For scalars and other objects: uses object broadcast @@ -983,8 +984,19 @@ def get_heap_bases(self): """ return self.heap_bases - def get_copy_engine_ctx(self) -> torch.Tensor | None: - """Return the device copy-engine context, or None when copy engines are disabled.""" + def get_copy_engine_ctx(self) -> "torch.Tensor | None": + """ + Get the device-side copy engine context for SDMA operations. + + The returned tensor holds the per-rank SDMA queue, read and write pointer, + and doorbell addresses that device-side kernels need in order to enqueue + copy packets. Pass it to ``iris.put`` as ``copy_engine_ctx`` when driving + the copy engine from inside a kernel. + + Returns: + torch.Tensor: The copy engine device context, or ``None`` when Iris was + not initialized with ``engines=("sdma",)``. + """ return self.copy_engines_device_ctx @staticmethod @@ -1040,6 +1052,7 @@ def put( Requires initialization with ``engines=("sdma",)``. Supports: + - Simple copy: put(src, to_rank) - Copy + signal: put(src, to_rank, signal_flag=flag) - Wait + copy: put(src, to_rank, wait_flag=flag, wait_value=N) @@ -1548,7 +1561,7 @@ class CCL: Collective Communication Library (CCL) interface for Iris. Provides collective operations that can be called as methods on the Iris instance. - Example usage: + Example: >>> ctx = iris.iris() >>> ctx.ccl.all_to_all(output_tensor, input_tensor) """ diff --git a/iris/mem/triton/context.py b/iris/mem/triton/context.py index cb6549313..63ee446e9 100644 --- a/iris/mem/triton/context.py +++ b/iris/mem/triton/context.py @@ -98,9 +98,11 @@ def initialize(context_tensor, rank, world_size, tracing: tl.constexpr = False): """ Initialize Context from the encoded context tensor. - The context tensor has the format: - - [cur_rank, num_ranks, heap_base_0, ..., heap_base_N, trace_info...] - - If tracing=True: extracts trace buffer pointers from context_tensor + The context tensor has the format:: + + [cur_rank, num_ranks, heap_base_0, ..., heap_base_N, trace_info...] + + If ``tracing=True``, trace buffer pointers are also extracted from it. Args: context_tensor: Pointer to encoded context data (from Iris.get_device_context()) diff --git a/iris/mem/triton/ops.py b/iris/mem/triton/ops.py index f56544758..7c355484f 100644 --- a/iris/mem/triton/ops.py +++ b/iris/mem/triton/ops.py @@ -604,7 +604,7 @@ def atomic_sub( This function performs an atomic subtraction operation by translating the pointer from the `from_rank`'s address space to the `to_rank`'s address space and atomically - subtracting the provided data from the `to_rank` memory location. If the `from_rank` and `to_rank` are the same, + subtracting the provided data from the `to_rank` memory location. If the `from_rank` and `to_rank` are the same, this function performs a local atomic subtraction operation. Args: diff --git a/iris/mem/triton/types.py b/iris/mem/triton/types.py index d04097824..4b0081ff8 100644 --- a/iris/mem/triton/types.py +++ b/iris/mem/triton/types.py @@ -109,16 +109,23 @@ class TileView: TileView storing BOTH runtime coordinates AND compile-time block sizes. This class uses the @constexpr_function pattern: + - Stores runtime coordinates (pid_m, pid_n) as tl.tensor (computed from tl.program_id) - Stores compile-time block sizes (block_m, block_n) as tl.constexpr - Constructor is decorated with @constexpr_function to execute at compile-time - Example usage: - pid = tl.program_id(0) - pid_m = pid // num_tiles_n - pid_n = pid % num_tiles_n - tile = TileView(pid_m, pid_n, BLOCK_M, BLOCK_N) - rm, rn, mask = tile.layout(M, N) + Attributes: + pid_m: Tile row index, computed from ``tl.program_id``. + pid_n: Tile column index, computed from ``tl.program_id``. + block_m: Tile height. Compile-time constant. + block_n: Tile width. Compile-time constant. + + Example: + >>> pid = tl.program_id(0) + >>> pid_m = pid // num_tiles_n + >>> pid_n = pid % num_tiles_n + >>> tile = TileView(pid_m, pid_n, BLOCK_M, BLOCK_N) + >>> rm, rn, mask = tile.layout(M, N) """ pid_m: tl.tensor @@ -154,10 +161,17 @@ class Tile: Extends TileView concept to include a data field for storing computed tile data (e.g., GEMM results in registers). - Example usage: - c = tl.dot(a, b) - tile = Tile(pid_m, pid_n, BLOCK_M, BLOCK_N, c) - ctx.all_reduce_atomic(tile, dst_view) + Attributes: + pid_m: Tile row index, computed from ``tl.program_id``. + pid_n: Tile column index, computed from ``tl.program_id``. + block_m: Tile height. Compile-time constant. + block_n: Tile width. Compile-time constant. + data: The tile's values, for example a GEMM result held in registers. + + Example: + >>> c = tl.dot(a, b) + >>> tile = Tile(pid_m, pid_n, BLOCK_M, BLOCK_N, c) + >>> ctx.all_reduce_atomic(tile, dst_view) """ pid_m: tl.tensor @@ -218,9 +232,16 @@ class TensorView: """ TensorView storing pointer and tensor metadata. - Example usage: - view = make_tensor_view(ptr, M, N, stride_m, stride_n) - tile_ptr, mask = view.tile_ptr(tile) + Attributes: + ptr: Pointer to the tensor's data. + M: Number of rows. + N: Number of columns. + stride_m: Stride in the M dimension. + stride_n: Stride in the N dimension. + + Example: + >>> view = make_tensor_view(ptr, M, N, stride_m, stride_n) + >>> tile_ptr, mask = view.tile_ptr(tile) """ ptr: tl.tensor @@ -287,7 +308,7 @@ class AllReduceConfig: """ Device-side configuration for all_reduce collective. - Fields: + Attributes: variant_code: Integer code for algorithm variant (0=atomic, 1=ring, 2=one_shot, 3=two_shot, 4=spinlock) locks_ptr: Pointer to locks array (ALWAYS required - pass dummy tensor if variant doesn't need locks) """