fuse_ops.cc 33.6 KB
Newer Older
1 2 3 4 5 6 7 8
/*
 * 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
9
 *
10
 *   http://www.apache.org/licenses/LICENSE-2.0
11
 *
12 13 14 15 16 17 18 19
 * 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.
 */

20
/*!
21
 * Copyright (c) 2019 by Contributors
22 23 24
 *
 * \file src/tvm/relay/pass/fuse_ops.cc
 *
25 26
 * \brief This is a backend-aware optimization pass.
 *   Fuse necessary ops into a single one.
27
 */
28
#include <tvm/expr_operator.h>
Zhi committed
29
#include <tvm/relay/analysis.h>
30
#include <tvm/relay/expr_functor.h>
31
#include <tvm/relay/op_attr_types.h>
32
#include <tvm/relay/transform.h>
33
#include "./pattern_util.h"
34 35
#include "../../common/arena.h"

36 37 38 39

namespace tvm {
namespace relay {

40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79
/*
  Note on Fusing algorithm:

  The main challenge of genenral fusor is to handle possible diamond shape branches,
  in the following graph, conv2d can be fused to elemwise add.

            conv2d
            /  |  \
           /   |   \
         op    op   op
          \    |    /
           \   |   /
          elemwise add
               |

  However, at the point of conv2d we do not necessarily know that all its future path
  will merge at the elemwise add. The new fusor algorithm applies post-dominator analysis.
  The immediate post-dominator of a node defined by the closest node where all the future path goes into.
  In the above case, the elemwise add is the post-dominator of conv2d. The general algorithm is as follows:

  - Construct a DAG of dataflow graph for dominator analysis
  - Construct a post-dominator tree which gives immediate post dominator of each node.
  - Run fusion algorithm with the given post-dominator information.

  Note that, because we run analysis on a DAG, we use a single pass post-dominator
  tree construction algorithm via LCA, which is simpler than the full version that handles cycles.

  The fusion algorithm traverses from each node and checks if it can be fused to its
  immediate post dominator. It has to check the following things:

  - CheckPath: check all the path between a node and its immediate post-dominator
               satiesfies the fuse condition.
  - Note that these intermediate node can already be fused with another nodes, the algorithm
      will still run correctly.
  - CommitFuse: mark all the nodes between source and post-dominator as the same group.
  - We use an Union-Find data structure to manage the groups.
*/
using common::LinkNode;
using common::LinkedList;

80 81
constexpr uint32_t kMaxFusedOps = 256;

82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186
/*!
 * \brief Indexed data flow graph in forward direction.
 *  This is a temporary data structure used for operator fusion analysis.
 *
 *  This data structure only captures the dataflow fragement and
 *  could ignore blocks like let by simply ordering each dataflow block
 *  and mark the output node as extern_ref;
 */
class IndexedForwardGraph {
 public:
  struct Node;
  /*!
   * The forward edge in the dataflow graph.
   */
  struct Edge {
    /*! \brief The corresponding node */
    Node* node{nullptr};
    /*! \brief The respective pattern of this op */
    OpPatternKind pattern{kOpaque};
  };
  /*! \brief A node in the graph. */
  struct Node {
    /*! \brief weak reference to the corresponding edge. */
    const tvm::Node* ref{nullptr};
    /*! \brief The index of the node in topological order. */
    size_t index{0};
    /*! \brief Whether this node is referenced by external source */
    bool extern_ref{false};
    /*! \brief The general pattern in the node */
    OpPatternKind pattern{kOpaque};
    /*! \brief The outputs of the node. */
    LinkedList<Edge> outputs;
  };
  /*! \brief The node map that maps node to graph */
  std::unordered_map<const tvm::Node*, Node*> node_map;
  /*! \brief All the nodes in post DFS order */
  std::vector<Node*> post_dfs_order;

  /*! \brief Dump the graph into string. */
  void DebugDump() {
    std::ostringstream os;
    for (size_t i = 0; i < post_dfs_order.size(); ++i) {
      Node* node = post_dfs_order[i];
      os << "node[" << i << "], "
         << GetRef<NodeRef>(node->ref)
         << " outputs=[";
      for (auto* link = node->outputs.head; link != nullptr; link = link->next) {
        os << link->value.node->index << ", ";
      }
      os << "]\n";
    }
    LOG(INFO) << os.str();
  }
  /*!
   * \brief create a indexed forward graph.
   * \param arena The arena used for data allocation.
   * \param body The body of the expression to create a graph.
   */
  static IndexedForwardGraph Create(common::Arena* arena, const Expr& body);

 private:
  class Creator;
};

// Creator of post dominator tree of the dataflow
class IndexedForwardGraph::Creator : private ExprVisitor {
 public:
  explicit Creator(common::Arena* arena)
      : arena_(arena) {}

