Fixed FP to FP8 downcast (#4748)

Fixes #4630

This implementation is based on an arithmetic approach (see below) and
provides an RTNE conversion from F32/16 to F8 types. Despite it is using
more expensive fp binary operations, I don't see any performance
degradation (you can find a simple bench commented out in the test). The
implementation is very simple and template-based. It replaces 4 existing
functions that are more complicated and adds a direct conversion from
F32 to F8 without intermediate F16.

The implementation is based on the following approach. To convert a
float from src to dst, we need to find such an exponent and mantissa
that src is equal or close to dst. For the normal numbers (exp != 0) it
means:
```python
2^(srcExp - SrcBias) * (1 + srcMan/2^SrcMBits) = 2^(dstExp - DstBias) * (1 + dstMan/2^DstMBits)
```
and the following for subnormals (dstExp == 0):
```python
src = 2^(1 - DstBias) * (dstMan/2^DstMBits)
```
The exponent is calculated as:
```python
dstExp = max(0, srcExp - SrcBias + DstBias)
```
Simplifying the first formula, we can get the following:
```python
dstMan = srcMan * 2^(DstMBits - SrcMBits)
```
If `SrcBias == DstBias` (in the case of FP16 to F8E5 conversion), this
formula also works for subnormals. In the general case, we can use the
following formula for subnormals:
```python
dstMan = src * 2^(DstMBits + DstBias - 1)
```
Thus, to get the mantissa, we are using a simple multiplication by a
constant and rounding the result to the nearest int. In case of the
mantissa overflow, we need to reset it to zero and increment the
exponent. It could be done in the following way:
```python
dst = max((dstExp << DstMBits) + dstMan, DST_MAX)
```
It gives us the required dst value without the sign bit.

Author: AndreyPavlenko

python/test/unit/intel/test_conversions.py

@@ -0,0 +1,54 @@
+import torch
+import pytest
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def type_convert(src, dst, rounding: tl.constexpr, BLOCK_SIZE: tl.constexpr):
+    idxs = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    x = tl.load(src + idxs)
+    y = x.to(dst.dtype.element_ty, fp_downcast_rounding=rounding)
+    tl.store(dst + idxs, y)
+
+
+@pytest.mark.parametrize("dst_dtype", [torch.float8_e4m3fn, torch.float8_e5m2], ids=["float8_e4m3fn", "float8_e5m2"])
+@pytest.mark.parametrize("src_dtype", [torch.float16, torch.bfloat16, torch.float32],
+                         ids=["float16", "bfloat16", "float32"])
+def test_convert_to_fp8(src_dtype, dst_dtype, device):
+    src_idtype = torch.int32 if src_dtype == torch.float32 else torch.int16
+    finfo = torch.finfo(dst_dtype)
+    min_val = torch.tensor(finfo.min, dtype=dst_dtype).view(torch.uint8).item()
+    max_val = torch.tensor(finfo.max, dtype=dst_dtype).view(torch.uint8).item()
+    SIZE = 2**16
+    BLOCK_SIZE = SIZE // 32
+    src = torch.arange(0, SIZE, dtype=src_idtype, device=device)
+    if src_dtype == torch.float32:
+        src = src << 16 | src
+    src = src.view(src_dtype)
+    dst = torch.empty_like(src, dtype=dst_dtype, device=device)
+    type_convert[(SIZE // BLOCK_SIZE, )](triton.reinterpret(src, src_dtype), triton.reinterpret(dst, dst_dtype), 'rtne',
+                                         BLOCK_SIZE)
+
+    dst = dst.view(torch.uint8)
+    expect = src.to(dtype=dst_dtype).view(torch.uint8)
+    diff_mask = dst != expect
+    src = src[diff_mask]
+    dst = dst[diff_mask]
+    expect = expect[diff_mask]
+
+    for s, si, e, d in zip(src, src.view(src_idtype), expect.view(torch.uint8), dst.view(torch.uint8)):
+        if torch.isnan(s):
+            e = 0b01111111
+        elif torch.isposinf(s) or (s >= 57344.) or (s >= 464. and dst_dtype == torch.float8_e4m3fn):
+            e = max_val
+        elif torch.isneginf(s) or (s <= -57344.) or (s <= -464. and dst_dtype == torch.float8_e4m3fn):
+            e = min_val
+        elif si == 0b1000000000000000:  # -0.0
+            e = 0b10000000
+
+        if d != e:
+            sfmt = "032b" if src_dtype == torch.float32 else "016b"
+            dfmt = "08b"
+            msg = f"Src={s}({format(si, sfmt)}). Expected={format(e, dfmt)}. Actual={format(d, dfmt)}."
+            pytest.fail(msg)

python/test/unit/language/test_conversions.py

@@ -234,17 +234,9 @@ def downcast_test(src_dtype, dst_dtype, rounding, exponent_bits, mantissa_bits,
 
     src = launch_exhaustive_populate(src_dtype, offset << 24, 2**24, False, src_dtype.primitive_bitwidth, max_repr, device)
     dst = launch_type_convert_triton(src, src_dtype, dst_dtype, device=device, rounding=rounding)
-    # Emulated cast always works on fp32. In XPU Triton kernels FP32 is casted to FP8 through FP16, which
-    # in some cases gives different results compared to direct FP32 to FP8 conversion (some precision might
-    # be lost due to two-step conversion). To get matching results, we convert FP32 source data to FP16 and
-    # back to FP32. This will need to be changed back when HW FP32->FP8 convertion is used for XPU.
-    if device=='xpu' and src_dtype.primitive_bitwidth == 32 and dst_dtype.primitive_bitwidth == 8:
-        src = launch_type_convert_triton(src, src_dtype, tl.float16, device=device, rounding=rounding)
-        src = launch_type_convert_triton(src, tl.float16, tl.float32, device=device)
-    else:
-        src = launch_type_convert_triton(src, src_dtype, tl.float32, device=device)
+    src = launch_type_convert_triton(src, src_dtype, tl.float32, device=device)
 
-    dst2 = launch_downcast_emulated(src, src_dtype, dst_dtype, rounding, exponent_bits, mantissa_bits, exponent_bias, device=device)
+    dst2 = launch_downcast_emulated(src, tl.float32, dst_dtype, rounding, exponent_bits, mantissa_bits, exponent_bias, device=device)
 
     dst = launch_upcast_emulated(dst, exponent_bits, mantissa_bits, exponent_bias, device=device)
     dst2 = launch_upcast_emulated(dst2, exponent_bits, mantissa_bits, exponent_bias, device=device)

scripts/skiplist/a770/language.txt

@@ -651,20 +651,6 @@ python/test/unit/language/test_matmul.py::test_lhs_in_tmem[float8e5-True-128-64-
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_store
 python/test/unit/language/test_tensor_descriptor.py::test_make_tensor_descriptor_matmul
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_rank_reducing_matmul
-# https://github.com/intel/intel-xpu-backend-for-triton/issues/4630
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e5--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e5--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-max]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-min]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e5--inf]
 # https://github.com/intel/intel-xpu-backend-for-triton/issues/4289
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[1-1024-host-1-float16-add]
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[1-1024-host-1-float32-add]

scripts/skiplist/arl-h/language.txt

@@ -524,20 +524,6 @@ python/test/unit/language/test_matmul.py::test_lhs_in_tmem[float8e5-True-128-64-
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_store
 python/test/unit/language/test_tensor_descriptor.py::test_make_tensor_descriptor_matmul
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_rank_reducing_matmul
-# https://github.com/intel/intel-xpu-backend-for-triton/issues/4630
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e5--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e5--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-max]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-min]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e5--inf]
 # https://github.com/intel/intel-xpu-backend-for-triton/issues/4289
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[1-1024-host-1-float16-add]
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[1-1024-host-1-float32-add]

