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_RELAY_EXPR_FUNCTOR_H_
#define TVM_RELAY_EXPR_FUNCTOR_H_

#include <tvm/node/functor.h>
#include <tvm/relay/adt.h>
#include <tvm/relay/error.h>
#include <tvm/relay/expr.h>
#include <tvm/relay/function.h>
#include <tvm/relay/op.h>

#include <deque>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

namespace tvm {
namespace relay {

template <typename FType>
class ExprFunctor;

// functions to be overriden.
#define EXPR_FUNCTOR_DEFAULT \
  { return VisitExprDefault_(op, std::forward<Args>(args)...); }

#define RELAY_EXPR_FUNCTOR_DISPATCH(OP)                                                    \
  vtable.template set_dispatch<OP>([](const ObjectRef& n, TSelf* self, Args... args) {     \
    return self->VisitExpr_(static_cast<const OP*>(n.get()), std::forward<Args>(args)...); \
  });

template <typename R, typename... Args>
class ExprFunctor<R(const Expr& n, Args...)> {
 private:
  using TSelf = ExprFunctor<R(const Expr& n, Args...)>;
  using FType = tvm::NodeFunctor<R(const 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) {
    ICHECK(n.defined()) << "Found null pointer node while traversing AST. The previous pass may "
                           "have generated invalid data.";
    static FType vtable = InitVTable();
    return vtable(n, this, std::forward<Args>(args)...);
  }
  // Functions that can be overriden by subclass
  virtual R VisitExpr_(const ConstantNode* 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 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 LetNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const IfNode* 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 RefCreateNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const RefReadNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const RefWriteNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const ConstructorNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExpr_(const MatchNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
  virtual R VisitExprDefault_(const Object* op, Args...) {
    LOG(FATAL) << "Do not have a default for " << op->GetTypeKey();
    throw;
  }

 private:
  // initialize the vtable.
  static FType InitVTable() {
    FType vtable;
    // Set dispatch
    RELAY_EXPR_FUNCTOR_DISPATCH(ConstantNode);
    RELAY_EXPR_FUNCTOR_DISPATCH(TupleNode);
    RELAY_EXPR_FUNCTOR_DISPATCH(VarNode);
    RELAY_EXPR_FUNCTOR_DISPATCH(GlobalVarNode);
    RELAY_EXPR_FUNCTOR_DISPATCH(FunctionNode);
    RELAY_EXPR_FUNCTOR_DISPATCH(CallNode);
    RELAY_EXPR_FUNCTOR_DISPATCH(LetNode);
    RELAY_EXPR_FUNCTOR_DISPATCH(IfNode);
    RELAY_EXPR_FUNCTOR_DISPATCH(OpNode);
    RELAY_EXPR_FUNCTOR_DISPATCH(TupleGetItemNode);
    RELAY_EXPR_FUNCTOR_DISPATCH(RefCreateNode);
    RELAY_EXPR_FUNCTOR_DISPATCH(RefReadNode);
    RELAY_EXPR_FUNCTOR_DISPATCH(RefWriteNode);
    RELAY_EXPR_FUNCTOR_DISPATCH(ConstructorNode);
    RELAY_EXPR_FUNCTOR_DISPATCH(MatchNode);
    return vtable;
  }
};

class ExprVisitor : public ::tvm::relay::ExprFunctor<void(const Expr& n)> {
 public:
  void VisitExpr(const Expr& expr) override;
  void VisitExpr_(const VarNode* op) override;
  void VisitExpr_(const GlobalVarNode* op) override;
  void VisitExpr_(const ConstantNode* op) override;
  void VisitExpr_(const TupleNode* op) override;
  void VisitExpr_(const FunctionNode* op) override;
  void VisitExpr_(const CallNode* op) override;
  void VisitExpr_(const LetNode* op) override;
  void VisitExpr_(const IfNode* op) override;
  void VisitExpr_(const OpNode* op) override;
  void VisitExpr_(const TupleGetItemNode* op) override;
  void VisitExpr_(const RefCreateNode* op) override;
  void VisitExpr_(const RefReadNode* op) override;
  void VisitExpr_(const RefWriteNode* op) override;
  void VisitExpr_(const ConstructorNode* op) override;
  void VisitExpr_(const MatchNode* op) override;
  virtual void VisitType(const Type& t);
  virtual void VisitClause(const Clause& c);
  virtual void VisitPattern(const Pattern& c);
  virtual void VisitSpan(const Span& span);

 protected:
  // Internal visiting counter
  std::unordered_map<const Object*, size_t> visit_counter_;
};

class ExprMutator : public ::tvm::relay::ExprFunctor<Expr(const Expr&)> {
 public:
  Expr Mutate(const Expr& expr) { return this->VisitExpr(expr); }
  Expr VisitExpr(const Expr& expr) override;
  Expr VisitExpr_(const VarNode* op) override;
  Expr VisitExpr_(const ConstantNode* op) override;
  Expr VisitExpr_(const GlobalVarNode* op) override;
  Expr VisitExpr_(const OpNode* op) override;
  Expr VisitExpr_(const TupleNode* op) override;
  Expr VisitExpr_(const FunctionNode* op) override;
  Expr VisitExpr_(const CallNode* call_node) override;
  Expr VisitExpr_(const LetNode* op) override;
  Expr VisitExpr_(const IfNode* op) override;
  Expr VisitExpr_(const TupleGetItemNode* op) override;
  Expr VisitExpr_(const RefCreateNode* op) override;
  Expr VisitExpr_(const RefReadNode* op) override;
  Expr VisitExpr_(const RefWriteNode* op) override;
  Expr VisitExpr_(const ConstructorNode* op) override;
  Expr VisitExpr_(const MatchNode* op) override;

  virtual Type VisitType(const Type& t);
  virtual Clause VisitClause(const Clause& c);
  virtual Pattern VisitPattern(const Pattern& c);

 protected:
  std::unordered_map<Expr, Expr, ObjectPtrHash, ObjectPtrEqual> memo_;
};

class MixedModeVisitor : public ::tvm::relay::ExprVisitor {
 public:
  using ::tvm::relay::ExprFunctor<void(const Expr& n)>::VisitExpr_;

  explicit MixedModeVisitor(int visit_limit = 1);

  using ExprVisitor::VisitExpr_;

  void VisitExpr(const Expr& expr) final;
  void VisitExpr_(const CallNode* op) override;
  void VisitExpr_(const TupleNode* op) override;
  void VisitExpr_(const TupleGetItemNode* op) override;

 protected:
  virtual void VisitLeaf(const Expr& expr);
  virtual bool CheckVisited(const Expr& expr);
  size_t visit_limit_;
};

class MixedModeMutator : public ::tvm::relay::ExprMutator {
 public:
  using ::tvm::relay::ExprFunctor<Expr(const Expr&)>::VisitExpr_;

  MixedModeMutator(bool pre = false) : pre_{pre} {};
  Expr VisitExpr(const Expr& expr) final;

  virtual Expr DispatchVisitExpr(const Expr& expr);
  Expr VisitExpr_(const TupleNode* op) final { return Rewrite(op); };
  Expr VisitExpr_(const CallNode* call_node) final { return Rewrite(call_node); };
  Expr VisitExpr_(const TupleGetItemNode* op) final { return Rewrite(op); };
  virtual Expr Rewrite_(const TupleNode* pre, const Expr& post) { return post; }
  virtual Expr Rewrite_(const CallNode* pre, const Expr& post) { return post; }
  virtual Expr Rewrite_(const TupleGetItemNode* pre, const Expr& post) { return post; }

 protected:
  bool pre_;
  template <typename T>
  Expr Rewrite(const T* op) {
    Expr post = ExprMutator::VisitExpr_(op);
    return Rewrite_(op, post);
  }

  virtual void VisitLeaf(const Expr& expr);
  virtual bool CheckVisited(const Expr& expr);
};

#define RELAY_EXPR_REWRITER_DISPATCH(OP)                                                   \
  vtable.template set_dispatch<OP>([](const ObjectRef& n, TSelf* self, const Expr& post) { \
    return self->Rewrite_(static_cast<const OP*>(n.get()), post);                          \
  });

#define EXPR_REWRITER_REWRITE_DEFAULT \
  { return post; }

class ExprRewriter {
 private:
  using TSelf = ExprRewriter;
  using FType = tvm::NodeFunctor<Expr(const ObjectRef& n, TSelf* self, const Expr& post)>;

 public:
  virtual ~ExprRewriter() {}
  Expr operator()(const Expr& pre, const Expr& post) { return Rewrite(pre, post); }
  virtual Expr Rewrite(const Expr& pre, const Expr& post) {
    ICHECK(pre.defined());
    static FType vtable = InitVTable();
    return vtable(pre, this, post);
  }
  // Functions that can be overriden by subclass, should not recurse
  virtual Expr Rewrite_(const VarNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
  virtual Expr Rewrite_(const GlobalVarNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
  virtual Expr Rewrite_(const ConstantNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
  virtual Expr Rewrite_(const TupleNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
  virtual Expr Rewrite_(const FunctionNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
  virtual Expr Rewrite_(const CallNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
  virtual Expr Rewrite_(const LetNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
  virtual Expr Rewrite_(const IfNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
  virtual Expr Rewrite_(const OpNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
  virtual Expr Rewrite_(const TupleGetItemNode* pre,
                        const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
  virtual Expr Rewrite_(const RefCreateNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
  virtual Expr Rewrite_(const RefReadNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
  virtual Expr Rewrite_(const RefWriteNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
  virtual Expr Rewrite_(const ConstructorNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;
  virtual Expr Rewrite_(const MatchNode* pre, const Expr& post) EXPR_REWRITER_REWRITE_DEFAULT;

 private:
  // initialize the vtable.
  static FType InitVTable() {
    FType vtable;
    // Set dispatch
    RELAY_EXPR_REWRITER_DISPATCH(ConstantNode);
    RELAY_EXPR_REWRITER_DISPATCH(TupleNode);
    RELAY_EXPR_REWRITER_DISPATCH(VarNode);
    RELAY_EXPR_REWRITER_DISPATCH(GlobalVarNode);
    RELAY_EXPR_REWRITER_DISPATCH(FunctionNode);
    RELAY_EXPR_REWRITER_DISPATCH(CallNode);
    RELAY_EXPR_REWRITER_DISPATCH(LetNode);
    RELAY_EXPR_REWRITER_DISPATCH(IfNode);
    RELAY_EXPR_REWRITER_DISPATCH(OpNode);
    RELAY_EXPR_REWRITER_DISPATCH(TupleGetItemNode);
    RELAY_EXPR_REWRITER_DISPATCH(RefCreateNode);
    RELAY_EXPR_REWRITER_DISPATCH(RefReadNode);
    RELAY_EXPR_REWRITER_DISPATCH(RefWriteNode);
    RELAY_EXPR_REWRITER_DISPATCH(ConstructorNode);
    RELAY_EXPR_REWRITER_DISPATCH(MatchNode);
    return vtable;
  }
};

Expr PostOrderRewrite(const Expr& expr, ExprRewriter* rewriter);

void PostOrderVisit(const Expr& node, std::function<void(const Expr&)> fvisit);

struct v_info {
  explicit v_info(Expr node_) : node{node_} {}
  v_info(Expr node_, bool children_expanded_)
      : node{node_}, children_expanded{children_expanded_} {};
  Expr node{};
  bool children_expanded{false};
};

template <typename FCheckVisited, typename FVisitLeaf, typename FExpandExpr>
void ExpandDataflow(Expr expr, FCheckVisited fcheck_visited, FVisitLeaf fvisit_leaf,
                    FExpandExpr fexpand_expr) {
  std::deque<v_info> stack;
  auto fpush_to_stack = [&fcheck_visited, &stack](const Expr& expr) {
    if (!fcheck_visited(expr)) {
      stack.emplace_front(v_info(expr));
    }
  };

  fpush_to_stack(expr);
  while (stack.size() > 0) {
    v_info* front = &stack.front();
    if (fcheck_visited(front->node)) {
      stack.pop_front();
    } else if (front->children_expanded) {
      fvisit_leaf(front->node);
      // TODO(d-smirnov): this is for compatibility with current implementation of MixedModeVisitor
      stack.pop_front();
    } else {
      front->children_expanded = true;
      for (auto e : fexpand_expr(front->node)) {
        fpush_to_stack(e);
      }
    }
  }
}

template <typename FCheckVisited, typename FVisitLeaf>
void ExpandDataflow(Expr expr, FCheckVisited fcheck_visited, FVisitLeaf fvisit_leaf) {
  auto fexpand_expr = [](const Expr& expr) {
    std::vector<Expr> result;
    if (const CallNode* op = expr.as<CallNode>()) {
      if (op->op == Op::Get("call_lowered")) {
        // Ignore the intermediate tuple since this is purely a calling-convention detail
        const auto* tuple_args = op->args[1].as<TupleNode>();
        ICHECK(tuple_args)
            << "Expected second arg to call_lowered to be a Tuple of input arguments.";
        for (auto it = tuple_args->fields.rbegin(); it != tuple_args->fields.rend(); ++it) {
          result.push_back(*it);
        }
        result.push_back(op->args[0]);
      } else {
        for (auto it = op->args.rbegin(); it != op->args.rend(); ++it) {
          result.push_back(*it);
        }
      }
      result.push_back(op->op);
    } else if (const TupleNode* op = expr.as<TupleNode>()) {
      for (auto it = op->fields.rbegin(); it != op->fields.rend(); ++it) {
        result.push_back(*it);
      }
    } else if (const TupleGetItemNode* op = expr.as<TupleGetItemNode>()) {
      result.push_back(op->tuple);
    }
    return result;
  };
  ExpandDataflow(expr, fcheck_visited, fvisit_leaf, fexpand_expr);
}

void ExpandANormalForm(const LetNode* op, std::function<void(const LetNode*)> pre_visit,
                       std::function<void(const LetNode*)> post_visit);

}  // namespace relay
}  // namespace tvm
#endif  // TVM_RELAY_EXPR_FUNCTOR_H_