Skip to content

Copy-engine support - #426

Open
dsidler wants to merge 70 commits into
mainfrom
dev/dasidler/sdma
Open

dsidler wants to merge 70 commits into
mainfrom
dev/dasidler/sdma

Conversation

@dsidler

@dsidler dsidler commented Mar 5, 2026

Copy link
Copy Markdown
Collaborator

No description provided.

Comment thread iris/host/memory/symmetric_heap.py
Comment thread iris/mem/triton/ops.py
Comment thread iris/mem/triton/ops.py
Comment thread tests/examples/test_message_passing.py
Comment thread tests/examples/test_message_passing.py
Comment thread tests/unittests/test_copy_engine_ops.py
@dsidler
dsidler marked this pull request as ready for review June 9, 2026 22:47
@dsidler

dsidler commented Jun 11, 2026

Copy link
Copy Markdown
Collaborator Author

@mawad-amd not sure how useful the changes and copy egine path for examples 07 and 10 are. I might revert those and only keep the changes for message-passing example

@mawad-amd

Copy link
Copy Markdown
Collaborator

Thanks, David! I think we will want to keep the benchmark you used for your results. The message passing one would be a great introductory example.

Will review this PR soon.

@mawad-amd mawad-amd left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Some minor comments. Thanks a lot for this David!

Please fix/resolve conflict, and just make sure CI passes and merge.

Comment thread iris/device/sdma_utils.py


@triton.jit
def wait_cnt():

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we move this to iris/mem/utils.py? Can we also guard for cdna3/4 and throw/return something for anything else?
Something like https://github.com/ROCm/iris/blob/main/iris/mem/utils.py#L69

Comment thread iris/host/iris.py
return ptr, bits

def put(
self,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Use similar naming scheme to the device side one

iris/iris/mem/triton/ops.py

Lines 281 to 299 in 9ff7048

def put(
from_ptr,
to_ptr,
from_rank,
to_rank,
heap_bases,
mask=None,
other=None,
load_cache_modifier=None,
store_cache_modifier=None,
hint: tl.constexpr = None,
copy_engine_ctx: tl.tensor = None,
src_row_stride: tl.constexpr = 0,
dst_row_stride: tl.constexpr = 0,
USE_COPY_ENGINE: tl.constexpr = False,
CONTIGUOUS_COPY: tl.constexpr = False,
from_base_ptr=None,
to_base_ptr=None,
):

Comment thread iris/host/iris.py
def put_tile(
self,
tile,
dst_rank: int,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same comment on dst_*

Comment thread iris/mem/triton/ops.py

Args:
from_ptr (triton.PointerType, or block of dtype=triton.PointerType): Pointer in the current rank's local memory from which to read data.
to_ptr (triton.PointerType, or block of dtype=triton.PointerType): Pointer in the current rank's address space that will be translated to the `to_rank`'s address space. Must be the current rank where the pointer is local.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why was this removed?

Comment thread iris/mem/triton/ops.py
copy_engine_ctx: tl.tensor = None,
src_row_stride: tl.constexpr = 0,
dst_row_stride: tl.constexpr = 0,
USE_COPY_ENGINE: tl.constexpr = False,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

let's just use lower case for all new arguments/constexpr

Comment thread iris/mem/triton/ops.py
scope=None,
hint: tl.constexpr = None,
copy_engine_ctx=None,
USE_COPY_ENGINE: tl.constexpr = False,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

case

Comment thread iris/mem/triton/ops.py
scope=None,
hint: tl.constexpr = None,
copy_engine_ctx=None,
USE_COPY_ENGINE: tl.constexpr = False,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

case

Comment thread pyproject.toml
"ruff",
"tritonblas @ git+https://github.com/ROCm/tritonBLAS.git@muhaawad/iris",

"rocm-xio @ git+https://github.com/ROCm/rocm-xio.git",

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This will likely break the simulation backend but let's merge it and see.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

in-progress We are working on it iris Iris project issue

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants