Hoist peer address calculations in Gluon all-gather

benchmark/ccl/bench_all_gather.py

@@ -16,7 +16,7 @@
 @bench.axis("M", M_VALUES)
 @bench.axis("N", N_VALUES)
 @bench.axis("dtype", DTYPES)
-@bench.axis("variant", ["persistent", "tdm"])
+@bench.axis("variant", ["persistent", "gluon", "tdm"])
 @bench.axis("comm_sms", [64, 128, 256])
 @bench.axis("backend", ["iris", "rccl"])
 def all_gather(state, ctx):
@@ -55,6 +55,8 @@ def all_gather(state, ctx):
             num_warps=16,
             comm_sms=comm_sms,
         )
+    elif variant == "gluon":
+        config = Config(use_gluon=True, block_size_m=8, block_size_n=256, comm_sms=comm_sms)
     else:
         config = Config(comm_sms=comm_sms)
 

iris/ccl/gluon/all_gather.py

@@ -13,7 +13,6 @@
 """
 
 try:
-    import triton.language as tl
     from triton.experimental import gluon
     from triton.experimental.gluon import language as gl
     from triton.language.extra.hip import utils as hip_utils
@@ -114,9 +113,15 @@ def persistent_all_gather_gluon(
     ELEMS_PER_THREAD: gl.constexpr = TOTAL_ELEMS // THREADS
     flat_layout: gl.constexpr = gl.BlockedLayout([ELEMS_PER_THREAD], [THREADS // WARPS], [WARPS], [0])
 
-    # Hoist local heap base outside the tile loop: eliminates redundant
-    # gl.load(heap_bases) calls in the inner store loop.
+    # Heap bases do not change between tiles.
     local_base = gl.load(ctx.heap_bases + iris_rank)
+    output_bases = (output_ptr,)
+    for rank_idx in gl.static_range(1, world_size):
+        dest_idx = (group_rank + rank_idx) % world_size
+        target_iris_rank = rank_start + dest_idx * rank_stride
+        ptr_delta = gl.load(ctx.heap_bases + target_iris_rank) - local_base
+        remote_base = gl.cast(gl.cast(output_ptr, gl.uint64) + ptr_delta, output_ptr.dtype)
+        output_bases += (remote_base,)
 
     for tile_id in range(pid, total_tiles, COMM_SMS):
         # Swizzled tile index computation for better L2 locality
@@ -146,26 +151,16 @@ def persistent_all_gather_gluon(
         # Output: this rank's data goes to output[group_rank * M + row, col]
         output_row = group_rank * M + row
         output_offsets = output_row * stride_out_m + col * stride_out_n
+        output_ptrs = output_ptr + output_offsets
 
         # Traffic-shaped stores to all ranks: stagger write order per rank
         # so each rank writes to a different target at any given moment,
         # avoiding memory controller contention on the receiver side.
-        for rank_idx in range(world_size):
-            dest_idx = (group_rank + rank_idx) % world_size
-            target_iris_rank = rank_start + dest_idx * rank_stride
-            output_ptrs = output_ptr + output_offsets
-
-            if dest_idx == group_rank:
+        for rank_idx in gl.static_range(world_size):
+            if rank_idx == 0:
                 gl.store(output_ptrs, data, mask=mask, cache_modifier=".wt")
             else:
-                # Hoisted translation: compute ptr_delta from pre-loaded
-                # local_base rather than calling ctx.store() which would
-                # do 2x gl.load(heap_bases) per call.
-                target_base = gl.load(ctx.heap_bases + target_iris_rank)
-                ptr_delta = target_base - local_base
-                output_ptrs_int = tl.cast(output_ptrs, gl.uint64)
-                remote_ptrs_int = output_ptrs_int + ptr_delta
-                remote_ptrs = tl.cast(remote_ptrs_int, output_ptrs.dtype)
+                remote_ptrs = output_bases[rank_idx] + output_offsets
                 gl.store(remote_ptrs, data, mask=mask)
 
 

tests/ccl/test_all_gather_gluon.py

@@ -32,18 +32,20 @@
     ],
 )
 @pytest.mark.parametrize(
-    "M, N, block_size_m, block_size_n",
+    "M, N, block_size_m, block_size_n, comm_sms",
     [
         # block_size_n must be a multiple of (threads_per_warp * num_warps).
         # With num_warps=4 the minimum is 256 on a 64-wide wavefront, 128 on a 32-wide one.
         # elems_per_thread = block_size_n / (threads_per_warp * num_warps): higher = wider
         # vector loads. The per-case counts below are for a 64-wide wavefront; double them at 32.
-        (256, 256, 32, 256),  # Small: elems_per_thread=1 (scalar loads)
-        (1024, 512, 32, 512),  # Medium: elems_per_thread=2 (dword loads)
-        (8192, 8192, 32, 1024),  # Large: elems_per_thread=4 (dwordx4, optimal)
+        (256, 256, 32, 256, 64),  # Small: elems_per_thread=1 (scalar loads)
+        (1024, 512, 32, 512, 64),  # Medium: elems_per_thread=2 (dword loads)
+        (8192, 8192, 32, 1024, 64),  # Large: elems_per_thread=4 (dwordx4, optimal)
+        (13, 300, 8, 256, 1),  # Reuse peer addresses across tiles with partial rows and columns.
+        (7, 64, 8, 256, 1),
     ],
 )
-def test_all_gather_gluon(dtype, M, N, block_size_m, block_size_n):
+def test_all_gather_gluon(dtype, M, N, block_size_m, block_size_n, comm_sms):
     """Test all-gather functionality using Gluon by comparing against PyTorch's implementation."""
     # Ensure torch.distributed is initialized (should be done by test runner)
     if not dist.is_initialized():
