You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
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, once115: fortile_idinrange(pid, total_tiles, COMM_SMS):
147: forrank_idxinrange(world_size):
150: output_ptrs=output_ptr+output_offsets# loop-invariant158: 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.
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.
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.
Activity
vishalshetagar commentedon Oct 18, 2026
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.