expr_functor.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
 
#ifndef TVM_RELAX_EXPR_FUNCTOR_H_
#define TVM_RELAX_EXPR_FUNCTOR_H_
 
#include <tvm/ir/object_functor.h>
#include <tvm/relax/block_builder.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/type.h>
#include <tvm/relax/type_functor.h>
#include <tvm/tirx/function.h>
 
#include <unordered_map>
#include <utility>
namespace tvm {
namespace relax {
 
template <typename FType>
class ExprFunctor;
 
// functions to be overriden.
#define EXPR_FUNCTOR_DEFAULT                                   \
  {                                                            \
    return VisitExprDefault_(op, std::forward<Args>(args)...); \
  }
 
#define EXPR_FUNCTOR_DISABLED                                                                    \
  final {                                                                                        \
    TVM_FFI_THROW(TypeError) << "Relax does not support " << op->GetTypeKey() << " expressions"; \
    throw;                                                                                       \
  }
 
#define RELAX_EXPR_FUNCTOR_DISPATCH(OP)                                                    \
  vtable.template SetDispatch<OP>([](const ffi::ObjectRef& n, TSelf* self, Args... args) { \
    return self->VisitExpr_(static_cast<const OP*>(n.get()), std::forward<Args>(args)...); \
  });
 
#define PY_EXPR_VISITOR_DEFAULT(N, PY_FUNC, DEFAULT_FUNC) \
  {                                                       \
    if (PY_FUNC != nullptr)                               \
      PY_FUNC(N);                                         \
    else                                                  \
      DEFAULT_FUNC;                                       \
  }
 
#define PY_EXPR_MUTATOR_DEFAULT(N, PY_FUNC, DEFAULT_FUNC, RET_TYPE) \
  {                                                                 \
    if (PY_FUNC != nullptr) {                                       \
      RET_TYPE ret = PY_FUNC(N).cast<RET_TYPE>();                   \
      return ret;                                                   \
    } else {                                                        \
      return DEFAULT_FUNC;                                          \
    }                                                               \
  }
 
#define PY_EXPR_VISITOR_DISPATCH(OP, PY_FUNC)                                \
  vtable.template SetDispatch<OP>([](const ffi::ObjectRef& n, TSelf* self) { \
    if (self->PY_FUNC != nullptr)                                            \
      self->PY_FUNC(n);                                                      \
    else                                                                     \
      self->VisitExpr_(static_cast<const OP*>(n.get()));                     \
  });
 
#define PY_EXPR_MUTATOR_DISPATCH(OP, PY_FUNC)                                \
  vtable.template SetDispatch<OP>([](const ffi::ObjectRef& n, TSelf* self) { \
    if (self->PY_FUNC != nullptr) {                                          \
      Expr expr = self->PY_FUNC(n).cast<Expr>();                             \
      return expr;                                                           \
    } else {                                                                 \
      return self->VisitExpr_(static_cast<const OP*>(n.get()));              \
    }                                                                        \
  });
 
#define PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(OP)                              \
  post_order_vtable.template SetDispatch<OP>([](const ffi::ObjectRef& n, TSelf* self) { \
    return self->VisitExprPostOrder_(static_cast<const OP*>(n.get()));                  \
  });
 
template <typename R, typename... Args>
class ExprFunctor<R(const Expr& n, Args...)> {
 private:
  using TSelf = ExprFunctor<R(const Expr& n, Args...)>;
  using FType = tvm::ObjectFunctor<R(const ffi::ObjectRef& n, TSelf* self, Args...)>;
 