@@ -61,8 +63,6 @@ def test_all_gather_gluon(dtype, M, N, block_size_m, block_size_n):
     # Each rank has an M x N input tensor
     # Output is (world_size * M, N) - concatenated along dimension 0
     pytorch_input_tensor = torch.randn(M, N, dtype=dtype, device=f"cuda:{rank}")
-    # Fill with deterministic values for easier debugging
-    pytorch_input_tensor.fill_(float(rank + 1))
 
     # Create output tensor for PyTorch: (world_size * M, N)
     pytorch_output_tensor = torch.zeros(world_size * M, N, dtype=dtype, device=f"cuda:{rank}")
@@ -80,19 +80,13 @@ def test_all_gather_gluon(dtype, M, N, block_size_m, block_size_n):
 
     # Run Iris Gluon all_gather
     shmem.barrier()
-    config = Config(use_gluon=True, block_size_m=block_size_m, block_size_n=block_size_n)
+    config = Config(use_gluon=True, block_size_m=block_size_m, block_size_n=block_size_n, comm_sms=comm_sms)
     shmem.ccl.all_gather(iris_output_tensor, iris_input_tensor, config=config)
     torch.cuda.synchronize()
 
     # Compare results
-    atol = 1e-3 if dtype == torch.float16 else 1e-5
-    max_diff = torch.abs(iris_output_tensor - pytorch_output_tensor).max().item()
-
     try:
-        assert torch.allclose(iris_output_tensor, pytorch_output_tensor, atol=atol), (
-            f"Max difference: {max_diff}, expected < {atol}\n"
-            f"Rank {rank}: Iris Gluon output doesn't match PyTorch's all_gather_into_tensor"
-        )
+        torch.testing.assert_close(iris_output_tensor, pytorch_output_tensor, rtol=0, atol=0)
     finally:
         # Final barrier to ensure all ranks complete before test cleanup
         # This helps with test isolation when running multiple tests

tests/ccl/test_process_groups.py

@@ -178,7 +178,7 @@ def test_all_reduce_with_groups(variant, group_type, dtype=torch.float32, M=256,
 # =============================================================================
 
 
-@pytest.mark.parametrize("variant", ["persistent", "tdm"])
+@pytest.mark.parametrize("variant", ["persistent", "gluon", "tdm"])
 @pytest.mark.parametrize("group_type", ["consecutive", "strided"])
 def test_all_gather_with_groups(group_type, variant, dtype=torch.float32, M=128, N=64):
     """Test all-gather with ProcessGroups."""
@@ -222,7 +222,7 @@ def test_all_gather_with_groups(group_type, variant, dtype=torch.float32, M=128,
         # N=64 fp32 is a 256-byte row: one 32 x 64 tile per 32 rows.
         config = Config(use_gluon=True, all_gather_variant="tdm", block_size_m=32, block_size_n=64)
     else:
-        config = Config()
+        config = Config(use_gluon=variant == "gluon")
     shmem.ccl.all_gather(iris_output_tensor, iris_input_tensor, group=my_group, config=config)
     torch.cuda.synchronize()