[Relay] Ignore Primitive functions in Visitors

include/tvm/relay/expr_functor.h

@@ -181,6 +181,7 @@ class ExprVisitor
  protected:
   // Internal visiting counter
   std::unordered_map<const Object*, size_t> visit_counter_;
+  bool enter_primitives_ = false;
 };
 
 /*!
@@ -231,6 +232,7 @@ class ExprMutator
  protected:
   /*! \brief Internal map used for memoization. */
   std::unordered_map<Expr, Expr, ObjectHash, ObjectEqual> memo_;
+  bool enter_primitives_ = false;
 };
 
 /*!

src/relay/ir/expr_functor.cc

@@ -85,6 +85,13 @@ Expr ExprMutator::VisitExpr_(const TupleNode* op) {
 Expr ExprMutator::VisitExpr_(const FunctionNode* op) {
   tvm::Array<TypeVar> ty_params;
   bool all_ty_params_unchanged = true;
+  // don't mutate primitive functions
+  const auto primitive = FunctionGetAttr(GetRef<Function>(op), attr::kPrimitive);
+  if (primitive->IsInstance<tir::IntImmNode>()) {
+    if (primitive.as<IntImmNode>()->value && !enter_primitives_) {
+        return GetRef<Expr>(op);
+    }
+  }
 
   for (auto ty_param : op->type_params) {
     TypeVar new_ty_param = Downcast<TypeVar>(VisitType(ty_param));
@@ -256,7 +263,12 @@ void ExprVisitor::ExprVisitor::VisitExpr_(const FunctionNode* op) {
   for (auto param : op->params) {
     this->VisitExpr(param);
   }
-
+  const auto primitive = FunctionGetAttr(GetRef<Function>(op), attr::kPrimitive);
+  if (primitive->IsInstance<tir::IntImmNode>()) {
+      if (primitive.as<IntImmNode>()->value && !enter_primitives_) {
+          return;
+      }
+  }
   this->VisitExpr(op->body);
 }
 

src/relay/pass/type_infer.cc

@@ -604,6 +604,7 @@ class TypeInferencer::Resolver : public ExprMutator, PatternMutator {
   Resolver(const std::unordered_map<Expr, ResolvedTypeInfo, ObjectHash, ObjectEqual>& tmap,
            TypeSolver* solver)
     : tmap_(tmap), solver_(solver) {
+    enter_primitives_ = true;
   }
 
   Expr VisitExpr_(const VarNode* op) final {