  IndexedForwardGraph Prepare(const Expr& body) {
    this->Update(body, nullptr, kOpaque);
    this->VisitExpr(body);
    return std::move(graph_);
  }

 private:
  /*! \brief allocator of all the internal node object */
  common::Arena* arena_;
  // The output.
  IndexedForwardGraph graph_;
  // attribute equal comparator
  AttrsEqual attr_equal_;
  // Update the message stored at the node.
  void Update(const Expr& node,
              IndexedForwardGraph::Node* parent,
              OpPatternKind pattern) {
    const tvm::Node* key = node.get();
    IndexedForwardGraph::Node* current;
    auto it = graph_.node_map.find(key);
    if (it != graph_.node_map.end()) {
      current = it->second;
    } else {
      current = arena_->make<IndexedForwardGraph::Node>();
      graph_.node_map[key] = current;
    }
    if (parent != nullptr) {
      auto* link = arena_->make<LinkNode<IndexedForwardGraph::Edge> >();
      link->value.node = parent;
      link->value.pattern = pattern;
      current->outputs.Push(link);
    } else {
      current->extern_ref = true;
    }
  }
187

188 189 190 191 192 193 194 195 196 197 198 199
  void AddNode(const tvm::Node* key) {
    auto it = graph_.node_map.find(key);
    CHECK(it != graph_.node_map.end())
        << "Cannot find node " << GetRef<NodeRef>(key);
    IndexedForwardGraph::Node* node = it->second;
    CHECK(node->ref == nullptr);
    node->ref = key;
    node->index = graph_.post_dfs_order.size();
    graph_.post_dfs_order.push_back(node);
  }

  // Post order tree
200
  void VisitExpr_(const FunctionNode* op) final {
201 202 203 204 205 206 207
    for (auto param : op->params) {
      this->Update(param, nullptr, kOpaque);
    }
    this->Update(op->body, nullptr, kOpaque);
    ExprVisitor::VisitExpr_(op);
  }

208
  void VisitExpr_(const ConstantNode* op) final {
209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227
    this->AddNode(op);
    Node* node = graph_.node_map.at(op);
    DataType dtype = TVMType2Type(op->data->dtype);
    // This rule must be consistent with code generator.
    bool is_simple_const = (
        dtype == Int(32) ||
        dtype == Int(64) ||
        dtype == Float(32) ||
        dtype == Float(64) ||
        dtype == Bool());
    if (op->is_scalar() && is_simple_const) {
      node->pattern = kElemWise;
    } else {
      // for now, mark non-scalar constant
      // as opaque, we will not choose to fuse it.
      node->pattern = kOpaque;
    }
  }

228
  void VisitExpr_(const CallNode* call) final {
229 230 231 232
    CHECK(graph_.node_map.count(call));
    Node* node = graph_.node_map.at(call);
    static auto fpattern =
        Op::GetAttr<TOpPattern>("TOpPattern");
233 234 235 236 237 238 239 240 241
    // Now we set the pattern of this call.
    //
    // If we see a call mentioning an operator we should mark it with its
    // annotated pattern.
    //
    // If the pattern is not annotated we will default to opaque.
    //
    // Finally if the operator position is not a call node we will
    // need to call Update, as it may be an arbitrary expression.
242 243 244
    OpPatternKind op_pattern = kOpaque;
    if (const OpNode* opnode = call->op.as<OpNode>()) {
      op_pattern = static_cast<OpPatternKind>(fpattern[GetRef<Op>(opnode)]);
245 246
    } else {
      this->Update(call->op, node, kOpaque);
247
    }
248

249
    node->pattern = op_pattern;
250
    this->Update(call->op, nullptr, kOpaque);
251
    const auto* rtype = call->checked_type().as<TensorTypeNode>();
252
    // pass the analysis back to all the children it references.
253 254 255
    for (size_t i = 0; i < call->args.size(); ++i) {
      const auto* arg_type =
          call->args[i]->checked_type().as<TensorTypeNode>();
256
      // specifically check if result type is the same as arguments type
257 258 259 260 261 262 263 264 265 266 267 268 269
      OpPatternKind edge_pattern = op_pattern;
      if (edge_pattern == kBroadcast &&
          arg_type != nullptr &&
          rtype != nullptr &&
          attr_equal_(rtype->shape, arg_type->shape)) {
        edge_pattern = kElemWise;
      }
      this->Update(call->args[i], node, edge_pattern);
    }
    ExprVisitor::VisitExpr_(call);
    this->AddNode(call);
  }

270
  void VisitExpr_(const TupleNode* op) final {
271 272
    CHECK(graph_.node_map.count(op));
    Node* tuple_node = graph_.node_map.at(op);
273
    tuple_node->pattern = kTuple;
274
    for (const Expr& field : op->fields) {
275 276 277 278 279
      if (field->checked_type().as<TensorTypeNode>()) {
        this->Update(field, tuple_node, kInjective);
      } else {
        this->Update(field, nullptr, kOpaque);
      }
280 281 282 283 284
    }
    ExprVisitor::VisitExpr_(op);
    this->AddNode(op);
  }

285
  void VisitExpr_(const TupleGetItemNode* op) final {
286 287
    auto tuple_type = op->tuple->checked_type().as<TupleTypeNode>();
    CHECK(tuple_type);
288 289 290 291 292 293
    // when TVM lowers a fused function, it expects all arguments to be a Tensor or
    // a tuple containing only Tensors. But this tuple may contain a reference or
    // another tuple. To avoid modifying codegen logic, we do not allow fusing through this node
    // if the tuple contains such non Tensor fields. However, all fields will be recursively
    // visited via call to ExprVisitor::VisitExpr_(op) below and corresponding visitor methods.
    bool has_non_tensor = false;
294
    for (auto ty : tuple_type->fields) {
295 296
      if (!ty.as<TensorTypeNode>()) {
        has_non_tensor = true;
297 298 299
        break;
      }
    }
300
    if (has_non_tensor) {
301 302 303 304 305 306 307
      this->Update(op->tuple, nullptr, kOpaque);
    } else {
      CHECK(graph_.node_map.count(op));
      Node* node = graph_.node_map.at(op);
      node->pattern = kInjective;
      this->Update(op->tuple, node, kInjective);
    }
308 309 310 311
    ExprVisitor::VisitExpr_(op);
    this->AddNode(op);
  }

312
  void VisitExpr_(const VarNode* op) final {
313 314 315
    this->AddNode(op);
  }

316
  void VisitExpr_(const LetNode* op) final {
317 318 319 320 321 322 323 324
    // do not fuse through let.
    this->Update(op->var, nullptr, kOpaque);
    this->Update(op->value, nullptr, kOpaque);
    this->Update(op->body, nullptr, kOpaque);
    ExprVisitor::VisitExpr_(op);
    this->AddNode(op);
  }

325
  void VisitExpr_(const IfNode* op) final {
326 327 328 329 330 331 332
    // do not fuse through if.
    this->Update(op->cond, nullptr, kOpaque);
    this->Update(op->true_branch, nullptr, kOpaque);
    this->Update(op->false_branch, nullptr, kOpaque);
    ExprVisitor::VisitExpr_(op);
    this->AddNode(op);
  }
333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351