 public:
  using result_type = R;
  virtual ~ExprFunctor() {}
  R operator()(const Expr& n, Args... args) { return VisitExpr(n, std::forward<Args>(args)...); }
  virtual R VisitExpr(const Expr& n, Args... args) {
    TVM_FFI_ICHECK(n.defined())
        << "Found null pointer node while traversing AST. The previous pass may "
           "have generated invalid data.";
    static FType vtable = InitVTable();
    if (vtable.CanDispatch(n)) {
      return vtable(n, this, std::forward<Args>(args)...);
    }
    return VisitExprFallback_(n.get(), std::forward<Args>(args)...);
  }
  // Functions that can be overriden by subclass
  // NOTE: cross dialect calls are invoked through global var
  // We do not expect inline tirx::Function to appear in relax IR.
  virtual R VisitExpr_(const GenericConstNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const TupleNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const VarNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const DataflowVarNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const ShapeExprNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const ExternFuncNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const GlobalVarNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const FunctionNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const CallNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const TensorLoadNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const tvm::IntImmNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const tvm::FloatImmNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const SeqExprNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const IfExprNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const OpNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const TupleGetItemNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const StringImmNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const DataTypeImmNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::LetNode* op, Args...) EXPR_FUNCTOR_DISABLED;
  virtual R VisitExpr_(const prim::AddNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::LShiftNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::RShiftNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::BitwiseAndNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::BitwiseOrNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::BitwiseXorNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::BitwiseNotNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::SubNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::MulNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::DivNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::ModNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::FloorDivNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::FloorModNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::MinNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::MaxNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::EQNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::NENode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::LTNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::LENode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::GTNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::GENode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::AndNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::OrNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::CastNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::NotNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::SelectNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::RampNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::BroadcastNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const prim::ShuffleNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExprFallback_(const ExprNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExprDefault_(const ffi::Object* op, Args...) {
    TVM_FFI_THROW(InternalError) << "Do not have a default for " << op->GetTypeKey();
    throw;
  }
 
 private:
  // initialize the vtable.
  static FType InitVTable() {
    FType vtable;
    // Set dispatch
    RELAX_EXPR_FUNCTOR_DISPATCH(GenericConstNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(TupleNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(VarNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(DataflowVarNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(ShapeExprNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(ExternFuncNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(GlobalVarNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(FunctionNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(CallNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::LetNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(TensorLoadNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::AddNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::LShiftNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::RShiftNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::BitwiseAndNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::BitwiseOrNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::BitwiseXorNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::BitwiseNotNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::SubNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::MulNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::DivNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::ModNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::FloorDivNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::FloorModNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::MinNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::MaxNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::EQNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::NENode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::LTNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::LENode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::GTNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::GENode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::AndNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::OrNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::CastNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::NotNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::SelectNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::RampNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::BroadcastNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(prim::ShuffleNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(tvm::IntImmNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(tvm::FloatImmNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(SeqExprNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(IfExprNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(OpNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(TupleGetItemNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(StringImmNode);
    RELAX_EXPR_FUNCTOR_DISPATCH(DataTypeImmNode);
    vtable.Finalize();
    return vtable;
  }
};
 
class ExprVisitor : public ExprFunctor<void(const Expr&)> {
 public:
  void VisitExpr(const Expr& expr) override;
  // specific leaf level visitor functions
  void VisitExpr_(const GenericConstNode* op) override;
  void VisitExpr_(const TupleNode* op) override;
  void VisitExpr_(const VarNode* op) override;
  void VisitExpr_(const DataflowVarNode* op) override;
  void VisitExpr_(const ShapeExprNode* op) override;
  void VisitExpr_(const ExternFuncNode* op) override;
  void VisitExpr_(const GlobalVarNode* op) override;
  void VisitExpr_(const FunctionNode* op) override;
  void VisitExpr_(const CallNode* op) override;
  void VisitExpr_(const TensorLoadNode* op) override;
  void VisitExpr_(const tvm::IntImmNode* op) override;
  void VisitExpr_(const tvm::FloatImmNode* op) override;
  void VisitExpr_(const SeqExprNode* op) override;
  void VisitExpr_(const IfExprNode* op) override;
  void VisitExpr_(const OpNode* op) override;
  void VisitExpr_(const TupleGetItemNode* op) override;
  void VisitExpr_(const StringImmNode* op) override;
  void VisitExpr_(const DataTypeImmNode* op) override;
  void VisitExpr_(const prim::AddNode* op) override;
  void VisitExpr_(const prim::LShiftNode* op) override;
  void VisitExpr_(const prim::RShiftNode* op) override;
  void VisitExpr_(const prim::BitwiseAndNode* op) override;
  void VisitExpr_(const prim::BitwiseOrNode* op) override;
  void VisitExpr_(const prim::BitwiseXorNode* op) override;
  void VisitExpr_(const prim::BitwiseNotNode* op) override;
  void VisitExpr_(const prim::SubNode* op) override;
  void VisitExpr_(const prim::MulNode* op) override;
  void VisitExpr_(const prim::DivNode* op) override;
  void VisitExpr_(const prim::ModNode* op) override;
  void VisitExpr_(const prim::FloorDivNode* op) override;
  void VisitExpr_(const prim::FloorModNode* op) override;
  void VisitExpr_(const prim::MinNode* op) override;
  void VisitExpr_(const prim::MaxNode* op) override;
  void VisitExpr_(const prim::EQNode* op) override;
  void VisitExpr_(const prim::NENode* op) override;
  void VisitExpr_(const prim::LTNode* op) override;
  void VisitExpr_(const prim::LENode* op) override;
  void VisitExpr_(const prim::GTNode* op) override;
  void VisitExpr_(const prim::GENode* op) override;
  void VisitExpr_(const prim::AndNode* op) override;
  void VisitExpr_(const prim::OrNode* op) override;
  void VisitExpr_(const prim::CastNode* op) override;
  void VisitExpr_(const prim::NotNode* op) override;
  void VisitExpr_(const prim::SelectNode* op) override;
  void VisitExpr_(const prim::RampNode* op) override;
  void VisitExpr_(const prim::BroadcastNode* op) override;
  void VisitExpr_(const prim::ShuffleNode* op) override;
  void VisitExprFallback_(const ExprNode* op) override;
 
  virtual void VisitBinding(const Binding& binding);
  // specific leaf level visitor functions
  virtual void VisitBinding_(const VarBindingNode* binding);
  virtual void VisitBinding_(const MatchCastNode* binding);
  // second level dispatching based on binding value type.
  // these dispatching functions get called from first-level dispatch on VarBinding
  virtual void VisitBinding_(const VarBindingNode* binding, const GenericConstNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const TupleNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const VarNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const DataflowVarNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const ShapeExprNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const ExternFuncNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const GlobalVarNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const FunctionNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const CallNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const SeqExprNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const IfExprNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const OpNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const TupleGetItemNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const ExprNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const StringImmNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const DataTypeImmNode* val);
  virtual void VisitBindingBlock(const BindingBlock& block);
  // specific leaf level visitor functions
  virtual void VisitBindingBlock_(const BindingBlockNode* block);
  virtual void VisitBindingBlock_(const DataflowBlockNode* block);
 
  virtual void VisitVarDef(const Var& var);
 
  virtual void VisitExprDepTypeField(const Type& ty);
 
  // specific leaf level visitor functions
  virtual void VisitVarDef_(const VarNode* var);
  virtual void VisitVarDef_(const DataflowVarNode* var);
 
  virtual void VisitLoc(const Location& loc);
  virtual void VisitTypePrimExprField(const PrimExpr& expr);
 
 private:
  using TSelf = ExprVisitor;
  using VisitBindingVTable = tvm::ObjectFunctor<void(const ffi::ObjectRef& n, ExprVisitor* self,
                                                     const VarBindingNode* binding)>;
  // initialize the vtable.
  static VisitBindingVTable InitVisitBindingVTable();
  class DefaultTypeFieldVisitor : public TypeVisitor {
   public:
    explicit DefaultTypeFieldVisitor(ExprVisitor* parent);
 
    // Override defaults in type visitor.
    void VisitTypeExprField(const Expr& expr) final;
    void VisitTypeExprField(const PrimExpr& expr) final;
    void VisitType_(const FuncTypeNode* op) final;
 
   private:
    ExprVisitor* parent_;
  };
  // This visitor is not visible to child classes and only
  // used to supported default visiting behavior.
  DefaultTypeFieldVisitor default_tyfield_visitor_{this};
};
 
void PostOrderVisit(const Expr& node, std::function<void(const Expr&)> fvisit);
 
class ExprMutatorBase : public ExprFunctor<Expr(const Expr&)> {
 public:
  Expr VisitExpr(const Expr& expr) override;
  Expr VisitExpr_(const GenericConstNode* op) override;
  Expr VisitExpr_(const TupleNode* op) override;
  Expr VisitExpr_(const VarNode* op) override;
  Expr VisitExpr_(const DataflowVarNode* op) override;
  Expr VisitExpr_(const ShapeExprNode* op) override;
  Expr VisitExpr_(const ExternFuncNode* op) override;
  Expr VisitExpr_(const GlobalVarNode* op) override;
  Expr VisitExpr_(const FunctionNode* op) override;
  Expr VisitExpr_(const CallNode* op) override;
  Expr VisitExpr_(const TensorLoadNode* op) override;
  Expr VisitExpr_(const tvm::IntImmNode* op) override;
  Expr VisitExpr_(const tvm::FloatImmNode* op) override;
  Expr VisitExpr_(const SeqExprNode* op) override;
  Expr VisitExpr_(const IfExprNode* op) override;
  Expr VisitExpr_(const OpNode* op) override;
  Expr VisitExpr_(const TupleGetItemNode* op) override;
  Expr VisitExpr_(const StringImmNode* op) override;
  Expr VisitExpr_(const DataTypeImmNode* op) override;
  Expr VisitExpr_(const prim::AddNode* op) override;
  Expr VisitExpr_(const prim::LShiftNode* op) override;
  Expr VisitExpr_(const prim::RShiftNode* op) override;
  Expr VisitExpr_(const prim::BitwiseAndNode* op) override;
  Expr VisitExpr_(const prim::BitwiseOrNode* op) override;
  Expr VisitExpr_(const prim::BitwiseXorNode* op) override;
  Expr VisitExpr_(const prim::BitwiseNotNode* op) override;
  Expr VisitExpr_(const prim::SubNode* op) override;
  Expr VisitExpr_(const prim::MulNode* op) override;
  Expr VisitExpr_(const prim::DivNode* op) override;
  Expr VisitExpr_(const prim::ModNode* op) override;
  Expr VisitExpr_(const prim::FloorDivNode* op) override;
  Expr VisitExpr_(const prim::FloorModNode* op) override;
  Expr VisitExpr_(const prim::MinNode* op) override;
  Expr VisitExpr_(const prim::MaxNode* op) override;
  Expr VisitExpr_(const prim::EQNode* op) override;
  Expr VisitExpr_(const prim::NENode* op) override;
  Expr VisitExpr_(const prim::LTNode* op) override;
  Expr VisitExpr_(const prim::LENode* op) override;
  Expr VisitExpr_(const prim::GTNode* op) override;
  Expr VisitExpr_(const prim::GENode* op) override;
  Expr VisitExpr_(const prim::AndNode* op) override;
  Expr VisitExpr_(const prim::OrNode* op) override;
  Expr VisitExpr_(const prim::CastNode* op) override;
  Expr VisitExpr_(const prim::NotNode* op) override;
  Expr VisitExpr_(const prim::SelectNode* op) override;
  Expr VisitExpr_(const prim::RampNode* op) override;
  Expr VisitExpr_(const prim::BroadcastNode* op) override;
  Expr VisitExpr_(const prim::ShuffleNode* op) override;
  Expr VisitExprFallback_(const ExprNode* op) override;
 
  virtual BindingBlock VisitBindingBlock(const BindingBlock& block);
 
  virtual PrimExpr VisitTypePrimExprField(const PrimExpr& expr);
 
  virtual Type VisitExprDepTypeField(const Type& ty);
 
 protected:
  bool VisitAndCheckTypeFieldUnchanged(const ffi::ObjectRef& ty) {
    if (const TypeNode* ty_node = ty.as<TypeNode>()) {
      Type type = ffi::GetRef<Type>(ty_node);
      return type.as<MissingType>().has_value() || this->VisitExprDepTypeField(type).same_as(ty);
    } else {
      return true;
    }
  }
 
 private:
  class DefaultTypeFieldMutator : public TypeMutator {
   public:
    explicit DefaultTypeFieldMutator(ExprMutatorBase* parent);
 
    // Override defaults in type visitor.
    Expr VisitTypeExprField(const Expr& expr) final;
    PrimExpr VisitTypeExprField(const PrimExpr& expr) final;
    Type VisitType_(const FuncTypeNode* op) final;
 
   private:
    ExprMutatorBase* parent_;
  };
  // This visitor is not visible to child classes and only
  // used to supported default visiting behavior.
  DefaultTypeFieldMutator default_tyfield_mutator_{this};
};
 
class ExprMutator : public ExprMutatorBase {
 public:
  using ExprMutatorBase::VisitExpr_;
 
  ExprMutator(ffi::Optional<IRModule> mod = std::nullopt) : builder_(BlockBuilder::Create(mod)) {}
  Expr VisitExpr(const Expr& expr) override;
  Expr VisitExpr_(const VarNode* op) override;
  Expr VisitExpr_(const DataflowVarNode* op) override;
  Expr VisitExpr_(const FunctionNode* op) override;
  Expr VisitExpr_(const SeqExprNode* op) override;
  Expr VisitExpr_(const IfExprNode* op) override;
 
  virtual void VisitBinding(const Binding& binding);
  // specific leaf level visitor functions
  virtual void VisitBinding_(const VarBindingNode* binding);
  virtual void VisitBinding_(const MatchCastNode* binding);
  // second level dispatching based on binding value type.
  // these dispatching functions get called from first-level dispatch on VarBinding
  virtual void VisitBinding_(const VarBindingNode* binding, const GenericConstNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const TupleNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const VarNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const DataflowVarNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const ShapeExprNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const ExternFuncNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const GlobalVarNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const FunctionNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const CallNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const SeqExprNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const IfExprNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const OpNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const TupleGetItemNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const ExprNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const StringImmNode* val);
  virtual void VisitBinding_(const VarBindingNode* binding, const DataTypeImmNode* val);
  virtual BindingBlock VisitBindingBlock(const BindingBlock& block) override;  // NOLINT(*)
  // specific leaf level visitor functions
  virtual BindingBlock VisitBindingBlock_(const BindingBlockNode* block);
  virtual BindingBlock VisitBindingBlock_(const DataflowBlockNode* block);
 
  virtual Var VisitVarDef(const Var& var);
  // specific leaf level visitor functions
  virtual Var VisitVarDef_(const VarNode* var);
  virtual Var VisitVarDef_(const DataflowVarNode* var);
 
 protected:
  void ReEmitBinding(const VarBindingNode* binding, Expr new_value);
 
  Expr VisitWithNewScope(const Expr& body_expr,
                         ffi::Optional<ffi::Array<Var>> params = std::nullopt);
 
  Expr VisitWithInnerScope(const Expr& body_expr);
 
  ffi::Optional<Expr> LookupBinding(const Var& var);
 
  template <typename T>
  Expr VisitExprPostOrder_(const T* op) {
    return builder_->Normalize(ExprMutator::VisitExpr_(op));
  }
 
  Var WithType(Var var, Type ty);
 
  BlockBuilder builder_;
 
  std::unordered_map<Var, Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> var_remap_;
 
 private:
  using TSelf = ExprMutator;
  using VisitBindingVTable = tvm::ObjectFunctor<void(const ffi::ObjectRef& n, ExprMutator* self,
                                                     const VarBindingNode* binding)>;
  // initialize the vtable.
  static VisitBindingVTable InitVisitBindingVTable();
};
 
}  // namespace relax
}  // namespace tvm
#endif  // TVM_RELAX_EXPR_FUNCTOR_H_