comaniac commented on a change in pull request #6655:
URL: https://github.com/apache/incubator-tvm/pull/6655#discussion_r528914573
##########
File path: src/relay/transforms/annotate_target.cc
##########
@@ -61,20 +67,27 @@ class AnnotateTargetRewriter : public ExprRewriter {
std::pair<std::string, Array<Expr>> AnnotateArgs(const Array<Expr>& args,
const std::string& target =
"") {
std::string ref_target = "";
+ Array<Expr> compiler_begins;
Array<Expr> compiler_ends;
for (auto arg : args) {
- std::string arg_target = "default";
+ std::string arg_target = default_target;
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.
ICHECK_EQ(call->args.size(), 1U);
+ // Do not alter existing annotation if not default
+ if (default_target != call->attrs.as<CompilerAttrs>()->compiler) {
+ compiler_begins.push_back(arg);
+ } else {
+ // Remove default
+ compiler_ends.push_back(call->args[0]);
+ }
Review comment:
My question was about the algorithm. I don't understand why we "remove
default (but preserve non-default) annotations in case if it is not first
invocation of the pass". Almost all my comments related to this question.
Meanwhile, PR derives another AnnotateTarget implementation to deal with the
case that `include_non_call_ops=False`, so you are more like making a new pass
and use `include_non_call_ops` to dispatch the one being used. Combining my
previous concern about keeping AnnotateTarget simple and straightforward, I
would suggest making it a separate pass. For example, we could name the new
pass "PruneSimpleRegions" and use it as following:
1. (MergeComposiste) -> AnnotateTarget -> PruneSimpleRegions ->
PartitionGraph
2. (MergeComposiste) -> AnnotateTarget -> MergeCompilerRegions ->
PruneSimpleRegions -> PartitionGraph
The first case could achieve your goal, and it could also serve the second
case to prune subgraphs (compiler regions) after merging. If we take one step
further, we could even let users provide a lambda function to determine if a
region is "simple", so that it can be customized for different backends such as
TensorRT. The implementation of this pass could also be much simpler and I
could even help with that if you prefer.
What do you think?
----------------------------------------------------------------
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]
