Rename iris/device/ to iris/mem/, unify Context class names - #520
Merged
Conversation
Move device-side remote memory access primitives from iris/device/ to iris/mem/ to better reflect what the module provides. Unify class names across backends: DeviceContext → Context, IrisDeviceCtx → Context, DeviceTracing → Tracing, GluonDeviceTracing → Tracing. iris/device/ files become thin backward-compat shims that re-export from iris.mem.*. The __translate dunder is explicitly re-exported since star imports exclude dunder names. All internal imports updated; examples left unchanged to exercise the shim path. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
github-actions
Bot
added
in-progress
Contributor
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
This PR reorganizes Iris device-side remote memory access code under iris.mem and standardizes backend class naming (Context, Tracing) while keeping iris.device.* as backward-compatible shims.
Changes:
- Introduces
iris/mem/(Triton + Gluon) withContext,Tracing, ops, and device utils. - Updates internal imports to canonical
iris.mem.*paths and re-exports oldiris.device.*modules as compatibility shims. - Updates top-level exports (
iris.Context,iris.DeviceContext, ops) and adjusts unit tests/import call sites accordingly.
Reviewed changes
Copilot reviewed 30 out of 30 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unittests/test_device_context.py | Updates import-path test to point at iris.mem.* |
| iris/x/reduce_scatter.py | Switches kernel context import to iris.mem.triton.context |
| iris/x/gather.py | Switches kernel context import to iris.mem.triton.context |
| iris/x/all_to_all.py | Switches kernel context import to iris.mem.triton.context |
| iris/x/all_reduce.py | Switches kernel context import to iris.mem.triton.context |
| iris/x/all_gather.py | Switches kernel context import to iris.mem.triton.context |
| iris/mem/utils.py | Adds device-side HIP/CDNA utility intrinsics under iris.mem |
| iris/mem/triton/tracing.py | Adds canonical Triton tracing aggregate under iris.mem.triton |
| iris/mem/triton/ops.py | Adds canonical Triton RMA functional API under iris.mem.triton |
| iris/mem/triton/context.py | Adds canonical Triton Context aggregate under iris.mem.triton |
| iris/mem/triton/init.py | Exposes canonical Triton API + backward-compat aliases |
| iris/mem/gluon/tracing.py | Adds canonical Gluon tracing aggregate under iris.mem.gluon |
| iris/mem/gluon/context.py | Adds canonical Gluon Context aggregate under iris.mem.gluon |
| iris/mem/gluon/init.py | Exposes canonical Gluon API + backward-compat aliases |
| iris/mem/init.py | Creates iris.mem package root |
| iris/host/iris.py | Re-points host re-export of device tracing to iris.mem.triton.tracing |
| iris/gluon.py | Re-points Gluon convenience imports to iris.mem.gluon.* |
| iris/experimental/init.py | Re-points experimental Gluon context import to iris.mem.gluon |
| iris/device/utils.py | Converts iris.device.utils into shim re-exporting iris.mem.utils |
| iris/device/triton/tracing.py | Converts Triton tracing module into shim re-exporting iris.mem.triton.tracing |
| iris/device/triton/ops.py | Converts Triton ops module into shim re-exporting iris.mem.triton.ops |
| iris/device/triton/context.py | Converts Triton context module into shim (explicitly re-exporting __translate) |
| iris/device/triton/init.py | Converts Triton package into shim re-exporting iris.mem.triton |
| iris/device/gluon/tracing.py | Converts Gluon tracing module into shim re-exporting iris.mem.gluon.tracing |
| iris/device/gluon/context.py | Converts Gluon context module into shim re-exporting iris.mem.gluon.context |
| iris/device/gluon/init.py | Converts Gluon package into shim re-exporting iris.mem.gluon |
| iris/device/init.py | Converts iris.device into shim re-exporting iris.mem |
| iris/ccl/all_to_all.py | Switches Gluon context import to iris.mem.gluon.context |
| iris/ccl/all_gather.py | Switches Gluon context import to iris.mem.gluon.context |
| iris/init.py | Exposes Context and repoints DeviceContext/ops to iris.mem.triton.* |
Address Copilot review: test_device_context_imports now also verifies that iris.device.triton.context.DeviceContext resolves to the same object as the canonical iris.mem.triton.context.Context path. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Fix stale documentation references caught by review agents: - iris/__init__.py and iris/experimental/__init__.py docstrings - docs/reference/gluon/device-functions.md automethod directives - docs/reference/triton/device-functions.md autofunction directives Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
iris/device/→iris/mem/to better describe what the module provides (remote memory access)DeviceContext→Context,IrisDeviceCtx→Context,DeviceTracing→Tracing,GluonDeviceTracing→Tracingiris/device/files become thin backward-compat shims (re-exports fromiris.mem.*). The__translatedunder is explicitly re-exported sinceimport *excludes dunder namesiris.mem.*paths; examples left unchanged to exercise shim pathWhy
mem?The Context provides remote memory access (load/store/atomics across GPUs). 3 chars, same energy as
x,ccl,ops. The hierarchy:mem(primitives) →ccl(collectives) →ops(fused). Users go deeper only when writing device kernels, and onlymemrequires explicit backend selection (iris.mem.tritonvsiris.mem.gluon).Test plan
from iris.device.triton.context import DeviceContextstill worksfrom iris.mem.triton.context import Contextworksfrom iris import DeviceContext, Context, load, storefrom iris.gluon import IrisDeviceCtxtorchrun --nproc_per_node=4 examples/25_ccl_all_gather/example.pytorchrun --nproc_per_node=4 examples/25_ccl_all_reduce/example.pypytest tests/ -x -v🤖 Generated with Claude Code