Make SDMA optional and explicitly enabled per Iris context by drprajap · Pull Request #568 · ROCm/iris

drprajap · 2026-10-02T22:54:26Z

Shader load/store, atomics, and collectives work without SDMA. import iris no longer imports rocm-xio, and an Iris context does not create SDMA queues unless it asks for them.

Consolidates #579. Queue setup stays conditional because each SDMA queue reserves 8 MiB of host/device memory, and the total grows with the number of local ranks. Shader-only callers, including TokenSpeed paths that are not on SDMA yet, should not pay that cost or take the rocm-xio dependency. On MI450 we also want SDMA off by default so CU and SDMA issues can be separated while it is brought up.

Behavior

Iris / iris.iris() take engines, default (). Pass engines=("sdma",) to create the host- and device-initiated SDMA queues. Any other name raises ValueError.
With SDMA off, get_copy_engine_ctx() returns None. Host put, put_tile, put_tiles, and quiet raise RuntimeError.
rocm-xio moves from the base dependencies to the sdma extra: pip install "iris[sdma]". It is imported only when "sdma" is in engines. If it is missing, that raises ModuleNotFoundError and points at the extra.
Installing the extra does not turn SDMA on. Each context still opts in through engines.

CI and docs

Scripts and workflows that run the SDMA tests or the SDMA examples install .[sdma].
docs/getting-started/installation.md documents the extra. docs/reference/host/class.md documents engines.

Test plan

tests/unittests/test_optional_sdma.py: import iris does not load xio; a default context does a cross-rank shader store, get_copy_engine_ctx() is None, and host put / quiet raise RuntimeError.
SDMA unit tests and the host-initiated message-passing examples pass engines=("sdma",) and skip when rocm-xio is not installed.

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

mawad-amd · 2026-10-09T12:28:39Z

@drprajap Can we address the comments and merge this? Also needs the branch updated with main.

drprajap · 2026-10-10T00:29:31Z

@drprajap Can we address the comments and merge this? Also needs the branch updated with main.

Addressed feedback and consolidated change from @efric 's #579

mawad-amd

My final three comments. Thanks!

efric

Overall LGTM % some comments!

efric · 2026-10-13T00:06:40Z

another nit:
the pr title seems slightly outdated, perhaps something like make SDMA optional and explicitly enabled per Iris context?

and can we update the description with some mention of why we'd like to keep copy engine initialization conditional (due to cost of queues as described in #579?

Shader load/store users no longer need rocm-xio. It moves from the required dependencies to the sdma extra, and import iris no longer imports xio. iris.iris() / Iris() take enable_copy_engine (True/False/None). None reads IRIS_ENABLE_COPY_ENGINE (1/0/auto, strictly parsed); in auto mode SDMA is enabled only if rocm-xio is installed. With SDMA off, get_copy_engine_ctx() returns None and the host SDMA APIs raise RuntimeError. CI installs .[sdma]; external validation stays on the base install. SDMA tests skip when the copy engine is disabled. Signed-off-by: Prajapati, Dimple <[email protected]>

…tra cover it - Remove the env var and its parser; None (default) enables SDMA only if rocm-xio is installed. - Bind rocm-xio in test_sdma_packets.py and skip it without rocm-xio. - Trim the flag tests to the no-rocm-xio cases and rename to test_optional_sdma.py. - Docs: use --force-reinstall --no-deps for a standalone rocm-xio install. Signed-off-by: Prajapati, Dimple <[email protected]>

- Bind rocm-xio lazily only when copy_engine="sdma"; ImportError points to iris[sdma]. - Skip SDMA tests when rocm-xio is not installed. - test_optional_sdma.py: no-xio import, fail-fast, and default shader-store checks. - Docs: installation covers the extra, class reference covers copy_engine. Signed-off-by: Prajapati, Dimple <[email protected]>

Signed-off-by: Prajapati, Dimple <[email protected]>

engines argument, sdma_utils owns sdma_ep, simpler install docs Signed-off-by: Prajapati, Dimple <[email protected]>

