Repository navigation
feat(rocSHMEM): add TDM in the IPC backend, #AIROCSHMEM-515 - #12568
Merged
Merged
Conversation
amirsojoodi
changed the title
✅ All Policy Checks Passed
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Optional registration can consume uninitialized LDS state, tile alignment can corrupt transfers, and the new tests currently exercise only fallback behavior.
Review effort: Balanced
Findings: 3
Open (10)
What changed in this PR
Adds optional gfx1250 Tensor Data Mover acceleration for rocSHMEM IPC work-group transfers.
Changes:
- Adds TDM descriptor, cache-policy, LDS staging, and configuration support.
- Integrates TDM with IPC operations and host-launched kernels.
- Adds functional integration and raw/structured transfer tests.
| File | Description |
|---|---|
projects/rocshmem/CMakeLists.txt |
Adds USE_TDM. |
projects/rocshmem/cmake/rocshmem_config.h.in |
Exposes the build option. |
projects/rocshmem/src/tdm.hpp |
Implements TDM operations and LDS registration. |
projects/rocshmem/src/util.hpp |
Adds TDM-staged work-group copying. |
projects/rocshmem/src/ipc_policy.hpp |
Routes IPC work-group copies through TDM. |
projects/rocshmem/src/constmem.hpp |
Adds TDM constant-memory configuration. |
projects/rocshmem/src/constmem.cpp |
Resolves runtime tile sizing. |
projects/rocshmem/src/envvar.hpp |
Declares the TDM environment category. |
projects/rocshmem/src/envvar.cpp |
Adds tile-size configuration. |
projects/rocshmem/src/rocshmem.cpp |
Implements the host LDS-size query. |
projects/rocshmem/src/rocshmem_gpu.cpp |
Implements device APIs and kernel opt-in. |
projects/rocshmem/src/host/host.cpp |
Allocates dynamic LDS for stream operations. |
projects/rocshmem/include/rocshmem/rocshmem.hpp |
Declares and documents public TDM APIs. |
projects/rocshmem/include/rocshmem/rocshmem_RMA_X.hpp |
Documents TDM-enabled RMA behavior. |
projects/rocshmem/tests/functional_tests/workgroup_primitives.cpp |
Enables TDM in work-group tests. |
projects/rocshmem/tests/functional_tests/team_reduction_tester.cpp |
Registers dynamic LDS in reduction tests. |
projects/rocshmem/tests/unit_tests/CMakeLists.txt |
Registers TDM unit tests. |
projects/rocshmem/tests/unit_tests/tdm_raw_put_gtest.hpp |
Implements raw-transfer fixtures. |
projects/rocshmem/tests/unit_tests/tdm_raw_put_gtest.cpp |
Defines raw-transfer cases. |
projects/rocshmem/tests/unit_tests/tdm_struct_put_gtest.hpp |
Implements structured-transfer fixtures. |
projects/rocshmem/tests/unit_tests/tdm_struct_put_gtest.cpp |
Defines structured-transfer cases. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
github-actions
Bot
added
the
documentation
amirsojoodi
force-pushed
the
users/asojoodi/rocshmem-tdm
branch
from
October 19, 2026 19:34
6982686 to
bc78308
Compare
amirsojoodi
force-pushed
the
users/asojoodi/rocshmem-tdm
branch
from
October 22, 2026 15:23
bc78308 to
8f00567
Compare
abouteiller
reviewed
Oct 22, 2026
abouteiller
approved these changes
Oct 23, 2026
Contributor
TheRock Submodule Bump ActivityNewest first
|
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
Adds an opt-in
USE_TDMbuild option that lets work-groupput/geton gfx1250 stage IPC transfers through LDS using the hardware Tensor Data Mover, instead of the existing lane-copy path. Disabled by default and a safe no-op everywhere it isn't applicable.Technical Details
rocshmem_set_tdm_lds()lets a kernel hand rocSHMEM a scratch LDS buffer;rocshmem_query_tdm_lds_bytes()reports the size to allocate. Without registration,put/getalways take the existing plain-copy path.ROCSHMEM_TDM_TILE_BYTES(default0= auto-size the tile from the device's available LDS per block).src/tdm.hppwraps thetensor_load_to_lds/tensor_store_from_ldsintrinsics (descriptor construction, cache-policy encoding, double-buffered tile-copy helpers); wired into the IPC work-group copy path inutil.hpp/ipc_policy.hpp.getmem/putmem/putmem_signalkernels auto-opt-in using their own dynamic LDS.tdm_raw_put_gtest,tdm_struct_put_gtest), built only whenUSE_IPC AND USE_TDM.Test Result
-DGPU_TARGETS="gfx1250"(noUSE_TDM) builds clean.-DUSE_TDM=1 -DGPU_TARGETS="gfx1250"builds and runs with no issues. Perf numbers will be reportedIssue Tracking
JIRA ID: AIROCSHMEM-515