[ROCDL] Add arch-aware universal s_waitcnt

python/flydsl/expr/rocdl/__init__.py

@@ -16,9 +16,13 @@
 
 from ..._mlir.dialects.rocdl import *  # noqa: F401,F403
 from ..meta import dsl_loc_tracing
+from . import cdna3 as cdna3
 from . import cdna4 as cdna4
 from . import cdna5 as cdna5
+from . import rdna3 as rdna3
+from . import rdna4 as rdna4
 from .enum import SyncScope as SyncScope
+from .universal import *
 
 # Keep references to ODS-generated builders so we can wrap them without losing access.
 _ods_wmma_scale_f32_16x16x128_f8f6f4 = globals().get("wmma_scale_f32_16x16x128_f8f6f4", None)
@@ -458,7 +462,6 @@ def lds_transpose_load(result_type, lds_memref, elem_offset, elem_bytes):
 
 
 # ── New high-level helpers from universal.py ──────────────────────────
-from .universal import *  # noqa: E402,F401,F403,I001
 from .cdna5 import *  # noqa: E402,F401,F403,I001
 from .inline_asm import *  # noqa: E402,F401,F403,I001
 

python/flydsl/expr/rocdl/cdna3.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 FlyDSL Project Contributors
+
+from ..._mlir.dialects import rocdl as mlir_rocdl
+from .utils import normalize_s_waitcnt_field
+
+
+def s_waitcnt(vmcnt=None, lgkmcnt=None, expcnt=None):
+    """
+    Emit a CDNA3/GFX942-encoded s_waitcnt operation.
+
+    vmcnt: split across [3:0] and [15:14]
+    expcnt: [6:4]
+    lgkmcnt: [11:8]
+    """
+    vmcnt = normalize_s_waitcnt_field("vmcnt", vmcnt, 63)
+    expcnt = normalize_s_waitcnt_field("expcnt", expcnt, 7)
+    lgkmcnt = normalize_s_waitcnt_field("lgkmcnt", lgkmcnt, 15)
+    encoded_vmcnt = (vmcnt & 0xF) | ((vmcnt & 0x30) << 10)
+    return mlir_rocdl.s_waitcnt(encoded_vmcnt | (lgkmcnt << 8) | (expcnt << 4))

python/flydsl/expr/rocdl/rdna3.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 FlyDSL Project Contributors
+
+from ..._mlir.dialects import rocdl as mlir_rocdl
+from .utils import normalize_s_waitcnt_field
+
+
+def s_waitcnt(vmcnt=None, lgkmcnt=None, expcnt=None):
+    """
+    Emit a RDNA3/GFX11-encoded s_waitcnt operation.
+
+    expcnt: [2:0]
+    unused: [3:3]
+    lgkmcnt: [9:4]
+    vmcnt: [15:10]
+    """
+    vmcnt = normalize_s_waitcnt_field("vmcnt", vmcnt, 63)
+    lgkmcnt = normalize_s_waitcnt_field("lgkmcnt", lgkmcnt, 63)
+    expcnt = normalize_s_waitcnt_field("expcnt", expcnt, 7)
+    return mlir_rocdl.s_waitcnt((vmcnt << 10) | (lgkmcnt << 4) | expcnt)

python/flydsl/expr/rocdl/rdna4.py

@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 FlyDSL Project Contributors
+
+from ..._mlir.dialects import rocdl as mlir_rocdl
+from .utils import normalize_s_waitcnt_field
+
+
+def s_waitcnt(vmcnt=None, lgkmcnt=None, expcnt=None):
+    """Emit a RDNA4/GFX120-encoded s_waitcnt operation.
+
+    expcnt: [2:0]
+    lgkmcnt: [9:4]
+    vmcnt: [15:10]
+    """
+    vmcnt = normalize_s_waitcnt_field("vmcnt", vmcnt, 63)
+    lgkmcnt = normalize_s_waitcnt_field("lgkmcnt", lgkmcnt, 63)
+    expcnt = normalize_s_waitcnt_field("expcnt", expcnt, 7)
+    return mlir_rocdl.s_waitcnt(vmcnt << 10 | lgkmcnt << 4 | expcnt)

python/flydsl/expr/rocdl/universal.py

@@ -7,6 +7,7 @@
     MmaOpGFX1250_WMMAType,
 )
 from ..._mlir.dialects import fly_rocdl
+from ..._mlir.dialects import rocdl as mlir_rocdl
 from ..._mlir.dialects.fly import AtomicOp, PointerType
 from ..._mlir.dialects.fly_rocdl import (
     CopyOpCDNA3BufferAtomicType,
@@ -16,6 +17,7 @@
     TargetAddressSpace,
 )
 from ..._mlir.extras import types as T
+from ...runtime.device import get_rocm_arch
 from ..meta import dsl_loc_tracing
 from ..primitive import cosize, get_iter, get_layout, get_scalar, make_ptr, make_view
 from ..typing import (
@@ -28,6 +30,28 @@
     is_generic_address_space,
     is_target_address_space,
 )