run_core_tests.sh runs every tests/<dir>/test_*.py, including the SDMA tests, which skip without rocm-xio. Match run_example_slurm.sh. Signed-off-by: Prajapati, Dimple <[email protected]>

drprajap · 2026-10-13T05:25:15Z

another nit: the pr title seems slightly outdated, perhaps something like make SDMA optional and explicitly enabled per Iris context?

and can we update the description with some mention of why we'd like to keep copy engine initialization conditional (due to cost of queues as described in #579?

Yes, I will align the PR description as per latest consolidation

drprajap · 2026-10-13T05:39:19Z

Updated the title and description to match the consolidated API (engines=("sdma",), sdma extra). The description now also notes why queue init stays conditional: each SDMA queue reserves 8 MiB of host/device memory, and that grows with the local rank count (#579).

Shader stores already run on a default Iris context. Keep the import check and the host put/quiet rejection, and drop the extra copy-engine assert from the device-context test. Signed-off-by: Prajapati, Dimple <[email protected]>

mawad-amd

Looks great. Thanks for addressing the comments, and for your patience.

Only conflict was get_copy_engine_ctx's docstring: ROCm#568 added the -> torch.Tensor | None annotation and shortened the prose to one line, while this branch had expanded it as part of documenting the 11 undocumented public APIs. Kept the annotation and the fuller docstring, reworded to name engines=("sdma",) since ROCm#568 replaced copy_engine=.

Quote the return annotation on get_copy_engine_ctx. ROCm#568 added `-> torch.Tensor | None`, but torch is in autodoc_mock_imports, so the annotation evaluates mock | None at class-definition time and raises TypeError: unsupported operand type(s) for |: 'Tensor' and 'NoneType'. That killed the import of iris.host.iris and took 45 autodoc directives with it -- the API reference went empty again, by a different route than the xio import did. Quoting defers evaluation; it is the only PEP 604 annotation on a mocked type in iris/, and nothing introspects annotations at runtime. Defer the get_copy_engine_ctx index entry to class.md. ROCm#568 documents the method in its new "Optional SDMA Copy Engine" section and this branch already documented it on the dedicated copy-engine page, so Sphinx warned about a duplicate object description. Added :no-index: to this branch's copy rather than touching the section ROCm#568 just added. Verified on 3.13: 39 pages, 0 autodoc import failures, 4 warnings (all external intersphinx), and every page rendered in headless Chrome over HTTP -- no dead links, no broken images, no page errors beyond two the rocm-docs-core theme also throws on rocm.github.io/iris. jQuery errors per page: 4 upstream, 0 here.

…ssing anchors atomic_sub and quiet are iris.mem.triton.ops device functions, but sat on reference/triton/ops.md with the iris.ops host namespace because both surfaces are called "ops". Moved to device-functions.md -- atomic_sub with the other atomics, quiet under a new Synchronization heading. ops.md is now only the four fused host ops. "Free function" is C++ for a non-member function and means nothing in Python. Replaced with module-level function vs method, shown rather than described. :noindex: was copied from the Gluon page onto types.md and device-context.md, where nothing else documents those objects, so it suppressed their index anchors for no reason. Removed, 29 lines; the site index goes from 60 to 114 API objects. copy-engine.md keeps its one: ROCm#568 documents get_copy_engine_ctx on class.md too, so without it Sphinx warns about a duplicate object description. Checked by removing each one and rebuilding. Built on 3.13: 34 pages, 0 autodoc import failures, 0 duplicate objects.

drprajap requested review from BKP, artulab, mawad-amd and neoblizz as code owners October 2, 2026 22:54

drprajap requested a review from nirvedhmeshram October 2, 2026 22:55

github-actions Bot added in-progress We are working on it iris Iris project issue labels Oct 2, 2026

mawad-amd reviewed Oct 2, 2026

View reviewed changes

Comment thread iris/host/iris.py Outdated

Comment thread docs/getting-started/installation.md Outdated

drprajap requested a review from efric as a code owner October 10, 2026 00:25

drprajap requested a review from mawad-amd October 12, 2026 05:52

mawad-amd reviewed Oct 12, 2026

View reviewed changes

Comment thread README.md Outdated

Comment thread docs/getting-started/installation.md Outdated

Comment thread iris/host/iris.py Outdated

efric reviewed Oct 12, 2026

View reviewed changes

Comment thread iris/device/sdma_utils.py Outdated

Comment thread iris/device/sdma_utils.py

Comment thread docs/getting-started/installation.md Outdated

Comment thread tests/unittests/test_optional_sdma.py Outdated

efric mentioned this pull request Oct 13, 2026

Make copy engine initialization optional #579

Closed

1 task done

drprajap and others added 7 commits October 13, 2026 02:09

Make SDMA initialization opt-in

dc139d6

Remove duplicate _require_copy_engine; trim repeated install docs

517bf00

Signed-off-by: Prajapati, Dimple <[email protected]>

Address review feedback

40a251d

engines argument, sdma_utils owns sdma_ep, simpler install docs Signed-off-by: Prajapati, Dimple <[email protected]>

Install the sdma extra in the SLURM core-test script

edd4d1d

run_core_tests.sh runs every tests/<dir>/test_*.py, including the SDMA tests, which skip without rocm-xio. Match run_example_slurm.sh. Signed-off-by: Prajapati, Dimple <[email protected]>

drprajap force-pushed the diprajap/optional-rocm-xio branch from fb94689 to edd4d1d Compare October 13, 2026 02:11

efric approved these changes Oct 13, 2026

View reviewed changes

drprajap changed the title ~~Make rocm-xio an optional sdma extra; add enable_copy_engine flag~~ make SDMA optional and explicitly enabled per Iris context Oct 13, 2026

drprajap changed the title ~~make SDMA optional and explicitly enabled per Iris context~~ Make SDMA optional and explicitly enabled per Iris context Oct 13, 2026

drprajap force-pushed the diprajap/optional-rocm-xio branch from e679481 to 638f3d0 Compare October 13, 2026 06:20

drprajap mentioned this pull request Oct 13, 2026

Stale - feat(comm): Iris across MI455X trays in one fabric domain lightseekorg/tokenspeed#2272

Closed

mawad-amd approved these changes Oct 13, 2026

View reviewed changes

drprajap merged commit e46ac4d into main Oct 13, 2026
24 checks passed

drprajap deleted the diprajap/optional-rocm-xio branch October 13, 2026 14:17

mawad-amd mentioned this pull request Oct 15, 2026

SDMA is Triton-only, and no collective uses it #593

Open

5 tasks

Navigation Menu

Repository navigation

Make SDMA optional and explicitly enabled per Iris context - #568

Make SDMA optional and explicitly enabled per Iris context#568
drprajap merged 8 commits into
mainfrom
diprajap/optional-rocm-xio

drprajap commented Oct 2, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

mawad-amd commented Oct 9, 2026 •

edited

Loading

Uh oh!

drprajap commented Oct 10, 2026

Uh oh!

mawad-amd left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

efric left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

efric commented Oct 13, 2026

Uh oh!

drprajap commented Oct 13, 2026

Uh oh!

drprajap commented Oct 13, 2026

Uh oh!

mawad-amd left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

drprajap commented Oct 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Behavior

CI and docs

Test plan

Submission Checklist

Uh oh!

Uh oh!

Uh oh!

mawad-amd commented Oct 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

drprajap commented Oct 10, 2026

Uh oh!

mawad-amd left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

efric left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

efric commented Oct 13, 2026

Uh oh!

drprajap commented Oct 13, 2026

Uh oh!

drprajap commented Oct 13, 2026

Uh oh!

mawad-amd left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

drprajap commented Oct 2, 2026 •

edited

Loading

mawad-amd commented Oct 9, 2026 •

edited

Loading