From 7cdd9cbdeb3456b4e152562ab09d5c2d0357b2c4 Mon Sep 17 00:00:00 2001 From: Christopher Milan Date: Tue, 19 May 2026 16:29:59 -0700 Subject: [PATCH] PYTHONREMU: V_CVT_PK_BF8_F32 saturation (#16268) --- test/amd/hw/test_cdna_vop3.py | 18 ++++++++++++++++++ test/mockgpu/amd/pcode.py | 4 ++-- tinygrad/uop/decompositions.py | 8 ++++---- 3 files changed, 24 insertions(+), 6 deletions(-) diff --git a/test/amd/hw/test_cdna_vop3.py b/test/amd/hw/test_cdna_vop3.py index 3bc4b960c8..6d634f6425 100644 --- a/test/amd/hw/test_cdna_vop3.py +++ b/test/amd/hw/test_cdna_vop3.py @@ -110,3 +110,21 @@ class TestCDNAVOP3(unittest.TestCase): cdna.v_cvt_pk_fp8_f32(cdna.v[2], cdna.v[0], cdna.v[1]), ]) self.assertEqual(out, 0xdead4038) + + def test_cvt_pk_bf8_f32_overflow_and_inf(self): + """V_CVT_PK_BF8_F32 converts finite overflow and infinities to E5M2 infinities.""" + for name, bits, expected in [ + ("finite_overflow", 0x47700000, 0x7c), + ("pos_inf", 0x7f800000, 0x7c), + ("neg_inf", 0xff800000, 0xfc), + ]: + with self.subTest(name=name): + out = run_cdna([ + cdna.s_mov_b32(cdna.s[0], 0xdeadbeef), + cdna.v_mov_b32_e32(cdna.v[2], cdna.s[0]), + cdna.s_mov_b32(cdna.s[0], bits), + cdna.v_mov_b32_e32(cdna.v[0], cdna.s[0]), + cdna.v_mov_b32_e32(cdna.v[1], 1.0), + cdna.v_cvt_pk_bf8_f32(cdna.v[2], cdna.v[0], cdna.v[1]), + ]) + self.assertEqual(out, 0xdead3c00 | expected) diff --git a/test/mockgpu/amd/pcode.py b/test/mockgpu/amd/pcode.py index 3a97f6885f..686f006233 100644 --- a/test/mockgpu/amd/pcode.py +++ b/test/mockgpu/amd/pcode.py @@ -105,9 +105,9 @@ def _bf8_to_f32(v: UOp) -> UOp: return is_sub.where(sub_f32.bitcast(dtypes.float32), normal) def _f32_to_fp8(v: UOp) -> UOp: - return f2f((v.bitcast(dtypes.float32) if v.dtype != dtypes.float32 else v).bitcast(dtypes.uint32), dtypes.float32, dtypes.fp8e4m3) + return f2f((v.bitcast(dtypes.float32) if v.dtype != dtypes.float32 else v).bitcast(dtypes.uint32), dtypes.float32, dtypes.fp8e4m3, sat=False) def _f32_to_bf8(v: UOp) -> UOp: - return f2f((v.bitcast(dtypes.float32) if v.dtype != dtypes.float32 else v).bitcast(dtypes.uint32), dtypes.float32, dtypes.fp8e5m2) + return f2f((v.bitcast(dtypes.float32) if v.dtype != dtypes.float32 else v).bitcast(dtypes.uint32), dtypes.float32, dtypes.fp8e5m2, sat=False) def _f32_to_bf16(v: UOp) -> UOp: """Convert f32 to bf16 with round-to-nearest-even. BF16 is the upper 16 bits of F32 with rounding.""" bits = (v.bitcast(dtypes.float32) if v.dtype != dtypes.float32 else v).bitcast(dtypes.uint32) diff --git a/tinygrad/uop/decompositions.py b/tinygrad/uop/decompositions.py index d468f1a467..170911eecc 100644 --- a/tinygrad/uop/decompositions.py +++ b/tinygrad/uop/decompositions.py @@ -384,7 +384,7 @@ f2f_dt = { f:getattr(dtypes, f"uint{f.bitsize}") for f in dtypes.floats } def rne(v: UOp, s) -> UOp: return shr(v, s) + ((shr(v, s - 1) & 1) & ((v & ((1 << (s - 1)) - 1)).ne(0).cast(v.dtype) | (shr(v, s) & 1))) -def f2f(v, fr:DType, to:DType): +def f2f(v, fr:DType, to:DType, sat=True): fs, fb, (fe, fm), ts, tb, (te, tm) = fr.bitsize, exponent_bias(fr), dtypes.finfo(fr), to.bitsize, exponent_bias(to), dtypes.finfo(to) # NB: denormals are zero! if fe <= te and fm < tm: @@ -399,7 +399,7 @@ def f2f(v, fr:DType, to:DType): is_nan = (nosign.eq(shl(1, fm + fe) - 1) if fr == dtypes.fp8e4m3 else exp.eq(shl(1, fe) - 1)) return (sign | exp.eq(0).where(0, is_nan.where(nan, norm))).bitcast(to) elif fe >= te and fm > tm: - v = f2f_clamp(v.bitcast(fr), to).bitcast(f2f_dt[fr]) + v = f2f_clamp(v.bitcast(fr), to, sat).bitcast(f2f_dt[fr]) sign, nosign = shr(v, fs - ts) & shl(1, ts - 1), v & (shl(1, fs - 1) - 1) norm = (rne(nosign, fm - tm) - shl(fb - tb, tm)).cast(f2f_dt[to]) underflow = (shr(v, fm) & (shl(1, fe) - 1)) < (1 + fb - tb) @@ -410,12 +410,12 @@ def f2f(v, fr:DType, to:DType): return is_nan.where(nan, sign.cast(f2f_dt[to]) | underflow.where(0, norm)) else: raise NotImplementedError(f"unsupported decomp {fr} -> {to}") -def f2f_clamp(val:UOp, dt:DType) -> UOp: +def f2f_clamp(val:UOp, dt:DType, sat=True) -> UOp: e, m = dtypes.finfo(dt) if dt in dtypes.fp8_fnuz: max_exp, max_man = (1 << e) - 1, (1 << m) - 1 else: max_exp, max_man = ((1 << e) - 1, (1 << m) - 2) if dt == dtypes.fp8e4m3 else ((1 << e) - 2, (1 << m) - 1) mx = val.const_like(2.0**(max_exp - exponent_bias(dt)) * (1.0 + max_man / (1 << m))) - sat = mx if dt in dtypes.fp8s else val.const_like(float('inf')) + sat = mx if dt in dtypes.fp8s and sat else val.const_like(float('inf')) # FIXME: CMPLT of nan is undefined return val.ne(val).where(val, (val < -mx).where(-sat, (mx < val).where(sat, val)))