+from . import cdna3, rdna3, rdna4
+from .utils import normalize_s_waitcnt_field
+
+
+@dsl_loc_tracing
+def s_waitcnt(bitfield=None, *, vmcnt=None, lgkmcnt=None, expcnt=None):
+    """Wait for named counters, or emit a legacy raw wait-counter bitfield."""
+    if bitfield is not None:
+        if vmcnt is not None or lgkmcnt is not None or expcnt is not None:
+            raise TypeError("legacy raw s_waitcnt bitfield cannot be combined with non-None keyword arguments")
+        return mlir_rocdl.s_waitcnt(normalize_s_waitcnt_field("bitfield", bitfield, 0xFFFF))
+
+    arch = get_rocm_arch()
+    if arch.startswith(("gfx942", "gfx950")):
+        return cdna3.s_waitcnt(vmcnt=vmcnt, lgkmcnt=lgkmcnt, expcnt=expcnt)
+    if arch.startswith("gfx11"):
+        return rdna3.s_waitcnt(vmcnt=vmcnt, lgkmcnt=lgkmcnt, expcnt=expcnt)
+    if arch.startswith("gfx120"):
+        return rdna4.s_waitcnt(vmcnt=vmcnt, lgkmcnt=lgkmcnt, expcnt=expcnt)
+    raise ValueError(
+        f"s_waitcnt is not supported on target arch {arch!r}; supported: gfx942 (CDNA3), gfx950 (CDNA4), gfx11xx (RDNA3 / RDNA3.5), and gfx120x (RDNA4). "
+    )
 
 
 def BufferCopy(bit_size, cache_modifier=0):
@@ -113,9 +137,8 @@ def WMMA(m, n, k, elem_ty_ab, elem_ty_acc=None, **kwargs):
     #   * RDNA4 (gfx12xx, e.g. gfx1201) and gfx1250 use the new v8-operand ABI;
     #     both route through MmaOpGFX1250_WMMAType via the gfx12 prefix below.
     #     (gfx1250 is its own arch, not RDNA4, but shares this WMMA atom.)
-    from ...runtime.device import get_rocm_arch
 
-    arch = (get_rocm_arch() or "").lower()
+    arch = get_rocm_arch() or ""
     if arch.startswith("gfx11"):
         return MmaOpGFX11_WMMAType.get(m, n, k, ty_ab, ty_ab, ty_acc, **kwargs)
     if arch.startswith("gfx12"):
@@ -153,7 +176,7 @@ def make_buffer_ptr(ptr: Pointer, num_records_bytes=None):
         # Coerce to i64: ROCDL make.buffer.rsrc requires an i64 num_records operand.
         num_records_bytes = Int64(num_records_bytes)
 
-    from ...runtime.device import get_rocm_arch, is_rdna_arch
+    from ...runtime.device import is_rdna_arch
 
     arch = get_rocm_arch()
     flags = (7 << 12) | (4 << 15)

python/flydsl/expr/rocdl/utils.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 FlyDSL Project Contributors
+
+from ..numeric import Numeric
+
+
+def normalize_s_waitcnt_field(name, value, maximum):
+    """Coerce a wait-counter argument to a static Python ``int``.
+
+    ``None`` means "do not wait on this counter" and maps to ``maximum``,
+    which is the encoding the hardware reads as "already satisfied".
+
+    Wait counters are encoded into an instruction's immediate field, so the
+    value has to be known at compile time; a run-time ``Integer`` is rejected
+    rather than silently materialised as a constant.
+    """
+    if value is None:
+        return maximum
+
+    if isinstance(value, Numeric):
+        if not value.is_static():
+            raise TypeError(f"{name} must be a static Python int or Integer, got a run-time value")
+        value = value.value
+
+    if not isinstance(value, int):
+        raise TypeError(f"{name} must be a static Python int or Integer, got {type(value).__name__}")
+
+    if not 0 <= value <= maximum:
+        raise ValueError(f"{name} must be in [0, {maximum}] on this target, got {value}")
+
+    return int(value)

python/flydsl/runtime/device.py

@@ -37,16 +37,22 @@ def _arch_from_hardware() -> str:
 
 
 def get_rocm_arch() -> str:
-    """Best-effort ROCm GPU arch string (e.g. 'gfx942')."""
+    """Best-effort ROCm GPU arch string, always lower-cased (e.g. 'gfx942').
+
+    Lower-casing happens here so every caller can compare against lower-case
+    literals without normalising first; ROCm itself only ever emits lower-case
+    names, so this only affects hand-set environment overrides.
+    """
     env = os.environ.get("FLYDSL_GPU_ARCH") or os.environ.get("HSA_OVERRIDE_GFX_VERSION")
     if env:
+        env = env.lower()
         if env.startswith("gfx"):
             return env
         if env.count(".") == 2:
             parts = env.split(".")
             return f"gfx{parts[0]}{parts[1]}{parts[2]}"
 
-    return _arch_from_hardware()
+    return _arch_from_hardware().lower()
 
 
 @functools.lru_cache(maxsize=None)