  void VisitExpr_(const RefCreateNode* op) final {
    this->Update(op->value, nullptr, kOpaque);
    ExprVisitor::VisitExpr_(op);
    this->AddNode(op);
  }

  void VisitExpr_(const RefReadNode* op) final {
    this->Update(op->ref, nullptr, kOpaque);
    ExprVisitor::VisitExpr_(op);
    this->AddNode(op);
  }

  void VisitExpr_(const RefWriteNode* op) final {
    this->Update(op->ref, nullptr, kOpaque);
    this->Update(op->value, nullptr, kOpaque);
    ExprVisitor::VisitExpr_(op);
    this->AddNode(op);
  }
352 353 354 355 356 357 358 359 360

  void VisitExpr_(const MatchNode* op) final {
    this->Update(op->data, nullptr, kOpaque);
    for (const Clause& c : op->clauses) {
      this->Update(c->rhs, nullptr, kOpaque);
    }
    ExprVisitor::VisitExpr_(op);
    this->AddNode(op);
  }
361 362 363 364 365 366 367 368 369 370 371 372
};

IndexedForwardGraph IndexedForwardGraph::Create(
    common::Arena* arena, const Expr& body) {
  return Creator(arena).Prepare(body);
}

/*!
 * \brief Dominator tree that represent domination or
 *  post domination relation of the node.
 */
class DominatorTree {
373
 public:
374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407
  /*!
   * \brief A node in the dominator tree.
   */
  struct Node {
    /*! \brief The node in the tree */
    IndexedForwardGraph::Node* gnode{nullptr};
    /*! \brief parent of the tree */
    Node* parent{nullptr};
    /*! \brief current depth*/
    int depth{0};
    /*! \brief aggregated pattern to parent */
    OpPatternKind pattern{kOpaque};
  };
  // index -> node.
  std::vector<Node*> nodes;
  /*!
   * \brief compute a post dominator relation for a given dataflow graph.
   * \param arena The arena used for node allocation.
   * \param graph The graph to be analyze.
   * \return The dominator tree of the graph.
   * \note This algorithm makes use of the fact that graph is DAG,
   *       and runs a single pass algorithm via LCA.
   */
  static DominatorTree PostDom(common::Arena* arena,
                               const IndexedForwardGraph& graph);

