diff --git a/kernels/gemm/rdna3_f16_gemm.py b/kernels/gemm/rdna3_f16_gemm.py index ef6a65f68..300ccda6f 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* @@ -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), diff --git a/kernels/gemm/rdna_f16_gemm.py b/kernels/gemm/rdna_f16_gemm.py index ca8217cc3..5ce2ec64f 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() @@ -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), 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/random/universal.py b/python/flydsl/expr/random/universal.py new file mode 100644 index 000000000..46b2aa2c9 --- /dev/null +++ b/python/flydsl/expr/random/universal.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +"""Target-neutral implementations of the random library.""" + +from ..numeric import BFloat16, Float32, Uint16, Uint32, Uint64 + +__all__ = [ + "cvt_f32_to_bf16_sr", + "philox_4x32", +] + + +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: + + * ``|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 = 10): + """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. + """ + _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) + + 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 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_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 new file mode 100644 index 000000000..6c8e90cdc --- /dev/null +++ b/tests/unit/test_stochastic_rounding.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +"""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. Device tests check the public paths, +special values, Philox reference output, reproducibility, and uniformity. +""" + +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 + + +_PHILOX_ZERO_VECTOR = (0x6627E8D5, 0xE169C58D, 0xBC57AC4C, 0x9B00DBD8) + + +@pytest.mark.l0_backend_agnostic +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 + 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_cvt_f32_to_bf16_sr_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.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=(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_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") +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.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)): + 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)