ROCm / ROCm/iris

all_reduce_one_shot / all_reduce_two_shot use hardcoded lock value 1

Open
#465 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug iris
Dominant language
Python
Stars
202
Forks
47
Avg merge
6d 11h
Merged PRs (30d)
4

Description

Bug

In iris/x/all_reduce.py, all_reduce_one_shot and all_reduce_two_shot use a hardcoded value of 1 to signal "tile ready":

  • Producers write: tl.atomic_xchg(lock_ptr, 1, sem="release")
  • Consumers spin: while iris.atomic_add(lock_ptr, 0, ...) != 1: pass

Between calls, the lock array must be zeroed back to 0 via a collective shmem.zeros + barrier, adding overhead to every kernel invocation. If the lock array is not properly zeroed (e.g., due to workspace reuse or error), consumers see lock == 1 from a previous call and read stale data.

Impact

  • Per-call overhead from mandatory lock zeroing + barrier between invocations
  • Fragile: skipping the zeroing step silently produces wrong results
  • Prevents efficient workspace reuse across calls

Fix

Add a call_number parameter to both functions:

  • Producers signal with: tl.atomic_xchg(lock_ptr, call_number, sem="release", scope="sys")
  • Consumers spin on: while iris.atomic_add(lock_ptr, 0, ...) != call_number: pass

Add a monotonically increasing call_counter field to FusedWorkspace, incremented on every matmul_all_reduce call. Each call uses a new version number, so stale locks from previous calls are automatically ignored without zeroing.

Component

iris/x/all_reduce.py, iris/ops/workspace.py

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in iris/x/all_reduce.py with all_reduce_one_shot and all_reduce_two_shot, then trace matmul_all_reduce and FusedWorkspace in iris/ops/workspace.py. Verify that each call receives a new call number, producers and consumers use it for lock signaling, and the previous lock-zeroing collective is no longer required.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Clearly specified
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.