 private:
  // Combine pattern together.
  static OpPatternKind CombinePattern(
      OpPatternKind lhs, OpPatternKind rhs) {
    if (lhs > rhs) return lhs;
    return rhs;
  }
  /*!
408
   * \brief Find the least common ancestor of the two nodes.
409 410 411 412
   * \param lhs The left node.
   * \param rhs The right node.
   * \param edge_pattern
   *        The combined edge pattern across all the parents.
413
   * \return The least common ancestor of the two.
414
   */
415
  static Node* LeastCommonAncestor(
416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434
      Node* lhs,
      Node* rhs,
      OpPatternKind* edge_pattern) {
    while (lhs != rhs) {
      if (lhs == nullptr) return nullptr;
      if (rhs == nullptr) return nullptr;
      if (lhs->depth < rhs->depth) {
        edge_pattern[0] = CombinePattern(
            edge_pattern[0], rhs->pattern);
        rhs = rhs->parent;
      } else if (rhs->depth < lhs->depth) {
        edge_pattern[0] = CombinePattern(
            edge_pattern[0], lhs->pattern);
        lhs = lhs->parent;
      } else {
        edge_pattern[0] = CombinePattern(
            edge_pattern[0], lhs->pattern);
        edge_pattern[0] = CombinePattern(
            edge_pattern[0], rhs->pattern);
435 436
        lhs = lhs->parent;
        rhs = rhs->parent;
437 438 439 440
      }
    }
    return lhs;
  }
441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476
  /*!
   * \brief Find the least common ancestor of a list of nodes.
   * \param nodes the nodes.
   * \param edge_pattern
   *        The combined edge pattern across all the parents.
   * \return The least common ancestor of all nodes.
   */
  Node* LeastCommonAncestor(const LinkedList<IndexedForwardGraph::Edge>& input_nodes,
                            OpPatternKind* edge_pattern) {
    auto link = input_nodes.head;
    if (link == nullptr) {
      return nullptr;
    }
    auto get_node = [&](const IndexedForwardGraph::Edge& edge) {
      size_t oindex = edge.node->index;
      CHECK_LT(oindex, nodes.size());
      Node* onode = nodes[oindex];
      CHECK(onode != nullptr);
      return onode;
    };
    Node* parent = get_node(link->value);
    *edge_pattern = CombinePattern(*edge_pattern, link->value.pattern);
    link = link->next;
    for (; link != nullptr; link = link->next) {
      parent = LeastCommonAncestor(parent, get_node(link->value), edge_pattern);
      *edge_pattern = CombinePattern(*edge_pattern, link->value.pattern);
    }
    return parent;
  }
  /*!
   * \brief Convert the Node from an IndexedForwardGraph Node into DomaintorTree Node.
   * \param arena The Arena.
   * \param gnode An IndexedForwardGraph Node.
   * \return The DominatorTree Node.
   */
  Node* GetNode(common::Arena* arena, IndexedForwardGraph::Node* gnode) {
477 478 479 480 481 482 483 484 485
    Node* tnode = arena->make<Node>();
    tnode->gnode = gnode;
    if (gnode->extern_ref) {
      tnode->depth = 1;
      tnode->parent = nullptr;
      tnode->pattern = kOpaque;
    } else {
      // find the LCAs of all outputs.
      OpPatternKind pattern = kElemWise;
486
      Node* parent = LeastCommonAncestor(gnode->outputs, &pattern);
487
      tnode->depth = parent ? parent->depth + 1 : 1;
488 489 490
      tnode->parent = parent;
      tnode->pattern = pattern;
    }
491 492 493 494 495 496 497 498 499 500 501 502 503
    return tnode;
  }
};


DominatorTree DominatorTree::PostDom(common::Arena* arena,
                                     const IndexedForwardGraph& graph) {
  DominatorTree tree;
  tree.nodes.resize(graph.post_dfs_order.size(), nullptr);
  // reverse topo order
  for (size_t i = graph.post_dfs_order.size(); i != 0; --i) {
    size_t index = i - 1;
    tree.nodes[index] = tree.GetNode(arena, graph.post_dfs_order[index]);
504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548
  }
  return tree;
}

/*!
 * \brief A partition of the graph marked by union find data structure.
 */
class GraphPartitioner {
 public:
  explicit GraphPartitioner(common::Arena* arena, int opt_level)
      : arena_(arena), opt_level_(opt_level) {}
  /*!
   * \brief Group as a union find data structure.
   */
  struct Group {
    /*! \brief The parent in the union find data structure. */
    Group* parent{nullptr};
    /*! \brief The pattern of the group */
    OpPatternKind pattern;
    /*! \brief reference to the root node. */
    const tvm::Node* root_ref{nullptr};
    /*!
     * \brief Reference to the master node,
     * this field is not nullptr only if pattern is kOutEWiseFusable.
     */
    const tvm::Node* master_ref{nullptr};
    /*!
     * \brief Find the group root, perform path compression
     * \return The root type node.
     */
    Group* FindRoot() {
      // fast path
      if (this->parent == nullptr) return this;
      // slow path with path compression.
      Group* root = this;
      while (root->parent != nullptr) {
        root = root->parent;
      }
      for (Group* p = this; p != root;) {
        Group* parent = p->parent;
        p->parent = root;
        p = parent;
      }
      return root;
    }
549 550 551 552 553

