[RELAY][PASS] CombineParallelConv2D

[vinx13]

This pass replace convolutions that share the same input node and the same arguments (except that the number of output channels can be different) with a single convolution. The weight of the new 2d convolution is the concatenation of the original weights.
The original conv2d nodes are replaced with strided_slice ops to get a slices of the output of the new conv2d.

This prevents launching multiple kernels in networks with multiple convolution branches, such as Inception block.

• use strided_slice instead of take
• fold following bias, relu, bn

Algorithm

1. Find parallel branches like

      data
    /      \
conv2d    conv2d
|            |
op           op
|            |

where op is elemwise or broadcast op.
2. Group branches by kernel shape and attrs of conv2d.
3. Combine conv2d in the same group and possibly combine subsequent ops.
4. Use strided_slice to split the output of the combined op.

Please review @tqchen @jroesch @ZihengJiang @masahi

[tqchen]

@vinx13 can you also elaborate a bit on possible use-cases of this pass?

[masahi]

I guess a bigger kernel is better for performance.

[merrymercy]

TensorRT does this optimization https://devblogs.nvidia.com/deploying-deep-learning-nvidia-tensorrt/#attachment_6827
It can benefit inception like networks.

[masahi]

@vinx13 What happens if convolutions are followed by elemwise ops? (bias, batch norm etc, which should be fused into a single op) Elemwise ops that follow a convolution can be different for each child convolution branch.

Can the folded convolution still be fused with elemwise ops?

[vinx13]

After this pass, the original convolutions will be replaced with Take. Elemwise ops will follow Take (each branch will follow a different Take).
Take is an injective op, so we need a way to fuse convolution and take.

[masahi]

hmm interesting. I understand that we can at least fuse elemwise ops into each Take. But even if we allow fusing convolution with Take, I don't see how we can fuse multiple Take with a single folded conv op.

Code itself looks good though. I understood how it works.

[merrymercy]

So we should also fold the following bias, relu and bn to keep the final graph easy for fusion. These elementwise ops should be placed before Take.

[masahi]

yeah, but element wise ops that follow convolution can be different for each child convolution branch (I don't know about inception network). So I think there will be less chance of folding?

Since there are multiple convolution branches, not everything can be fused anyway. If we fold element wise ops as well, we are left with multiple Take ops that cannot be fused with anything.

I think we might as well stop fusing (or "Realize" in a NNVM term) at the folded convolution op, and let Take ops be fused with element wise ops.

[tqchen]

I am working on a new fusion pass in the new relay IR and hopefully we can followup on this topic there

[tqchen]

Some specific followup comments:

• Let us also think a bit more about naming, usually fold means removing(e.g. constant folding). Maybe a more appropriate name is CombineParallelConv2D?
• Let use avoid using take, but instead use slice operator to slice the outputs
• Let us also fold followup bias, relu and bn when possible as suggested by @merrymercy

[tqchen]

Let us use slice instead of take to get the final result

[tqchen]

The new set of changes LGTM, @vinx13 can you also add support for fusing followup ops and a testcase?

[tqchen]

strided_slice is added in #2094

[tqchen]

@MarisaKirisame @merrymercy can you please take another look and https://docs.tvm.ai/contribute/code_review.html#approve-and-request-changes-explicitly

[tqchen]

@vinx13 please rebase against the master

[tqchen]

Thanks, @MarisaKirisame @masahi @vinx13 this is merged!