diff --git a/python/flydsl/expr/rocdl/__init__.py b/python/flydsl/expr/rocdl/__init__.py index 4b18f15cc..d06b04220 100644 --- a/python/flydsl/expr/rocdl/__init__.py +++ b/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 diff --git a/python/flydsl/expr/rocdl/cdna3.py b/python/flydsl/expr/rocdl/cdna3.py new file mode 100644 index 000000000..40c42a892 --- /dev/null +++ b/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)) diff --git a/python/flydsl/expr/rocdl/rdna3.py b/python/flydsl/expr/rocdl/rdna3.py new file mode 100644 index 000000000..c23269e65 --- /dev/null +++ b/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) diff --git a/python/flydsl/expr/rocdl/rdna4.py b/python/flydsl/expr/rocdl/rdna4.py new file mode 100644 index 000000000..071c6a390 --- /dev/null +++ b/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) diff --git a/python/flydsl/expr/rocdl/universal.py b/python/flydsl/expr/rocdl/universal.py index e55579588..06b6086cc 100644 --- a/python/flydsl/expr/rocdl/universal.py +++ b/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) diff --git a/python/flydsl/expr/rocdl/utils.py b/python/flydsl/expr/rocdl/utils.py new file mode 100644 index 000000000..80195abb3 --- /dev/null +++ b/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) diff --git a/python/flydsl/runtime/device.py b/python/flydsl/runtime/device.py index ed5fd47b2..7383d03fd 100644 --- a/python/flydsl/runtime/device.py +++ b/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)