    /*!
     * \brief The number of nodes belonging to this group
     */
    uint32_t num_nodes{1};
554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587
  };
  /*!
   * \brief Partition a graph.
   * \return group assignments of each node.
   */
  std::vector<Group*> Partition(const IndexedForwardGraph& graph);

 private:
  /*! \brief The internal arena for temporary space. */
  common::Arena* arena_;
  /*! \brief optimization level for fuse operation. */
  int opt_level_;
  /*! \brief The internal groups. */
  std::vector<Group*> groups_;
  /*! \brief internal field used for deduplication */
  std::unordered_set<IndexedForwardGraph::Node*> visited_;
  // Internal implelementation of CheckPath
  template<typename F>
  bool CheckPath_(IndexedForwardGraph::Node* src,
                  IndexedForwardGraph::Node* sink,
                  F fcond) {
    if (visited_.count(src)) return true;
    visited_.insert(src);
    Group* gnode =  groups_[src->index];
    CHECK(gnode != nullptr);
    gnode = gnode->FindRoot();
    if (!fcond(gnode->pattern, src == sink)) return false;
    if (src == sink) return true;
    for (auto link = src->outputs.head; link != nullptr; link = link->next) {
      if (!CheckPath_(link->value.node, sink, fcond)) return false;
    }
    return true;
  }
  /*!
588 589
   * \brief Check all the node and edge pattern
   *  between src and sink satisfies fcond.
590
   *
591
   * src is not checked.
592 593 594 595
   *
   * \param src The source node.
   * \param sink The termination node.
   * \param fcond The condition to be checked.
596
   * \tparam F the condition function, with signature
597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625
   * \note sink must be a post-dominator of src.
   */
  template<typename F>
  bool CheckPath(IndexedForwardGraph::Node* src,
                 IndexedForwardGraph::Node* sink,
                 F fcond) {
    CHECK(!src->extern_ref);
    visited_.clear();
    CHECK(src != sink);
    for (auto link = src->outputs.head; link != nullptr; link = link->next) {
      if (!CheckPath_(link->value.node, sink, fcond)) return false;
    }
    return true;
  }
  // Combine two patterns together.
  static OpPatternKind CombinePattern(
      OpPatternKind lhs, OpPatternKind rhs) {
    if (lhs > kBroadcast && rhs > kBroadcast) {
      LOG(FATAL) << "Cannot merge two complex group together";
    }
    if (lhs > rhs) return lhs;
    return rhs;
  }
  /*!
   * \brief Merge the child group to the parent.
   * \param child The child group.
   * \param parent The parent group.
   */
  void MergeFromTo(Group* child, Group* parent) {
626
    // update the number of nodes of the parent group
627
    parent->num_nodes += child->num_nodes;
628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651
    child = child->FindRoot();
    parent = parent->FindRoot();
    if (child == parent) return;
    child->parent = parent;
    // update master ref and pattern
    if (child->master_ref != nullptr) {
      CHECK(parent->master_ref == nullptr);
      parent->master_ref = child->master_ref;
      parent->pattern = CombinePattern(
          child->pattern, parent->pattern);
    }
  }
  // Internal implelementation of CommitFuse
  void CommitFuse_(IndexedForwardGraph::Node* src,
                   IndexedForwardGraph::Node* sink,
                   Group* target) {
    if (src == sink) return;
    if (visited_.count(src)) return;
    visited_.insert(src);
    Group* gnode = groups_[src->index];
    CHECK(gnode != nullptr);
    // merge the current group to the parent if possible.
    MergeFromTo(gnode, target);
    for (auto link = src->outputs.head; link != nullptr; link = link->next) {
652
      CommitFuse_(link->value.node, sink, target);
653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700
    }
  }
  /*!
   * \brief Commit fusion operation.
   * \param src The source node.
   * \param sink The termination node.
   * \note sink must be a post-dominator of src.
   */
  void CommitFuse(IndexedForwardGraph::Node* src,
                  IndexedForwardGraph::Node* sink) {
    Group* target = groups_[sink->index];
    visited_.clear();
    CHECK(src != sink);
    CommitFuse_(src, sink, target);
  }

