manupa-arm commented on a change in pull request #6655:
URL: https://github.com/apache/incubator-tvm/pull/6655#discussion_r503299324
##########
File path: src/relay/transforms/annotate_target.cc
##########
@@ -40,11 +40,14 @@ static const PackedFunc* make_begin_op =
static const PackedFunc* make_end_op =
runtime::Registry::Get("relay.op.annotation._make.compiler_end");
+#define DEFAULT_TARGET_NAME ("default")
Review comment:
Can we use constexpr or const specifier instead ? -- to scope the
parameter.
##########
File path: src/relay/transforms/annotate_target.cc
##########
@@ -113,18 +118,26 @@ class AnnotateTargetRewriter : public ExprRewriter {
std::vector<std::string> supported_targets;
auto op_node = pre->op.as<OpNode>();
-
// This graph has annotations, meaning that this is not the first time
running this pass.
if (op_node && pre->op == CompilerBeginOp()) {
// Bypass compiler begin due to lack of target information. It will be
processed
// when the following op handling arguments.
CHECK_EQ(pre->args.size(), 1U);
- return post.as<CallNode>()->args[0];
+
+ std::string begin_target = pre->attrs.as<CompilerAttrs>()->compiler;
+ if ("" == begin_target || begin_target == DEFAULT_TARGET_NAME)
Review comment:
brackets
##########
File path: src/relay/transforms/annotate_target.cc
##########
@@ -62,19 +65,22 @@ class AnnotateTargetRewriter : public ExprRewriter {
const std::string& target =
"") {
std::string ref_target = "";
Array<Expr> compiler_ends;
+ Array<Expr> compiler_begins;
for (auto arg : args) {
- std::string arg_target = "default";
+ std::string arg_target = DEFAULT_TARGET_NAME;
const CallNode* call = arg.as<CallNode>();
-
if (call && call->op == CompilerBeginOp()) {
// Argument is already compiler begin node meaning that this is not
the first time
// running this pass, so we simply remove it and will add a new one
later.
CHECK_EQ(call->args.size(), 1U);
const CallNode* end = call->args[0].as<CallNode>();
- if (end->op == CompilerEndOp()) {
+ if (end && end->op == CompilerEndOp()) {
arg_target = end->attrs.as<CompilerAttrs>()->compiler;
}
- compiler_ends.push_back(call->args[0]);
+ if (arg_target == "" || arg_target == DEFAULT_TARGET_NAME)
Review comment:
brackets (here and else as well) ?, did this pass through clang-format ?
##########
File path: src/relay/transforms/annotate_target.cc
##########
@@ -62,19 +65,22 @@ class AnnotateTargetRewriter : public ExprRewriter {
const std::string& target =
"") {
std::string ref_target = "";
Array<Expr> compiler_ends;
+ Array<Expr> compiler_begins;
for (auto arg : args) {
- std::string arg_target = "default";
+ std::string arg_target = DEFAULT_TARGET_NAME;
const CallNode* call = arg.as<CallNode>();
-
if (call && call->op == CompilerBeginOp()) {
// Argument is already compiler begin node meaning that this is not
the first time
// running this pass, so we simply remove it and will add a new one
later.
CHECK_EQ(call->args.size(), 1U);
const CallNode* end = call->args[0].as<CallNode>();
- if (end->op == CompilerEndOp()) {
+ if (end && end->op == CompilerEndOp()) {
arg_target = end->attrs.as<CompilerAttrs>()->compiler;
}
- compiler_ends.push_back(call->args[0]);
+ if (arg_target == "" || arg_target == DEFAULT_TARGET_NAME)
+ compiler_ends.push_back(call->args[0]);
+ else
+ compiler_begins.push_back(arg);
Review comment:
Maybe its worth adding a comment, this is to not touch (non-empty and
non-default) annotation. Took sometime figure out the unusual compiler_begins
append at this location of the code. Just a suggestion.
##########
File path: tests/python/relay/test_pass_annotate_target.py
##########
@@ -327,6 +327,42 @@ def after():
assert tvm.ir.structural_equal(expected, result)
+def test_tuple_two_targets():
+ """Tests whether the TupleNode is promoted to previously annotatated
operation or is excluded."""
+ target_relu = "relu_target"
+ target_maximum = "maximum_target"
+ target_default = "default"
+
+ @tvm.ir.register_op_attr("nn.relu", "target." + target_relu)
+ def relu(attrs, args): # pylint: disable=unused-variable
+ return True
+
+ @tvm.ir.register_op_attr("maximum", "target." + target_maximum)
+ def maximum(attrs, args): # pylint: disable=unused-variable
+ return True
+
+ def before():
+ a = relay.var("a", shape=(10, 5))
+ b = relay.var("b", shape=(10, 5))
+ r = relay.nn.relu(b)
+ t1 = relay.Tuple((r, r))
+ r2 = relay.nn.relu(t1)
+ m = relay.maximum(a, b)
+ t2 = relay.Tuple((m, r2))
+ f = relay.Function([a, b], t2)
+ return tvm.IRModule.from_expr(f)
+
+ for default_tuples, parts in [(True, 3), (False, 2)]:
+ result = before()
+ result = transform.AnnotateTarget([target_relu],
default_tuples)(result)
+ result = transform.AnnotateTarget([target_maximum], True)(result)
+ result = transform.MergeCompilerRegions()(result)
+ result = transform.PartitionGraph()(result)
+ assert parts == len(
+ list(filter(lambda _: "target" in _.name_hint,
result.get_global_vars()))
Review comment:
While the number of partitions gives an indication whether tuples are
going to into partition, don't you think checking explicitly for sub-graphs of
tuples make sense?
##########
File path: src/relay/transforms/annotate_target.cc
##########
@@ -113,18 +118,26 @@ class AnnotateTargetRewriter : public ExprRewriter {
std::vector<std::string> supported_targets;
auto op_node = pre->op.as<OpNode>();
-
// This graph has annotations, meaning that this is not the first time
running this pass.
if (op_node && pre->op == CompilerBeginOp()) {
// Bypass compiler begin due to lack of target information. It will be
processed
// when the following op handling arguments.
CHECK_EQ(pre->args.size(), 1U);
- return post.as<CallNode>()->args[0];
+
+ std::string begin_target = pre->attrs.as<CompilerAttrs>()->compiler;
+ if ("" == begin_target || begin_target == DEFAULT_TARGET_NAME)
Review comment:
This leaves the above comment out of place; Now it seems we only bypass
compiler begins if its default or empty. Thus, move the comment here and change
accordingly. Anyway, is there a reason why we are bypassing "default"
compiler_begins here and adding them back in AnnotateArgs ? -- If we did not
bypass them, we dont need to add them back (in line 80) ?
##########
File path: src/relay/transforms/annotate_target.cc
##########
@@ -113,18 +118,26 @@ class AnnotateTargetRewriter : public ExprRewriter {
std::vector<std::string> supported_targets;
auto op_node = pre->op.as<OpNode>();
-
// This graph has annotations, meaning that this is not the first time
running this pass.
if (op_node && pre->op == CompilerBeginOp()) {
// Bypass compiler begin due to lack of target information. It will be
processed
// when the following op handling arguments.
CHECK_EQ(pre->args.size(), 1U);
- return post.as<CallNode>()->args[0];
+
+ std::string begin_target = pre->attrs.as<CompilerAttrs>()->compiler;
+ if ("" == begin_target || begin_target == DEFAULT_TARGET_NAME)
+ return post.as<CallNode>()->args[0];
+
+ return post;
} else if (op_node && pre->op == CompilerEndOp()) {
// Override compiler end with the new target.
CHECK_EQ(pre->args.size(), 1U);
auto input_expr = post.as<CallNode>()->args[0];
CHECK(op_expr_to_target_.find(input_expr) != op_expr_to_target_.end());
+
+ std::string begin_target = pre->attrs.as<CompilerAttrs>()->compiler;
+ if ("" != begin_target && begin_target != DEFAULT_TARGET_NAME) return
post;
Review comment:
Same question here as well; reason for treating already put "default"
annotation as they were empty ?
----------------------------------------------------------------
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:
[email protected]
