Take the wavefront width from the device, not an assumed 64

examples/06_message_passing/message_passing_gluon.py

@@ -17,6 +17,7 @@
 import torch.multiprocessing as mp
 from triton.experimental import gluon
 from triton.experimental.gluon import language as gl
+from triton.language.extra.hip import utils as hip_utils
 import triton
 
 import iris
@@ -42,8 +43,8 @@ def producer_kernel(
 
     # Compute start index of this block
     block_start = pid * BLOCK_SIZE
-    # Create a simple 1D layout for the arange operation (64 threads per warp for AMD)
-    layout: gl.constexpr = gl.BlockedLayout([1], [64], [1], [0])
+    # Create a simple 1D layout for the arange operation: one wavefront (64 lanes on CDNA3/4, 32 on CDNA5)
+    layout: gl.constexpr = gl.BlockedLayout([1], [hip_utils.num_threads()], [1], [0])
     offsets = block_start + gl.arange(0, BLOCK_SIZE, layout=layout)
 
     # Guard for out-of-bounds accesses
@@ -80,8 +81,8 @@ def consumer_kernel(
     pid = gl.program_id(0)
 
     block_start = pid * BLOCK_SIZE
-    # Create a simple 1D layout for the arange operation (64 threads per warp for AMD)
-    layout: gl.constexpr = gl.BlockedLayout([1], [64], [1], [0])
+    # Create a simple 1D layout for the arange operation: one wavefront (64 lanes on CDNA3/4, 32 on CDNA5)
+    layout: gl.constexpr = gl.BlockedLayout([1], [hip_utils.num_threads()], [1], [0])
     offsets = block_start + gl.arange(0, BLOCK_SIZE, layout=layout)
     mask = offsets < buffer_size
 

examples/32_gluon_all_gather_tracing/all_gather_tracing.py

@@ -41,6 +41,7 @@
 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
 
 import iris
 from iris.gluon import IrisDeviceCtx
@@ -92,11 +93,12 @@ def all_gather_put_kernel(
 
     pid = gl.program_id(0)
 
-    # AMD GPUs have 64 threads per warp (wavefront size 64).
-    # Total threads per CTA = NUM_WARPS * 64.
-    # Each thread handles SPT = BLOCK_SIZE // (NUM_WARPS * 64) elements.
-    SPT: gl.constexpr = BLOCK_SIZE // (NUM_WARPS * 64)
-    layout: gl.constexpr = gl.BlockedLayout([SPT], [64], [NUM_WARPS], [0])
+    # The wavefront is 64 wide on CDNA3/4 and 32 on CDNA5, so read the CTA's
+    # thread count from the launch options. Each thread handles
+    # SPT = BLOCK_SIZE // (NUM_WARPS * wavefront) elements.
+    THREADS: gl.constexpr = hip_utils.num_threads()
+    SPT: gl.constexpr = BLOCK_SIZE // THREADS
+    layout: gl.constexpr = gl.BlockedLayout([SPT], [THREADS // NUM_WARPS], [NUM_WARPS], [0])
     offsets = pid * BLOCK_SIZE + gl.arange(0, BLOCK_SIZE, layout=layout)
     mask = offsets < num_elements
 
@@ -128,13 +130,13 @@ def all_gather_put_kernel(
 def _launch(shmem, local_buf, global_buf, context_tensor, enable_tracing: bool):
     """Launch one iteration of the all-gather kernel.
 
-    Layout constraints for AMD GPUs (warp size = 64):
-      BLOCK_SIZE = sizePerThread * 64 * NUM_WARPS
-    Here NUM_WARPS=4, sizePerThread=1 → BLOCK_SIZE = 1 * 64 * 4 = 256.
+    Layout constraint: BLOCK_SIZE = sizePerThread * wavefront * NUM_WARPS.
+    Here NUM_WARPS=4 and BLOCK_SIZE=256, so sizePerThread is 1 on a 64-wide
+    wavefront (CDNA3/4) and 2 on a 32-wide one (CDNA5).
     """
     num_elements = local_buf.numel()
     NUM_WARPS = 4
-    BLOCK_SIZE = 64 * NUM_WARPS  # 256 elements per tile (1 el/thread, 64 threads/warp, 4 warps)
+    BLOCK_SIZE = 256  # elements per tile
     grid = ((num_elements + BLOCK_SIZE - 1) // BLOCK_SIZE,)
     all_gather_put_kernel[grid](
         IrisDeviceCtx,

iris/ccl/config.py

@@ -46,13 +46,10 @@ class Config:
                                 Only "two_shot" is supported
         num_stages: Number of pipeline stages for the kernel (default: 1)
         num_warps: Number of warps per workgroup (default: 4). For gluon kernels,
-                   this also sets WARPS_PER_CTA in the BlockedLayout. The product
-                   threads_per_warp * num_warps determines the minimum tile size
+                   this also sets the warps of the BlockedLayout. The product of
+                   the wavefront width (64 on CDNA3/4, 32 on CDNA5, read in-kernel)
+                   and num_warps determines the minimum tile size
                    (block_size_m * block_size_n for flat-2D, or block_size_n for 1D).
-        threads_per_warp: Threads per warp/wavefront (default: 64). Must match the
-                          hardware wavefront size: 64 for AMD GPUs, 32 for NVIDIA.
-                          Used by gluon kernels to construct BlockedLayout for
-                          vectorized memory access.
         waves_per_eu: Waves per execution unit hint for occupancy (default: 0, auto)
 
     Example:
@@ -92,7 +89,6 @@ class Config:
     reduce_scatter_variant: str = "two_shot"
     num_stages: int = 1
     num_warps: int = 4
-    threads_per_warp: int = 64
     waves_per_eu: int = 0
 
     def __post_init__(self):
@@ -144,7 +140,5 @@ def __post_init__(self):
         if self.reduce_scatter_variant != "two_shot":
             raise ValueError(f"reduce_scatter_variant must be 'two_shot', got '{self.reduce_scatter_variant}'")
 
-        if self.threads_per_warp not in (32, 64):
-            raise ValueError(f"threads_per_warp must be 32 (NVIDIA) or 64 (AMD), got {self.threads_per_warp}")
         if self.num_warps <= 0:
             raise ValueError(f"num_warps must be positive, got {self.num_warps}")

iris/ccl/gluon/all_gather.py

@@ -16,6 +16,7 @@
     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
 except ImportError as e:
     raise ValueError("Gluon is not available. Install Triton with Gluon support or set use_gluon=False.") from e
 
@@ -44,8 +45,6 @@ def persistent_all_gather_gluon(
     BLOCK_SIZE_N: gl.constexpr,
     GROUP_SIZE_M: gl.constexpr,
     COMM_SMS: gl.constexpr,
-    THREADS_PER_WARP: gl.constexpr,
-    WARPS_PER_CTA: gl.constexpr,
     TRACING: gl.constexpr = False,
 ):
     """
@@ -60,7 +59,7 @@ def persistent_all_gather_gluon(
     Memory layout (BlockedLayout):
         A 1D BlockedLayout distributes TOTAL_ELEMS = BLOCK_SIZE_M * BLOCK_SIZE_N
         elements across the thread hierarchy:
-            ELEMS_PER_THREAD = TOTAL_ELEMS // (THREADS_PER_WARP * WARPS_PER_CTA)
+            ELEMS_PER_THREAD = TOTAL_ELEMS // (threads per warp * num_warps)
 
         Each thread handles ELEMS_PER_THREAD contiguous elements in the
         flattened row-major order. Row/col are recovered via:
@@ -69,7 +68,7 @@ def persistent_all_gather_gluon(
 
     Constraints:
         - BLOCK_SIZE_M * BLOCK_SIZE_N must be a multiple of
-          (THREADS_PER_WARP * WARPS_PER_CTA).
+          (threads per warp * num_warps), checked at compile time.
         - Optimal tile: 2048-4096 total elements (8-16 per thread).
           Larger tiles cause register spilling and performance collapse.
         - Recommended: BLOCK_SIZE_M=8, BLOCK_SIZE_N=256 (2048 elems, 8/thread).
@@ -92,8 +91,6 @@ def persistent_all_gather_gluon(
         BLOCK_SIZE_N: Number of columns per tile.
         GROUP_SIZE_M: Swizzle group size for M-dimension tiling.
         COMM_SMS: Number of CUs used for persistent scheduling.
-        THREADS_PER_WARP: Threads per warp/wavefront (64 for AMD, 32 for NVIDIA).
-        WARPS_PER_CTA: Number of warps per workgroup. Must match num_warps.
     """
     ctx = IrisDeviceCtx.initialize(context_tensor, tracing=TRACING)
 
@@ -104,9 +101,18 @@ def persistent_all_gather_gluon(
     total_tiles = num_pid_m * num_pid_n
 
     # Flat 1D layout covering BLOCK_SIZE_M * BLOCK_SIZE_N elements
+    # Lanes and warps come from the compiler's launch options: the wavefront is
+    # 64 wide on CDNA3/4 but 32 on CDNA5, and the warp count is config.num_warps.
     TOTAL_ELEMS: gl.constexpr = BLOCK_SIZE_M * BLOCK_SIZE_N
-    ELEMS_PER_THREAD: gl.constexpr = TOTAL_ELEMS // (THREADS_PER_WARP * WARPS_PER_CTA)
-    flat_layout: gl.constexpr = gl.BlockedLayout([ELEMS_PER_THREAD], [THREADS_PER_WARP], [WARPS_PER_CTA], [0])
+    THREADS: gl.constexpr = hip_utils.num_threads()
+    WARPS: gl.constexpr = hip_utils.num_warps()
+    gl.static_assert(
+        TOTAL_ELEMS % THREADS == 0,
+        "Gluon all-gather requires block_size_m * block_size_n to be a positive multiple of "
+        "threads per warp * num_warps. Recommended: block_size_m=8, block_size_n=256.",
+    )
+    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.
@@ -194,26 +200,6 @@ def launch(
         block_size_m = 8
         block_size_n = 256
 
-    # Validate flat-2D layout constraints.
-    # TOTAL_ELEMS = BLOCK_SIZE_M * BLOCK_SIZE_N must be a multiple of
-    # THREADS_PER_WARP * WARPS_PER_CTA so each thread gets a whole
-    # number of elements.
-    total_elems = block_size_m * block_size_n
-    threads_per_cta = config.threads_per_warp * config.num_warps
-    if total_elems < threads_per_cta:
-        raise ValueError(
-            f"Gluon all-gather requires block_size_m * block_size_n >= "
-            f"threads_per_warp * num_warps ({threads_per_cta}), "
-            f"got {block_size_m} * {block_size_n} = {total_elems}."
-        )
-    if total_elems % threads_per_cta != 0:
-        raise ValueError(
-            f"Gluon all-gather requires block_size_m * block_size_n to be a "
-            f"multiple of threads_per_warp * num_warps ({threads_per_cta}), "
-            f"got {block_size_m} * {block_size_n} = {total_elems}. "
-            f"Recommended: block_size_m=8, block_size_n=256."
-        )
-
     context_tensor = ctx.get_device_context()
     tracing = getattr(ctx, "tracing", None)
     tracing_enabled = bool(tracing and getattr(tracing, "enabled", False))
@@ -240,8 +226,6 @@ def launch(
         block_size_n,
         config.swizzle_size,
         config.comm_sms,
-        config.threads_per_warp,
-        config.num_warps,
         tracing_enabled,
         num_stages=config.num_stages,
         num_warps=config.num_warps,

iris/ccl/gluon/all_to_all.py

@@ -11,6 +11,7 @@
 try:
     from triton.experimental import gluon
     from triton.experimental.gluon import language as gl
+    from triton.language.extra.hip import utils as hip_utils
 except ImportError as e:
     raise ValueError("Gluon is not available. Install Triton with Gluon support or set use_gluon=False.") from e
 
@@ -83,12 +84,14 @@ def persistent_all_to_all_gluon(
         pid_m = first_pid_m + ((tile_id % num_pid_in_group) % group_size_m)
         pid_n = (tile_id % num_pid_in_group) // group_size_m
 
-        # Optimized layout for maximum VGPR usage and dwordx4 vectorization
-        # Use layout that maximizes register utilization and enables wider loads
-        # For AMD: 64 threads/warp, 4 warps = 256 threads total
-        # BlockedLayout: [size_per_thread], [threads_per_warp], [warps_per_cta], [order]
-        layout_col: gl.constexpr = gl.BlockedLayout([1], [64], [4], [0])  # Column access
-        layout_row: gl.constexpr = gl.BlockedLayout([1], [64], [4], [0])  # Row indices
+        # BlockedLayout: [size_per_thread], [threads_per_warp], [warps_per_cta], [order].
+        # Lanes and warps come from the compiler's launch options rather than literals:
+        # the wavefront is 64 wide on CDNA3/4 but 32 on CDNA5, and the launch's warp
+        # count is config.num_warps.
+        WARPS: gl.constexpr = hip_utils.num_warps()
+        LANES: gl.constexpr = hip_utils.num_threads() // WARPS
+        layout_col: gl.constexpr = gl.BlockedLayout([1], [LANES], [WARPS], [0])  # Column access
+        layout_row: gl.constexpr = gl.BlockedLayout([1], [LANES], [WARPS], [0])  # Row indices
 
         rm = (pid_m * BLOCK_SIZE_M + gl.arange(0, BLOCK_SIZE_M, layout=layout_row)) % M
         rn = (pid_n * BLOCK_SIZE_N + gl.arange(0, BLOCK_SIZE_N, layout=layout_col)) % N

tests/ccl/test_all_gather_gluon.py

@@ -35,8 +35,9 @@
     "M, N, block_size_m, block_size_n",
     [
         # block_size_n must be a multiple of (threads_per_warp * num_warps).
-        # With defaults (threads_per_warp=64, num_warps=4), minimum is 256.
-        # elems_per_thread = block_size_n / 256: higher = wider vector loads.
+        # 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)

tests/unittests/test_atomic_add_gluon.py

@@ -5,6 +5,7 @@
 import pytest
 from triton.experimental import gluon
 from triton.experimental.gluon import language as gl
+from triton.language.extra.hip import utils as hip_utils
 import iris
 from iris.gluon import IrisDeviceCtx
 
@@ -23,7 +24,7 @@ def atomic_add_kernel(
     ctx = IrisDeviceCtx.initialize(context_tensor)
     pid = gl.program_id(0)
     block_start = pid * BLOCK_SIZE
-    layout: gl.constexpr = gl.BlockedLayout([1], [64], [1], [0])
+    layout: gl.constexpr = gl.BlockedLayout([1], [hip_utils.num_threads()], [1], [0])
     offsets = block_start + gl.arange(0, BLOCK_SIZE, layout=layout)
     mask = offsets < BLOCK_SIZE
 

tests/unittests/test_atomic_and_gluon.py

@@ -5,6 +5,7 @@
 import pytest
 from triton.experimental import gluon
 from triton.experimental.gluon import language as gl
+from triton.language.extra.hip import utils as hip_utils
 import iris
 from iris.gluon import IrisDeviceCtx
 
@@ -23,7 +24,7 @@ def atomic_and_kernel(
     ctx = IrisDeviceCtx.initialize(context_tensor)
     pid = gl.program_id(0)
     block_start = pid * BLOCK_SIZE
-    layout: gl.constexpr = gl.BlockedLayout([1], [64], [1], [0])
+    layout: gl.constexpr = gl.BlockedLayout([1], [hip_utils.num_threads()], [1], [0])
     offsets = block_start + gl.arange(0, BLOCK_SIZE, layout=layout)
     mask = offsets < BLOCK_SIZE
 

tests/unittests/test_atomic_max_gluon.py

@@ -5,6 +5,7 @@
 import pytest
 from triton.experimental import gluon
 from triton.experimental.gluon import language as gl
+from triton.language.extra.hip import utils as hip_utils
 import iris
 from iris.gluon import IrisDeviceCtx
 
@@ -23,7 +24,7 @@ def atomic_max_kernel(
     ctx = IrisDeviceCtx.initialize(context_tensor)
     pid = gl.program_id(0)
     block_start = pid * BLOCK_SIZE
-    layout: gl.constexpr = gl.BlockedLayout([1], [64], [1], [0])
+    layout: gl.constexpr = gl.BlockedLayout([1], [hip_utils.num_threads()], [1], [0])
     offsets = block_start + gl.arange(0, BLOCK_SIZE, layout=layout)
     mask = offsets < BLOCK_SIZE
 

tests/unittests/test_atomic_min_gluon.py

@@ -5,6 +5,7 @@
 import pytest
 from triton.experimental import gluon
 from triton.experimental.gluon import language as gl
+from triton.language.extra.hip import utils as hip_utils
 import iris
 from iris.gluon import IrisDeviceCtx
 
@@ -23,7 +24,7 @@ def atomic_min_kernel(
     ctx = IrisDeviceCtx.initialize(context_tensor)
     pid = gl.program_id(0)
     block_start = pid * BLOCK_SIZE
-    layout: gl.constexpr = gl.BlockedLayout([1], [64], [1], [0])
+    layout: gl.constexpr = gl.BlockedLayout([1], [hip_utils.num_threads()], [1], [0])
     offsets = block_start + gl.arange(0, BLOCK_SIZE, layout=layout)
     mask = offsets < BLOCK_SIZE
 

tests/unittests/test_atomic_or_gluon.py

@@ -5,6 +5,7 @@
 import pytest
 from triton.experimental import gluon
 from triton.experimental.gluon import language as gl
+from triton.language.extra.hip import utils as hip_utils
 import iris
 from iris.gluon import IrisDeviceCtx
 
@@ -23,7 +24,7 @@ def atomic_or_kernel(
     ctx = IrisDeviceCtx.initialize(context_tensor)
     pid = gl.program_id(0)
     block_start = pid * BLOCK_SIZE
-    layout: gl.constexpr = gl.BlockedLayout([1], [64], [1], [0])
+    layout: gl.constexpr = gl.BlockedLayout([1], [hip_utils.num_threads()], [1], [0])
     offsets = block_start + gl.arange(0, BLOCK_SIZE, layout=layout)
     mask = offsets < BLOCK_SIZE
 

tests/unittests/test_atomic_xor_gluon.py

@@ -5,6 +5,7 @@
 import pytest
 from triton.experimental import gluon
 from triton.experimental.gluon import language as gl
+from triton.language.extra.hip import utils as hip_utils
 import iris
 from iris.gluon import IrisDeviceCtx
 
@@ -23,7 +24,7 @@ def atomic_xor_kernel(
     ctx = IrisDeviceCtx.initialize(context_tensor)
     pid = gl.program_id(0)
     block_start = pid * BLOCK_SIZE
-    layout: gl.constexpr = gl.BlockedLayout([1], [64], [1], [0])
+    layout: gl.constexpr = gl.BlockedLayout([1], [hip_utils.num_threads()], [1], [0])
     offsets = block_start + gl.arange(0, BLOCK_SIZE, layout=layout)
     mask = offsets < BLOCK_SIZE
 

tests/unittests/test_copy_gluon.py

@@ -5,6 +5,7 @@
 import pytest
 from triton.experimental import gluon
 from triton.experimental.gluon import language as gl
+from triton.language.extra.hip import utils as hip_utils
 import iris
 from iris.gluon import IrisDeviceCtx
 
@@ -23,7 +24,7 @@ def copy_get_kernel(
     ctx = IrisDeviceCtx.initialize(context_tensor)
     pid = gl.program_id(0)
     block_start = pid * BLOCK_SIZE
-    layout: gl.constexpr = gl.BlockedLayout([1], [64], [1], [0])
+    layout: gl.constexpr = gl.BlockedLayout([1], [hip_utils.num_threads()], [1], [0])
     offsets = block_start + gl.arange(0, BLOCK_SIZE, layout=layout)
     mask = offsets < BLOCK_SIZE
 
@@ -47,7 +48,7 @@ def copy_put_kernel(
     ctx = IrisDeviceCtx.initialize(context_tensor)
     pid = gl.program_id(0)
     block_start = pid * BLOCK_SIZE
-    layout: gl.constexpr = gl.BlockedLayout([1], [64], [1], [0])
+    layout: gl.constexpr = gl.BlockedLayout([1], [hip_utils.num_threads()], [1], [0])
     offsets = block_start + gl.arange(0, BLOCK_SIZE, layout=layout)
     mask = offsets < BLOCK_SIZE
 
@@ -71,7 +72,7 @@ def copy_local_kernel(
     ctx = IrisDeviceCtx.initialize(context_tensor)
     pid = gl.program_id(0)
     block_start = pid * BLOCK_SIZE
-    layout: gl.constexpr = gl.BlockedLayout([1], [64], [1], [0])
+    layout: gl.constexpr = gl.BlockedLayout([1], [hip_utils.num_threads()], [1], [0])
     offsets = block_start + gl.arange(0, BLOCK_SIZE, layout=layout)
     mask = offsets < BLOCK_SIZE