From 4d9c0d2545ddbc92fbd11fec7086546d917e574c Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram <96096277+nirvedhmeshram@users.noreply.github.com> Date: Thu, 15 Oct 2026 16:45:19 +0000 Subject: [PATCH 1/5] Take the wavefront width from the device, not an assumed 64 gfx1250 is an AMD part with a 32-wide wavefront. Two places assumed AMD means 64 lanes, and on gfx1250 both Gluon store-path collectives failed to compile ("Result has an invalid layout ... threadsPerWarp = [64]"): * Config.threads_per_warp defaulted to 64, and its validator described 32 as NVIDIA and 64 as AMD. It now defaults to the current device's warp_size; an explicit value still wins. * gluon/all_to_all.py built its BlockedLayout from literal [64] lanes and [4] warps. It now reads both from the launch (hip_utils.num_threads and num_warps), so it also follows config.num_warps. On gfx1250 at world size 4, tests/ccl/test_all_gather_gluon.py goes from 0/9 to 9/9 and test_all_to_all_gluon.py from 0/9 to 9/9; a new test checks the default follows the device. The rest of tests/ccl is unchanged. From #566 (Muhammad Awad), which found and fixed this alongside its TDM all-gather. Co-authored-by: Muhammad Awad Co-Authored-By: Claude Opus 5.5 (1M context) --- iris/ccl/config.py | 18 ++++++++++++------ iris/ccl/gluon/all_gather.py | 2 +- iris/ccl/gluon/all_to_all.py | 15 +++++++++------ tests/ccl/test_all_gather_gluon.py | 15 +++++++++++++-- 4 files changed, 35 insertions(+), 15 deletions(-) diff --git a/iris/ccl/config.py b/iris/ccl/config.py index 81188c5fb..9ebd052da 100644 --- a/iris/ccl/config.py +++ b/iris/ccl/config.py @@ -6,6 +6,9 @@ """ from dataclasses import dataclass + +import torch + import iris @@ -49,10 +52,11 @@ class Config: this also sets WARPS_PER_CTA in the BlockedLayout. The product threads_per_warp * 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. + threads_per_warp: Threads per warp/wavefront (default: None, the current + device's warp size). Must match the hardware wavefront + size; both 32 and 64 occur on AMD GPUs (64 on CDNA, 32 on + gfx1250). 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 +96,7 @@ class Config: reduce_scatter_variant: str = "two_shot" num_stages: int = 1 num_warps: int = 4 - threads_per_warp: int = 64 + threads_per_warp: int | None = None waves_per_eu: int = 0 def __post_init__(self): @@ -144,7 +148,9 @@ 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 is None: + self.threads_per_warp = torch.cuda.get_device_properties(torch.cuda.current_device()).warp_size 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}") + raise ValueError(f"threads_per_warp must be 32 or 64 (the wavefront width), got {self.threads_per_warp}") if self.num_warps <= 0: raise ValueError(f"num_warps must be positive, got {self.num_warps}") diff --git a/iris/ccl/gluon/all_gather.py b/iris/ccl/gluon/all_gather.py index 6befb9fcf..4a5b2ebb7 100644 --- a/iris/ccl/gluon/all_gather.py +++ b/iris/ccl/gluon/all_gather.py @@ -92,7 +92,7 @@ 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). + THREADS_PER_WARP: Threads per warp/wavefront: 64 on CDNA, 32 on gfx1250 and NVIDIA. WARPS_PER_CTA: Number of warps per workgroup. Must match num_warps. """ ctx = IrisDeviceCtx.initialize(context_tensor, tracing=TRACING) diff --git a/iris/ccl/gluon/all_to_all.py b/iris/ccl/gluon/all_to_all.py index 315b9b0d8..6c6596491 100644 --- a/iris/ccl/gluon/all_to_all.py +++ b/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 CDNA but 32 on gfx1250, and the launch's + # num_warps 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 diff --git a/tests/ccl/test_all_gather_gluon.py b/tests/ccl/test_all_gather_gluon.py index 3c57baead..e9666ed2c 100644 --- a/tests/ccl/test_all_gather_gluon.py +++ b/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) @@ -103,3 +104,13 @@ def test_all_gather_gluon(dtype, M, N, block_size_m, block_size_n): import gc gc.collect() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_config_threads_per_warp_follows_device(): + """The default wavefront width is the device's, not an assumed 64: gfx1250 is 32 wide.""" + expected = torch.cuda.get_device_properties(torch.cuda.current_device()).warp_size + assert Config().threads_per_warp == expected + assert Config(threads_per_warp=64).threads_per_warp == 64 + with pytest.raises(ValueError, match="threads_per_warp must be 32 or 64"): + Config(threads_per_warp=16) From 05e62255b7a6575be547e9bc7ed0ca2556e76c59 Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram <96096277+nirvedhmeshram@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:34:25 +0000 Subject: [PATCH 2/5] Address review: device-neutral all-to-all layout, cached warp-size default - Config: cache the per-device warp_size query and fall back to 64 when no GPU is available, so Config() stays constructible on CPU-only hosts. - all_to_all (gluon): drop the HIP-only num_threads() helper; take THREADS_PER_WARP from config.threads_per_warp and the warp count from gl.num_warps(), matching the gluon all-gather. - Test the no-GPU fallback. Co-authored-by: Muhammad Awad Co-Authored-By: Claude Opus 5.5 (1M context) --- iris/ccl/config.py | 17 +++++++++++++++-- iris/ccl/gluon/all_to_all.py | 16 ++++++++-------- tests/ccl/test_all_gather_gluon.py | 8 ++++++++ 3 files changed, 31 insertions(+), 10 deletions(-) diff --git a/iris/ccl/config.py b/iris/ccl/config.py index 9ebd052da..55cb170ff 100644 --- a/iris/ccl/config.py +++ b/iris/ccl/config.py @@ -5,6 +5,7 @@ Configuration structures for iris-ccl collective operations. """ +import functools from dataclasses import dataclass import torch @@ -12,6 +13,18 @@ import iris +@functools.lru_cache(maxsize=None) +def _device_warp_size(device_index: int) -> int: + return torch.cuda.get_device_properties(device_index).warp_size + + +def _default_threads_per_warp() -> int: + """The current device's wavefront width, or 64 (the previous default) without a GPU.""" + if not torch.cuda.is_available(): + return 64 + return _device_warp_size(torch.cuda.current_device()) + + @dataclass class Config: """ @@ -53,7 +66,7 @@ class Config: threads_per_warp * 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: None, the current - device's warp size). Must match the hardware wavefront + device's warp size; 64 when no GPU is available). Must match the hardware wavefront size; both 32 and 64 occur on AMD GPUs (64 on CDNA, 32 on gfx1250). Used by gluon kernels to construct BlockedLayout for vectorized memory access. @@ -149,7 +162,7 @@ def __post_init__(self): raise ValueError(f"reduce_scatter_variant must be 'two_shot', got '{self.reduce_scatter_variant}'") if self.threads_per_warp is None: - self.threads_per_warp = torch.cuda.get_device_properties(torch.cuda.current_device()).warp_size + self.threads_per_warp = _default_threads_per_warp() if self.threads_per_warp not in (32, 64): raise ValueError(f"threads_per_warp must be 32 or 64 (the wavefront width), got {self.threads_per_warp}") if self.num_warps <= 0: diff --git a/iris/ccl/gluon/all_to_all.py b/iris/ccl/gluon/all_to_all.py index 6c6596491..5ac32f26c 100644 --- a/iris/ccl/gluon/all_to_all.py +++ b/iris/ccl/gluon/all_to_all.py @@ -11,7 +11,6 @@ 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 @@ -58,6 +57,7 @@ def persistent_all_to_all_gluon( COMM_SMS: gl.constexpr, NUM_XCDS: gl.constexpr, CHUNK_SIZE: gl.constexpr, + THREADS_PER_WARP: gl.constexpr, ): """ Persistent all-to-all kernel using Gluon. @@ -85,13 +85,12 @@ def persistent_all_to_all_gluon( pid_n = (tile_id % num_pid_in_group) // group_size_m # 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 CDNA but 32 on gfx1250, and the launch's - # num_warps 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 + # No literals: the wavefront is 64 wide on CDNA but 32 on gfx1250 + # (THREADS_PER_WARP = config.threads_per_warp, which defaults to the device's), + # and the launch's warp count is config.num_warps. + WARPS: gl.constexpr = gl.num_warps() + layout_col: gl.constexpr = gl.BlockedLayout([1], [THREADS_PER_WARP], [WARPS], [0]) # Column access + layout_row: gl.constexpr = gl.BlockedLayout([1], [THREADS_PER_WARP], [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 @@ -198,6 +197,7 @@ def launch( config.comm_sms, config.num_xcds, config.chunk_size, + config.threads_per_warp, num_stages=config.num_stages, num_warps=config.num_warps, waves_per_eu=config.waves_per_eu, diff --git a/tests/ccl/test_all_gather_gluon.py b/tests/ccl/test_all_gather_gluon.py index e9666ed2c..ad35e8fce 100644 --- a/tests/ccl/test_all_gather_gluon.py +++ b/tests/ccl/test_all_gather_gluon.py @@ -114,3 +114,11 @@ def test_config_threads_per_warp_follows_device(): assert Config(threads_per_warp=64).threads_per_warp == 64 with pytest.raises(ValueError, match="threads_per_warp must be 32 or 64"): Config(threads_per_warp=16) + + +def test_default_threads_per_warp_without_gpu(monkeypatch): + """Without a GPU the default falls back to 64 instead of raising.""" + from iris.ccl import config as config_module + + monkeypatch.setattr(config_module.torch.cuda, "is_available", lambda: False) + assert config_module._default_threads_per_warp() == 64 From 85f3442777db37c2f327be0e283a4c35f73cb708 Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram <96096277+nirvedhmeshram@users.noreply.github.com> Date: Thu, 15 Oct 2026 18:16:13 +0000 Subject: [PATCH 3/5] Address review: name architectures, drop warp-size cache and CPU fallback - Describe wavefront width by architecture (64 on CDNA3/4, 32 on CDNA5), matching upstream Triton, and drop the NVIDIA mention. - Read the device warp size directly in Config.__post_init__: the query is sub-microsecond, so the lru_cache bought nothing, and the tests already require a GPU, so the no-GPU fallback and its test are removed. Co-authored-by: Muhammad Awad Co-Authored-By: Claude Opus 5.5 (1M context) --- iris/ccl/config.py | 22 ++++------------------ iris/ccl/gluon/all_gather.py | 2 +- iris/ccl/gluon/all_to_all.py | 2 +- tests/ccl/test_all_gather_gluon.py | 10 +--------- 4 files changed, 7 insertions(+), 29 deletions(-) diff --git a/iris/ccl/config.py b/iris/ccl/config.py index 55cb170ff..05ecf0035 100644 --- a/iris/ccl/config.py +++ b/iris/ccl/config.py @@ -5,7 +5,6 @@ Configuration structures for iris-ccl collective operations. """ -import functools from dataclasses import dataclass import torch @@ -13,18 +12,6 @@ import iris -@functools.lru_cache(maxsize=None) -def _device_warp_size(device_index: int) -> int: - return torch.cuda.get_device_properties(device_index).warp_size - - -def _default_threads_per_warp() -> int: - """The current device's wavefront width, or 64 (the previous default) without a GPU.""" - if not torch.cuda.is_available(): - return 64 - return _device_warp_size(torch.cuda.current_device()) - - @dataclass class Config: """ @@ -66,10 +53,9 @@ class Config: threads_per_warp * 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: None, the current - device's warp size; 64 when no GPU is available). Must match the hardware wavefront - size; both 32 and 64 occur on AMD GPUs (64 on CDNA, 32 on - gfx1250). Used by gluon kernels to construct BlockedLayout - for vectorized memory access. + device's warp size). Must match the hardware wavefront + size: 64 on CDNA3/4, 32 on CDNA5. 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: @@ -162,7 +148,7 @@ def __post_init__(self): raise ValueError(f"reduce_scatter_variant must be 'two_shot', got '{self.reduce_scatter_variant}'") if self.threads_per_warp is None: - self.threads_per_warp = _default_threads_per_warp() + self.threads_per_warp = torch.cuda.get_device_properties(torch.cuda.current_device()).warp_size if self.threads_per_warp not in (32, 64): raise ValueError(f"threads_per_warp must be 32 or 64 (the wavefront width), got {self.threads_per_warp}") if self.num_warps <= 0: diff --git a/iris/ccl/gluon/all_gather.py b/iris/ccl/gluon/all_gather.py index 4a5b2ebb7..eac91d66a 100644 --- a/iris/ccl/gluon/all_gather.py +++ b/iris/ccl/gluon/all_gather.py @@ -92,7 +92,7 @@ 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 on CDNA, 32 on gfx1250 and NVIDIA. + THREADS_PER_WARP: Threads per warp/wavefront: 64 on CDNA3/4, 32 on CDNA5. WARPS_PER_CTA: Number of warps per workgroup. Must match num_warps. """ ctx = IrisDeviceCtx.initialize(context_tensor, tracing=TRACING) diff --git a/iris/ccl/gluon/all_to_all.py b/iris/ccl/gluon/all_to_all.py index 5ac32f26c..1798afe4b 100644 --- a/iris/ccl/gluon/all_to_all.py +++ b/iris/ccl/gluon/all_to_all.py @@ -85,7 +85,7 @@ def persistent_all_to_all_gluon( pid_n = (tile_id % num_pid_in_group) // group_size_m # BlockedLayout: [size_per_thread], [threads_per_warp], [warps_per_cta], [order]. - # No literals: the wavefront is 64 wide on CDNA but 32 on gfx1250 + # No literals: the wavefront is 64 wide on CDNA3/4 but 32 on CDNA5 # (THREADS_PER_WARP = config.threads_per_warp, which defaults to the device's), # and the launch's warp count is config.num_warps. WARPS: gl.constexpr = gl.num_warps() diff --git a/tests/ccl/test_all_gather_gluon.py b/tests/ccl/test_all_gather_gluon.py index ad35e8fce..d09db1b26 100644 --- a/tests/ccl/test_all_gather_gluon.py +++ b/tests/ccl/test_all_gather_gluon.py @@ -108,17 +108,9 @@ def test_all_gather_gluon(dtype, M, N, block_size_m, block_size_n): @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") def test_config_threads_per_warp_follows_device(): - """The default wavefront width is the device's, not an assumed 64: gfx1250 is 32 wide.""" + """The default wavefront width is the device's, not an assumed 64: CDNA5 is 32 wide.""" expected = torch.cuda.get_device_properties(torch.cuda.current_device()).warp_size assert Config().threads_per_warp == expected assert Config(threads_per_warp=64).threads_per_warp == 64 with pytest.raises(ValueError, match="threads_per_warp must be 32 or 64"): Config(threads_per_warp=16) - - -def test_default_threads_per_warp_without_gpu(monkeypatch): - """Without a GPU the default falls back to 64 instead of raising.""" - from iris.ccl import config as config_module - - monkeypatch.setattr(config_module.torch.cuda, "is_available", lambda: False) - assert config_module._default_threads_per_warp() == 64 From 51e13b4447695c3594bc1aced3f9b4964aeaf3c2 Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram <96096277+nirvedhmeshram@users.noreply.github.com> Date: Thu, 15 Oct 2026 13:36:22 -0500 Subject: [PATCH 4/5] Address review: read all-to-all lanes and warps from hip_utils Go back to hip_utils.num_warps()/num_threads() in the all-to-all layout instead of passing THREADS_PER_WARP from the host. Iris is AMD-only, and these builtins (triton#9604) are in the Triton CI runs. Config keeps its device default because all-gather checks its tile shape against threads_per_warp * num_warps on the host before launch. Co-Authored-By: Claude Opus 5.5 (1M context) --- iris/ccl/gluon/all_to_all.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/iris/ccl/gluon/all_to_all.py b/iris/ccl/gluon/all_to_all.py index 1798afe4b..f2b5221c6 100644 --- a/iris/ccl/gluon/all_to_all.py +++ b/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 @@ -57,7 +58,6 @@ def persistent_all_to_all_gluon( COMM_SMS: gl.constexpr, NUM_XCDS: gl.constexpr, CHUNK_SIZE: gl.constexpr, - THREADS_PER_WARP: gl.constexpr, ): """ Persistent all-to-all kernel using Gluon. @@ -85,12 +85,13 @@ def persistent_all_to_all_gluon( pid_n = (tile_id % num_pid_in_group) // group_size_m # BlockedLayout: [size_per_thread], [threads_per_warp], [warps_per_cta], [order]. - # No literals: the wavefront is 64 wide on CDNA3/4 but 32 on CDNA5 - # (THREADS_PER_WARP = config.threads_per_warp, which defaults to the device's), - # and the launch's warp count is config.num_warps. - WARPS: gl.constexpr = gl.num_warps() - layout_col: gl.constexpr = gl.BlockedLayout([1], [THREADS_PER_WARP], [WARPS], [0]) # Column access - layout_row: gl.constexpr = gl.BlockedLayout([1], [THREADS_PER_WARP], [WARPS], [0]) # Row indices + # 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 @@ -197,7 +198,6 @@ def launch( config.comm_sms, config.num_xcds, config.chunk_size, - config.threads_per_warp, num_stages=config.num_stages, num_warps=config.num_warps, waves_per_eu=config.waves_per_eu, From 8a912a908e0fb11c7505b482d9d91bb4fefa495d Mon Sep 17 00:00:00 2001 From: Nirvedh Meshram <96096277+nirvedhmeshram@users.noreply.github.com> Date: Thu, 15 Oct 2026 16:46:20 -0500 Subject: [PATCH 5/5] Address review: read the wavefront width in-kernel everywhere All-gather now reads lanes and warps from hip_utils like all-to-all, and its tile-shape check moves from the host into a gl.static_assert. With no host-side user left, Config.threads_per_warp and its torch device query are removed. The gluon unit tests and examples 06 and 32 hard-coded 64-lane layouts (or a lane count equal to BLOCK_SIZE), which fail to compile on a 32-wide wavefront. They now take the lane count from hip_utils.num_threads(). Co-Authored-By: Claude Opus 5.5 (1M context) --- .../message_passing_gluon.py | 9 ++-- .../all_gather_tracing.py | 20 +++++---- iris/ccl/config.py | 17 ++----- iris/ccl/gluon/all_gather.py | 44 ++++++------------- tests/ccl/test_all_gather_gluon.py | 10 ----- tests/unittests/test_atomic_add_gluon.py | 3 +- tests/unittests/test_atomic_and_gluon.py | 3 +- tests/unittests/test_atomic_max_gluon.py | 3 +- tests/unittests/test_atomic_min_gluon.py | 3 +- tests/unittests/test_atomic_or_gluon.py | 3 +- tests/unittests/test_atomic_xor_gluon.py | 3 +- tests/unittests/test_copy_gluon.py | 7 +-- tests/unittests/test_device_context_gluon.py | 5 ++- tests/unittests/test_get_gluon.py | 3 +- tests/unittests/test_gluon_cache_modifiers.py | 13 +++--- tests/unittests/test_load_gluon.py | 3 +- tests/unittests/test_put_gluon.py | 3 +- tests/unittests/test_store_gluon.py | 3 +- 18 files changed, 67 insertions(+), 88 deletions(-) diff --git a/examples/06_message_passing/message_passing_gluon.py b/examples/06_message_passing/message_passing_gluon.py index 63b7f0568..154618c5d 100644 --- a/examples/06_message_passing/message_passing_gluon.py +++ b/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 diff --git a/examples/32_gluon_all_gather_tracing/all_gather_tracing.py b/examples/32_gluon_all_gather_tracing/all_gather_tracing.py index 000336e5d..4bf44a27f 100644 --- a/examples/32_gluon_all_gather_tracing/all_gather_tracing.py +++ b/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, diff --git a/iris/ccl/config.py b/iris/ccl/config.py index 05ecf0035..bdceaf1b9 100644 --- a/iris/ccl/config.py +++ b/iris/ccl/config.py @@ -6,9 +6,6 @@ """ from dataclasses import dataclass - -import torch - import iris @@ -49,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: None, the current - device's warp size). Must match the hardware wavefront - size: 64 on CDNA3/4, 32 on CDNA5. 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: @@ -95,7 +89,6 @@ class Config: reduce_scatter_variant: str = "two_shot" num_stages: int = 1 num_warps: int = 4 - threads_per_warp: int | None = None waves_per_eu: int = 0 def __post_init__(self): @@ -147,9 +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 is None: - self.threads_per_warp = torch.cuda.get_device_properties(torch.cuda.current_device()).warp_size - if self.threads_per_warp not in (32, 64): - raise ValueError(f"threads_per_warp must be 32 or 64 (the wavefront width), got {self.threads_per_warp}") if self.num_warps <= 0: raise ValueError(f"num_warps must be positive, got {self.num_warps}") diff --git a/iris/ccl/gluon/all_gather.py b/iris/ccl/gluon/all_gather.py index eac91d66a..5ae3fa129 100644 --- a/iris/ccl/gluon/all_gather.py +++ b/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 on CDNA3/4, 32 on CDNA5. - 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, diff --git a/tests/ccl/test_all_gather_gluon.py b/tests/ccl/test_all_gather_gluon.py index d09db1b26..438f01cfe 100644 --- a/tests/ccl/test_all_gather_gluon.py +++ b/tests/ccl/test_all_gather_gluon.py @@ -104,13 +104,3 @@ def test_all_gather_gluon(dtype, M, N, block_size_m, block_size_n): import gc gc.collect() - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") -def test_config_threads_per_warp_follows_device(): - """The default wavefront width is the device's, not an assumed 64: CDNA5 is 32 wide.""" - expected = torch.cuda.get_device_properties(torch.cuda.current_device()).warp_size - assert Config().threads_per_warp == expected - assert Config(threads_per_warp=64).threads_per_warp == 64 - with pytest.raises(ValueError, match="threads_per_warp must be 32 or 64"): - Config(threads_per_warp=16) diff --git a/tests/unittests/test_atomic_add_gluon.py b/tests/unittests/test_atomic_add_gluon.py index 3c398291d..cb9c43dbb 100644 --- a/tests/unittests/test_atomic_add_gluon.py +++ b/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 diff --git a/tests/unittests/test_atomic_and_gluon.py b/tests/unittests/test_atomic_and_gluon.py index f27057906..3f53c9fe1 100644 --- a/tests/unittests/test_atomic_and_gluon.py +++ b/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 diff --git a/tests/unittests/test_atomic_max_gluon.py b/tests/unittests/test_atomic_max_gluon.py index 4662a5180..5b04a37a5 100644 --- a/tests/unittests/test_atomic_max_gluon.py +++ b/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 diff --git a/tests/unittests/test_atomic_min_gluon.py b/tests/unittests/test_atomic_min_gluon.py index 8e91b264c..a5c0f743f 100644 --- a/tests/unittests/test_atomic_min_gluon.py +++ b/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 diff --git a/tests/unittests/test_atomic_or_gluon.py b/tests/unittests/test_atomic_or_gluon.py index 0de585142..7466c0520 100644 --- a/tests/unittests/test_atomic_or_gluon.py +++ b/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 diff --git a/tests/unittests/test_atomic_xor_gluon.py b/tests/unittests/test_atomic_xor_gluon.py index 271cc223e..e2ed7a692 100644 --- a/tests/unittests/test_atomic_xor_gluon.py +++ b/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 diff --git a/tests/unittests/test_copy_gluon.py b/tests/unittests/test_copy_gluon.py index 549025b2f..33d500b02 100644 --- a/tests/unittests/test_copy_gluon.py +++ b/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 diff --git a/tests/unittests/test_device_context_gluon.py b/tests/unittests/test_device_context_gluon.py index 889695740..b51004864 100644 --- a/tests/unittests/test_device_context_gluon.py +++ b/tests/unittests/test_device_context_gluon.py @@ -4,6 +4,7 @@ import torch 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 from iris.host.tracing.events import TraceEvent @@ -21,7 +22,7 @@ def device_context_tracing_1d_address_kernel( """Test ctx.tracing.record_event_start/end with a 1D address (block of pointers).""" ctx = IrisDeviceCtx.initialize(context_tensor, tracing=True) - layout: gl.constexpr = gl.BlockedLayout([1], [BLOCK_SIZE], [1], [0]) + layout: gl.constexpr = gl.BlockedLayout([1], [hip_utils.num_threads()], [1], [0]) offsets = gl.arange(0, BLOCK_SIZE, layout=layout) address_1d = dummy_buffer + offsets @@ -46,7 +47,7 @@ def test_device_context_gluon_tracing_1d_address(): source_rank = shmem.get_rank() num_ranks = shmem.get_num_ranks() - BLOCK_SIZE = 64 # AMD wavefront size (64 threads per warp) + BLOCK_SIZE = 64 # Dummy buffer only to form 1D pointer block; never read/write dummy_buffer = shmem.zeros((BLOCK_SIZE,), dtype=torch.int64) diff --git a/tests/unittests/test_get_gluon.py b/tests/unittests/test_get_gluon.py index a203a703d..48359e246 100644 --- a/tests/unittests/test_get_gluon.py +++ b/tests/unittests/test_get_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 @@ -26,7 +27,7 @@ def 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 diff --git a/tests/unittests/test_gluon_cache_modifiers.py b/tests/unittests/test_gluon_cache_modifiers.py index fe40eca25..054e545cf 100644 --- a/tests/unittests/test_gluon_cache_modifiers.py +++ b/tests/unittests/test_gluon_cache_modifiers.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 from itertools import product @@ -30,7 +31,7 @@ def load_cache_modifier_kernel( partner = int((source_rank + num_ranks // 2) % num_ranks) 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 @@ -53,7 +54,7 @@ def store_cache_modifier_kernel( 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 @@ -79,7 +80,7 @@ def get_cache_modifier_kernel( 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 @@ -115,7 +116,7 @@ def put_cache_modifier_kernel( 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 @@ -145,7 +146,7 @@ def copy_local_read_remote_write_cache_modifier_kernel( 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 @@ -179,7 +180,7 @@ def copy_remote_read_local_write_cache_modifier_kernel( 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 diff --git a/tests/unittests/test_load_gluon.py b/tests/unittests/test_load_gluon.py index aa0e895fd..f37a02d0c 100644 --- a/tests/unittests/test_load_gluon.py +++ b/tests/unittests/test_load_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 @@ -25,7 +26,7 @@ def load_kernel( partner = int((source_rank + num_ranks // 2) % num_ranks) # Compute start index of this block 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) # Guard for out-of-bounds accesses diff --git a/tests/unittests/test_put_gluon.py b/tests/unittests/test_put_gluon.py index 24d9c074b..bca335dfa 100644 --- a/tests/unittests/test_put_gluon.py +++ b/tests/unittests/test_put_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 @@ -26,7 +27,7 @@ def 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 diff --git a/tests/unittests/test_store_gluon.py b/tests/unittests/test_store_gluon.py index 3abc82a3c..3174fa971 100644 --- a/tests/unittests/test_store_gluon.py +++ b/tests/unittests/test_store_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 store_kernel( 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