junrushao commented on code in PR #668:
URL: https://github.com/apache/tvm-ffi/pull/668#discussion_r3590433213
##########
python/tvm_ffi/_dtype.py:
##########
@@ -288,6 +289,72 @@ def lanes(self) -> int:
"""
return self._tvm_ffi_dtype.lanes
+ @property
+ def is_bool(self) -> builtins.bool:
+ """Whether this dtype stores boolean values."""
+ return self.type_code == DataTypeCode.BOOL
+
+ @property
+ def is_integer(self) -> builtins.bool:
+ """Whether this dtype stores signed or unsigned integer values."""
+ return self.type_code in {DataTypeCode.INT, DataTypeCode.UINT}
Review Comment:
Addressed in 8cef0767. Ruff flags the suggested direct or expression with
PLR1714, so I hoisted the integer codes into the class-level _INTEGER_CODES set
and use membership against that shared set. This keeps lint clean while
eliminating the per-access allocation.
##########
include/tvm/ffi/dtype.h:
##########
@@ -138,6 +138,160 @@ inline String DLDataTypeToString(DLDataType dtype) {
return TypeTraits<String>::MoveFromAnyAfterCheck(&out);
}
+/*!
+ * \brief Check whether a DLDataType is a floating-point type.
+ * \param dtype The DLDataType to check.
+ * \return True if dtype is a floating-point type, false otherwise.
+ */
+inline bool DTypeIsFloat(DLDataType dtype) {
+ switch (static_cast<int>(dtype.code)) {
+ case kDLFloat:
+ case kDLBfloat:
+ case kDLFloat8_e3m4:
+ case kDLFloat8_e4m3:
+ case kDLFloat8_e4m3b11fnuz:
+ case kDLFloat8_e4m3fn:
+ case kDLFloat8_e4m3fnuz:
+ case kDLFloat8_e5m2:
+ case kDLFloat8_e5m2fnuz:
+ case kDLFloat8_e8m0fnu:
+ case kDLFloat6_e2m3fn:
+ case kDLFloat6_e3m2fn:
+ case kDLFloat4_e2m1fn:
+ return true;
+ default:
+ return false;
+ }
+ return false;
+}
Review Comment:
Addressed in 8cef0767. I removed the unreachable return after the
DTypeIsFloat switch. The C++ test target was rebuilt and the 9 focused
DType/DataType tests pass.
##########
python/tvm_ffi/_dtype.py:
##########
@@ -288,6 +289,72 @@ def lanes(self) -> int:
"""
return self._tvm_ffi_dtype.lanes
+ @property
+ def is_bool(self) -> builtins.bool:
+ """Whether this dtype stores boolean values."""
+ return self.type_code == DataTypeCode.BOOL
+
+ @property
+ def is_integer(self) -> builtins.bool:
+ """Whether this dtype stores signed or unsigned integer values."""
+ return self.type_code in {DataTypeCode.INT, DataTypeCode.UINT}
+
+ @property
+ def is_float(self) -> builtins.bool:
+ """Whether this dtype stores floating-point values."""
+ return self.type_code in {
+ DataTypeCode.FLOAT,
+ DataTypeCode.BFLOAT,
+ DataTypeCode.Float8E3M4,
+ DataTypeCode.Float8E4M3,
+ DataTypeCode.Float8E4M3B11FNUZ,
+ DataTypeCode.Float8E4M3FN,
+ DataTypeCode.Float8E4M3FNUZ,
+ DataTypeCode.Float8E5M2,
+ DataTypeCode.Float8E5M2FNUZ,
+ DataTypeCode.Float8E8M0FNU,
+ DataTypeCode.Float6E2M3FN,
+ DataTypeCode.Float6E3M2FN,
+ DataTypeCode.Float4E2M1FN,
+ }
Review Comment:
Addressed in 8cef0767. The floating-point codes now live in the class-level
_FLOAT_CODES set, and is_float reuses it instead of allocating a new set on
every access.
##########
include/tvm/ffi/dtype.h:
##########
@@ -138,6 +138,160 @@ inline String DLDataTypeToString(DLDataType dtype) {
return TypeTraits<String>::MoveFromAnyAfterCheck(&out);
}
+/*!
+ * \brief Check whether a DLDataType is a floating-point type.
+ * \param dtype The DLDataType to check.
+ * \return True if dtype is a floating-point type, false otherwise.
+ */
+inline bool DTypeIsFloat(DLDataType dtype) {
+ switch (static_cast<int>(dtype.code)) {
+ case kDLFloat:
+ case kDLBfloat:
+ case kDLFloat8_e3m4:
+ case kDLFloat8_e4m3:
+ case kDLFloat8_e4m3b11fnuz:
+ case kDLFloat8_e4m3fn:
+ case kDLFloat8_e4m3fnuz:
+ case kDLFloat8_e5m2:
+ case kDLFloat8_e5m2fnuz:
+ case kDLFloat8_e8m0fnu:
+ case kDLFloat6_e2m3fn:
+ case kDLFloat6_e3m2fn:
+ case kDLFloat4_e2m1fn:
+ return true;
+ default:
+ return false;
+ }
+ return false;
+}
+
+/*!
+ * \brief Check whether a DLDataType is an integer type.
+ * \param dtype The DLDataType to check.
+ * \return True if dtype is a signed or unsigned integer type, false otherwise.
+ */
+inline bool DTypeIsInt(DLDataType dtype) {
+ switch (static_cast<int>(dtype.code)) {
+ case kDLInt:
+ case kDLUInt:
+ return true;
+ default:
+ return false;
+ }
+ return false;
+}
Review Comment:
Addressed in 8cef0767. I removed the unreachable return after the DTypeIsInt
switch. The C++ test target was rebuilt and the 9 focused DType/DataType tests
pass.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
