[TVMScript] Add parser and printer support for e4m3/e5m2 fp8

include/tvm/script/ir_builder/tir/ir.h

@@ -489,6 +489,18 @@ TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES(Int, DataType::Int);
 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(Float, DataType::Float);
 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(UInt, DataType::UInt);
 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(Int, DataType::Int);
+
+#define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(DType, FDType) \
+  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType, FDType(1));                    \
+  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x4, FDType(4));                \
+  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x8, FDType(8));                \
+  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x16, FDType(16));              \
+  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x32, FDType(32));              \
+  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x64, FDType(64));
+
+TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(E4M3Float8, DataType::NVFloat8E4M3);
+TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(E5M2Float8, DataType::NVFloat8E5M2);
+
 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(Boolean, DataType::Bool());
 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(Void, DataType::Void());
 

python/tvm/script/ir_builder/tir/ir.py

@@ -1408,30 +1408,39 @@ def func(
 uint32x64 = func_gen(("UInt32x64"))
 uint64x64 = func_gen(("UInt64x64"))
 
-float8 = func_gen(("Float8"))
 float16 = func_gen(("Float16"))
 float32 = func_gen(("Float32"))
 float64 = func_gen(("Float64"))
-float8x4 = func_gen(("Float8x4"))
 float16x4 = func_gen(("Float16x4"))
 float32x4 = func_gen(("Float32x4"))
 float64x4 = func_gen(("Float64x4"))
-float8x8 = func_gen(("Float8x8"))
 float16x8 = func_gen(("Float16x8"))
 float32x8 = func_gen(("Float32x8"))
 float64x8 = func_gen(("Float64x8"))
-float8x16 = func_gen(("Float8x16"))
 float16x16 = func_gen(("Float16x16"))
 float32x16 = func_gen(("Float32x16"))
 float64x16 = func_gen(("Float64x16"))
-float8x32 = func_gen(("Float8x32"))
 float16x32 = func_gen(("Float16x32"))
 float32x32 = func_gen(("Float32x32"))
 float64x32 = func_gen(("Float64x32"))
-float8x64 = func_gen(("Float8x64"))
 float16x64 = func_gen(("Float16x64"))
 float32x64 = func_gen(("Float32x64"))
 float64x64 = func_gen(("Float64x64"))
+
+e4m3_float8 = func_gen(("E4M3Float8"))
+e4m3_float8x4 = func_gen(("E4M3Float8x4"))
+e4m3_float8x8 = func_gen(("E4M3Float8x8"))
+e4m3_float8x16 = func_gen(("E4M3Float8x16"))
+e4m3_float8x32 = func_gen(("E4M3Float8x32"))
+e4m3_float8x64 = func_gen(("E4M3Float8x64"))
+
+e5m2_float8 = func_gen(("E5M2Float8"))
+e5m2_float8x4 = func_gen(("E5M2Float8x4"))
+e5m2_float8x8 = func_gen(("E5M2Float8x8"))
+e5m2_float8x16 = func_gen(("E5M2Float8x16"))
+e5m2_float8x32 = func_gen(("E5M2Float8x32"))
+e5m2_float8x64 = func_gen(("E5M2Float8x64"))
+
 # pylint: enable=invalid-name
 
 
@@ -1954,27 +1963,33 @@ def wrapped(*args, **kwargs):
     "uint16x64",
     "uint32x64",
     "uint64x64",
-    "float8",
+    "e4m3_float8",
+    "e5m2_float8",
     "float16",
     "float32",
     "float64",
-    "float8x4",
+    "e4m3_float8x4",
+    "e5m2_float8x4",
     "float16x4",
     "float32x4",
     "float64x4",
-    "float8x8",
+    "e4m3_float8x8",
+    "e5m2_float8x8",
     "float16x8",
     "float32x8",
     "float64x8",
-    "float8x16",
+    "e4m3_float8x16",
+    "e5m2_float8x16",
     "float16x16",
     "float32x16",
     "float64x16",
-    "float8x32",
+    "e4m3_float8x32",
+    "e5m2_float8x32",
     "float16x32",
     "float32x32",
     "float64x32",
-    "float8x64",
+    "e4m3_float8x64",
+    "e5m2_float8x64",
     "float16x64",
     "float32x64",
     "float64x64",

src/script/ir_builder/tir/ir.cc

@@ -751,6 +751,11 @@ TVM_REGISTER_GLOBAL_SIZES_LANES("script.ir_builder.tir.Float", Float);
 TVM_REGISTER_GLOBAL_SIZES_LANES("script.ir_builder.tir.UInt", UInt);
 TVM_REGISTER_GLOBAL_SIZES_LANES("script.ir_builder.tir.Int", Int);
 
+TVM_REGISTER_GLOBAL("script.ir_builder.tir.E4M3Float8").set_body_typed(E4M3Float8);
+TVM_REGISTER_GLOBAL("script.ir_builder.tir.E5M2Float8").set_body_typed(E5M2Float8);
+TVM_REGISTER_GLOBAL_LANES("script.ir_builder.tir.E4M3Float8", E4M3Float8);
+TVM_REGISTER_GLOBAL_LANES("script.ir_builder.tir.E5M2Float8", E5M2Float8);
+
 TVM_REGISTER_GLOBAL("script.ir_builder.tir.Boolean").set_body_typed(Boolean);
 TVM_REGISTER_GLOBAL("script.ir_builder.tir.Handle").set_body_typed(Handle);
 TVM_REGISTER_GLOBAL("script.ir_builder.tir.Void").set_body_typed(Void);

tests/python/tvmscript/test_tvmscript_printer_tir.py

@@ -917,5 +917,36 @@ def func():
     _assert_print(func, expected_output)
 
 
+@pytest.mark.parametrize("dtype", ["e4m3_float8", "e5m2_float8"])
+def test_float8(dtype):
+    from tvm.script import tir as T
+
+    def get_func(dtype):
+        if dtype == "e4m3_float8":
+
+            @T.prim_func
+            def func():
+                T.evaluate(T.e4m3_float8(0.0))
+
+            return func
+        elif dtype == "e5m2_float8":
+
+            @T.prim_func
+            def func():
+                T.evaluate(T.e5m2_float8(0.0))
+
+            return func
+
+    expected_output = f"""
+# from tvm.script import tir as T
+
+@T.prim_func
+def func():
+    T.evaluate(T.{dtype}(0))
+    """
+    func = get_func(dtype)
+    _assert_print(func, expected_output)
+
+
 if __name__ == "__main__":
     tvm.testing.main()