scripts/skiplist/arl-s/language.txt

@@ -524,20 +524,6 @@ python/test/unit/language/test_matmul.py::test_lhs_in_tmem[float8e5-True-128-64-
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_store
 python/test/unit/language/test_tensor_descriptor.py::test_make_tensor_descriptor_matmul
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_rank_reducing_matmul
-# https://github.com/intel/intel-xpu-backend-for-triton/issues/4630
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e5--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e5--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-max]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-min]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e5--inf]
 # https://github.com/intel/intel-xpu-backend-for-triton/issues/4289
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[1-1024-host-1-float16-add]
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[1-1024-host-1-float32-add]

scripts/skiplist/default/language.txt

@@ -1,20 +1,6 @@
 # https://github.com/intel/intel-xpu-backend-for-triton/issues/4665
 python/test/unit/language/test_core.py::test_dot3d[8-1-32-32-32-32-32-float64-float64]
 python/test/unit/language/test_core.py::test_dot3d[4-1-64-64-64-32-32-float64-float64]
-# https://github.com/intel/intel-xpu-backend-for-triton/issues/4630
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e5--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e5--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-max]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-min]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e5--inf]
 # https://github.com/intel/intel-xpu-backend-for-triton/issues/4289
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[1-1024-host-1-float16-add]
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[1-1024-host-1-float32-add]

scripts/skiplist/lts/language.txt

@@ -252,20 +252,6 @@ python/test/unit/language/test_core.py::test_convert_mma2mma[mma_pair0-float16-2
 python/test/unit/language/test_matmul.py::test_lhs_in_tmem
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_batched_gemm_2d_tma
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_batched_gemm_3d_tma
-# https://github.com/intel/intel-xpu-backend-for-triton/issues/4630
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e5--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e5--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-max]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-min]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e5--inf]
 # https://github.com/intel/intel-xpu-backend-for-triton/issues/4289
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[1-1024-host-1-float16-add]
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[1-1024-host-1-float32-add]

scripts/skiplist/mtl/language.txt

@@ -301,20 +301,6 @@ python/test/unit/language/test_core.py::test_dot[1-64-128-128-2-False-False-none
 python/test/unit/language/test_core.py::test_dot[1-64-128-128-2-False-False-none-tf32-float32-float32-1-None1]
 python/test/unit/language/test_core.py::test_dot[1-128-128-64-4-False-False-chain-dot-ieee-float8e5-float32-1-None]
 python/test/unit/language/test_core.py::test_dot[1-128-128-64-4-False-False-chain-dot-ieee-float8e4nv-float32-1-None]
-# https://github.com/intel/intel-xpu-backend-for-triton/issues/4630
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e5--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e5--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-max]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-min]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e5--inf]
 # https://github.com/intel/intel-xpu-backend-for-triton/issues/4289
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[1-1024-host-1-float16-add]
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[1-1024-host-1-float32-add]

scripts/skiplist/xe2/language.txt

@@ -1,17 +1,3 @@
-# https://github.com/intel/intel-xpu-backend-for-triton/issues/4630
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float32-float8e5--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[float16-float8e5--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-max]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-min]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv--inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e4nv-nan]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e5-inf]
-python/test/unit/language/test_conversions.py::test_typeconvert_downcast_clamping[bfloat16-float8e5--inf]
 # https://github.com/intel/intel-xpu-backend-for-triton/issues/4289
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[1-1024-host-1-float16-add]
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[1-1024-host-1-float32-add]