From 7fcfa698c0e63ee6a00313e2c6b25d737612ee23 Mon Sep 17 00:00:00 2001 From: Feng Shijie Date: Tue, 28 Jul 2026 03:54:07 +0000 Subject: [PATCH 1/3] [ROCDL] Add arch-aware universal s_waitcnt --- python/flydsl/expr/rocdl/__init__.py | 5 ++- python/flydsl/expr/rocdl/cdna3.py | 20 ++++++++++ python/flydsl/expr/rocdl/rdna3.py | 20 ++++++++++ python/flydsl/expr/rocdl/rdna4.py | 18 +++++++++ python/flydsl/expr/rocdl/universal.py | 55 +++++++++++++++++++++++++-- python/flydsl/expr/rocdl/utils.py | 31 +++++++++++++++ python/flydsl/runtime/device.py | 10 ++++- 7 files changed, 152 insertions(+), 7 deletions(-) create mode 100644 python/flydsl/expr/rocdl/cdna3.py create mode 100644 python/flydsl/expr/rocdl/rdna3.py create mode 100644 python/flydsl/expr/rocdl/rdna4.py create mode 100644 python/flydsl/expr/rocdl/utils.py 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..738462b76 --- /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) + 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..ad2ff04a5 --- /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) + 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..47c2b9fcc --- /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) + 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..dba6418fa 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,52 @@ is_generic_address_space, is_target_address_space, ) +from . import cdna3, rdna3, rdna4 +from .utils import normalize_s_waitcnt_field + +__all__ = [ + "s_waitcnt", + "BufferCopy", + "BufferCopy8b", + "BufferCopy16b", + "BufferCopy32b", + "BufferCopy64b", + "BufferCopy128b", + "BufferCopyLDS", + "BufferCopyLDS32b", + "BufferCopyLDS64b", + "BufferCopyLDS128b", + "BufferAtomic", + "BufferAtomicAdd", + "BufferAtomicMax", + "BufferAtomicMin", + "BufferAtomicPkAdd", + "MFMA", + "WMMA", + "make_buffer_ptr", + "make_buffer_tensor", + "get_buffer_rsrc", +] + + +@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("gfx9"): # 942, 950 + 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 gfx1200 (RDNA4). " + ) def BufferCopy(bit_size, cache_modifier=0): @@ -113,12 +161,11 @@ 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"): + if arch.startswith("gfx125"): return MmaOpGFX1250_WMMAType.get( m, n, @@ -153,7 +200,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) From c865ab91ff8d728b15d72b9b957791398dea354b Mon Sep 17 00:00:00 2001 From: Feng Shijie Date: Tue, 28 Jul 2026 04:08:31 +0000 Subject: [PATCH 2/3] fix comments --- python/flydsl/expr/rocdl/cdna3.py | 2 +- python/flydsl/expr/rocdl/rdna3.py | 2 +- python/flydsl/expr/rocdl/rdna4.py | 2 +- python/flydsl/expr/rocdl/universal.py | 6 +++--- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/python/flydsl/expr/rocdl/cdna3.py b/python/flydsl/expr/rocdl/cdna3.py index 738462b76..40c42a892 100644 --- a/python/flydsl/expr/rocdl/cdna3.py +++ b/python/flydsl/expr/rocdl/cdna3.py @@ -17,4 +17,4 @@ def s_waitcnt(vmcnt=None, lgkmcnt=None, expcnt=None): expcnt = normalize_s_waitcnt_field("expcnt", expcnt, 7) lgkmcnt = normalize_s_waitcnt_field("lgkmcnt", lgkmcnt, 15) encoded_vmcnt = (vmcnt & 0xF) | ((vmcnt & 0x30) << 10) - mlir_rocdl.s_waitcnt(encoded_vmcnt | (lgkmcnt << 8) | (expcnt << 4)) + 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 index ad2ff04a5..c23269e65 100644 --- a/python/flydsl/expr/rocdl/rdna3.py +++ b/python/flydsl/expr/rocdl/rdna3.py @@ -17,4 +17,4 @@ def s_waitcnt(vmcnt=None, lgkmcnt=None, expcnt=None): vmcnt = normalize_s_waitcnt_field("vmcnt", vmcnt, 63) lgkmcnt = normalize_s_waitcnt_field("lgkmcnt", lgkmcnt, 63) expcnt = normalize_s_waitcnt_field("expcnt", expcnt, 7) - mlir_rocdl.s_waitcnt((vmcnt << 10) | (lgkmcnt << 4) | expcnt) + 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 index 47c2b9fcc..071c6a390 100644 --- a/python/flydsl/expr/rocdl/rdna4.py +++ b/python/flydsl/expr/rocdl/rdna4.py @@ -15,4 +15,4 @@ def s_waitcnt(vmcnt=None, lgkmcnt=None, expcnt=None): vmcnt = normalize_s_waitcnt_field("vmcnt", vmcnt, 63) lgkmcnt = normalize_s_waitcnt_field("lgkmcnt", lgkmcnt, 63) expcnt = normalize_s_waitcnt_field("expcnt", expcnt, 7) - mlir_rocdl.s_waitcnt(vmcnt << 10 | lgkmcnt << 4 | expcnt) + 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 dba6418fa..b8caefe42 100644 --- a/python/flydsl/expr/rocdl/universal.py +++ b/python/flydsl/expr/rocdl/universal.py @@ -67,14 +67,14 @@ def s_waitcnt(bitfield=None, *, vmcnt=None, lgkmcnt=None, expcnt=None): return mlir_rocdl.s_waitcnt(normalize_s_waitcnt_field("bitfield", bitfield, 0xFFFF)) arch = get_rocm_arch() - if arch.startswith("gfx9"): # 942, 950 + 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 gfx1200 (RDNA4). " + f"s_waitcnt is not supported on target arch {arch!r}; supported: gfx942 (CDNA3), gfx950 (CDNA4), gfx11xx (RDNA3 / RDNA3.5), and gfx120x (RDNA4). " ) @@ -165,7 +165,7 @@ def WMMA(m, n, k, elem_ty_ab, elem_ty_acc=None, **kwargs): 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("gfx125"): + if arch.startswith("gfx12"): return MmaOpGFX1250_WMMAType.get( m, n, From 83da9a1da4f227eded986ed65b103bfd0ff52073 Mon Sep 17 00:00:00 2001 From: Feng Shijie Date: Tue, 28 Jul 2026 07:16:33 +0000 Subject: [PATCH 3/3] [Fix] Restore universal ROCDL exports Avoid hiding existing atom types behind a restrictive export list and dispatch CDNA wait counters correctly. Co-authored-by: Cursor --- python/flydsl/expr/rocdl/universal.py | 26 +------------------------- 1 file changed, 1 insertion(+), 25 deletions(-) diff --git a/python/flydsl/expr/rocdl/universal.py b/python/flydsl/expr/rocdl/universal.py index b8caefe42..06b6086cc 100644 --- a/python/flydsl/expr/rocdl/universal.py +++ b/python/flydsl/expr/rocdl/universal.py @@ -33,30 +33,6 @@ from . import cdna3, rdna3, rdna4 from .utils import normalize_s_waitcnt_field -__all__ = [ - "s_waitcnt", - "BufferCopy", - "BufferCopy8b", - "BufferCopy16b", - "BufferCopy32b", - "BufferCopy64b", - "BufferCopy128b", - "BufferCopyLDS", - "BufferCopyLDS32b", - "BufferCopyLDS64b", - "BufferCopyLDS128b", - "BufferAtomic", - "BufferAtomicAdd", - "BufferAtomicMax", - "BufferAtomicMin", - "BufferAtomicPkAdd", - "MFMA", - "WMMA", - "make_buffer_ptr", - "make_buffer_tensor", - "get_buffer_rsrc", -] - @dsl_loc_tracing def s_waitcnt(bitfield=None, *, vmcnt=None, lgkmcnt=None, expcnt=None): @@ -67,7 +43,7 @@ def s_waitcnt(bitfield=None, *, vmcnt=None, lgkmcnt=None, expcnt=None): return mlir_rocdl.s_waitcnt(normalize_s_waitcnt_field("bitfield", bitfield, 0xFFFF)) arch = get_rocm_arch() - if arch.startswith("gfx942", "gfx950"): + 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)