  // Initialize the groups.
  void InitGroups(const IndexedForwardGraph& graph) {
    groups_.resize(graph.post_dfs_order.size());
    for (size_t nid = 0; nid < groups_.size(); ++nid) {
      const auto* graph_node = graph.post_dfs_order[nid];
      auto* group_node = arena_->make<Group>();
      group_node->pattern = graph_node->pattern;
      group_node->root_ref = graph_node->ref;
      // set master ref if necessary.
      if (group_node->pattern == kOutEWiseFusable) {
        group_node->master_ref = graph_node->ref;
      }
      groups_[nid] = group_node;
    }
  }

  // execute the fusion algorithm.
  void RunFuse(const IndexedForwardGraph& graph,
               const DominatorTree& post_dom_tree,
               int phase) {
    for (size_t nid = 0; nid < groups_.size(); ++nid) {
      // the group of current node has been specified already.
      auto* graph_node = graph.post_dfs_order[nid];
      auto* dom_node = post_dom_tree.nodes[nid];
      Group* group_node = groups_[nid];
      CHECK(group_node != nullptr);
      // no actions for opaque nodes
      if (group_node->pattern == kOpaque) continue;
      // no actions needed if the current node have no dominator
      if (dom_node->parent == nullptr) continue;
      CHECK(!graph_node->extern_ref);
      size_t dom_parent_gindex = dom_node->parent->gnode->index;
701

702 703 704 705
      // refuse the fusion if too many ops are going to be fused together
      if (groups_[dom_parent_gindex]->num_nodes + group_node->num_nodes > kMaxFusedOps)
        continue;

706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727
      if (phase == 2) {
        // Fuse injective ops into intermediate tuples, if any
        if (group_node->pattern > kInjective) continue;
        Group* dom_parent_group = groups_[dom_parent_gindex];
        Group* dom_root_group = dom_parent_group->FindRoot();
        // If dom node group has a tuple as its root, we do not fuse tuple fields into it
        if (dom_root_group->pattern == kTuple) continue;
        if (dom_parent_group->pattern == kTuple && dom_root_group->pattern <= kInjective) {
          // Now we know the tuple has been fused into subsequent injective ops
          auto fcond = [](OpPatternKind kind, bool is_sink) {
            return kind <= kInjective;
          };
          // dom_root_group can also be tuple, as in inception layers
          // CheckPath is needed to avoid fusing two intermediate tuples
          if (CheckPath(graph_node, dom_node->parent->gnode, fcond)) {
            CommitFuse(graph_node, dom_node->parent->gnode);
          }
        }
        continue;
      }

      // Skip if current node is already fused to the parent.
728 729 730 731
      if (groups_[dom_parent_gindex] != nullptr &&
          group_node->FindRoot() == groups_[dom_parent_gindex]->FindRoot()) {
        continue;
      }
732 733
      // Do not fuse into tuple for now
      if (groups_[dom_parent_gindex]->pattern == kTuple) continue;
734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749
      // Try to fuse current node to its post-dominator.
      if (group_node->pattern == kOutEWiseFusable) {
        if (phase != 0) continue;
        // Path for OutEWiseFusable: conv2d
        // Check if the dominator relation is elemwise.
        if (dom_node->parent != nullptr && dom_node->pattern == kElemWise) {
          CHECK(dom_node->parent->gnode != nullptr);
          // The fuse can be executed if all the intermediate ops are still broadcast.
          auto fcond = [](OpPatternKind kind, bool is_sink) {
            return kind <= kBroadcast;
          };
          if (CheckPath(graph_node, dom_node->parent->gnode, fcond)) {
            CommitFuse(graph_node, dom_node->parent->gnode);
          }
        }
      } else if (group_node->pattern <= kBroadcast) {
750 751 752 753 754 755 756 757
        // Pre-condition: can only be fused to parent which is injective or reduction.
        if (dom_node->parent != nullptr &&
            (dom_node->pattern <= kInjective ||
             dom_node->pattern == kCommReduce)) {
          // Check if all the intermediate ops are still broadcast.
          // The final terminal node can already be fused to a OutEWiseFusable group.
          auto fcond = [](OpPatternKind kind, bool is_sink) {
            if (!is_sink) {
758 759 760
              // Elemwise, broadcast, and injective ops on the parallel branches
              // are allowed be fused to the elemwise/broadcast master.
              return kind <= kInjective;
761 762 763
            } else {
              return (kind <= kBroadcast ||
                      kind == kCommReduce ||
764
                      kind == kInjective ||
765 766 767 768 769
                      kind == kOutEWiseFusable);
            }
          };
          if (CheckPath(graph_node, dom_node->parent->gnode, fcond)) {
            CommitFuse(graph_node, dom_node->parent->gnode);
770 771
          }
        }
772
      } else if (group_node->pattern == kInjective || group_node->pattern == kTuple) {
773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797
        // defer injective fusion to second phase.
        // so conv2d always finishes fusing.
        if (phase != 1) continue;
        // Check if all path are injective.
        auto fcond = [](OpPatternKind kind, bool is_sink) {
          return kind <= kInjective;
        };
        if (CheckPath(graph_node, dom_node->parent->gnode, fcond)) {
          CommitFuse(graph_node, dom_node->parent->gnode);
        }
      } else {
        // do nothing.
        CHECK(group_node->pattern == kCommReduce);
      }
    }
  }
};

std::vector<GraphPartitioner::Group*>
GraphPartitioner::Partition(const IndexedForwardGraph& graph) {
  this->InitGroups(graph);
  if (opt_level_ == 0) return std::move(groups_);
  // get post dominator tree
  auto post_dom_tree = DominatorTree::PostDom(arena_, graph);
  // run fusion algorithm.
798
  for (int phase = 0; phase < 3; ++phase) {
799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850
    this->RunFuse(graph, post_dom_tree, phase);
  }
  return std::move(groups_);
}

class FuseMutator : private ExprMutator {
 public:
  // Run the transform
  Expr Transform(const Expr& body, int fuse_opt_level) {
    // setup the group map.
    auto graph = IndexedForwardGraph::Create(&arena_, body);
    auto groups = GraphPartitioner(&arena_, fuse_opt_level).Partition(
        graph);
    for (size_t nid = 0; nid < graph.post_dfs_order.size(); ++nid) {
      CHECK(graph.post_dfs_order[nid]->ref != nullptr);
      gmap_[graph.post_dfs_order[nid]->ref] = groups[nid];
    }
    // The following line can be used for debug.
    // this->DebugDumpGroup(body);
    return this->Mutate(body);
  }


