From 726a3a9ed499ca6ff54190a7ee1e03fb93a08a1f Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Tue, 28 Jul 2026 12:29:13 +0200 Subject: [PATCH 1/3] Add stochastic rounding and philox RNG under rocdl --- kernels/gemm/rdna3_f16_gemm.py | 18 ++++- kernels/gemm/rdna_f16_gemm.py | 18 ++++- python/flydsl/expr/rocdl/__init__.py | 1 + python/flydsl/expr/rocdl/random.py | 71 +++++++++++++++++ tests/kernels/test_rdna_gemm.py | 35 +++++++++ tests/unit/test_stochastic_rounding.py | 102 +++++++++++++++++++++++++ 6 files changed, 239 insertions(+), 6 deletions(-) create mode 100644 python/flydsl/expr/rocdl/random.py create mode 100644 tests/unit/test_stochastic_rounding.py diff --git a/kernels/gemm/rdna3_f16_gemm.py b/kernels/gemm/rdna3_f16_gemm.py index ef6a65f68..1e8397e87 100644 --- a/kernels/gemm/rdna3_f16_gemm.py +++ b/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* @@ -331,11 +336,17 @@ def _do_compute_rk(accs_in, rk, buf_offset): 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"): + # Stochastic rounding: perturb the discarded bits with a + # per-element Philox draw keyed by the element index, so + # the f32 -> bf16 store is unbiased in expectation. + rbits = fx.rocdl.philox_4x32(fx.Uint32(elem_off), fx.Uint32(sr_seed))[0] + val = fx.rocdl.stochastic_round_bf16(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 +355,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), diff --git a/kernels/gemm/rdna_f16_gemm.py b/kernels/gemm/rdna_f16_gemm.py index ca8217cc3..cd28a0dae 100644 --- a/kernels/gemm/rdna_f16_gemm.py +++ b/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() @@ -319,11 +324,17 @@ def _load_a_single_from_lds(rk, rm_val, buf_offset): 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"): + # Stochastic rounding: perturb the discarded bits with a + # per-element Philox draw keyed by the element index, so + # the f32 -> bf16 store is unbiased in expectation. + rbits = fx.rocdl.philox_4x32(fx.Uint32(elem_off), fx.Uint32(sr_seed))[0] + val = fx.rocdl.stochastic_round_bf16(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 +344,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), diff --git a/python/flydsl/expr/rocdl/__init__.py b/python/flydsl/expr/rocdl/__init__.py index e35d4b90a..b85f5604f 100644 --- a/python/flydsl/expr/rocdl/__init__.py +++ b/python/flydsl/expr/rocdl/__init__.py @@ -502,6 +502,7 @@ def lds_transpose_load(result_type, lds_memref, elem_offset, elem_bytes): # ── New high-level helpers from universal.py ────────────────────────── from .cdna5 import * # noqa: E402,F401,F403,I001 from .inline_asm import * # noqa: E402,F401,F403,I001 +from .random import * # noqa: E402,F401,F403,I001 # ── Wrappers: accept DSL Numeric args (fx.Int32, fx.Float32, etc.) ───────── # ODS-generated ops require raw ir.Value. These wrappers auto-convert. diff --git a/python/flydsl/expr/rocdl/random.py b/python/flydsl/expr/rocdl/random.py new file mode 100644 index 000000000..8e8474308 --- /dev/null +++ b/python/flydsl/expr/rocdl/random.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +"""Counter-based RNG and stochastic rounding for AMD targets. + +Stochastic rounding does not fit the target-neutral IEEE rounding modes +(``fx.RoundingMode`` on ``.to()``, lowered through ``arith.truncf``): it needs +entropy, and the hardware form is a target-specific instruction (e.g. the CDNA +``v_cvt_*_sr`` conversions, PTX ``cvt.rs`` on NVIDIA). So it lives under the +rocdl target package rather than as a neutral rounding mode. The f32 -> bf16 +path here is a portable software emulation; a CDNA hardware ``sr`` variant can +slot in alongside it later. + +``philox_4x32`` is co-located as the entropy source that feeds it. +""" + +from ..numeric import BFloat16, Float32, Uint16, Uint32, Uint64 + +__all__ = ["philox_4x32", "stochastic_round_bf16"] + +_PHILOX_ROUND_A = 0xD2511F53 +_PHILOX_ROUND_B = 0xCD9E8D57 +_PHILOX_KEY_A = 0x9E3779B9 +_PHILOX_KEY_B = 0xBB67AE85 + + +def philox_4x32(counter, key, rounds: int = 7): + """Philox 4x32 counter-based PRNG. Returns four ``Uint32`` random words. + + ``counter`` and ``key`` are ``Uint32`` (e.g. a global element index and a + seed). The same ``(counter, key)`` always yields the same words, so no RNG + state has to be threaded through a kernel. Each round does two widening + 32x32 to 64 bit multiplies. + """ + c0 = Uint32(counter) + c1 = Uint32(0) + c2 = Uint32(0) + c3 = Uint32(0) + k0 = Uint32(key) + k1 = Uint32(0) + # Materialize the round multipliers, key increments and shift once instead + # of rebuilding identical constants on every unrolled round. + round_a = Uint64(_PHILOX_ROUND_A) + round_b = Uint64(_PHILOX_ROUND_B) + key_a = Uint32(_PHILOX_KEY_A) + key_b = Uint32(_PHILOX_KEY_B) + shift32 = Uint64(32) + for _ in range(rounds): + prod_b = Uint64(c2) * round_b + prod_a = Uint64(c0) * round_a + c0 = Uint32(prod_b >> shift32) ^ c1 ^ k0 + c2 = Uint32(prod_a >> shift32) ^ c3 ^ k1 + c1 = Uint32(prod_b) + c3 = Uint32(prod_a) + k0 = k0 + key_a + k1 = k1 + key_b + return c0, c1, c2, c3 + + +def stochastic_round_bf16(x, rand): + """Round a ``Float32`` to ``BFloat16`` with stochastic rounding. + + ``rand`` is a ``Uint32`` supplying entropy; its low 16 bits are used. + bf16 keeps the top 16 bits of the f32, so adding the random value to the + 16 discarded mantissa bits turns truncation into a round-up with + probability equal to the fractional distance to the next bf16 value. + bf16 shares the f32 exponent field, so Inf and NaN pass through unchanged + and no special-case handling is needed. + """ + u = Float32(x).bitcast(Uint32) + (Uint32(rand) & Uint32(0xFFFF)) + return Uint16(u >> Uint32(16)).bitcast(BFloat16) diff --git a/tests/kernels/test_rdna_gemm.py b/tests/kernels/test_rdna_gemm.py index 9744075c9..b6bbf5db4 100644 --- a/tests/kernels/test_rdna_gemm.py +++ b/tests/kernels/test_rdna_gemm.py @@ -103,6 +103,41 @@ def test_f16_gemm_correctness(M, N, K, in_dtype, out_dtype): assert verify_output(C.float(), C_ref, atol=0.05, rtol=0.05) +def test_f16_gemm_stochastic_rounding(): + """BF16 GEMM with the stochastic-rounding epilogue: bounded, seed-varying, reproducible. + + Unbiasedness of the rounding itself is proven in tests/unit/test_stochastic_rounding.py; + here we only check the GEMM wiring. + """ + _requires_rdna_wmma() + + M = N = K = 256 + torch.manual_seed(0) + A = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") * 0.1 + B_T = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") * 0.1 + C_ref = A.float() @ B_T.float().T + stream = torch.cuda.current_stream() + + rn_gemm, _, _, _ = create_wmma_gemm_module(M, N, K, in_dtype="bf16", out_dtype="bf16", rounding="rn") + rs_gemm, _, _, _ = create_wmma_gemm_module(M, N, K, in_dtype="bf16", out_dtype="bf16", rounding="rs") + + def run(launch_fn, *seed): + C = torch.zeros(M, N, dtype=torch.bfloat16, device="cuda") + launch_fn(C, A, B_T, stream, *seed) + torch.cuda.synchronize() + return C.float() + + C_rn = run(rn_gemm) # default 4-arg launch: seed is unused for round-to-nearest + C_rs1, C_rs2 = run(rs_gemm, 1), run(rs_gemm, 2) + + # SR stays within about an ULP of round-to-nearest against the f32 reference + assert (C_rs1 - C_ref).abs().max() < 3 * (C_rn - C_ref).abs().max() + 5e-3 + # the seed is a runtime argument: one compiled kernel, different seeds vary + # the rounding, and a repeated seed is reproducible + assert not torch.equal(C_rs1, C_rs2) + assert torch.equal(C_rs1, run(rs_gemm, 1)) + + @pytest.mark.parametrize( "M, N, K", [ diff --git a/tests/unit/test_stochastic_rounding.py b/tests/unit/test_stochastic_rounding.py new file mode 100644 index 000000000..20391c1ac --- /dev/null +++ b/tests/unit/test_stochastic_rounding.py @@ -0,0 +1,102 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +"""Tests for fx.rocdl stochastic rounding + philox RNG. + +The rounding math is checked exactly in pure Python (L0): over the full 16-bit +random range the number of round-ups equals the discarded low 16 bits, which is +what makes stochastic rounding unbiased. The device path and the RNG are checked +on GPU (L2): outputs land on the two nearest bf16 values with the right up +probability, and the empirical mean matches the input. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +import flydsl.compiler as flyc +import flydsl.expr as fx + +try: + import torch +except ImportError: + torch = None + + +@pytest.mark.l0_backend_agnostic +@pytest.mark.parametrize("x", [1.0, 1.001953125, 3.14159, 1e-3, float("inf"), float("-inf")]) +def test_stochastic_round_bf16_unbiased(x): + u = int(np.array([x], dtype=np.float32).view(np.uint32)[0]) + low16 = u & 0xFFFF + base = u >> 16 + r = np.arange(0, 1 << 16, dtype=np.uint64) + truncated = ((np.uint64(u) + r) >> np.uint64(16)).astype(np.uint64) + # exactly `low16` of the 65536 random draws round up + up_count = int(np.sum(truncated != np.uint64(base))) + assert up_count == low16 + # every result is one of the two nearest bf16 values + assert set(truncated.tolist()) <= {base, base + 1} + + +@pytest.mark.l2_device +@pytest.mark.rocm_lower +@pytest.mark.skipif(torch is None or not torch.cuda.is_available(), reason="requires GPU") +@pytest.mark.parametrize("frac", [0.0, 0.5, 0.9]) +def test_stochastic_round_bf16_device(frac): + BLOCK = 256 + NBLOCKS = 4096 + P = BLOCK * NBLOCKS + + @flyc.kernel(known_block_size=[BLOCK, 1, 1]) + def kernel(Out: fx.Tensor, x: fx.Float32, seed: fx.Int32): + g = fx.block_idx.x * BLOCK + fx.thread_idx.x + r = fx.rocdl.philox_4x32(fx.Uint32(g), fx.Uint32(seed))[0] + fx.memref_store(fx.rocdl.stochastic_round_bf16(x, r), Out, g) + + @flyc.jit + def launch(Out: fx.Tensor, x: fx.Float32, seed: fx.Int32, stream: fx.Stream = fx.Stream(None)): + kernel(Out, x, seed).launch(grid=(NBLOCKS, 1, 1), block=(BLOCK, 1, 1), stream=stream) + + step = 2.0**-7 # bf16 step above 1.0 + x = 1.0 + frac * step + lo, hi = 1.0, 1.0 + step + + out = torch.zeros(P, dtype=torch.bfloat16, device="cuda") + launch(out, x, 1234, stream=torch.cuda.Stream()) + torch.cuda.synchronize() + f = out.float() + + assert torch.all((f == lo) | (f == hi)), "outputs must be one of the two nearest bf16 values" + up_prob = (f == hi).float().mean().item() + assert abs(up_prob - frac) < 0.01, f"up_prob {up_prob} far from frac {frac}" + assert abs(f.mean().item() - x) < 1e-4 + + +@pytest.mark.l2_device +@pytest.mark.rocm_lower +@pytest.mark.skipif(torch is None or not torch.cuda.is_available(), reason="requires GPU") +def test_philox_4x32_uniform(): + BLOCK = 256 + NBLOCKS = 1024 + P = BLOCK * NBLOCKS + + @flyc.kernel(known_block_size=[BLOCK, 1, 1]) + def kernel(Out: fx.Tensor, seed: fx.Int32): + g = fx.block_idx.x * BLOCK + fx.thread_idx.x + fx.memref_store(fx.Int32(fx.rocdl.philox_4x32(fx.Uint32(g), fx.Uint32(seed))[0]), Out, g) + + @flyc.jit + def launch(Out: fx.Tensor, seed: fx.Int32, stream: fx.Stream = fx.Stream(None)): + kernel(Out, seed).launch(grid=(NBLOCKS, 1, 1), block=(BLOCK, 1, 1), stream=stream) + + out = torch.zeros(P, dtype=torch.int32, device="cuda") + launch(out, 7, stream=torch.cuda.Stream()) + torch.cuda.synchronize() + + u = out.cpu().numpy().astype(np.uint32).astype(np.float64) / 2**32 + assert abs(u.mean() - 0.5) < 0.01, f"mean {u.mean()} not uniform" + # counter-based PRNG on distinct counters must not collide in bulk + assert np.unique(out.cpu().numpy()).size > int(0.999 * P) From 4d1431b87833ba3092f54b4e4c98794af4a6abec Mon Sep 17 00:00:00 2001 From: Feng Shijie Date: Fri, 31 Jul 2026 09:47:18 +0000 Subject: [PATCH 2/3] [Feat] Add target-neutral fx.random library --- kernels/gemm/rdna3_f16_gemm.py | 4 +- kernels/gemm/rdna_f16_gemm.py | 4 +- python/flydsl/compiler/jit_function.py | 25 ++---- python/flydsl/expr/__init__.py | 5 +- python/flydsl/expr/_library.py | 62 +++++++++++++ python/flydsl/expr/random/__init__.py | 19 ++++ python/flydsl/expr/random/rocdl.py | 41 +++++++++ .../{rocdl/random.py => random/universal.py} | 77 +++++++++------- python/flydsl/expr/rocdl/__init__.py | 1 - tests/unit/test_expr_optional_rocdl.py | 23 ++++- tests/unit/test_stochastic_rounding.py | 90 ++++++++++++++++--- 11 files changed, 284 insertions(+), 67 deletions(-) create mode 100644 python/flydsl/expr/_library.py create mode 100644 python/flydsl/expr/random/__init__.py create mode 100644 python/flydsl/expr/random/rocdl.py rename python/flydsl/expr/{rocdl/random.py => random/universal.py} (50%) diff --git a/kernels/gemm/rdna3_f16_gemm.py b/kernels/gemm/rdna3_f16_gemm.py index 1e8397e87..e3db715fa 100644 --- a/kernels/gemm/rdna3_f16_gemm.py +++ b/kernels/gemm/rdna3_f16_gemm.py @@ -341,8 +341,8 @@ def _do_compute_rk(accs_in, rk, buf_offset): # Stochastic rounding: perturb the discarded bits with a # per-element Philox draw keyed by the element index, so # the f32 -> bf16 store is unbiased in expectation. - rbits = fx.rocdl.philox_4x32(fx.Uint32(elem_off), fx.Uint32(sr_seed))[0] - val = fx.rocdl.stochastic_round_bf16(val, rbits) + rbits = fx.random.philox_4x32(fx.Uint32(elem_off), fx.Uint32(sr_seed))[0] + 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"): diff --git a/kernels/gemm/rdna_f16_gemm.py b/kernels/gemm/rdna_f16_gemm.py index cd28a0dae..a7c8de263 100644 --- a/kernels/gemm/rdna_f16_gemm.py +++ b/kernels/gemm/rdna_f16_gemm.py @@ -329,8 +329,8 @@ def _load_a_single_from_lds(rk, rm_val, buf_offset): # Stochastic rounding: perturb the discarded bits with a # per-element Philox draw keyed by the element index, so # the f32 -> bf16 store is unbiased in expectation. - rbits = fx.rocdl.philox_4x32(fx.Uint32(elem_off), fx.Uint32(sr_seed))[0] - val = fx.rocdl.stochastic_round_bf16(val, rbits) + rbits = fx.random.philox_4x32(fx.Uint32(elem_off), fx.Uint32(sr_seed))[0] + 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"): diff --git a/python/flydsl/compiler/jit_function.py b/python/flydsl/compiler/jit_function.py index 9274e4657..fbf8c4d98 100644 --- a/python/flydsl/compiler/jit_function.py +++ b/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(): diff --git a/python/flydsl/expr/__init__.py b/python/flydsl/expr/__init__.py index a76223a53..53d9b37e3 100644 --- a/python/flydsl/expr/__init__.py +++ b/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}") diff --git a/python/flydsl/expr/_library.py b/python/flydsl/expr/_library.py new file mode 100644 index 000000000..d3da71d77 --- /dev/null +++ b/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 ``.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) diff --git a/python/flydsl/expr/random/__init__.py b/python/flydsl/expr/random/__init__.py new file mode 100644 index 000000000..98eea8724 --- /dev/null +++ b/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.`` runs the implementation chosen for the compilation +target; ``fx.random.universal.`` 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__) diff --git a/python/flydsl/expr/random/rocdl.py b/python/flydsl/expr/random/rocdl.py new file mode 100644 index 000000000..a402d2f62 --- /dev/null +++ b/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] diff --git a/python/flydsl/expr/rocdl/random.py b/python/flydsl/expr/random/universal.py similarity index 50% rename from python/flydsl/expr/rocdl/random.py rename to python/flydsl/expr/random/universal.py index 8e8474308..46b2aa2c9 100644 --- a/python/flydsl/expr/rocdl/random.py +++ b/python/flydsl/expr/random/universal.py @@ -1,30 +1,49 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 FlyDSL Project Contributors -"""Counter-based RNG and stochastic rounding for AMD targets. +"""Target-neutral implementations of the random library.""" -Stochastic rounding does not fit the target-neutral IEEE rounding modes -(``fx.RoundingMode`` on ``.to()``, lowered through ``arith.truncf``): it needs -entropy, and the hardware form is a target-specific instruction (e.g. the CDNA -``v_cvt_*_sr`` conversions, PTX ``cvt.rs`` on NVIDIA). So it lives under the -rocdl target package rather than as a neutral rounding mode. The f32 -> bf16 -path here is a portable software emulation; a CDNA hardware ``sr`` variant can -slot in alongside it later. +from ..numeric import BFloat16, Float32, Uint16, Uint32, Uint64 -``philox_4x32`` is co-located as the entropy source that feeds it. -""" +__all__ = [ + "cvt_f32_to_bf16_sr", + "philox_4x32", +] -from ..numeric import BFloat16, Float32, Uint16, Uint32, Uint64 -__all__ = ["philox_4x32", "stochastic_round_bf16"] +def cvt_f32_to_bf16_sr(x, rand): + """Round a ``Float32`` to ``BFloat16`` with stochastic rounding. + + ``rand`` is a ``Uint32`` supplying entropy. Pass a value that is uniformly + random over all 32 bits: this implementation reads only the low 16, but a + target override may consume different ones, so identical rounding + decisions are only reproducible for a fixed compilation target. + + bf16 keeps the top 16 bits of the f32, so adding the random value to the + 16 discarded mantissa bits turns truncation into a round-up with + probability equal to the fractional distance to the next bf16 value. The + result is one of the two bf16 neighbours of ``x`` and unbiased in + expectation. The add works on the raw sign-magnitude bits, so a round-up + moves away from zero for both signs; zero and any exactly representable + value are returned unchanged. + + Everything below follows from that single add-then-truncate, and differs + from an IEEE conversion: -_PHILOX_ROUND_A = 0xD2511F53 -_PHILOX_ROUND_B = 0xCD9E8D57 -_PHILOX_KEY_A = 0x9E3779B9 -_PHILOX_KEY_B = 0xBB67AE85 + * ``|x|`` above the largest bf16 finite (``0x1.FEp+127``) can carry into + the exponent and become Inf, with probability equal to the same + fractional distance. + * Inf is preserved: a zero mantissa leaves nothing for the add to carry. + * A NaN keeps only its top 7 payload bits, so every quiet NaN stays NaN. + A NaN carrying payload solely in the discarded low 16 bits (such as the + signaling NaN ``0x7F800001``) instead becomes Inf, unless the add + happens to carry. + """ + u = Float32(x).bitcast(Uint32) + (Uint32(rand) & Uint32(0xFFFF)) + return Uint16(u >> Uint32(16)).bitcast(BFloat16) -def philox_4x32(counter, key, rounds: int = 7): +def philox_4x32(counter, key, rounds: int = 10): """Philox 4x32 counter-based PRNG. Returns four ``Uint32`` random words. ``counter`` and ``key`` are ``Uint32`` (e.g. a global element index and a @@ -32,19 +51,24 @@ def philox_4x32(counter, key, rounds: int = 7): state has to be threaded through a kernel. Each round does two widening 32x32 to 64 bit multiplies. """ + _PHILOX_ROUND_A = 0xD2511F53 + _PHILOX_ROUND_B = 0xCD9E8D57 + _PHILOX_KEY_A = 0x9E3779B9 + _PHILOX_KEY_B = 0xBB67AE85 + c0 = Uint32(counter) c1 = Uint32(0) c2 = Uint32(0) c3 = Uint32(0) k0 = Uint32(key) k1 = Uint32(0) - # Materialize the round multipliers, key increments and shift once instead - # of rebuilding identical constants on every unrolled round. + round_a = Uint64(_PHILOX_ROUND_A) round_b = Uint64(_PHILOX_ROUND_B) key_a = Uint32(_PHILOX_KEY_A) key_b = Uint32(_PHILOX_KEY_B) shift32 = Uint64(32) + for _ in range(rounds): prod_b = Uint64(c2) * round_b prod_a = Uint64(c0) * round_a @@ -54,18 +78,5 @@ def philox_4x32(counter, key, rounds: int = 7): c3 = Uint32(prod_a) k0 = k0 + key_a k1 = k1 + key_b - return c0, c1, c2, c3 - - -def stochastic_round_bf16(x, rand): - """Round a ``Float32`` to ``BFloat16`` with stochastic rounding. - ``rand`` is a ``Uint32`` supplying entropy; its low 16 bits are used. - bf16 keeps the top 16 bits of the f32, so adding the random value to the - 16 discarded mantissa bits turns truncation into a round-up with - probability equal to the fractional distance to the next bf16 value. - bf16 shares the f32 exponent field, so Inf and NaN pass through unchanged - and no special-case handling is needed. - """ - u = Float32(x).bitcast(Uint32) + (Uint32(rand) & Uint32(0xFFFF)) - return Uint16(u >> Uint32(16)).bitcast(BFloat16) + return c0, c1, c2, c3 diff --git a/python/flydsl/expr/rocdl/__init__.py b/python/flydsl/expr/rocdl/__init__.py index b85f5604f..e35d4b90a 100644 --- a/python/flydsl/expr/rocdl/__init__.py +++ b/python/flydsl/expr/rocdl/__init__.py @@ -502,7 +502,6 @@ def lds_transpose_load(result_type, lds_memref, elem_offset, elem_bytes): # ── New high-level helpers from universal.py ────────────────────────── from .cdna5 import * # noqa: E402,F401,F403,I001 from .inline_asm import * # noqa: E402,F401,F403,I001 -from .random import * # noqa: E402,F401,F403,I001 # ── Wrappers: accept DSL Numeric args (fx.Int32, fx.Float32, etc.) ───────── # ODS-generated ops require raw ir.Value. These wrappers auto-convert. diff --git a/tests/unit/test_expr_optional_rocdl.py b/tests/unit/test_expr_optional_rocdl.py index d8f67c943..dd9cc67c1 100644 --- a/tests/unit/test_expr_optional_rocdl.py +++ b/tests/unit/test_expr_optional_rocdl.py @@ -59,6 +59,17 @@ def find_spec(self, fullname, path, target=None): assert "rocdl" not in expr.__dict__ assert "tdm_ops" not in expr.__dict__ assert "cluster_barrier" not in expr.gpu.__dict__ +assert "random" not in expr.__dict__ +assert expr.random.philox_4x32 is not None +assert expr.random.cvt_f32_to_bf16_sr is not None +assert "rocdl" not in expr.random.__dict__ +assert "flydsl.expr.random.rocdl" not in sys.modules + +from flydsl.compiler.jit_function import _flydsl_key_cached + +_flydsl_key_cached(False, "", ()) +assert "flydsl.expr.rocdl" not in sys.modules +assert "flydsl.expr.random.rocdl" not in sys.modules def _assert_rocdl_import_error(alias, exc): @@ -77,18 +88,23 @@ def _assert_rocdl_import_error(alias, exc): _LAZY_ALIAS_CHECK = """ import importlib +import sys expr = importlib.import_module("flydsl.expr") assert "buffer_ops" not in expr.__dict__ assert "rocdl" not in expr.__dict__ assert "tdm_ops" not in expr.__dict__ +assert "flydsl.expr.random.rocdl" not in sys.modules from flydsl.expr import rocdl, tdm_ops from flydsl.expr.rocdl import cluster +from flydsl.expr.random import rocdl as random_rocdl assert rocdl.__name__ == "flydsl.expr.rocdl" assert tdm_ops.__name__ == "flydsl.expr.rocdl.tdm_ops" assert cluster.__name__ == "flydsl.expr.rocdl.cluster" +assert random_rocdl.__name__ == "flydsl.expr.random.rocdl" +assert expr.random.rocdl is random_rocdl assert cluster.CLUSTER_BARRIER_ID == -3 assert expr.rocdl is rocdl assert expr.tdm_ops is tdm_ops @@ -96,7 +112,12 @@ def _assert_rocdl_import_error(alias, exc): def _build_env(): - pkg = _REPO_ROOT / "build-fly" / "python_packages" / "flydsl" + pkg = Path( + os.environ.get( + "FLYDSL_TEST_BUILD_FLYDSL_PKG", + _REPO_ROOT / "build-fly" / "python_packages" / "flydsl", + ) + ) if not pkg.is_dir(): pytest.skip("build-fly python_packages not found (run scripts/build.sh)") diff --git a/tests/unit/test_stochastic_rounding.py b/tests/unit/test_stochastic_rounding.py index 20391c1ac..6c8e90cdc 100644 --- a/tests/unit/test_stochastic_rounding.py +++ b/tests/unit/test_stochastic_rounding.py @@ -3,13 +3,12 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 FlyDSL Project Contributors -"""Tests for fx.rocdl stochastic rounding + philox RNG. +"""Tests for target-neutral RNG and stochastic BF16 conversion. The rounding math is checked exactly in pure Python (L0): over the full 16-bit random range the number of round-ups equals the discarded low 16 bits, which is -what makes stochastic rounding unbiased. The device path and the RNG are checked -on GPU (L2): outputs land on the two nearest bf16 values with the right up -probability, and the empirical mean matches the input. +what makes stochastic rounding unbiased. Device tests check the public paths, +special values, Philox reference output, reproducibility, and uniformity. """ from __future__ import annotations @@ -26,9 +25,25 @@ torch = None +_PHILOX_ZERO_VECTOR = (0x6627E8D5, 0xE169C58D, 0xBC57AC4C, 0x9B00DBD8) + + @pytest.mark.l0_backend_agnostic -@pytest.mark.parametrize("x", [1.0, 1.001953125, 3.14159, 1e-3, float("inf"), float("-inf")]) -def test_stochastic_round_bf16_unbiased(x): +def test_philox_4x32_reference_vector(): + words = fx.random.philox_4x32(fx.Uint32(0), fx.Uint32(0)) + assert tuple(int(word.value) for word in words) == _PHILOX_ZERO_VECTOR + repeated = fx.random.philox_4x32(fx.Uint32(0), fx.Uint32(0)) + different_seed = fx.random.philox_4x32(fx.Uint32(0), fx.Uint32(1)) + assert tuple(int(word.value) for word in repeated) == _PHILOX_ZERO_VECTOR + assert tuple(int(word.value) for word in different_seed) != _PHILOX_ZERO_VECTOR + + +@pytest.mark.l0_backend_agnostic +@pytest.mark.parametrize( + "x", + [1.0, 1.001953125, 3.14159, 1e-3, -1.001953125, -3.14159, float("inf"), float("-inf")], +) +def test_cvt_f32_to_bf16_sr_unbiased(x): u = int(np.array([x], dtype=np.float32).view(np.uint32)[0]) low16 = u & 0xFFFF base = u >> 16 @@ -45,7 +60,7 @@ def test_stochastic_round_bf16_unbiased(x): @pytest.mark.rocm_lower @pytest.mark.skipif(torch is None or not torch.cuda.is_available(), reason="requires GPU") @pytest.mark.parametrize("frac", [0.0, 0.5, 0.9]) -def test_stochastic_round_bf16_device(frac): +def test_cvt_f32_to_bf16_sr_device(frac): BLOCK = 256 NBLOCKS = 4096 P = BLOCK * NBLOCKS @@ -53,8 +68,8 @@ def test_stochastic_round_bf16_device(frac): @flyc.kernel(known_block_size=[BLOCK, 1, 1]) def kernel(Out: fx.Tensor, x: fx.Float32, seed: fx.Int32): g = fx.block_idx.x * BLOCK + fx.thread_idx.x - r = fx.rocdl.philox_4x32(fx.Uint32(g), fx.Uint32(seed))[0] - fx.memref_store(fx.rocdl.stochastic_round_bf16(x, r), Out, g) + r = fx.Uint32(g) * fx.Uint32(1664525) + fx.Uint32(seed) + fx.memref_store(fx.random.cvt_f32_to_bf16_sr(x, r), Out, g) @flyc.jit def launch(Out: fx.Tensor, x: fx.Float32, seed: fx.Int32, stream: fx.Stream = fx.Stream(None)): @@ -75,6 +90,61 @@ def launch(Out: fx.Tensor, x: fx.Float32, seed: fx.Int32, stream: fx.Stream = fx assert abs(f.mean().item() - x) < 1e-4 +@pytest.mark.l2_device +@pytest.mark.rocm_lower +@pytest.mark.skipif(torch is None or not torch.cuda.is_available(), reason="requires GPU") +def test_cvt_f32_to_bf16_sr_special_values_device(): + BLOCK = 1 + + @flyc.kernel(known_block_size=[BLOCK, 1, 1]) + def kernel(Out: fx.Tensor, x: fx.Float32, seed: fx.Int32): + g = fx.block_idx.x * BLOCK + fx.thread_idx.x + r = fx.Uint32(g) * fx.Uint32(1664525) + fx.Uint32(seed) + fx.memref_store(fx.random.cvt_f32_to_bf16_sr(x, r), Out, g) + + @flyc.jit + def launch(Out: fx.Tensor, x: fx.Float32, seed: fx.Int32, stream: fx.Stream = fx.Stream(None)): + kernel(Out, x, seed).launch(grid=(1, 1, 1), block=(BLOCK, 1, 1), stream=stream) + + def run(x, seed): + out = torch.empty(BLOCK, dtype=torch.bfloat16, device="cuda") + launch(out, x, seed, stream=torch.cuda.Stream()) + torch.cuda.synchronize() + return out + + for x in (float("inf"), float("-inf"), float("nan")): + first = run(x, 1) + second = run(x, 2) + assert torch.equal(first.view(torch.int16), second.view(torch.int16)) + if np.isnan(x): + assert torch.isnan(first).all() + else: + assert torch.isinf(first).all() + assert (first < 0).all().item() == (x < 0) + + +@pytest.mark.l2_device +@pytest.mark.rocm_lower +@pytest.mark.skipif(torch is None or not torch.cuda.is_available(), reason="requires GPU") +def test_philox_4x32_reference_vector_device(): + @flyc.kernel(known_block_size=[1, 1, 1]) + def kernel(Out: fx.Tensor): + words = fx.random.philox_4x32(fx.Uint32(0), fx.Uint32(0)) + fx.memref_store(fx.Int32(words[0]), Out, 0) + fx.memref_store(fx.Int32(words[1]), Out, 1) + fx.memref_store(fx.Int32(words[2]), Out, 2) + fx.memref_store(fx.Int32(words[3]), Out, 3) + + @flyc.jit + def launch(Out: fx.Tensor, stream: fx.Stream = fx.Stream(None)): + kernel(Out).launch(grid=(1, 1, 1), block=(1, 1, 1), stream=stream) + + out = torch.empty(4, dtype=torch.int32, device="cuda") + launch(out, stream=torch.cuda.Stream()) + torch.cuda.synchronize() + assert tuple(int(value) & 0xFFFFFFFF for value in out.cpu().tolist()) == _PHILOX_ZERO_VECTOR + + @pytest.mark.l2_device @pytest.mark.rocm_lower @pytest.mark.skipif(torch is None or not torch.cuda.is_available(), reason="requires GPU") @@ -86,7 +156,7 @@ def test_philox_4x32_uniform(): @flyc.kernel(known_block_size=[BLOCK, 1, 1]) def kernel(Out: fx.Tensor, seed: fx.Int32): g = fx.block_idx.x * BLOCK + fx.thread_idx.x - fx.memref_store(fx.Int32(fx.rocdl.philox_4x32(fx.Uint32(g), fx.Uint32(seed))[0]), Out, g) + fx.memref_store(fx.Int32(fx.random.philox_4x32(fx.Uint32(g), fx.Uint32(seed))[0]), Out, g) @flyc.jit def launch(Out: fx.Tensor, seed: fx.Int32, stream: fx.Stream = fx.Stream(None)): From 37bfa9931ead17131dc3ec10df6bfc2f73e08cd6 Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Fri, 31 Jul 2026 14:18:46 +0200 Subject: [PATCH 3/3] Amortize philox draws across the SR gemm epilogue block --- kernels/gemm/rdna3_f16_gemm.py | 14 ++++++++++---- kernels/gemm/rdna_f16_gemm.py | 14 ++++++++++---- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/kernels/gemm/rdna3_f16_gemm.py b/kernels/gemm/rdna3_f16_gemm.py index e3db715fa..300ccda6f 100644 --- a/kernels/gemm/rdna3_f16_gemm.py +++ b/kernels/gemm/rdna3_f16_gemm.py @@ -332,16 +332,22 @@ 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] elem_off = g_row * N + g_col if const_expr(rounding == "rs"): - # Stochastic rounding: perturb the discarded bits with a - # per-element Philox draw keyed by the element index, so - # the f32 -> bf16 store is unbiased in expectation. - rbits = fx.random.philox_4x32(fx.Uint32(elem_off), fx.Uint32(sr_seed))[0] + 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) diff --git a/kernels/gemm/rdna_f16_gemm.py b/kernels/gemm/rdna_f16_gemm.py index a7c8de263..5ce2ec64f 100644 --- a/kernels/gemm/rdna_f16_gemm.py +++ b/kernels/gemm/rdna_f16_gemm.py @@ -320,16 +320,22 @@ 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] elem_off = g_row * N + g_col if const_expr(rounding == "rs"): - # Stochastic rounding: perturb the discarded bits with a - # per-element Philox draw keyed by the element index, so - # the f32 -> bf16 store is unbiased in expectation. - rbits = fx.random.philox_4x32(fx.Uint32(elem_off), fx.Uint32(sr_seed))[0] + 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)