Add stochastic rounding and philox RNG under rocdl

kernels/gemm/rdna3_f16_gemm.py

@@ -45,6 +45,7 @@ def create_wmma_gemm_module(
     in_dtype="bf16",
     out_dtype="bf16",
     *,
+    rounding="rn",  # "rn" (round to nearest) or "rs" (stochastic rounding)
     reg_m=4,
     reg_n=4,
     reg_k=2,
@@ -69,6 +70,9 @@ def create_wmma_gemm_module(
     THREADS_PER_BLOCK = NUM_WAVES * WAVE_SIZE  # 128
 
     assert reg_k >= 2 and reg_k % 2 == 0
+    assert rounding in ("rn", "rs"), f"rounding must be 'rn' or 'rs', got {rounding!r}"
+    if rounding == "rs":
+        assert out_dtype == "bf16", "stochastic rounding currently supports bf16 output only"
 
     LOAD_VEC = 8  # 8 bf16 = 128-bit GMEM/LDS load
     A_TILE_ELEMS = BLOCK_M * BLOCK_K
@@ -119,6 +123,7 @@ def wmma_gemm_kernel(
         arg_c: fx.Tensor,
         arg_a: fx.Tensor,
         arg_bt: fx.Tensor,
+        sr_seed: fx.Int32,  # runtime seed; only read on the stochastic-rounding path
     ):
         lds_storage = fx.SharedAllocator().allocate(_SharedStorage).peek()
         lds_ptr = lds_storage.lds.ptr  # i8-base aliased as elem_dtype*
@@ -327,15 +332,27 @@ def _do_compute_rk(accs_in, rk, buf_offset):
                 idx = rm * reg_n + rn
                 wmma_m_off = wave_m * (reg_m * WMMA_M) + 16 * rm
                 wmma_n_off = wave_n * (reg_n * WMMA_N) + 16 * rn
+                if const_expr(rounding == "rs"):
+                    # One Philox draw per 8-element block, keyed on the block's
+                    # base output element: the 4 random words cover all 8
+                    # stores, each taking a distinct 16-bit slice (low/high of a
+                    # word), so the f32 -> bf16 store is unbiased in expectation
+                    # without a per-element draw.
+                    base_off = (tile_m0 + wmma_m_off + klane) * N + (tile_n0 + wmma_n_off + lane16)
+                    rand_words = fx.random.philox_4x32(fx.Uint32(base_off), fx.Uint32(sr_seed))
                 for si in range_constexpr(8):
                     g_row = tile_m0 + wmma_m_off + 2 * si + klane
                     g_col = tile_n0 + wmma_n_off + lane16
                     val = accs[idx][si]
-                    if const_expr(out_dtype == "bf16"):
+                    elem_off = g_row * N + g_col
+                    if const_expr(rounding == "rs"):
+                        word = rand_words[si // 2]
+                        rbits = word if si % 2 == 0 else (word >> fx.Uint32(16))
+                        val = fx.random.cvt_f32_to_bf16_sr(val, rbits)
+                    elif const_expr(out_dtype == "bf16"):
                         val = val.to(fx.BFloat16)
                     elif const_expr(out_dtype == "f16"):
                         val = val.to(fx.Float16)
-                    elem_off = g_row * N + g_col
                     buffer_ops.buffer_store(val, c_rsrc, elem_off)
 
     @flyc.jit
@@ -344,12 +361,13 @@ def launch_gemm(
         arg_a: fx.Tensor,
         arg_bt: fx.Tensor,
         stream: fx.Stream,
+        sr_seed: fx.Int32 = 0,
     ):
         c1 = 1
         total_blocks = grid_m * grid_n
         bk = THREADS_PER_BLOCK
 
-        launcher = wmma_gemm_kernel(arg_c, arg_a, arg_bt)
+        launcher = wmma_gemm_kernel(arg_c, arg_a, arg_bt, sr_seed)
         launcher.launch(
             grid=(total_blocks, c1, c1),
             block=(bk, c1, c1),

kernels/gemm/rdna_f16_gemm.py

@@ -39,6 +39,7 @@ def create_wmma_gemm_module(
     in_dtype="bf16",
     out_dtype="bf16",
     *,
+    rounding="rn",  # "rn" (round to nearest) or "rs" (stochastic rounding)
     reg_m=4,  # M-repeats per warp
     reg_n=4,  # N-repeats per warp
     reg_k=2,  # K-steps per tile (32/16=2)
@@ -56,6 +57,9 @@ def create_wmma_gemm_module(
     THREADS_PER_BLOCK = NUM_WAVES * WAVE_SIZE  # 128
 
     assert reg_k >= 2 and reg_k % 2 == 0
+    assert rounding in ("rn", "rs"), f"rounding must be 'rn' or 'rs', got {rounding!r}"
+    if rounding == "rs":
+        assert out_dtype == "bf16", "stochastic rounding currently supports bf16 output only"
 
     # Loading: each thread loads 8 bf16 elements per load (128 bits = buffer_load_b128)
     LOAD_VEC = 8
@@ -96,6 +100,7 @@ def wmma_gemm_kernel(
         arg_c: fx.Tensor,
         arg_a: fx.Tensor,
         arg_bt: fx.Tensor,
+        sr_seed: fx.Int32,  # runtime seed; only read on the stochastic-rounding path
     ):
         lds = fx.SharedAllocator(static=False).allocate(fx.Array[lds_elem_dtype, LDS_TOTAL, 16]).peek()
 
@@ -315,15 +320,27 @@ def _load_a_single_from_lds(rk, rm_val, buf_offset):
                 idx = rm * reg_n + rn
                 wmma_m_off = wave_m * (reg_m * WMMA_M) + 16 * rm
                 wmma_n_off = wave_n * (reg_n * WMMA_N) + 16 * rn
+                if const_expr(rounding == "rs"):
+                    # One Philox draw per 8-element block, keyed on the block's
+                    # base output element: the 4 random words cover all 8
+                    # stores, each taking a distinct 16-bit slice (low/high of a
+                    # word), so the f32 -> bf16 store is unbiased in expectation
+                    # without a per-element draw.
+                    base_off = (tile_m0 + wmma_m_off + base8) * N + (tile_n0 + wmma_n_off + lane16)
+                    rand_words = fx.random.philox_4x32(fx.Uint32(base_off), fx.Uint32(sr_seed))
                 for si in range_constexpr(8):
                     g_row = tile_m0 + wmma_m_off + base8 + si
                     g_col = tile_n0 + wmma_n_off + lane16
                     val = accs[idx][si]
-                    if const_expr(out_dtype == "bf16"):
+                    elem_off = g_row * N + g_col
+                    if const_expr(rounding == "rs"):
+                        word = rand_words[si // 2]
+                        rbits = word if si % 2 == 0 else (word >> fx.Uint32(16))
+                        val = fx.random.cvt_f32_to_bf16_sr(val, rbits)
+                    elif const_expr(out_dtype == "bf16"):
                         val = val.to(fx.BFloat16)
                     elif const_expr(out_dtype == "f16"):
                         val = val.to(fx.Float16)
-                    elem_off = g_row * N + g_col
                     buffer_ops.buffer_store(val, c_rsrc, elem_off)
 
     # ── Host launcher ──────────────────────────────────────────────────────
@@ -333,12 +350,13 @@ def launch_gemm(
         arg_a: fx.Tensor,
         arg_bt: fx.Tensor,
         stream: fx.Stream,
+        sr_seed: fx.Int32 = 0,
     ):
         c1 = 1
         total_blocks = grid_m * grid_n
         bk = THREADS_PER_BLOCK
 
-        launcher = wmma_gemm_kernel(arg_c, arg_a, arg_bt)
+        launcher = wmma_gemm_kernel(arg_c, arg_a, arg_bt, sr_seed)
         launcher.launch(
             grid=(total_blocks, c1, c1),
             block=(bk, c1, c1),

python/flydsl/compiler/jit_function.py

@@ -7,7 +7,6 @@
 import inspect
 import os
 import pickle
-import pkgutil
 import threading
 import time
 import types
@@ -321,23 +320,15 @@ def _flydsl_key_cached(use_external_binary: bool, llvm_dir: str, extra_source_di
     flydsl_root = Path(flydsl.__file__).resolve().parent
 
     # 1) Hash all Python source files in key sub-packages.
-    pkg_prefixes = [
-        (str(flydsl_root / "compiler"), "flydsl.compiler."),
-        (str(flydsl_root / "expr"), "flydsl.expr."),
-        (str(flydsl_root / "runtime"), "flydsl.runtime."),
-        (str(flydsl_root / "utils"), "flydsl.utils."),
-    ]
-    for pkg_path, prefix in pkg_prefixes:
-        if not os.path.isdir(pkg_path):
+    for package_name in ("compiler", "expr", "runtime", "utils"):
+        pkg_path = flydsl_root / package_name
+        if not pkg_path.is_dir():
             continue
-        for lib in pkgutil.walk_packages([pkg_path], prefix=prefix):
-            try:
-                spec = lib.module_finder.find_spec(lib.name)
-                if spec and spec.origin and os.path.isfile(spec.origin):
-                    with open(spec.origin, "rb") as f:
-                        contents.append(hashlib.sha256(f.read()).hexdigest())
-            except Exception:
-                pass
+        for py_file in sorted(pkg_path.rglob("*.py")):
+            with open(py_file, "rb") as f:
+                source_hash = hashlib.sha256(f.read()).hexdigest()
+            relative_path = py_file.relative_to(flydsl_root).as_posix()
+            contents.append(f"{relative_path}:{source_hash}")
 
     p = flydsl_root / "__init__.py"
     if p.is_file():

python/flydsl/expr/__init__.py

@@ -21,11 +21,14 @@
     "rocdl": ".rocdl",
     "tdm_ops": ".rocdl.tdm_ops",  # deprecated, use .rocdl.tdm_ops instead
 }
+_LIBRARY_MODULES = {
+    "random": ".random",
+}
 
 
 # lazy load backend subpackages
 def __getattr__(name: str):
-    module_name = _BACKEND_MODULES.get(name)
+    module_name = _BACKEND_MODULES.get(name) or _LIBRARY_MODULES.get(name)
     if module_name is None:
         raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 

python/flydsl/expr/_library.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 FlyDSL Project Contributors
+
+"""Lazy target overrides for target-neutral expression libraries."""
+
+import importlib
+from functools import wraps
+from types import ModuleType
+from typing import Callable, Mapping
+
+from ..compiler.backends import compile_backend_name
+
+
+class Library:
+    """Dispatch universal library functions to optional target modules.
+
+    ``targets`` maps compiler backend IDs to public target namespaces. For
+    example, ``{"rocm": "rocdl"}`` maps the ROCm compiler backend to the
+    sibling ``<library>.rocdl`` module.
+    """
+
+    def __init__(self, package: str, *, targets: Mapping[str, str]):
+        self.package = package
+        self.targets = dict(targets)
+        self._resolved_impls = {}
+
+    def _resolve(self, backend: str, name: str, universal_impl: Callable) -> Callable:
+        key = (backend, name)
+        if key not in self._resolved_impls:
+            implementation = universal_impl
+            target_name = self.targets.get(backend)
+            if target_name is not None:
+                module = self.load_target(target_name)
+                if name in getattr(module, "__all__", ()):
+                    implementation = getattr(module, name)
+            self._resolved_impls[key] = implementation
+        return self._resolved_impls[key]
+
+    def dispatch(self, universal_impl: Callable) -> Callable:
+        """Wrap a universal implementation with cached target override lookup."""
+
+        name = universal_impl.__name__
+
+        @wraps(universal_impl)
+        def wrapper(*args, **kwargs):
+            implementation = self._resolve(compile_backend_name(), name, universal_impl)
+            return implementation(*args, **kwargs)
+
+        return wrapper
+
+    def dispatch_all(self, namespace: dict, names) -> None:
+        """Wrap every name in ``names`` with target override lookup."""
+
+        for name in names:
+            namespace[name] = self.dispatch(namespace[name])
+
+    def load_target(self, name: str) -> ModuleType:
+        """Load one explicitly declared target namespace."""
+
+        if name not in self.targets.values():
+            raise AttributeError(f"module {self.package!r} has no attribute {name!r}")
+        return importlib.import_module(f".{name}", self.package)

python/flydsl/expr/random/__init__.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 FlyDSL Project Contributors
+
+"""Target-neutral random-number algorithms with optional target overrides.
+
+``fx.random.<name>`` runs the implementation chosen for the compilation
+target; ``fx.random.universal.<name>`` always runs the portable one.
+"""
+
+from .._library import Library
+from . import universal
+from .universal import *  # noqa: F401,F403
+
+__all__ = list(universal.__all__)
+
+_library = Library(__name__, targets={"rocm": "rocdl"})
+__getattr__ = _library.load_target
+
+_library.dispatch_all(globals(), __all__)

python/flydsl/expr/random/rocdl.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 FlyDSL Project Contributors
+
+"""ROCDL-specific overrides for random algorithms."""
+
+from ..._mlir.dialects import llvm
+from ...compiler.backends import get_backend
+from ..numeric import Boolean, Float32, Uint32
+from ..typing import BFloat16x2
+from . import universal
+
+__all__ = [
+    "cvt_f32_to_bf16_sr",
+]
+
+
+def cvt_f32_to_bf16_sr(x, rand):
+    """Convert with one ``v_cvt_sr_bf16_f32`` where the architecture has it.
+
+    Only gfx950 does this.
+    """
+    if get_backend().target.arch not in ("gfx950",):
+        return universal.cvt_f32_to_bf16_sr(x, rand)
+
+    # The result is a tied 2xbf16: word_sel picks the half that gets written and
+    # the other half keeps the old value, undefined here because only half 0 is
+    # read back. An undef old value also keeps the tie free of an extra v_mov.
+    return BFloat16x2(
+        llvm.call_intrinsic(
+            BFloat16x2.ir_type,
+            "llvm.amdgcn.cvt.sr.bf16.f32",
+            [
+                llvm.mlir_undef(BFloat16x2.ir_type),
+                Float32(x).ir_value(),
+                Uint32(rand).ir_value(),
+                Boolean(0).ir_value(),  # word_sel
+            ],
+            [],
+            [],
+        )
+    )[0]