 private:
  /*! \brief Temporary information from each group. */
  struct GroupInfo {
   public:
    // The parameters of the function.
    Array<Var> params;
    // The arguments to call the functions.
    Array<Expr> arguments;
    // Get a new parameter or allocate an old one
    Var GetOrAllocParam(const Expr& expr, const Type& type) {
      // run linear scan as most fused groups contain only a few inputs.
      for (size_t i = 0; i < arguments.size(); ++i) {
        if (expr.same_as(arguments[i])) return params[i];
      }
      // create a new parameter.
      std::ostringstream os;
      os << "p" << params.size();
      auto var = VarNode::make(os.str(), type);
      params.push_back(var);
      arguments.push_back(expr);
      return var;
    }
  };
  /*! \brief Internal arena. */
  common::Arena arena_;
  /*! \brief The group assignment map. */
  std::unordered_map<const Node*, GraphPartitioner::Group*> gmap_;
  /* \brief Internal group information map. */
  std::unordered_map<GraphPartitioner::Group*, GroupInfo> ginfo_;
851

852 853
  // Skip primitive function.
  Expr VisitExpr_(const FunctionNode* fn_node) {
854
    if (fn_node->IsPrimitive()) {
855 856 857 858 859
      return GetRef<Expr>(fn_node);
    } else {
      return ExprMutator::VisitExpr_(fn_node);
    }
  }
860

861
  // Transform calls.
862
  Expr VisitExpr_(const CallNode* call) {
863
    static const Op& stop_fusion = Op::Get("annotation.stop_fusion");
864
    if (call->op.as<OpNode>()) {
865 866 867 868 869 870 871
      static auto fnoncomputational =
        Op::GetAttr<TNonComputational>("TNonComputational");

      if (fnoncomputational.get(Downcast<Op>(call->op), false)) {
        return ExprMutator::VisitExpr_(call);
      }

872 873 874
      // If it is a primitive op call
      // then we must have a group assignment for it already.
      CHECK(gmap_.count(call));
875 876 877
      if (call->op.same_as(stop_fusion)) {
        return ExprMutator::VisitExpr(call->args[0]);
      }
878
      auto* ret_group = gmap_.at(call)->FindRoot();
879 880
      Array<Expr> new_args = GetNewArguments(call->args, ret_group);

881 882 883 884 885 886
      auto new_call = CallNode::make(
          call->op, new_args, call->attrs, call->type_args);

      if (ret_group->root_ref == call) {
        // This is the root of the group
        // create the new call node.
887
        return MakeNewFunction(ret_group, call->checked_type(), new_call);
888 889 890
      } else {
        // This is an intermediate node of a fused function
        // simply return the new call.
891
        return std::move(new_call);
892 893 894 895 896
      }
    } else {
      return ExprMutator::VisitExpr_(call);
    }
  }
897 898 899

