tqchen commented on code in PR #19910:
URL: https://github.com/apache/tvm/pull/19910#discussion_r3499060489


##########
src/relax/op/op.cc:
##########
@@ -420,9 +420,24 @@ static ffi::Optional<Type> 
InferCallTIROutputTypeFromArguments(
     return dummy_args;
   }();
 
-  auto derived_ret_ty =
-      DeriveCallRetType(dummy_callee_ty, Call(Var("dummy_callee", 
dummy_callee_ty), dummy_args),
-                        BlockBuilder::Create(std::nullopt));
+  Type derived_ret_ty = DeriveCallRetType(
+      dummy_callee_ty, Call(Type::Missing(), Var("dummy_callee", 
dummy_callee_ty), dummy_args),
+      BlockBuilder::Create(std::nullopt));

Review Comment:
   Implemented. The DeriveCallRetType result is now checked with IsMissing(), 
and output inference returns std::nullopt before TIRVarsInType can visit 
Type::Missing().



##########
include/tvm/ir/base_expr.h:
##########
@@ -314,6 +322,87 @@ class Expr : public ffi::ObjectRef {
   TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Expr, ffi::ObjectRef, ExprNode);
 };
 
+/*!
+ * \brief Typed reference/view over an expression whose result type is a
+ * specific Type subtype.
+ * \tparam ExpectedType The expected expression result type.
+ */
+template <typename ExpectedType>
+class TypedExpr : public Expr {
+ public:
+  /*! \return the typed result of this expression. */
+  ExpectedType ty() const {
+    const auto* node = get();
+    TVM_FFI_DCHECK(node != nullptr);
+    const Type& ty = node->ExprNode::ty;
+    const auto* ty_node = ty.get();
+    TVM_FFI_DCHECK(ty_node != nullptr);
+    TVM_FFI_DCHECK(ty_node->template IsInstance<typename 
ExpectedType::ContainerType>());
+    return ffi::details::ObjectUnsafe::ObjectRefFromObjectPtr<ExpectedType>(
+        ffi::details::ObjectUnsafe::ObjectPtrFromUnowned<ffi::Object>(
+            const_cast<ffi::Object*>(static_cast<const 
ffi::Object*>(ty_node))));
+  }

Review Comment:
   Implemented with the checked public API available in this tvm-ffi revision. 
TypedExpr::ty() uses Type::as<ExpectedType::ContainerType>() followed by 
ffi::GetRef<ExpectedType>(); there is no free ffi::downcast helper in this 
checkout, and the old const_cast/ObjectUnsafe path is gone.



##########
python/tvm/ir/expr.py:
##########
@@ -77,7 +73,10 @@ def __call__(self, *args: Expr) -> Expr:
         """
         # pylint: disable=import-outside-toplevel
 
-        if args and all(isinstance(x, Number | PrimExpr) for x in args):
+        if args and all(
+            isinstance(x, Number) or (isinstance(x, Expr) and isinstance(x.ty, 
tvm.ir.PrimType))
+            for x in args
+        ):

Review Comment:
   Applied in incremental commit bb25d071fa3e397b68b8867b6d1e723349a3595b. 
GlobalVar.__call__ now uses is_prim_expr(x); the same canonical predicate was 
generalized across every equivalent changed Python boundary found by the 
whole-diff audit.



##########
python/tvm/relax/expr.py:
##########
@@ -45,24 +44,24 @@
 GlobalVar = tvm.ir.GlobalVar
 
 
-def prim_value(value: PrimExpr | int | float, dtype: str | None = None) -> 
PrimExpr:
-    """Convert a Python scalar or primitive expression to ``PrimExpr``.
+def prim_value(value: Expr | int | float, dtype: str | None = None) -> Expr:
+    """Convert a Python scalar or primitive expression to ``Expr``.
 
     Parameters
     ----------
-    value : PrimExpr | int | float
+    value : Expr | int | float
         The value to convert.
 
     dtype : Optional[str]
         The dtype to use when converting Python numeric values.
 
     Returns
     -------
-    result : PrimExpr
-        The converted primitive expression.  Existing ``PrimExpr`` inputs are
+    result : Expr
+        The converted primitive expression.  Existing ``Expr`` inputs are
         returned unchanged.
     """
-    if isinstance(value, PrimExpr):
+    if isinstance(value, Expr) and isinstance(value.ty, tvm.ir.PrimType):
         return value

Review Comment:
   Applied in incremental commit bb25d071fa3e397b68b8867b6d1e723349a3595b. 
prim_value now uses tvm.ir.is_prim_expr(value), matching the shared Expr plus 
PrimType category check without repeating its implementation.



##########
python/tvm/relax/expr.py:
##########
@@ -71,9 +70,9 @@ def prim_value(value: PrimExpr | int | float, dtype: str | 
None = None) -> PrimE
     if isinstance(value, Real):
         return tvm.tirx.FloatImm(dtype or "float64", float(value))
     tvm_value = tvm_ffi.convert(value)
-    if isinstance(tvm_value, PrimExpr):
+    if isinstance(tvm_value, Expr) and isinstance(tvm_value.ty, 
tvm.ir.PrimType):
         return tvm_value

Review Comment:
   Applied in incremental commit bb25d071fa3e397b68b8867b6d1e723349a3595b. The 
converted-value path now uses tvm.ir.is_prim_expr(tvm_value), consistent with 
the direct-input path and the other audited Python boundaries.



-- 
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]

Reply via email to