[TIR] Implement max/min_value for fp8 data types (#16723)

vinx13 · web-flow · commit 9ec72494cf71 · 2024-03-15T16:05:53.000+08:00
diff --git a/src/tir/op/op.cc b/src/tir/op/op.cc
@@ -262,6 +262,12 @@ PrimExpr max_value(const DataType& dtype, Span span) {
     }
   } else if (dtype.is_bfloat16()) {
     return FloatImm(dtype, std::numeric_limits<float>::max(), span);
+  } else if (dtype.is_float8()) {
+    if (dtype.code() == DataType::TypeCode::kE5M2Float) {
+      return FloatImm(dtype, 57344.0, span);
+    } else if (dtype.code() == DataType::TypeCode::kE4M3Float) {
+      return FloatImm(dtype, 448.0, span);
+    }
   }
   LOG(FATAL) << "Cannot decide max_value for type" << dtype;
 }
@@ -296,6 +302,12 @@ PrimExpr min_value(const DataType& dtype, Span span) {
     }
   } else if (dtype.is_bfloat16()) {
     return FloatImm(dtype, std::numeric_limits<float>::lowest(), span);
+  } else if (dtype.is_float8()) {
+    if (dtype.code() == DataType::TypeCode::kE5M2Float) {
+      return FloatImm(dtype, -57344.0, span);
+    } else if (dtype.code() == DataType::TypeCode::kE4M3Float) {
+      return FloatImm(dtype, -448.0, span);
+    }
   }
   LOG(FATAL) << "Cannot decide min_value for type" << dtype;
 }

Original file line number	Diff line number	Diff line change
`@@ -262,6 +262,12 @@ PrimExpr max_value(const DataType& dtype, Span span) {`
`262`	`262`	`}`
`263`	`263`	`} else if (dtype.is_bfloat16()) {`
`264`	`264`	`return FloatImm(dtype, std::numeric_limits<float>::max(), span);`
	`265`	`+ } else if (dtype.is_float8()) {`
	`266`	`+ if (dtype.code() == DataType::TypeCode::kE5M2Float) {`
	`267`	`+ return FloatImm(dtype, 57344.0, span);`
	`268`	`+ } else if (dtype.code() == DataType::TypeCode::kE4M3Float) {`
	`269`	`+ return FloatImm(dtype, 448.0, span);`
	`270`	`+ }`
`265`	`271`	`}`
`266`	`272`	`LOG(FATAL) << "Cannot decide max_value for type" << dtype;`
`267`	`273`	`}`
`@@ -296,6 +302,12 @@ PrimExpr min_value(const DataType& dtype, Span span) {`
`296`	`302`	`}`
`297`	`303`	`} else if (dtype.is_bfloat16()) {`
`298`	`304`	`return FloatImm(dtype, std::numeric_limits<float>::lowest(), span);`
	`305`	`+ } else if (dtype.is_float8()) {`
	`306`	`+ if (dtype.code() == DataType::TypeCode::kE5M2Float) {`
	`307`	`+ return FloatImm(dtype, -57344.0, span);`
	`308`	`+ } else if (dtype.code() == DataType::TypeCode::kE4M3Float) {`
	`309`	`+ return FloatImm(dtype, -448.0, span);`
	`310`	`+ }`
`299`	`311`	`}`
`300`	`312`	`LOG(FATAL) << "Cannot decide min_value for type" << dtype;`
`301`	`313`	`}`