Address comment and update doc

include/tvm/relax/transform.h

@@ -102,9 +102,11 @@ TVM_DLL Pass ToNonDataflow();
 TVM_DLL Pass CallTIRRewrite();
 
 /*!
- * \brief Convert all reshape-like call_tir to VM reshape operator call.
- * The VM reshape operator calls will be further lowered to a CreateView
- * operation at runtime, instead of doing real data copy.
+ * \brief Convert all reshape-like call_tir whose corresponding binding
+ * vars are DataflowVars to relax.reshape operator calls. The relax.reshape
+ * calls will be lowered an external builtin function call in a subsequent
+ * pass, where the external builtin function does a CreateView operation
+ * at runtime, instead of doing real data copy.
  * Here "reshape-like" includes reshape, expand_dims, flatten, etc.
  *
  * \return The Pass.

python/tvm/relax/transform/transform.py

@@ -102,9 +102,12 @@ def CallTIRRewrite() -> tvm.ir.transform.Pass:
 
 
 def RewriteDataflowReshape() -> tvm.ir.transform.Pass:
-    """Convert all reshape-like call_tir to VM reshape operator call.
-    The VM reshape operator calls will be further lowered to a CreateView
-    operation at runtime, instead of doing real data copy.
+    """Convert all reshape-like call_tir whose corresponding binding
+    vars are DataflowVars to relax.reshape operator calls. The relax.reshape
+    calls will be lowered an external builtin function call in a subsequent
+    pass, where the external builtin function does a CreateView operation
+    at runtime, instead of doing real data copy.
+
     Here "reshape-like" includes reshape, expand_dims, flatten, etc.
 
     Returns

src/relax/transform/rewrite_dataflow_reshape.cc

@@ -18,7 +18,7 @@
  */
 /*!
  * \file src/relax/transform/rewrite_dataflow_reshape.cc
- * \brief Transform all reshape within dataflow block to a specialized reshape operator
+ * \brief Transform all reshape within dataflow block to a relax.reshape operator
  */
 #include <tvm/relax/analysis.h>
 #include <tvm/relax/expr_functor.h>
@@ -75,8 +75,11 @@ class DataflowReshapeRewriter : public ExprMutator {
     if (call->op != call_tir_op) {
       return false;
     }
-    GlobalVar gv = Downcast<GlobalVar>(call->args[0]);
-    const auto* func = mod_->functions.Get(gv).as<tir::PrimFuncNode>();
+    const auto* gv = call->args[0].as<GlobalVarNode>();
+    if (gv == nullptr) {
+      return false;
+    }
+    const auto* func = mod_->functions.Get(GetRef<GlobalVar>(gv)).as<tir::PrimFuncNode>();
     ICHECK_NOTNULL(func);
     return HasReshapePattern(GetRef<tir::PrimFunc>(func));
   }