You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2020/03/24 17:43:57 UTC
[GitHub] [incubator-tvm] manupa-arm opened a new pull request #5143: [RELAY]
Re-wrote the Graph Partitioner to support multiple outputs
manupa-arm opened a new pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143
This PR aims to **include support for multiple outputs** in the regions that get outlined as functions for different compiler backends in the existing Partition Graph Pass. Such regions are annotated and bounded by compiler_begin and compiler_end annotation ops.
This is a required step as step 4 in [RFC - BYOC](https://discuss.tvm.ai/t/rfc-byoc-an-extended-graph-partitioning-flow/6028). Moreover, this feature is requested in prior discussions such as [Improved graph partitioning algorithm](https://discuss.tvm.ai/t/relay-improved-graph-partitioning-algorithm/5830)
This PR uses the utility functions provided by the [AnnotationRegionSet PR](https://github.com/apache/incubator-tvm/pull/5030).
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] zhiics commented on a change in pull request #5143:
[RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
zhiics commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r397980929
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -364,36 +388,114 @@ class Partitioner : public ExprMutator {
IRModule Partition() {
auto glob_funcs = module_->functions;
- for (const auto& pair : glob_funcs) {
- if (auto* fn = pair.second.as<FunctionNode>()) {
+ for (const auto &pair : glob_funcs) {
+ if (auto *fn = pair.second.as<FunctionNode>()) {
auto func = GetRef<Function>(fn);
func = Function(func->params,
- VisitExpr(func->body),
- func->ret_type,
- func->type_params,
- func->attrs);
+ VisitExpr(func->body),
+ func->ret_type,
+ func->type_params,
+ func->attrs);
module_->Update(pair.first, func);
}
}
return module_;
}
private:
- int var_id_{0};
- int subgraph_id_{0};
- std::unordered_set<std::shared_ptr<Subgraph>> subgraphs_;
+ /*!
+ * \brief Get the region an expression belongs to
Review comment:
align "*", same to the other functions below.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] zhiics commented on a change in pull request #5143:
[RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
zhiics commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r397975848
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -102,61 +89,71 @@ class AnnotationChecker : public ExprVisitor {
bool found_end_{false};
};
-/*! \brief This class partitions the expr labeled with begin and end annoations
+/*! \brief This class partitions the expr labeled with begin and end annotations
* into function containing multiple regions. Each region is labeled with
* a compiler attribute so that it will be handled by any compilers that are not
* in the TVM stack.
*
- * TODO(@zhiics) This following algorithm is not adequate to handle all cases,
- * i.e. multiple `compiler_end` nodes.
+ * Input : A Relay module that have functions with disjoint annotated regions
+ * using compiler_begin and compiler_end. There could be multiple outputs.
+ *
+ * Output : A Relay module with global functions for such disjoint annotated regions
+ * with calls inserted at the respective location
+ *
+ * Dependencies : RegionSet Utility class.
+ *
+ * Methodology :
+ * 1) The RegionSet utility class is able to construct a collection of
+ * nodes that are bound by a given annotation -- here we use compiler_begin
+ * and compiler_end
+ * 2) Initially, for each function in the module RegionSets are populated.
+ * 3) Then, Vistor pass is traversed until a compiler_end node is encountered
+ * that belongs to a "region".
+ * 4) When the first compiler_end of a given annotated region is found, a function is
+ * formed and inserted.
+ * a) if the region has multiple outputs, a Tuple node (capturing all outputs)
+ * is returned.
+ * 5) Thereafter, if we encounter an another output of the same annotated region,
+ * it is important to note that the function is already formed. Therefore, it will
+ * lookup the function and add a TupleGetItemNode.
+ * a) We will use the location index of "rets" of each "Region" of RegionSet
+ * as TupleGetItemNode index.
+ * 6) Therefore, functions will be created for all annotated regions. The name for each
+ * global function is created using "Region" id and the compiler name.
+ *
+ * Expected Usecase :
+ * This pass is intended to run as the last pass in a series of passes as follows :
+ * 1) Annotate Supported Single Ops - annotated each single op with supported backends.
Review comment:
1) and 2) are inaccurate, we don't have these two annotations
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] zhiics commented on a change in pull request #5143:
[RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
zhiics commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r397982123
##########
File path: tests/python/relay/test_pass_partition_graph.py
##########
@@ -678,6 +678,81 @@ def expected():
check_result(mod, {"y": y_data}, (8, 8), np.log(np_add))
+def test_multiple_outputs():
Review comment:
Can you add a test case for something like `conv2d + batch_norm`?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] manupa-arm commented on a change in pull request
#5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r399295536
##########
File path: tests/python/relay/test_pass_partition_graph.py
##########
@@ -678,6 +678,81 @@ def expected():
check_result(mod, {"y": y_data}, (8, 8), np.log(np_add))
+def test_multiple_outputs():
Review comment:
Added the testcase
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] zhiics commented on a change in pull request #5143:
[RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
zhiics commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r398760547
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -365,35 +374,113 @@ class Partitioner : public ExprMutator {
IRModule Partition() {
auto glob_funcs = module_->functions;
for (const auto& pair : glob_funcs) {
- if (auto* fn = pair.second.as<FunctionNode>()) {
+ if (auto *fn = pair.second.as<FunctionNode>()) {
auto func = GetRef<Function>(fn);
func = Function(func->params,
- VisitExpr(func->body),
- func->ret_type,
- func->type_params,
- func->attrs);
+ VisitExpr(func->body),
+ func->ret_type,
+ func->type_params,
+ func->attrs);
module_->Update(pair.first, func);
}
}
return module_;
}
private:
- int var_id_{0};
- int subgraph_id_{0};
- std::unordered_set<std::shared_ptr<Subgraph>> subgraphs_;
+ /*!
+ * \brief Get the region an expression belongs to
+ * if its in a region.
+ */
+ AnnotatedRegion GetRegion(const Expr& e) {
+ for (auto sg_set_it : regions_sets_) {
+ auto sg_set = sg_set_it.first;
+ AnnotatedRegion sg = sg_set->GetRegion(e);
+ if (sg.defined()) {
+ return sg;
+ }
+ }
+ return AnnotatedRegion(nullptr);
+ }
+
+ /*!
+ * \brief Get the function an expression belongs to
+ * if its in a region.
+ */
+ BaseFunc GetFunc(const Expr& e) {
+ for (auto sg_set_it : regions_sets_) {
+ auto sg_set = sg_set_it.first;
+ auto func = sg_set_it.second;
+
+ AnnotatedRegion sg = sg_set->GetRegion(e);
+ if (sg.defined()) {
+ return func;
+ }
+ }
+ return BaseFunc(nullptr);
+ }
+
+ /*!
Review comment:
align
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] zhiics commented on a change in pull request #5143:
[RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
zhiics commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r397973058
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -36,36 +36,23 @@
#include <tvm/relay/expr_functor.h>
#include <tvm/relay/transform.h>
-#include <string>
+#include <utility>
#include <unordered_map>
#include <unordered_set>
-#include <utility>
#include <vector>
#include "../backend/utils.h"
+#include "../analysis/annotated_region_set.h"
+
namespace tvm {
namespace relay {
namespace partitioning {
// Cache compiler_begin and compiler_end annotation ops for equivalence check to
// reduce registry lookup overhead.
-static const Op& compiler_begin_op = Op::Get("annotation.compiler_begin");
-static const Op& compiler_end_op = Op::Get("annotation.compiler_end");
-
-/*!
- * \brief The subgraph properties for partitioning.
- */
-struct Subgraph {
- /*! \brief The subgraph ID. */
- int id;
-
- /*! \brief The input arguments of this subgraph. */
- std::vector<std::pair<Var, Expr>> args;
-
- /*! \brief Nodes in this subgraph. */
- std::unordered_set<Expr, ObjectHash, ObjectEqual> nodes;
-};
+static const Op &compiler_begin_op = Op::Get("annotation.compiler_begin");
Review comment:
keep the original code style: `static const Op&`
Same for all the other places in the rest of PR
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] manupa-arm commented on a change in pull request
#5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r398520226
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -36,36 +36,23 @@
#include <tvm/relay/expr_functor.h>
#include <tvm/relay/transform.h>
-#include <string>
+#include <utility>
#include <unordered_map>
#include <unordered_set>
-#include <utility>
#include <vector>
#include "../backend/utils.h"
+#include "../analysis/annotated_region_set.h"
+
namespace tvm {
namespace relay {
namespace partitioning {
// Cache compiler_begin and compiler_end annotation ops for equivalence check to
// reduce registry lookup overhead.
-static const Op& compiler_begin_op = Op::Get("annotation.compiler_begin");
-static const Op& compiler_end_op = Op::Get("annotation.compiler_end");
-
-/*!
- * \brief The subgraph properties for partitioning.
- */
-struct Subgraph {
- /*! \brief The subgraph ID. */
- int id;
-
- /*! \brief The input arguments of this subgraph. */
- std::vector<std::pair<Var, Expr>> args;
-
- /*! \brief Nodes in this subgraph. */
- std::unordered_set<Expr, ObjectHash, ObjectEqual> nodes;
-};
+static const Op &compiler_begin_op = Op::Get("annotation.compiler_begin");
Review comment:
Fixed. Prolly feedback this to lint?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] manupa-arm commented on a change in pull request
#5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r399294950
##########
File path: tests/python/relay/test_pass_partition_graph.py
##########
@@ -678,6 +678,81 @@ def expected():
check_result(mod, {"y": y_data}, (8, 8), np.log(np_add))
+def test_multiple_outputs():
+ def create_merged_graph():
+ data = relay.var('data', shape=(10, 10))
+
+ cb_1 = compiler_begin(data, 'test_target')
+ O_1 = relay.abs(cb_1)
+ ce_2 = compiler_end(O_1, 'test_target')
+ O_2 = relay.nn.relu(O_1)
+ ce_3 = compiler_end(O_2, 'test_target')
+
+ X = relay.tanh(ce_2)
+
+ cb_3 = compiler_begin(ce_3, 'test_target')
+ cb_4 = compiler_begin(X, 'test_target')
+ O_3 = relay.add(cb_3, cb_4)
+ ce_4 = compiler_end(O_3, 'test_target')
+
+ func = relay.Function([data], ce_4)
+ return func
+
+ def expected():
+ mod = tvm.IRModule()
+
+ # function 1
+ f1_cb1 = relay.var('test_target_1_i0', shape=(10, 10))
+ f1_O_1 = relay.abs(f1_cb1)
+ f1_O_2 = relay.nn.relu(f1_O_1)
+ f1_out = relay.Tuple((f1_O_2,f1_O_1))
+ func1 = relay.Function([f1_cb1], f1_out)
+
+ func1 = func1.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
+ func1 = func1.with_attr("Inline", tvm.tir.IntImm("int32", 1))
+ func1 = func1.with_attr("Compiler",
+ tvm.tir.StringImm("test_target"))
+ func1 = func1.with_attr("ExternalSymbol",
+ tvm.tir.StringImm("test_target_1"))
+ gv1 = relay.GlobalVar("test_target_1")
+ mod[gv1] = func1
+
+ # function 0
+ f2_cb3 = relay.var('test_target_0_i0', shape=(10, 10))
+ f2_cb4 = relay.var('test_target_0_i1', shape=(10, 10))
+ f2_O_3 = relay.add(f2_cb3, f2_cb4)
+ func0 = relay.Function([f2_cb3, f2_cb4], f2_O_3)
+
+ func0 = func0.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
+ func0 = func0.with_attr("Inline", tvm.tir.IntImm("int32", 1))
+ func0 = func0.with_attr("Compiler",
+ tvm.tir.StringImm("test_target"))
+ func0 = func0.with_attr("ExternalSymbol",
+ tvm.tir.StringImm("test_target_0"))
+ gv0 = relay.GlobalVar("test_target_0")
+ mod[gv0] = func0
+
+ # body
+ data = relay.var('data', shape=(10, 10))
+ tuple_out = gv1(data)
+ ce_2 = relay.TupleGetItem(tuple_out, 1)
+ ce_3 = relay.TupleGetItem(tuple_out, 0)
+
+ X = relay.tanh(ce_2)
+ ce_4 = gv0(ce_3, X)
+ func = relay.Function([data], ce_4)
+ mod["main"] = func
+
+ return mod
+
+ # print(create_merged_graph())
Review comment:
done
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] manupa-arm commented on a change in pull request
#5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r398761014
##########
File path: tests/python/relay/test_pass_partition_graph.py
##########
@@ -678,6 +678,81 @@ def expected():
check_result(mod, {"y": y_data}, (8, 8), np.log(np_add))
+def test_multiple_outputs():
Review comment:
But isnt it returned as a tuple already ? Looks like a single output to me. Unless, you do some procesing afterwards in the region. Is that what you imply?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] comaniac commented on a change in pull request
#5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r398768598
##########
File path: tests/python/relay/test_pass_partition_graph.py
##########
@@ -678,6 +678,81 @@ def expected():
check_result(mod, {"y": y_data}, (8, 8), np.log(np_add))
+def test_multiple_outputs():
Review comment:
Oh yes. The case in my mind is `conv2d->BN->conv2d`. In this case, the region output would be a new tuple of the rest 2 BN outputs and 1 conv2d output.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] zhiics commented on a change in pull request #5143:
[RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
zhiics commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r398760603
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -365,35 +374,113 @@ class Partitioner : public ExprMutator {
IRModule Partition() {
auto glob_funcs = module_->functions;
for (const auto& pair : glob_funcs) {
- if (auto* fn = pair.second.as<FunctionNode>()) {
+ if (auto *fn = pair.second.as<FunctionNode>()) {
auto func = GetRef<Function>(fn);
func = Function(func->params,
- VisitExpr(func->body),
- func->ret_type,
- func->type_params,
- func->attrs);
+ VisitExpr(func->body),
+ func->ret_type,
+ func->type_params,
+ func->attrs);
module_->Update(pair.first, func);
}
}
return module_;
}
private:
- int var_id_{0};
- int subgraph_id_{0};
- std::unordered_set<std::shared_ptr<Subgraph>> subgraphs_;
+ /*!
+ * \brief Get the region an expression belongs to
+ * if its in a region.
+ */
+ AnnotatedRegion GetRegion(const Expr& e) {
+ for (auto sg_set_it : regions_sets_) {
+ auto sg_set = sg_set_it.first;
+ AnnotatedRegion sg = sg_set->GetRegion(e);
+ if (sg.defined()) {
+ return sg;
+ }
+ }
+ return AnnotatedRegion(nullptr);
+ }
+
+ /*!
Review comment:
align
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] zhiics commented on a change in pull request #5143:
[RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
zhiics commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r398761325
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -364,36 +388,114 @@ class Partitioner : public ExprMutator {
IRModule Partition() {
auto glob_funcs = module_->functions;
- for (const auto& pair : glob_funcs) {
- if (auto* fn = pair.second.as<FunctionNode>()) {
+ for (const auto &pair : glob_funcs) {
+ if (auto *fn = pair.second.as<FunctionNode>()) {
auto func = GetRef<Function>(fn);
func = Function(func->params,
- VisitExpr(func->body),
- func->ret_type,
- func->type_params,
- func->attrs);
+ VisitExpr(func->body),
+ func->ret_type,
+ func->type_params,
+ func->attrs);
module_->Update(pair.first, func);
}
}
return module_;
}
private:
- int var_id_{0};
- int subgraph_id_{0};
- std::unordered_set<std::shared_ptr<Subgraph>> subgraphs_;
+ /*!
+ * \brief Get the region an expression belongs to
Review comment:
still not aligned, you want something like:
```c++
/*!
* \brief
*/
```
Same to the other places as well
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] manupa-arm commented on a change in pull request
#5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r398590002
##########
File path: tests/python/relay/test_pass_partition_graph.py
##########
@@ -678,6 +678,81 @@ def expected():
check_result(mod, {"y": y_data}, (8, 8), np.log(np_add))
+def test_multiple_outputs():
Review comment:
Not sure I get you; Can you elaborate more as to how multiple outputs in regions are exercised in conv2d + batch_norm?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] manupa-arm commented on a change in pull request
#5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r399295343
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -365,35 +374,113 @@ class Partitioner : public ExprMutator {
IRModule Partition() {
auto glob_funcs = module_->functions;
for (const auto& pair : glob_funcs) {
- if (auto* fn = pair.second.as<FunctionNode>()) {
+ if (auto *fn = pair.second.as<FunctionNode>()) {
auto func = GetRef<Function>(fn);
func = Function(func->params,
- VisitExpr(func->body),
- func->ret_type,
- func->type_params,
- func->attrs);
+ VisitExpr(func->body),
+ func->ret_type,
+ func->type_params,
+ func->attrs);
module_->Update(pair.first, func);
}
}
return module_;
}
private:
- int var_id_{0};
- int subgraph_id_{0};
- std::unordered_set<std::shared_ptr<Subgraph>> subgraphs_;
+ /*!
+ * \brief Get the region an expression belongs to
+ * if its in a region.
+ */
+ AnnotatedRegion GetRegion(const Expr& e) {
+ for (auto sg_set_it : regions_sets_) {
+ auto sg_set = sg_set_it.first;
+ AnnotatedRegion sg = sg_set->GetRegion(e);
+ if (sg.defined()) {
+ return sg;
+ }
+ }
+ return AnnotatedRegion(nullptr);
+ }
+
+ /*!
Review comment:
Hopefully done!
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] manupa-arm commented on a change in pull request
#5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r398522710
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -102,61 +89,71 @@ class AnnotationChecker : public ExprVisitor {
bool found_end_{false};
};
-/*! \brief This class partitions the expr labeled with begin and end annoations
+/*! \brief This class partitions the expr labeled with begin and end annotations
* into function containing multiple regions. Each region is labeled with
* a compiler attribute so that it will be handled by any compilers that are not
* in the TVM stack.
*
- * TODO(@zhiics) This following algorithm is not adequate to handle all cases,
- * i.e. multiple `compiler_end` nodes.
+ * Input : A Relay module that have functions with disjoint annotated regions
+ * using compiler_begin and compiler_end. There could be multiple outputs.
+ *
+ * Output : A Relay module with global functions for such disjoint annotated regions
+ * with calls inserted at the respective location
+ *
+ * Dependencies : RegionSet Utility class.
+ *
+ * Methodology :
+ * 1) The RegionSet utility class is able to construct a collection of
+ * nodes that are bound by a given annotation -- here we use compiler_begin
+ * and compiler_end
+ * 2) Initially, for each function in the module RegionSets are populated.
+ * 3) Then, Vistor pass is traversed until a compiler_end node is encountered
+ * that belongs to a "region".
+ * 4) When the first compiler_end of a given annotated region is found, a function is
+ * formed and inserted.
+ * a) if the region has multiple outputs, a Tuple node (capturing all outputs)
+ * is returned.
+ * 5) Thereafter, if we encounter an another output of the same annotated region,
+ * it is important to note that the function is already formed. Therefore, it will
+ * lookup the function and add a TupleGetItemNode.
+ * a) We will use the location index of "rets" of each "Region" of RegionSet
+ * as TupleGetItemNode index.
+ * 6) Therefore, functions will be created for all annotated regions. The name for each
+ * global function is created using "Region" id and the compiler name.
+ *
+ * Expected Usecase :
+ * This pass is intended to run as the last pass in a series of passes as follows :
+ * 1) Annotate Supported Single Ops - annotated each single op with supported backends.
Review comment:
Removed the use-case segment. This is not needed now; as things are still in discussion as how this pass would potentially fit in the bigger picture.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] comaniac commented on a change in pull request
#5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r398747412
##########
File path: tests/python/relay/test_pass_partition_graph.py
##########
@@ -678,6 +678,81 @@ def expected():
check_result(mod, {"y": y_data}, (8, 8), np.log(np_add))
+def test_multiple_outputs():
Review comment:
batch_norm has 3 outputs.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] zhiics commented on a change in pull request #5143:
[RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
zhiics commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r397979319
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -165,101 +162,142 @@ class Partitioner : public ExprMutator {
// Traverse the rest graph.
auto input_expr = VisitExpr(call->args[0]);
- // Replace the begin annotation with an external call input variable.
- auto compiler_attrs = call->attrs.as<CompilerAttrs>();
+ AnnotatedRegion sg = GetRegion(GetRef<Call>(call));
+ int index = GetArgIdx(sg, GetRef<Call>(call));
+ CHECK_NE(index, -1);
// The type of the created variable is the same as the compiler_begin
// node.
- auto var = Var(compiler_attrs->compiler + "_input" + std::to_string(var_id_++),
- call->checked_type_);
-
- // Find the corresponding subgraph and add the argument.
- auto subgraph = GetSubgraph(GetRef<Call>(call));
- if (!subgraph) {
- throw Error(ErrorBuilder()
- << "Cannot find the corresponding subgraph for start annotation:\n"
- << AsText(GetRef<Call>(call), false));
- }
- subgraph->args.push_back({var, input_expr});
+ std::string target = call->attrs.as<CompilerAttrs>()->compiler;
+ std::string varname = target + "_" + std::to_string(sg->GetID())
+ + "_i" + std::to_string(index);
+ auto var = Var(varname, GetRef<Call>(call)->checked_type_);
+
+ auto cand = std::make_pair(var, input_expr);
+ if (std::find(region_args[sg].begin(),
+ region_args[sg].end(), cand) == region_args[sg].end()) {
+ region_args[sg].push_back(cand);
+ }
+
return std::move(var);
} else {
CHECK_EQ(call->op, compiler_end_op);
// The annotation node is inserted on edge so it must have only one argument.
CHECK_EQ(call->args.size(), 1U);
- auto compiler_attrs = call->attrs.as<CompilerAttrs>();
+ AnnotatedRegion region = GetRegion(GetRef<Call>(call));
- // Check if the argument already belongs to an existing subgraph
- auto subgraph = GetSubgraph(call->args[0]);
- if (!subgraph) {
- auto ret = this->subgraphs_.emplace(std::make_shared<Subgraph>());
- subgraph = *ret.first;
- subgraph->nodes.insert(call->args[0]);
- subgraph->id = this->subgraph_id_++;
- }
- subgraph->nodes.insert(GetRef<Call>(call));
+ // TODO(@manupa-arm) : need to use the parent function (to which region
+ // belongs to) name/key for the funtions that are created
+ BaseFunc f = GetFunc(GetRef<Call>(call));
// Traverse subgraph inputs.
auto input = VisitExpr(call->args[0]);
- Array<Var> params;
- Array<Expr> args;
- std::unordered_map<std::string, runtime::NDArray> params_bind;
-
- // The subgraph may be merged so we need to update it again.
- subgraph = GetSubgraph(GetRef<Call>(call));
- CHECK(subgraph);
-
- // Record the constants for propagation.
- for (auto pair : subgraph->args) {
- params.push_back(pair.first);
- if (const auto* cn = pair.second.as<ConstantNode>()) {
- params_bind[pair.first->name_hint()] = cn->data;
+ CHECK(region.defined()) << "Region not defined for " << GetRef<Call>(call);
+ // functions are created for each annotated regions,
+ // when their first output is encountered.
+ // If multiple outputs are there, a tuple node is inserted at the end.
+ // region_function_calls is map that maintains
+ // (each annotated regions) --> created function
+
+ if (region_function_calls.find(region) != region_function_calls.end()) {
+ // This section is executed only if there are multiple outputs in the region
+ // Thus, the function is always created and at the end there would be a tuple node
+ // Therefore, we insert a tuple get item node.
+
+ // Use the already created tuple node
+ auto sg_call = region_function_calls[region];
+ int index = GetRetIdx(region, GetRef<Call>(call));
+ CHECK_NE(index, -1);
+
+ auto tuple_get_item_ = TupleGetItem(sg_call, index);
+ tuple_get_item_->checked_type_ = GetRef<Call>(call)->args[0]->checked_type_;
+ return std::move(tuple_get_item_);
+ } else {
+ // First time this region is encountered in the traversal
+ // Creating the function
+
+ Array<Expr> fields;
+
+ for (auto ret : region->GetOutputs()) {
+ auto ret_expr = VisitExpr(Downcast<Call>(ret)->args[0]);
+ fields.push_back(ret_expr);
+ }
+ int index = GetRetIdx(region, GetRef<Call>(call));
+ CHECK_NE(index, -1);
+
+ Array<Var> params;
+ Array<Expr> param_expr;
+ std::unordered_map<std::string, runtime::NDArray> params_bind;
+
+ for (auto pair : region_args[region]) {
+ params.push_back(pair.first);
+ if (const auto* cn = pair.second.as<ConstantNode>()) {
+ params_bind[pair.first->name_hint()] = cn->data;
+ } else {
+ param_expr.push_back(pair.second);
+ }
+ }
+
+ Function global_region_func;
+ if (region->GetOutputs().size() == 1) {
+ // If there are only a single output; no need to add a tuple
+ global_region_func = Function(params, fields[0],
+ call->args[0]->checked_type_, {}, DictAttrs());
} else {
- args.push_back(pair.second);
+ auto tuple = Tuple(fields);
+ global_region_func = Function(params, tuple, tuple->checked_type_, {}, DictAttrs());
+ }
+
+ std::string target = call->attrs.as<CompilerAttrs>()->compiler;
+ std::string name = target + "_" + std::to_string(region->GetID());
+
+ global_region_func = WithAttr(std::move(global_region_func), attr::kExternalSymbol,
+ tir::StringImmNode::make(name));
+ global_region_func = WithAttr(std::move(global_region_func), attr::kPrimitive,
+ tvm::Integer(1));
+ global_region_func = WithAttr(std::move(global_region_func), attr::kCompiler,
+ tvm::tir::StringImmNode::make(target));
+ global_region_func = WithAttr(std::move(global_region_func), attr::kInline,
+ tvm::Integer(1));
+
+ // Constant propagation
+ if (!params_bind.empty()) {
+ global_region_func = backend::BindParamsByName(global_region_func, params_bind);
}
- }
- auto subgraph_func =
- Function(params, input, call->checked_type_, {});
-
- std::string name = compiler_attrs->compiler + "_" + std::to_string(subgraph->id);
- subgraph_func =
- WithAttr(std::move(subgraph_func), attr::kExternalSymbol, tir::StringImmNode::make(name));
- subgraph_func =
- WithAttr(std::move(subgraph_func), attr::kPrimitive, tvm::Integer(1));
- subgraph_func =
- WithAttr(std::move(subgraph_func), attr::kCompiler,
- tvm::tir::StringImmNode::make(compiler_attrs->compiler));
- subgraph_func =
- WithAttr(std::move(subgraph_func), attr::kInline, tvm::Integer(1));
-
- // Constant propagation
- if (!params_bind.empty()) {
- subgraph_func = backend::BindParamsByName(subgraph_func, params_bind);
+ std::string fname = name;
+ CHECK(!module_->ContainGlobalVar(fname))
+ << "Global function " << fname << " already exists";
+ // Create a global function and add it to the IRModule for the region.
+ // This way we lift the functions that should be handled by external
+ // codegen to the module scope and rely on the pass manager to prevent relay
+ // function level passes (i.e. simplify inference and fusion) optimizing it.
+ GlobalVar glob_func(fname);
+ module_->Add(glob_func, global_region_func);
+
+ // The return type of callnode is the same as the type of the
+ // compiler_end node.
+ auto ret = Call(glob_func, param_expr);
+ region_function_calls[region] = ret;
+
+ if (region->GetOutputs().size() == 1) {
+ // If there is only a single output; no need to add a tuplegetitem node
+ return Call(glob_func, param_expr);
Review comment:
no need to create another call, just `return ret`
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] zhiics commented on issue #5143: [RELAY] Re-wrote
the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
zhiics commented on issue #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#issuecomment-606795653
Thanks @manupa-arm @comaniac
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on a change in pull request #5143:
[RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r401568941
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -36,36 +36,23 @@
#include <tvm/relay/expr_functor.h>
#include <tvm/relay/transform.h>
-#include <string>
+#include <utility>
#include <unordered_map>
#include <unordered_set>
-#include <utility>
#include <vector>
#include "../backend/utils.h"
+#include "../analysis/annotated_region_set.h"
+
namespace tvm {
namespace relay {
namespace partitioning {
// Cache compiler_begin and compiler_end annotation ops for equivalence check to
// reduce registry lookup overhead.
-static const Op& compiler_begin_op = Op::Get("annotation.compiler_begin");
-static const Op& compiler_end_op = Op::Get("annotation.compiler_end");
-
-/*!
- * \brief The subgraph properties for partitioning.
- */
-struct Subgraph {
- /*! \brief The subgraph ID. */
- int id;
-
- /*! \brief The input arguments of this subgraph. */
- std::vector<std::pair<Var, Expr>> args;
-
- /*! \brief Nodes in this subgraph. */
- std::unordered_set<Expr, ObjectHash, ObjectEqual> nodes;
-};
+static const Op &compiler_begin_op = Op::Get("annotation.compiler_begin");
Review comment:
There are still many places where the style is violated. I recommend running clang-format from your editor. We have our own .clang-format file at the root dir.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] manupa-arm commented on a change in pull request
#5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r399295612
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -364,36 +388,114 @@ class Partitioner : public ExprMutator {
IRModule Partition() {
auto glob_funcs = module_->functions;
- for (const auto& pair : glob_funcs) {
- if (auto* fn = pair.second.as<FunctionNode>()) {
+ for (const auto &pair : glob_funcs) {
+ if (auto *fn = pair.second.as<FunctionNode>()) {
auto func = GetRef<Function>(fn);
func = Function(func->params,
- VisitExpr(func->body),
- func->ret_type,
- func->type_params,
- func->attrs);
+ VisitExpr(func->body),
+ func->ret_type,
+ func->type_params,
+ func->attrs);
module_->Update(pair.first, func);
}
}
return module_;
}
private:
- int var_id_{0};
- int subgraph_id_{0};
- std::unordered_set<std::shared_ptr<Subgraph>> subgraphs_;
+ /*!
+ * \brief Get the region an expression belongs to
Review comment:
Hopefully done!
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on a change in pull request #5143:
[RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r401568941
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -36,36 +36,23 @@
#include <tvm/relay/expr_functor.h>
#include <tvm/relay/transform.h>
-#include <string>
+#include <utility>
#include <unordered_map>
#include <unordered_set>
-#include <utility>
#include <vector>
#include "../backend/utils.h"
+#include "../analysis/annotated_region_set.h"
+
namespace tvm {
namespace relay {
namespace partitioning {
// Cache compiler_begin and compiler_end annotation ops for equivalence check to
// reduce registry lookup overhead.
-static const Op& compiler_begin_op = Op::Get("annotation.compiler_begin");
-static const Op& compiler_end_op = Op::Get("annotation.compiler_end");
-
-/*!
- * \brief The subgraph properties for partitioning.
- */
-struct Subgraph {
- /*! \brief The subgraph ID. */
- int id;
-
- /*! \brief The input arguments of this subgraph. */
- std::vector<std::pair<Var, Expr>> args;
-
- /*! \brief Nodes in this subgraph. */
- std::unordered_set<Expr, ObjectHash, ObjectEqual> nodes;
-};
+static const Op &compiler_begin_op = Op::Get("annotation.compiler_begin");
Review comment:
There are still many places the style is violated. I recommend running clang-format from your editor. We have our own .clang-format file at the root dir.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] manupa-arm commented on a change in pull request
#5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r399295386
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -365,35 +374,113 @@ class Partitioner : public ExprMutator {
IRModule Partition() {
auto glob_funcs = module_->functions;
for (const auto& pair : glob_funcs) {
- if (auto* fn = pair.second.as<FunctionNode>()) {
+ if (auto *fn = pair.second.as<FunctionNode>()) {
auto func = GetRef<Function>(fn);
func = Function(func->params,
- VisitExpr(func->body),
- func->ret_type,
- func->type_params,
- func->attrs);
+ VisitExpr(func->body),
+ func->ret_type,
+ func->type_params,
+ func->attrs);
module_->Update(pair.first, func);
}
}
return module_;
}
private:
- int var_id_{0};
- int subgraph_id_{0};
- std::unordered_set<std::shared_ptr<Subgraph>> subgraphs_;
+ /*!
+ * \brief Get the region an expression belongs to
+ * if its in a region.
+ */
+ AnnotatedRegion GetRegion(const Expr& e) {
+ for (auto sg_set_it : regions_sets_) {
+ auto sg_set = sg_set_it.first;
+ AnnotatedRegion sg = sg_set->GetRegion(e);
+ if (sg.defined()) {
+ return sg;
+ }
+ }
+ return AnnotatedRegion(nullptr);
+ }
+
+ /*!
+ * \brief Get the function an expression belongs to
+ * if its in a region.
+ */
+ BaseFunc GetFunc(const Expr& e) {
+ for (auto sg_set_it : regions_sets_) {
+ auto sg_set = sg_set_it.first;
+ auto func = sg_set_it.second;
+
+ AnnotatedRegion sg = sg_set->GetRegion(e);
+ if (sg.defined()) {
+ return func;
+ }
+ }
+ return BaseFunc(nullptr);
+ }
+
+ /*!
Review comment:
Hopefully done!
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] manupa-arm commented on a change in pull request
#5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r401638617
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -36,36 +36,23 @@
#include <tvm/relay/expr_functor.h>
#include <tvm/relay/transform.h>
-#include <string>
+#include <utility>
#include <unordered_map>
#include <unordered_set>
-#include <utility>
#include <vector>
#include "../backend/utils.h"
+#include "../analysis/annotated_region_set.h"
+
namespace tvm {
namespace relay {
namespace partitioning {
// Cache compiler_begin and compiler_end annotation ops for equivalence check to
// reduce registry lookup overhead.
-static const Op& compiler_begin_op = Op::Get("annotation.compiler_begin");
-static const Op& compiler_end_op = Op::Get("annotation.compiler_end");
-
-/*!
- * \brief The subgraph properties for partitioning.
- */
-struct Subgraph {
- /*! \brief The subgraph ID. */
- int id;
-
- /*! \brief The input arguments of this subgraph. */
- std::vector<std::pair<Var, Expr>> args;
-
- /*! \brief Nodes in this subgraph. */
- std::unordered_set<Expr, ObjectHash, ObjectEqual> nodes;
-};
+static const Op &compiler_begin_op = Op::Get("annotation.compiler_begin");
Review comment:
@masahi thanks for tip! fixed this in this follow-up [PR](https://github.com/apache/incubator-tvm/pull/5202)
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] zhiics merged pull request #5143: [RELAY] Re-wrote
the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
zhiics merged pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] masahi commented on a change in pull request #5143:
[RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r401568941
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -36,36 +36,23 @@
#include <tvm/relay/expr_functor.h>
#include <tvm/relay/transform.h>
-#include <string>
+#include <utility>
#include <unordered_map>
#include <unordered_set>
-#include <utility>
#include <vector>
#include "../backend/utils.h"
+#include "../analysis/annotated_region_set.h"
+
namespace tvm {
namespace relay {
namespace partitioning {
// Cache compiler_begin and compiler_end annotation ops for equivalence check to
// reduce registry lookup overhead.
-static const Op& compiler_begin_op = Op::Get("annotation.compiler_begin");
-static const Op& compiler_end_op = Op::Get("annotation.compiler_end");
-
-/*!
- * \brief The subgraph properties for partitioning.
- */
-struct Subgraph {
- /*! \brief The subgraph ID. */
- int id;
-
- /*! \brief The input arguments of this subgraph. */
- std::vector<std::pair<Var, Expr>> args;
-
- /*! \brief Nodes in this subgraph. */
- std::unordered_set<Expr, ObjectHash, ObjectEqual> nodes;
-};
+static const Op &compiler_begin_op = Op::Get("annotation.compiler_begin");
Review comment:
There are still some places the style is violated. I recommend running clang-format from your editor. We have our own .clang-format file at the root dir.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] manupa-arm commented on a change in pull request
#5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r398523020
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -364,36 +388,114 @@ class Partitioner : public ExprMutator {
IRModule Partition() {
auto glob_funcs = module_->functions;
- for (const auto& pair : glob_funcs) {
- if (auto* fn = pair.second.as<FunctionNode>()) {
+ for (const auto &pair : glob_funcs) {
+ if (auto *fn = pair.second.as<FunctionNode>()) {
auto func = GetRef<Function>(fn);
func = Function(func->params,
- VisitExpr(func->body),
- func->ret_type,
- func->type_params,
- func->attrs);
+ VisitExpr(func->body),
+ func->ret_type,
+ func->type_params,
+ func->attrs);
module_->Update(pair.first, func);
}
}
return module_;
}
private:
- int var_id_{0};
- int subgraph_id_{0};
- std::unordered_set<std::shared_ptr<Subgraph>> subgraphs_;
+ /*!
+ * \brief Get the region an expression belongs to
Review comment:
fixed
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] manupa-arm commented on a change in pull request
#5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r398590002
##########
File path: tests/python/relay/test_pass_partition_graph.py
##########
@@ -678,6 +678,81 @@ def expected():
check_result(mod, {"y": y_data}, (8, 8), np.log(np_add))
+def test_multiple_outputs():
Review comment:
Not sure that I get you; Can you elaborate more as to how multiple outputs in regions are exercised in conv2d + batch_norm?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] zhiics commented on a change in pull request #5143:
[RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
zhiics commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r398772992
##########
File path: tests/python/relay/test_pass_partition_graph.py
##########
@@ -678,6 +678,81 @@ def expected():
check_result(mod, {"y": y_data}, (8, 8), np.log(np_add))
+def test_multiple_outputs():
+ def create_merged_graph():
+ data = relay.var('data', shape=(10, 10))
+
+ cb_1 = compiler_begin(data, 'test_target')
+ O_1 = relay.abs(cb_1)
+ ce_2 = compiler_end(O_1, 'test_target')
+ O_2 = relay.nn.relu(O_1)
+ ce_3 = compiler_end(O_2, 'test_target')
+
+ X = relay.tanh(ce_2)
+
+ cb_3 = compiler_begin(ce_3, 'test_target')
+ cb_4 = compiler_begin(X, 'test_target')
+ O_3 = relay.add(cb_3, cb_4)
+ ce_4 = compiler_end(O_3, 'test_target')
+
+ func = relay.Function([data], ce_4)
+ return func
+
+ def expected():
+ mod = tvm.IRModule()
+
+ # function 1
+ f1_cb1 = relay.var('test_target_1_i0', shape=(10, 10))
+ f1_O_1 = relay.abs(f1_cb1)
+ f1_O_2 = relay.nn.relu(f1_O_1)
+ f1_out = relay.Tuple((f1_O_2,f1_O_1))
+ func1 = relay.Function([f1_cb1], f1_out)
+
+ func1 = func1.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
+ func1 = func1.with_attr("Inline", tvm.tir.IntImm("int32", 1))
+ func1 = func1.with_attr("Compiler",
+ tvm.tir.StringImm("test_target"))
+ func1 = func1.with_attr("ExternalSymbol",
+ tvm.tir.StringImm("test_target_1"))
+ gv1 = relay.GlobalVar("test_target_1")
+ mod[gv1] = func1
+
+ # function 0
+ f2_cb3 = relay.var('test_target_0_i0', shape=(10, 10))
+ f2_cb4 = relay.var('test_target_0_i1', shape=(10, 10))
+ f2_O_3 = relay.add(f2_cb3, f2_cb4)
+ func0 = relay.Function([f2_cb3, f2_cb4], f2_O_3)
+
+ func0 = func0.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
+ func0 = func0.with_attr("Inline", tvm.tir.IntImm("int32", 1))
+ func0 = func0.with_attr("Compiler",
+ tvm.tir.StringImm("test_target"))
+ func0 = func0.with_attr("ExternalSymbol",
+ tvm.tir.StringImm("test_target_0"))
+ gv0 = relay.GlobalVar("test_target_0")
+ mod[gv0] = func0
+
+ # body
+ data = relay.var('data', shape=(10, 10))
+ tuple_out = gv1(data)
+ ce_2 = relay.TupleGetItem(tuple_out, 1)
+ ce_3 = relay.TupleGetItem(tuple_out, 0)
+
+ X = relay.tanh(ce_2)
+ ce_4 = gv0(ce_3, X)
+ func = relay.Function([data], ce_4)
+ mod["main"] = func
+
+ return mod
+
+ # print(create_merged_graph())
Review comment:
remove the debugging stmt as well
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] manupa-arm commented on a change in pull request
#5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r398522928
##########
File path: src/relay/transforms/partition_graph.cc
##########
@@ -165,101 +162,142 @@ class Partitioner : public ExprMutator {
// Traverse the rest graph.
auto input_expr = VisitExpr(call->args[0]);
- // Replace the begin annotation with an external call input variable.
- auto compiler_attrs = call->attrs.as<CompilerAttrs>();
+ AnnotatedRegion sg = GetRegion(GetRef<Call>(call));
+ int index = GetArgIdx(sg, GetRef<Call>(call));
+ CHECK_NE(index, -1);
// The type of the created variable is the same as the compiler_begin
// node.
- auto var = Var(compiler_attrs->compiler + "_input" + std::to_string(var_id_++),
- call->checked_type_);
-
- // Find the corresponding subgraph and add the argument.
- auto subgraph = GetSubgraph(GetRef<Call>(call));
- if (!subgraph) {
- throw Error(ErrorBuilder()
- << "Cannot find the corresponding subgraph for start annotation:\n"
- << AsText(GetRef<Call>(call), false));
- }
- subgraph->args.push_back({var, input_expr});
+ std::string target = call->attrs.as<CompilerAttrs>()->compiler;
+ std::string varname = target + "_" + std::to_string(sg->GetID())
+ + "_i" + std::to_string(index);
+ auto var = Var(varname, GetRef<Call>(call)->checked_type_);
+
+ auto cand = std::make_pair(var, input_expr);
+ if (std::find(region_args[sg].begin(),
+ region_args[sg].end(), cand) == region_args[sg].end()) {
+ region_args[sg].push_back(cand);
+ }
+
return std::move(var);
} else {
CHECK_EQ(call->op, compiler_end_op);
// The annotation node is inserted on edge so it must have only one argument.
CHECK_EQ(call->args.size(), 1U);
- auto compiler_attrs = call->attrs.as<CompilerAttrs>();
+ AnnotatedRegion region = GetRegion(GetRef<Call>(call));
- // Check if the argument already belongs to an existing subgraph
- auto subgraph = GetSubgraph(call->args[0]);
- if (!subgraph) {
- auto ret = this->subgraphs_.emplace(std::make_shared<Subgraph>());
- subgraph = *ret.first;
- subgraph->nodes.insert(call->args[0]);
- subgraph->id = this->subgraph_id_++;
- }
- subgraph->nodes.insert(GetRef<Call>(call));
+ // TODO(@manupa-arm) : need to use the parent function (to which region
+ // belongs to) name/key for the funtions that are created
+ BaseFunc f = GetFunc(GetRef<Call>(call));
// Traverse subgraph inputs.
auto input = VisitExpr(call->args[0]);
- Array<Var> params;
- Array<Expr> args;
- std::unordered_map<std::string, runtime::NDArray> params_bind;
-
- // The subgraph may be merged so we need to update it again.
- subgraph = GetSubgraph(GetRef<Call>(call));
- CHECK(subgraph);
-
- // Record the constants for propagation.
- for (auto pair : subgraph->args) {
- params.push_back(pair.first);
- if (const auto* cn = pair.second.as<ConstantNode>()) {
- params_bind[pair.first->name_hint()] = cn->data;
+ CHECK(region.defined()) << "Region not defined for " << GetRef<Call>(call);
+ // functions are created for each annotated regions,
+ // when their first output is encountered.
+ // If multiple outputs are there, a tuple node is inserted at the end.
+ // region_function_calls is map that maintains
+ // (each annotated regions) --> created function
+
+ if (region_function_calls.find(region) != region_function_calls.end()) {
+ // This section is executed only if there are multiple outputs in the region
+ // Thus, the function is always created and at the end there would be a tuple node
+ // Therefore, we insert a tuple get item node.
+
+ // Use the already created tuple node
+ auto sg_call = region_function_calls[region];
+ int index = GetRetIdx(region, GetRef<Call>(call));
+ CHECK_NE(index, -1);
+
+ auto tuple_get_item_ = TupleGetItem(sg_call, index);
+ tuple_get_item_->checked_type_ = GetRef<Call>(call)->args[0]->checked_type_;
+ return std::move(tuple_get_item_);
+ } else {
+ // First time this region is encountered in the traversal
+ // Creating the function
+
+ Array<Expr> fields;
+
+ for (auto ret : region->GetOutputs()) {
+ auto ret_expr = VisitExpr(Downcast<Call>(ret)->args[0]);
+ fields.push_back(ret_expr);
+ }
+ int index = GetRetIdx(region, GetRef<Call>(call));
+ CHECK_NE(index, -1);
+
+ Array<Var> params;
+ Array<Expr> param_expr;
+ std::unordered_map<std::string, runtime::NDArray> params_bind;
+
+ for (auto pair : region_args[region]) {
+ params.push_back(pair.first);
+ if (const auto* cn = pair.second.as<ConstantNode>()) {
+ params_bind[pair.first->name_hint()] = cn->data;
+ } else {
+ param_expr.push_back(pair.second);
+ }
+ }
+
+ Function global_region_func;
+ if (region->GetOutputs().size() == 1) {
+ // If there are only a single output; no need to add a tuple
+ global_region_func = Function(params, fields[0],
+ call->args[0]->checked_type_, {}, DictAttrs());
} else {
- args.push_back(pair.second);
+ auto tuple = Tuple(fields);
+ global_region_func = Function(params, tuple, tuple->checked_type_, {}, DictAttrs());
+ }
+
+ std::string target = call->attrs.as<CompilerAttrs>()->compiler;
+ std::string name = target + "_" + std::to_string(region->GetID());
+
+ global_region_func = WithAttr(std::move(global_region_func), attr::kExternalSymbol,
+ tir::StringImmNode::make(name));
+ global_region_func = WithAttr(std::move(global_region_func), attr::kPrimitive,
+ tvm::Integer(1));
+ global_region_func = WithAttr(std::move(global_region_func), attr::kCompiler,
+ tvm::tir::StringImmNode::make(target));
+ global_region_func = WithAttr(std::move(global_region_func), attr::kInline,
+ tvm::Integer(1));
+
+ // Constant propagation
+ if (!params_bind.empty()) {
+ global_region_func = backend::BindParamsByName(global_region_func, params_bind);
}
- }
- auto subgraph_func =
- Function(params, input, call->checked_type_, {});
-
- std::string name = compiler_attrs->compiler + "_" + std::to_string(subgraph->id);
- subgraph_func =
- WithAttr(std::move(subgraph_func), attr::kExternalSymbol, tir::StringImmNode::make(name));
- subgraph_func =
- WithAttr(std::move(subgraph_func), attr::kPrimitive, tvm::Integer(1));
- subgraph_func =
- WithAttr(std::move(subgraph_func), attr::kCompiler,
- tvm::tir::StringImmNode::make(compiler_attrs->compiler));
- subgraph_func =
- WithAttr(std::move(subgraph_func), attr::kInline, tvm::Integer(1));
-
- // Constant propagation
- if (!params_bind.empty()) {
- subgraph_func = backend::BindParamsByName(subgraph_func, params_bind);
+ std::string fname = name;
+ CHECK(!module_->ContainGlobalVar(fname))
+ << "Global function " << fname << " already exists";
+ // Create a global function and add it to the IRModule for the region.
+ // This way we lift the functions that should be handled by external
+ // codegen to the module scope and rely on the pass manager to prevent relay
+ // function level passes (i.e. simplify inference and fusion) optimizing it.
+ GlobalVar glob_func(fname);
+ module_->Add(glob_func, global_region_func);
+
+ // The return type of callnode is the same as the type of the
+ // compiler_end node.
+ auto ret = Call(glob_func, param_expr);
+ region_function_calls[region] = ret;
+
+ if (region->GetOutputs().size() == 1) {
+ // If there is only a single output; no need to add a tuplegetitem node
+ return Call(glob_func, param_expr);
Review comment:
fixed
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] comaniac commented on a change in pull request
#5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#discussion_r398768598
##########
File path: tests/python/relay/test_pass_partition_graph.py
##########
@@ -678,6 +678,81 @@ def expected():
check_result(mod, {"y": y_data}, (8, 8), np.log(np_add))
+def test_multiple_outputs():
Review comment:
Oh yes. The case in my mind is `conv2d->BN->relu`. In this case, the region output would be a new tuple of the rest 2 BN outputs and 1 relu output (relu could be any op tho).
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services
[GitHub] [incubator-tvm] manupa-arm commented on issue #5143: [RELAY]
Re-wrote the Graph Partitioner to support multiple outputs
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on issue #5143: [RELAY] Re-wrote the Graph Partitioner to support multiple outputs
URL: https://github.com/apache/incubator-tvm/pull/5143#issuecomment-603402680
cc : @zhiics @comaniac @tqchen
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services