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 81188c5fb..bdceaf1b9 100644 --- a/iris/ccl/config.py +++ b/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}") diff --git a/iris/ccl/gluon/all_gather.py b/iris/ccl/gluon/all_gather.py index 6befb9fcf..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 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, diff --git a/iris/ccl/gluon/all_to_all.py b/iris/ccl/gluon/all_to_all.py index 315b9b0d8..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 @@ -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 diff --git a/tests/ccl/test_all_gather_gluon.py b/tests/ccl/test_all_gather_gluon.py index 3c57baead..438f01cfe 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) 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