  Expr VisitExpr_(const TupleNode* tuple) {
    auto* ret_group = gmap_.at(tuple)->FindRoot();
900
    if (ret_group->root_ref == tuple) {
901
      return ExprMutator::VisitExpr_(tuple);
902 903
    }
    // This tuple is an intermediate node in the group
904
    Array<Expr> new_fields = GetNewArguments(tuple->fields, ret_group);
905
    return TupleNode::make(new_fields);
906 907
  }

908 909 910 911
  Expr VisitExpr_(const TupleGetItemNode* tuple_get) {
    auto* ret_group = gmap_.at(tuple_get)->FindRoot();
    auto new_tuple = GetNewArguments({tuple_get->tuple}, ret_group)[0];
    auto new_node = TupleGetItemNode::make(new_tuple, tuple_get->index);
912
    if (ret_group->root_ref == tuple_get) {
913 914 915 916 917 918 919 920 921
      if (gmap_.at(tuple_get->tuple.get())->FindRoot() != ret_group) {
        // Isolated. This case occurs when tuple is created by an Opaque op
        // e.g. multibox_transform_loc
        return ExprMutator::VisitExpr_(tuple_get);
      }
      // A new function whose output is a tuple field access
      return MakeNewFunction(ret_group, tuple_get->checked_type(), new_node);
    }
    // This is an intermediate node in the group
922
    return std::move(new_node);
923 924
  }

925
  Expr MakeNewFunction(GraphPartitioner::Group* group, Type ret_type, Expr body) {
926 927 928 929 930 931 932 933
    // If the function has no call, it is not a primitive function.
    struct HasCallVisitor : ExprVisitor {
      bool has_call = false;
      void VisitExpr_(const CallNode* op) final {
        has_call = true;
      }
    } visitor;
    visitor(body);
934 935
    const GroupInfo& ginfo = ginfo_[group];
    auto func = FunctionNode::make(ginfo.params, body, ret_type, {});
936
    func = FunctionSetAttr(func, "Primitive", tvm::Integer(visitor.has_call));
937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956
    return CallNode::make(func, ginfo.arguments, Attrs());
  }

  Array<Expr> GetNewArguments(const tvm::Array<Expr>& args,
                              GraphPartitioner::Group* current_group) {
    Array<Expr> new_args;
    for (auto arg : args) {
      auto* arg_group = gmap_.at(arg.get())->FindRoot();
      auto type = arg->checked_type();
      Expr new_arg = this->Mutate(arg);
      if (current_group != arg_group) {
        Var param = ginfo_[current_group].GetOrAllocParam(new_arg, type);
        new_args.push_back(param);
      } else {
        new_args.push_back(new_arg);
      }
    }
    return new_args;
  }

957 958
  // Debug function, dump the group assignment in text.
  void DebugDumpGroup(const Expr& body) {
959
    std::string text = AsText(body, false, [this](const Expr& expr) -> std::string {
960 961 962 963
        auto it = gmap_.find(expr.get());
        if (it == gmap_.end()) return "";
        std::ostringstream os;
        auto *group = it->second->FindRoot();
964
        os << " /* group=" << group << " */";
965 966 967 968
        return os.str();
      });
    LOG(INFO) << "Dump of group info:\n" << text;
  }
969 970
};

971
Expr FuseOps(const Expr& expr, int fuse_opt_level, const Module& module) {
972
  return FuseMutator().Transform(expr, fuse_opt_level);
973 974
}

975 976 977 978 979 980 981 982
namespace transform {

Pass FuseOps(int fuse_opt_level) {
  runtime::TypedPackedFunc<Function(Function, Module, PassContext)> pass_func =
    [=](Function f, Module m, PassContext pc) {
    int opt_level = fuse_opt_level == -1 ? pc->opt_level : fuse_opt_level;
    return Downcast<Function>(FuseOps(f, opt_level, m));
  };
983 984
  return CreateFunctionPass(pass_func, 1, "FuseOps",
                            {ir::StringImm::make("InferType")});
985 986
}

987 988 989
TVM_REGISTER_API("relay._transform.FuseOps")
.set_body_typed(FuseOps);

990 991
}  // namespace transform

992 993
}  // namespace relay
}  // namespace tvm