track

python/tvm/relay/transform/memory_alloc.py

@@ -40,18 +40,25 @@ def is_primitive(call):
 
 
 class CheckReshapeOnly(ExprVisitor):
-    """A pass to check if the fused op contains only reshape ops."""
+    """
+    A pass to check if the fused op contains only reshape ops.
+    TODO(@Jared) this is capturing the case where there is no any ops at all.
+    I had put a quick hack and require that it must have >= 1 reshape, but this must be masking some bigger problem.
+    Please fix - if you want to collapse reshape on reshape, perhaps you should do fusion as such.
+    """
     def __init__(self):
         super().__init__()
         self._reshape_ops = [op.get("reshape"), op.get("contrib_reverse_reshape"),
                              op.get("dyn.reshape")]
         self.reshape_only = True
+        self.has_reshape = False
 
     def visit_call(self, call):
         if not self.reshape_only:
             return
         if call.op not in self._reshape_ops:
             self.reshape_only = False
+        self.has_reshape = True
         for arg in call.args:
             self.visit(arg)
 
@@ -60,7 +67,7 @@ def is_reshape_only(func):
     """Check if the primitive function contains only reshape ops."""
     check = CheckReshapeOnly()
     check.visit(func)
-    return check.reshape_only
+    return check.reshape_only and check.has_reshape
 
 
 class ManifestAllocPass(ExprMutator):

src/relay/backend/vm/compiler.cc

@@ -923,6 +923,31 @@ transform::Sequential MemoryOpt(tvm::Target host_target) {
   return transform::Sequential(pass_seqs);
 }
 
+Pass CheckPrimeFunc() {
+  runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
+    [=](Function f, IRModule m, PassContext pc) {
+      struct CheckPrimeFuncVisitor : ExprVisitor {
+        bool inside_primitive = false;
+        void VisitExpr_(const ConstantNode* op) override {
+          CHECK_EQ(inside_primitive, false);
+        }
+        void VisitExpr_(const FunctionNode* op) override {
+          if (op->HasNonzeroAttr(attr::kPrimitive)) {
+            CHECK_EQ(inside_primitive, false);
+            inside_primitive = true;
+            VisitExpr(op->body);
+            inside_primitive = false;
+          } else {
+            VisitExpr(op->body);
+          }
+        }
+      } vis;
+      vis(f);
+      return f;
+    };
+  return CreateFunctionPass(pass_func, 1, "CheckPrimeFunc", {});
+}
+
 IRModule VMCompiler::OptimizeModule(const IRModule& mod, const TargetsMap& targets,
                                     const Target& target_host) {
   Array<Pass> pass_seqs;
@@ -956,14 +981,19 @@ IRModule VMCompiler::OptimizeModule(const IRModule& mod, const TargetsMap& targe
     }
     *rv = false;
   });
+
   pass_seqs.push_back(transform::EliminateCommonSubexpr(fskip));
   pass_seqs.push_back(transform::SimplifyExpr());
   pass_seqs.push_back(transform::InlinePrimitives());
 
+
   pass_seqs.push_back(transform::CombineParallelConv2D(3));
   pass_seqs.push_back(transform::CombineParallelDense(3));
   pass_seqs.push_back(transform::CombineParallelBatchMatmul(3));
+
   pass_seqs.push_back(transform::FoldConstant());
+  //pass_seqs.push_back(tvm::transform::PrintIR());
+  //pass_seqs.push_back(CheckPrimeFunc());
   pass_seqs.push_back(transform::FoldScaleAxis());
   pass_seqs.push_back(transform::CanonicalizeCast());
   pass_seqs.push_back(transform::CanonicalizeOps());
@@ -1004,7 +1034,6 @@ IRModule VMCompiler::OptimizeModule(const IRModule& mod, const TargetsMap& targe
 }
 
 void VMCompiler::PopulateGlobalMap() {
-  // First we populate global map.
   size_t global_index = 0;
   for (auto named_func : context_.module->functions) {
     auto gvar = named_func.first;

src/relay/transforms/fold_constant.cc

@@ -104,7 +104,23 @@ class ConstantFolder : public ExprMutator {
     }
   }
 
+  bool inside_primitive = false;
+  Expr VisitExpr_(const FunctionNode* op) final {
+    if (op->HasNonzeroAttr(attr::kPrimitive)) {
+      CHECK_EQ(inside_primitive, false);
+      inside_primitive = true;
+      auto ret = ExprMutator::VisitExpr_(op);
+      inside_primitive = false;
+      return ret;
+    } else {
+      return ExprMutator::VisitExpr_(op);
+    }
+  }
+
   Expr VisitExpr_(const CallNode* call) final {
+    if (inside_primitive) {
+      return GetRef<Expr>(call);
+    }
     static auto op_stateful = Op::GetAttrMap<TOpIsStateful>("TOpIsStateful");
 
     std::unordered_set<std::string> skip_list{"zeros_like", "ones_like", "full_like", "full"};

src/relay/transforms/partial_eval.cc

@@ -91,6 +91,7 @@
  */
 #include <tvm/ir/type_functor.h>
 #include <tvm/relay/analysis.h>
+#include <tvm/relay/feature.h>
 #include <tvm/relay/expr_functor.h>
 #include <tvm/relay/feature.h>
 #include <tvm/relay/interpreter.h>
@@ -776,8 +777,9 @@ class PartialEvaluator : public ExprFunctor<PStatic(const Expr& e, LetList* ll)>
 
   Func VisitFuncStatic(const Function& func, const Expr& var) {
     CHECK(IsAtomic(var));
+    // todo: figure out primitive semantic
     if (func->HasNonzeroAttr(attr::kPrimitive)) {
-      return ConstEvaluateFunc(func);
+      // return ConstEvaluateFunc(func);
     }
     std::vector<std::pair<Var, PStatic> > free_vars;
     for (const auto& v : FreeVars(func)) {
@@ -1200,7 +1202,7 @@ namespace transform {
 Pass PartialEval() {
   runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> pass_func =
       [=](IRModule m, PassContext pc) { return relay::PartialEval(m); };
-  return CreateModulePass(pass_func, 1, "PartialEval", {});
+  return CreateModulePass(pass_func, 1, "PartialEvaluate", {});
 }
 
 TVM_REGISTER_GLOBAL("relay._transform.PartialEvaluate").set_body_typed(PartialEval);