[RELAY] Add MergeCompilerRegions pass #5134

mbaret · 2020-03-23T17:05:32Z

This is part of the infrastructure necessary to support multiple outputs in the BYOC graph partitioning framework. It corresponds to step 3 of the flow outlined here and implements the algorithm described in this RFC. In summary, it merges together consecutive operators supported by the same target while ensuring that no data dependency issues are created.

This PR depends upon the utility functions provided by the AnnotationRegionSet PR and the Improved graph partitioner.

mbaret · 2020-03-23T17:05:58Z

cc @zhiics @comaniac @masahi @tqchen

src/relay/transforms/merge_compiler_regions.cc

src/relay/transforms/annotate_target.cc

masahi · 2020-03-24T21:13:24Z

@mbaret Sorry I haven't been following the development of a new partitioning algorithm. Is this related to annotator support for composite?

mbaret · 2020-03-25T08:18:06Z

It isn't (although I hope to handle that next). This is to help support the case where we partition into a region that has multiple outputs.

zhiics

@comaniac please review the algorithm part as you also worked on it before.

src/relay/transforms/annotate_target.cc

src/relay/transforms/merge_compiler_regions.cc

comaniac

The algorithm behavior seems fine to me.

src/relay/transforms/merge_compiler_regions.cc

comaniac · 2020-03-25T17:48:24Z

src/relay/transforms/merge_compiler_regions.cc

+  // By now, all the nodes have some sort of annotation.
+  // Region merger is ExprVisitor that will update the
+  // AnnotatedRegionSet.
+  RegionMerger merger(regions);


This name is confusing. RegionMerger seems like a mutator but it is actually a visitor. Maybe something like MergableRegionAnalyzer. Or it's better to combine this one with MergeAnnotation. The logic of MergeAnnotation is pretty simple and it can be done in this pass as well, making this pass a real "merger".

I don't think I can easily combine the two passes because of the visit order. I agree the name is slightly confusing, but it does merge 'Regions' in that you pass it an AnnotatedRegionSet and it returns a merged version of that set.

The regions it merges is not a Relay graph but a separate data structure. It looks weird for a "merger" to only refer the given Relay graph and manipulate a separate data structure.

Another way I can think of to make it more reasonable is moving RegionMerger to the AnnotateRegionSet and make it optional, like:

AnnotatedRegionSet regions = AnnotatedRegionSet::Create( expr_all_annotated, compiler_begin_op, compiler_end_op, true);

where the last argument indicates if the created regions should be merged.

Anyway for the time being, we could revisit this part in the future.

trevor-m · 2020-03-27T00:21:03Z

Hello, I am trying to use this (in combination with #5143) but I am getting a segfault every time I call MergeCompilerRegions. Here is a small repro script. I expect the entire model to be partitioned into a single subgraph.

from tvm import relay
import mxnet as mx
from mxnet.gluon.model_zoo.vision import get_model
from tvm.relay import op as reg

def _register_external_op_helper(op_name, supported=True):
    @reg.register(op_name, "target.test")
    def _func_wrapper(attrs, args):
        return supported
    return _func_wrapper

_register_external_op_helper("nn.conv2d")
_register_external_op_helper("nn.dense")
_register_external_op_helper("nn.relu")
_register_external_op_helper("add")
_register_external_op_helper("multiply")
_register_external_op_helper("nn.bias_add")
_register_external_op_helper("nn.batch_flatten")
_register_external_op_helper("nn.max_pool2d")
_register_external_op_helper("nn.dropout")
_register_external_op_helper("nn.batch_norm")
_register_external_op_helper("nn.global_avg_pool2d")

def test_resnet():
    block = get_model('resnet18_v1', pretrained=True)
    mod, params = relay.frontend.from_mxnet(block, shape={'data': (1, 3, 224, 224)}, dtype='float32')
    mod = relay.transform.AnnotateTarget("test")(mod)
    mod = relay.transform.MergeCompilerRegions()(mod)
    mod = relay.transform.PartitionGraph()(mod)
    print(mod)

test_resnet()

comaniac · 2020-03-27T00:30:13Z

@mbaret could you reproduce the failure and try to resolve it? It crashed before partition graph so it should not be the issue of multiple output support.

mbaret · 2020-03-27T15:17:05Z

@mbaret could you reproduce the failure and try to resolve it? It crashed before partition graph so it should not be the issue of multiple output support.

I've found the issue (was to do with changes I made to the merging during code review on the AnnotatedRegionSet patch). I've made the necessary change in this PR, but we'll also need this fix PR. I've also added in a more extensive test (based on the one in your PR) to verify the issue is resolved.

zhiics · 2020-03-27T15:27:22Z

Can you just incorporate that change as it is really short?

zhiics · 2020-03-27T15:32:44Z

also please fix @comaniac's comments

mbaret · 2020-03-27T15:59:11Z

I've incorporated that fix into this PR now.

This pass is part of the flow to support creating compiler regions with multiple outputs. It should be called after AnnotateTarget and will merge together regions that share the same target to create larger compiler regions that can be off-loaded to external codegens. This pass implements an algorithm to ensure that during the merging, no data dependency issues are created. See the tests for an example of this case. Co-authored-by: Ramana Radhakrishnan <[email protected]> Co-authored-by: Manupa Karunaratne <[email protected]> Change-Id: Ibd99083564608d888482f57c5080109f3eefec88

This alters the behaviour of the AnnotateTarget pass to enforce the property that all compiler annotations exist along a single data flow edge. Specifically, this means they should have exactly one parent and one child. Change-Id: I0e74803a77767f4f377d17755a13a74a30909797

Deleting items from a list while iterating it seems to result in undefined behaviour which sometimes segfaults. This makes sure all the item deletion happens separately.

trevor-m · 2020-03-27T16:50:18Z

@mbaret Thanks for investigating and making the fixes! However, I'm still getting the segfault on call to MergeCompilerRegions in that example I posted earlier using your fixes. Are you able to run that example successfully?

mbaret · 2020-03-27T20:30:55Z

There's still another issue to resolve with this (in AnnotateRestDefault). I'll take a look at it on Monday.

tqchen · 2020-03-30T20:59:31Z

Thanks @mbaret @trevor-m @zhiics @comaniac @siju-samuel

* [RELAY] Add MergeCompilerRegions pass This pass is part of the flow to support creating compiler regions with multiple outputs. It should be called after AnnotateTarget and will merge together regions that share the same target to create larger compiler regions that can be off-loaded to external codegens. This pass implements an algorithm to ensure that during the merging, no data dependency issues are created. See the tests for an example of this case. Co-authored-by: Ramana Radhakrishnan <[email protected]> Co-authored-by: Manupa Karunaratne <[email protected]> Change-Id: Ibd99083564608d888482f57c5080109f3eefec88 * [RELAY] Annotate compiler_ends on each edge This alters the behaviour of the AnnotateTarget pass to enforce the property that all compiler annotations exist along a single data flow edge. Specifically, this means they should have exactly one parent and one child. Change-Id: I0e74803a77767f4f377d17755a13a74a30909797 * Fix comment * Rebase *Node::make * Moved block outside for loop * Code style * Update make API * Remove comment * Remove redundant 'else's * Make one line * Fix comment * RefWrite * Fix merge ordering * Add the RFC example as a test * [FIX] Fixed merging behaviour in AnnotateRegionSet Deleting items from a list while iterating it seems to result in undefined behaviour which sometimes segfaults. This makes sure all the item deletion happens separately. * Added checks * Move comment * Update comments

siju-samuel reviewed Mar 24, 2020

View reviewed changes

src/relay/transforms/merge_compiler_regions.cc Outdated Show resolved Hide resolved

src/relay/transforms/merge_compiler_regions.cc Outdated Show resolved Hide resolved

src/relay/transforms/annotate_target.cc Outdated Show resolved Hide resolved

mbaret force-pushed the merge-compiler branch from 451a031 to 52dce66 Compare March 24, 2020 15:14

zhiics reviewed Mar 25, 2020

View reviewed changes

comaniac requested changes Mar 25, 2020

View reviewed changes

mbaret force-pushed the merge-compiler branch from 52dce66 to 9cc0a5f Compare March 26, 2020 14:03

mbaret force-pushed the merge-compiler branch from aabd94a to 3c1cfbd Compare March 27, 2020 15:58

mbaret and others added 15 commits March 27, 2020 16:29

Fix comment

f1a3e75

Rebase *Node::make

3b9abca

Moved block outside for loop

c98a9eb

Code style

6b56aa3

Update make API

17a2ede

Remove comment

7297be2

Remove redundant 'else's

b1cf5d5

Make one line

df4beb8

Fix comment

648b40d

RefWrite

b9b52e4

Fix merge ordering

cf759c8

Add the RFC example as a test

f89fe7d

Added checks

451f8d1

mbaret added 3 commits March 27, 2020 16:30

[FIX] Fixed merging behaviour in AnnotateRegionSet

3c12183

Deleting items from a list while iterating it seems to result in undefined behaviour which sometimes segfaults. This makes sure all the item deletion happens separately.

Move comment

d9b1c71

Update comments

a3848b0

mbaret force-pushed the merge-compiler branch from 3c1cfbd to a3848b0 Compare March 27, 2020 16:34

comaniac approved these changes Mar 27, 2020

View reviewed changes

tqchen assigned tqchen and unassigned tqchen Mar 30, 2020

tqchen approved these changes Mar 30, 2020

View reviewed changes

tqchen merged commit 0212138 into apache:master Mar 30, 2020

tqchen added the status: accepted label Mar 30, 2020

tqchen mentioned this pull request Mar 30, 2020

[TEST] Various CI fixes for the VTA and Relay #5181

Merged

mbaret mentioned this pull request May 12, 2020

Pattern Language, Matcher, Rewriter, and Function Paritioner #5231

Merged

ZihengJiang mentioned this pull request Sep 25, 2020

TVM v0.7 Release Note Candidate #6486

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[RELAY] Add MergeCompilerRegions pass #5134

[RELAY] Add MergeCompilerRegions pass #5134

mbaret commented Mar 23, 2020 •

edited

Loading

mbaret commented Mar 23, 2020

masahi commented Mar 24, 2020

mbaret commented Mar 25, 2020

zhiics left a comment

comaniac left a comment

comaniac Mar 25, 2020

mbaret Mar 27, 2020

comaniac Mar 27, 2020

trevor-m commented Mar 27, 2020 •

edited

Loading

comaniac commented Mar 27, 2020

mbaret commented Mar 27, 2020 •

edited

Loading

zhiics commented Mar 27, 2020 •

edited

Loading

zhiics commented Mar 27, 2020

mbaret commented Mar 27, 2020

trevor-m commented Mar 27, 2020 •

edited

Loading

mbaret commented Mar 27, 2020

tqchen commented Mar 30, 2020 •

edited

Loading

[RELAY] Add MergeCompilerRegions pass #5134

[RELAY] Add MergeCompilerRegions pass #5134

Conversation

mbaret commented Mar 23, 2020 • edited Loading

mbaret commented Mar 23, 2020

masahi commented Mar 24, 2020

mbaret commented Mar 25, 2020

zhiics left a comment

Choose a reason for hiding this comment

comaniac left a comment

Choose a reason for hiding this comment

comaniac Mar 25, 2020

Choose a reason for hiding this comment

mbaret Mar 27, 2020

Choose a reason for hiding this comment

comaniac Mar 27, 2020

Choose a reason for hiding this comment

trevor-m commented Mar 27, 2020 • edited Loading

comaniac commented Mar 27, 2020

mbaret commented Mar 27, 2020 • edited Loading

zhiics commented Mar 27, 2020 • edited Loading

zhiics commented Mar 27, 2020

mbaret commented Mar 27, 2020

trevor-m commented Mar 27, 2020 • edited Loading

mbaret commented Mar 27, 2020

tqchen commented Mar 30, 2020 • edited Loading

mbaret commented Mar 23, 2020 •

edited

Loading

trevor-m commented Mar 27, 2020 •

edited

Loading

mbaret commented Mar 27, 2020 •

edited

Loading

zhiics commented Mar 27, 2020 •

edited

Loading

trevor-m commented Mar 27, 2020 •

edited

Loading

tqchen commented Mar 30, 2020 •

edited

Loading