Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[RELAY] Fixes to MergeCompilerRegions #5195

Merged
merged 9 commits into from
Apr 1, 2020
Merged

[RELAY] Fixes to MergeCompilerRegions #5195

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
2563f1f
[RELAY] Fixed issues with MergeCompilerRegions
mbaret Mar 30, 2020
3166c8d
Fixed issue using pre-merged regions
mbaret Mar 31, 2020
c7ca0c1
Removed some debugging logic
mbaret Apr 1, 2020
c7a51ec
Remove default annotations
mbaret Apr 1, 2020
f8e378e
Annotate default 'if's
mbaret Apr 1, 2020
79b9c13
Clang format
mbaret Apr 1, 2020
a293997
Use src/dest in merge
mbaret Apr 1, 2020
763bca6
Fixed partition test
mbaret Apr 1, 2020
6c02abe
Removed comments
mbaret Apr 1, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions src/relay/analysis/annotated_region_set.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,12 +70,12 @@ void AnnotatedRegionSetNode::MergeRegions(AnnotatedRegion src,
regions_.erase(src);
}

void AnnotatedRegionSetNode::AddToRegion(AnnotatedRegion region, const Expr& expr) {
auto region2 = GetRegion(expr);
if (region2.defined()) {
MergeRegions(region, region2);
void AnnotatedRegionSetNode::AddToRegion(AnnotatedRegion dest, const Expr& expr) {
auto src = GetRegion(expr);
if (src.defined()) {
MergeRegions(src, dest);
} else {
region->nodes.insert(expr);
dest->nodes.insert(expr);
}
}

Expand Down
4 changes: 2 additions & 2 deletions src/relay/analysis/annotated_region_set.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,10 +178,10 @@ class AnnotatedRegionSetNode : public Object {
/*!
* \brief Add an expression to a region.
*
* \param region The region to add the expression to.
* \param dest The region to add the expression to.
* \param expr The expression.
*/
void AddToRegion(AnnotatedRegion region, const Expr& expr);
void AddToRegion(AnnotatedRegion dest, const Expr& expr);

/*!
* \brief Make a new region.
Expand Down
21 changes: 18 additions & 3 deletions src/relay/transforms/annotate_target.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,9 @@ namespace tvm {
namespace relay {
namespace annotate_target {

// Cache compiler_begin op for equivalence check.
static const Op& compiler_begin_op = Op::Get("annotation.compiler_begin");

// A helper class to insert annotation boundaries for a program region that will
// be handled by a specific compiler.
class AnnotateTargetWrapper : public ExprMutator {
Expand All @@ -52,6 +55,13 @@ class AnnotateTargetWrapper : public ExprMutator {
return fannotate[op](call->attrs, call->args);
}
}
if (expr->IsInstance<TupleGetItemNode>()) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have a concern about this. It seems like you have to list all possible nodes here and it's easy to miss something. For now I would suggest adding an exception at least to indicate that we need to add more here.

Copy link
Contributor

@trevor-m trevor-m Mar 31, 2020 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think that many nodes will need to be supported here. TupleGetItemNode is needed for batchnorm. I have a PR ready to add support for TupleNode which is needed for concatenate. I haven't encountered any other nodes yet that would need to be added here.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure yet what we plan on doing with Let nodes (or Ref* nodes). I think just saying they're not supported is sufficient for now though without also throwing an error.

TupleGetItem get = Downcast<TupleGetItem>(expr);
if (get->tuple->IsInstance<CallNode>() &&
get->tuple.as<CallNode>()->op == compiler_begin_op) {
return true;
}
}
return false;
}

Expand Down Expand Up @@ -110,9 +120,14 @@ class AnnotateTargetWrapper : public ExprMutator {
auto new_e = ExprMutator::VisitExpr_(op);

auto get = Downcast<TupleGetItem>(new_e);
return TupleGetItem(
InsertEnd(get->tuple),
get->index);
if (IsSupported(get->tuple)) {
const auto* begin_op =
runtime::Registry::Get("relay.op.annotation._make.compiler_begin");
CHECK(begin_op);
return TupleGetItem((*begin_op)(InsertEnd(get->tuple), target_), get->index);
} else {
return TupleGetItem(InsertEnd(get->tuple), get->index);
}
}

Expr VisitExpr_(const FunctionNode* op) {
Expand Down
Loading