Add context-free Call reinference and binding type propagation (#20480)

- Separate native context-free Call type inference from
builder-dependent Relax and distributed inference, with strict read-only
reinference from current Call inputs.
- Infer `tirx.buffer_data` from its BufferVar argument and reinfer after
pointer storage-scope rewrites.
- Migrate TVM structural visit and mutation hooks to native tvm-ffi
callbacks and remove the structural compatibility header.
- Propagate Bind and Let value type changes through their binder
definitions and later uses while preserving unchanged node identity.
- Reuse existing Call result Type objects when inference establishes an
equivalent result cheaply.

Author: tqchen

3rdparty/tvm-ffi

@@ -26,10 +26,10 @@
 
 #include <tvm/ffi/cast.h>
 #include <tvm/ffi/dtype.h>
+#include <tvm/ffi/extra/structural_mutate.h>
+#include <tvm/ffi/extra/structural_visit.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ffi/string.h>
-// Keep raw and typed structural hook return paths available to all TVM IR nodes.
-#include <tvm/ir/ffi_structural_compat.h>
 #include <tvm/ir/source_map.h>
 
 #include <cstddef>

include/tvm/ir/base_expr.h

@@ -385,6 +385,9 @@ class Var : public Expr {
   /*! \brief Return a fresh ordinary Var with a new primitive type. */
   TVM_DLL Var CopyWithDType(PrimType dtype) const;
 
+  /*! \brief Return a fresh ordinary Var with a new type, retaining its metadata. */
+  TVM_DLL Var CopyWithType(Type type) const;
+
   TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Var, Expr, VarNode);
 };
 
@@ -513,6 +516,13 @@ class Call : public Expr {
   /*! \brief Construct a provisional Call without invoking its Op validator. */
   TVM_DLL static Call Unchecked(Type ret_ty, Expr op, ffi::Array<Expr> args, Attrs attrs = Attrs(),
                                 ffi::Array<Type> ty_args = ffi::Array<Type>(), Span span = Span());
+  /*! \brief Recompute a result type from the Call's current explicit inputs.
+   *
+   * This ignores the Call's stored result type and does not mutate the Call.
+   * A context-free inference hook must be registered for the operator. Missing
+   * or invalid inputs are reported by the hook rather than screened here.
+   */
+  TVM_DLL static Type ReinferType(const CallNode* call);
 
   TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Call, Expr, CallNode);
   TVM_DEFINE_OBJECT_REF_COW_METHOD(CallNode);

include/tvm/ir/expr.h

@@ -510,6 +510,19 @@ class TVM_DLL ExprMutator : public ObjectMutator {
   virtual UnchangedOr<PrimExpr> Mutate_(const prim::ShuffleNode* node, InplaceMode inplace_mode);
 
  protected:
+  /*!
+   * \brief Reinfer the result type of a Call after ordinary mutation.
+   * \param mutated The ordinary mutation result, or Unchanged for the original Call.
+   * \param original The borrowed Call before ordinary mutation.
+   * \param inplace_mode Permission to modify the original Call in place.
+   * \return The combined mutation result, preserving Unchanged when possible.
+   * \note A unique replacement Call may be updated in place even when the original
+   *       could not be. Inference errors propagate to the caller.
+   */
+  static UnchangedOr<Expr> ReinferMutatedCallType(UnchangedOr<Expr> mutated,
+                                                  const CallNode* original,
+                                                  InplaceMode inplace_mode);
+
   /*!
    * \brief Construct a mutator with an extended native dispatch table.
    * \param vtable The finalized table, which must outlive this mutator.

include/tvm/ir/expr_functor.h

@@ -43,6 +43,9 @@ namespace tvm {
 template <typename>
 class OpAttrMap;
 
+/*! \brief Infer a Call's result type from its explicit inputs without builder state. */
+using FInferType = ffi::reflection::NativeFunctionView<ffi::Expected<Type>(const CallNode* call)>;
+
 /*! \brief An operator argument's name and documentation. */
 class ArgumentInfoNode : public ffi::Object {
  public:

include/tvm/ir/ffi_structural_compat.h

@@ -66,7 +66,7 @@ using FCallPacked = ffi::String;
  * \param call The call expression to be derived.
  * \param ctx The builder context.
  */
-using FInferType = ffi::TypedFunction<Type(const Call& call, const BlockBuilder& ctx)>;
+using FInferTypeWithBuilder = ffi::TypedFunction<Type(const Call& call, const BlockBuilder& ctx)>;
 
 /*!
  * \brief The function type of a normalization function.

include/tvm/ir/op.h

@@ -53,6 +53,7 @@
     Var,
     is_prim_expr,
     is_prim_var,
+    reinfer_type,
 )
 from . import prim
 from .function import BaseFunc, CallingConv

include/tvm/relax/op_attr_types.h

@@ -593,6 +593,15 @@ def unchecked(
         )
 
 
+def reinfer_type(call: Call) -> "tvm.ir.Type":
+    """Derive a Call's result type from its current inputs without changing the Call.
+
+    The operator must have a context-free inference rule. Invalid inputs and
+    missing rules raise errors.
+    """
+    return _ffi_api.reinfer_type(call)
+
+
 @tvm_ffi.register_object("ir.TensorRegion")
 class TensorRegion(Expr, Scriptable):
     """A region of an arbitrary tensor expression.

python/tvm/ir/__init__.py

@@ -52,7 +52,7 @@ def te_matmul(a: te.Tensor, b: te.Tensor) -> te.Tensor:
 
         a_relax = relax.Var("a", relax.TensorType(a.shape))
         b_relax = relax.Var("b", relax.TensorType(b.shape))
-        f_infer_ty = call.op.get_attr("FInferType")
+        f_infer_ty = call.op.get_attr("relax.FInferTypeWithBuilder")
         output_shape = f_infer_ty(relax.op.matmul(a_relax, b_relax), bb).shape
         if isinstance(a_shape[-1], tirx.IntImm) and a_shape[-1] == 0:
             return te.compute(