PYTHONREMU: V_CVT_PK_BF8_F32 saturation (#16268)

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)

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)

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)))