Skip to content

Gluon all-gather reloads peer heap bases inside the tile loop #591

Open

Description

@mawad-amd

Gluon all-gather reloads peer heap bases inside the tile loop

Follow-up to #585, where this was raised and deferred.

1. Hoist the peer heap bases

iris/ccl/gluon/all_gather.py:

113:  local_base = gl.load(ctx.heap_bases + iris_rank)      # hoisted, once
115:  for tile_id in range(pid, total_tiles, COMM_SMS):
147:      for rank_idx in range(world_size):
150:          output_ptrs = output_ptr + output_offsets      # loop-invariant
158:              target_base = gl.load(ctx.heap_bases + target_iris_rank)   # per dest, per tile

The comment at line 111 says the hoist "eliminates redundant gl.load(heap_bases)
calls in the inner store loop", but only the local base was hoisted. At 8 ranks
the loop still does 7 peer loads per tile. heap_bases, rank_start, rank_stride
and world_size are all tile-invariant, so each destination's ptr_delta can be
computed once before the tile loop — the same fix #585 applied to the TDM path in
d626261b. output_ptrs at line 150 is invariant across rank_idx too.

2. INLINE grouping, deferred from #585

From that review thread:

The grouping idea seems most useful for the register store path (gl.store),
where it bounds live values; I'd leave that to a later PR.

That is this file. These stores hold data live in registers across the whole
destination loop, and the docstring already warns the path is register-sensitive:
"Larger tiles cause register spilling and performance collapse."

Work

  • Hoist the per-destination ptr_delta out of the tile loop
  • Hoist output_ptrs out of the rank_idx loop
  • Confirm in the ISA whether range(world_size) at line 147 unrolls — it is
    range, not tl.static_range, and the grouping argument rests on it
  • Try INLINE grouping at 4 and 8; measure VGPR count and spill, not just time
  • Check iris/ccl/gluon/all_to_all.py:130 for the same pattern
  • Correctness vs torch.distributed on 2/4/8 ranks before timing
  • Before/after through benchmark/ccl/bench_all_gather.py

We expect a performance improvement here; it needs measuring.

Activity

  1. added
    irisIris project issue
    pythonPull requests that update python code
    on Oct 15, 2026
  2. vishalshetagar commented on Oct 18, 2026

    @vishalshetagar
    Contributor

    I’d like to work on this issue. I’ll start with hoisting the address calculations, then investigate INLINE grouping and performance. I’ll open a draft PR for feedback once the initial changes are ready.

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    gluongood first issueGood for newcomershelp wantedExtra attention is neededirisIris project issuepythonPull requests that update python code

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions