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.

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

[Relay] Add support for tuple node in operator fusion #2187

Merged
merged 10 commits into from
Nov 30, 2018

Conversation

masahi
Copy link
Member

@masahi masahi commented Nov 28, 2018

Fixes #2183. It enables fusing a concat node with other ops before it.
@tqchen please review.

Example: Upsampling + Concat

fn (%x: Tensor[(1, 16, 64, 64), float32])
    -> Tensor[(1, 32, 64, 64), float32] {
  %0 = fn(%p0: Tensor[(1, 16, 64, 64), float32])
          -> Tensor[(1, 16, 32, 32), float32] {
    %1 = nn.max_pool2d(%p0, pool_size=[2, 2], strides=[2, 2]) # ty=Tensor[(1, 16, 32, 32), float32]
    %1
  }
  %2 = %0(%x) # ty=Tensor[(1, 16, 32, 32), float32]
  %3 = fn(%p01: Tensor[(1, 16, 32, 32), float32],
          %p1: Tensor[(1, 16, 64, 64), float32])
          -> Tensor[(1, 32, 64, 64), float32] {
    %4 = nn.upsampling(%p01, scale=2) # ty=Tensor[(1, 16, 64, 64), float32]
    %5 = (%4, %p1)
    %6 = concatenate(%5, axis=1) # ty=Tensor[(1, 32, 64, 64), float32]
    %6
  }
  %7 = %3(%2, %x) # ty=Tensor[(1, 32, 64, 64), float32]
  %7
}

@tqchen
Copy link
Member

tqchen commented Nov 28, 2018

This looks right, @masahi We also need to fix for opt_level=0, which will break every op into a single function. But we dont want to do that for a single Tuple node and isolate that into a function. A solution can either be:

  • Force tuple into followup concat when possible in the post processing of grouping phase.
  • In the post processing, don't isolate Tuple into a single function.
    • If the function is only contains MakeTuple, just inline that back.

src/relay/pass/fuse_ops.cc Show resolved Hide resolved
@tqchen tqchen added the status: need update need update based on feedbacks label Nov 28, 2018
@tqchen
Copy link
Member

tqchen commented Nov 28, 2018

I would recommend handling the Tuple as root case, so that this can be used for future purposes

@masahi
Copy link
Member Author

masahi commented Nov 29, 2018

@tqchen I handled the cases for opt level = 0 and when tuple is the root of the group. But I'm not sure if I handled them correctly.

When tuple is a isolated, it is not put into a function. But if it is fused with other ops before it, I make a function that returns a tuple. Please have a look at the new test case.

@tqchen
Copy link
Member

tqchen commented Nov 29, 2018

@masahi your implementation is correct, However, it is likely we don't need to maintain a node count.What you can do instead is to quickly check if all the fields in the tuple are params of that GroupInfo and if they are, it means it is an isolated function. Good job in making this happen

@masahi
Copy link
Member Author

masahi commented Nov 29, 2018

@tqchen let me know if there is a better way to test equality of two Array<Expr>. Don't know why std::equal won't compile.

@jroesch
Copy link
Member

jroesch commented Nov 29, 2018

👍 looks like a good change, this might of been one of the reasons fusion was being blocked for me and @MarisaKirisame on one of our examples for the paper.

if (ret_group == gmap_.at(tuple)) {
bool isolated = true;
for (size_t i = 0; i < new_fields.size(); ++i) {
isolated &= (new_fields[i] == ginfo_[ret_group].params[i]);
Copy link
Member

Choose a reason for hiding this comment

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

change to new_fields[i].same_as(ginfo_[ret_group].params[i]), in case == get overloaded in the future

@tqchen tqchen added status: accepted and removed status: need update need update based on feedbacks labels Nov 29, 2018
@tqchen tqchen merged commit 04f7a18 into apache:master Nov 30, 2018
@tqchen
Copy link
Member

tqchen commented Nov 30, 2018

Thanks, @masahi @jroesch , this is merged

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants