(tvm) branch main updated: [FIX][IR] Share typed expression mismatch diagnostics (#20523)

[tqchen]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new f508277458 [FIX][IR] Share typed expression mismatch diagnostics 
(#20523)
f508277458 is described below

commit f50827745838df7b2eaf8d36cf10fef657342027
Author: Tianqi Chen <tqc...@users.noreply.github.com>
AuthorDate: Fri Oct 2 20:11:00 2026 -0400

    [FIX][IR] Share typed expression mismatch diagnostics (#20523)
    
    Use one TypedExpr mismatch formatter for PrimExpr and IntExpr so
    rejected expressions consistently report their type with ty=, including
    primitive dtype details. Retain the existing strict scalar integer
    checks and conversions.
---
 include/tvm/ir/base_expr.h                  | 38 +++++++++++++----------------
 tests/python/tirx-base/test_tir_op_types.py |  2 +-
 2 files changed, 18 insertions(+), 22 deletions(-)

diff --git a/include/tvm/ir/base_expr.h b/include/tvm/ir/base_expr.h
index 4302a570b2..c61a09c4e5 100644
--- a/include/tvm/ir/base_expr.h
+++ b/include/tvm/ir/base_expr.h
@@ -489,6 +489,20 @@ struct TypeTraits<TypedExpr<ExpectedType>>
     const auto* expr = 
details::ObjectUnsafe::RawObjectPtrFromUnowned<ExprNode>(src->v_obj);
     return details::AnyUnsafe::CheckAnyViewStrict<ExpectedType>(expr->ty);
   }
+
+  TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* src) {
+    if (src->type_index >= TypeIndex::kTVMFFIStaticObjectBegin && src->v_obj 
!= nullptr &&
+        details::IsObjectInstance<ExprNode>(src->type_index)) {
+      const auto* expr = 
details::ObjectUnsafe::RawObjectPtrFromUnowned<ExprNode>(src->v_obj);
+      if (expr->ty.defined()) {
+        const auto* prim_ty = expr->ty.as<PrimTypeNode>();
+        std::string ty =
+            prim_ty ? std::string(ffi::DLDataTypeToString(prim_ty->dtype)) : 
expr->ty.GetTypeKey();
+        return TypeIndexToTypeKey(src->type_index) + "[ty=" + ty + "]";
+      }
+    }
+    return TypeTraitsBase::GetMismatchTypeInfo(src);
+  }
 };
 
 template <>
@@ -499,9 +513,6 @@ inline constexpr bool use_default_type_traits_v<PrimExpr> = 
false;
 template <>
 struct TypeTraits<PrimExpr> : public ObjectRefWithFallbackTraitsBase<PrimExpr, 
StrictBool, int64_t,
                                                                      double, 
PrimExprConvertible> {
-  using Base =
-      ObjectRefWithFallbackTraitsBase<PrimExpr, StrictBool, int64_t, double, 
PrimExprConvertible>;
-
   TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) {
     if (src->type_index == TypeIndex::kTVMFFINone) return 
PrimExpr::_type_is_nullable;
     return TypeTraits<TypedExpr<PrimType>>::CheckAnyStrict(src);
@@ -509,15 +520,8 @@ struct TypeTraits<PrimExpr> : public 
ObjectRefWithFallbackTraitsBase<PrimExpr, S
 
   TVM_FFI_INLINE static std::string TypeStr() { return "ir.PrimExpr"; }
 
-  TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* 
source) {
-    if (source->type_index >= TypeIndex::kTVMFFIStaticObjectBegin && 
source->v_obj != nullptr &&
-        details::IsObjectInstance<ExprNode>(source->type_index)) {
-      const auto* expr = 
details::ObjectUnsafe::RawObjectPtrFromUnowned<ExprNode>(source->v_obj);
-      if (expr->ty.defined() && 
!details::AnyUnsafe::CheckAnyViewStrict<PrimType>(expr->ty)) {
-        return TypeIndexToTypeKey(source->type_index) + "[ty=" + 
expr->ty.GetTypeKey() + "]";
-      }
-    }
-    return Base::GetMismatchTypeInfo(source);
+  TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* src) {
+    return TypeTraits<TypedExpr<PrimType>>::GetMismatchTypeInfo(src);
   }
 
   TVM_DLL static PrimExpr ConvertFallbackValue(StrictBool value);
@@ -551,15 +555,7 @@ struct TypeTraits<IntExpr> : public 
ObjectRefWithFallbackTraitsBase<IntExpr, int
   TVM_FFI_INLINE static std::string TypeStr() { return "ir.IntExpr"; }
 
   TVM_FFI_INLINE static std::string GetMismatchTypeInfo(const TVMFFIAny* src) {
-    if (src->type_index >= TypeIndex::kTVMFFIStaticObjectBegin &&
-        details::IsObjectInstance<ExprNode>(src->type_index)) {
-      const auto* expr = 
details::ObjectUnsafe::RawObjectPtrFromUnowned<ExprNode>(src->v_obj);
-      if (const auto* ty = expr->ty.as<PrimTypeNode>()) {
-        return TypeIndexToTypeKey(src->type_index) +
-               "[dtype=" + ffi::DLDataTypeToString(ty->dtype) + "]";
-      }
-    }
-    return TypeTraits<PrimExpr>::GetMismatchTypeInfo(src);
+    return TypeTraits<TypedExpr<PrimType>>::GetMismatchTypeInfo(src);
   }
 };
 
diff --git a/tests/python/tirx-base/test_tir_op_types.py 
b/tests/python/tirx-base/test_tir_op_types.py
index 38c5d83e36..fad780f8cc 100644
--- a/tests/python/tirx-base/test_tir_op_types.py
+++ b/tests/python/tirx-base/test_tir_op_types.py
@@ -36,7 +36,7 @@ def test_scalar_integer_signature():
         assert tirx.tvm_struct_get(ptr, index, 2, 
dtype="int32").args[1].same_as(index)
     assert tirx.tvm_struct_get(ptr, 1 << 40, 2, 
dtype="int32").args[1].ty.dtype == "int64"
     for dtype in ("float32", "bool", "int32x4", "int32xvscalex4"):
-        with pytest.raises(TypeError, match=r"index.*expected 
`ir.IntExpr`.*dtype="):
+        with pytest.raises(TypeError, match=rf"index.*expected 
`ir.IntExpr`.*\[ty={dtype}\]"):
             tirx.tvm_struct_get(ptr, tirx.Var("index", dtype), 2, 
dtype="int32")