[Relay] Added Merge Composite pass

[mbaret]

This adds a pass, MergeComposite, which merges patterns (expressed as Relay expressions). The detected patterns are wrapped in a function call which is marked with a 'Composite' attribute that names the pattern.

This is primarily for use with the external codegen infrastructure in the case where a combination of Relay ops map to a single external codegen op. It is not the same as fusion.

Further discussion on this PR can be found at https://discuss.tvm.ai/t/rfc-external-codegen-defining-composite-relay-operators/5470.

[masahi]

ok reading your original proposal carefully I understand that I can look up 'Composite' attribute to know if a particular function has a pattern I'm looking for. Of course I still need to traverse the arguments, but I can remove the "detection" logic from "traversal" logic in DetectFusedConv2DBiasReLU.

It seems I have to make only minimal change to make use of this feature in my PR. Also it enables removing my custom annotator. Sounds great!

[masahi]

ok I already have a working test case for detecting conv + bias + relu (Note the "Composite" attribute). Now I really want to hook this up with "Compiler" and "ExternalSymbol" attributes to enable codegen.

def @main(%data: Tensor[(1, 3, 224, 224), float32]) -> Tensor[(1, 16, 224, 224), float32] {
  %2 = fn (%data1: Tensor[(1, 3, 224, 224), float32], %weight: Tensor[(16, 3, 3, 3), float32], %bias: Tensor[(16, 1, 1), float32], Primitive=1, Composite="dnnl.conv_bias_relu") -> Tensor[(1, 16, 224, 224), float32] {
    %0 = nn.conv2d(%data1, %weight, padding=[1, 1], channels=16, kernel_size=[3, 3]) /* ty=Tensor[(1, 16, 224, 224), float32] */;
    %1 = add(%0, %bias) /* ty=Tensor[(1, 16, 224, 224), float32] */;
    nn.relu(%1) /* ty=Tensor[(1, 16, 224, 224), float32] */
  };
  %3 = %2(%data, meta[relay.Constant][0] /* ty=Tensor[(16, 3, 3, 3), float32] */ /* ty=Tensor[(16, 3, 3, 3), float32] */, meta[relay.Constant][1] /* ty=Tensor[(16, 1, 1), float32] */ /* ty=Tensor[(16, 1, 1), float32] */) /* ty=Tensor[(1, 16, 224, 224), float32] */;
  %6 = fn (%data2: Tensor[(1, 16, 224, 224), float32], %weight1: Tensor[(16, 16, 3, 3), float32], %bias1: Tensor[(16, 1, 1), float32], Primitive=1, Composite="dnnl.conv_bias_relu") -> Tensor[(1, 16, 224, 224), float32] {
    %4 = nn.conv2d(%data2, %weight1, padding=[1, 1], channels=16, kernel_size=[3, 3]) /* ty=Tensor[(1, 16, 224, 224), float32] */;
    %5 = add(%4, %bias1) /* ty=Tensor[(1, 16, 224, 224), float32] */;
    nn.relu(%5) /* ty=Tensor[(1, 16, 224, 224), float32] */
  };
  %6(%3, meta[relay.Constant][2] /* ty=Tensor[(16, 16, 3, 3), float32] */ /* ty=Tensor[(16, 16, 3, 3), float32] */, meta[relay.Constant][3] /* ty=Tensor[(16, 1, 1), float32] */ /* ty=Tensor[(16, 1, 1), float32] */) /* ty=Tensor[(1, 16, 224, 224), float32] */
}

[mbaret]

To hook it up end-to-end will, for now, require writing a custom annotation pass that can recognise composite functions and understand that they need to be treated in the same way as operators. I've got a generic pass that does this, but it builds on top of @zhiics annotation PR which has since been withdrawn. If we agree that this merge_composite pass is useful then I can discuss with zhiics how we want to move forward with the annotator.

[masahi]

If there are multiple patterns to detect, can a composite function detected by one pattern be used as a part of a match for another patterns?

Not sure if this is useful though. Maybe it enables breaking up a big pattern into chunks or pattern reuse.

[comaniac]

Thanks for the PR. We may start from this PR and gradually figure out useful patterns for op-level annotation.

One suggestion is that could we limit this pass to only work for external functions? From the unit tests I suppose this pass works for an entire Relay program, but in practice we want it to work on the subgraph in compiler_begin and compiler_end.

[mbaret]

It's the intention that this can be called on the entire Relay graph so that it can be used to help implement a generic annotation pass (one that is aware of composite functions). That way we can define functions similar to the 'Is Supported?' mechanism in the original annotation PR (since taken down) where you could declare in Python whether an operator was supported. That could be extended to say whether a composite function is supported without having to add pattern matching code to the annotator.

The problem there is the case where composite functions do not end up in an external function after partitioning. My thinking is to have some legalize pass after the partitioning that removes the composite functions from sections of the graph not marked as external.

[comaniac]

So you want the flow to be:

CompositeMerge -> Annotation -> Partitioning. 

I agree that this would make annotation generic and straightforward, although it seems like we don't need annotation anymore if we specify all patterns including single ops. While there are lots of approaches to do so, maybe we could accept this solution first and consider the further steps. What you do think? @zhiics

Also, please help clarify the question asked by @masahi and me about multiple matching. Thanks.

[mbaret]

I think we probably can get rid of the annotation pass. If we expose the logic that determines whether an operator is supported by a given external compiler to the partitioning pass, then this will also mean we don't need to define a pattern for every single operator. Just define patterns for composite operators and have a common way of expressing whether or not a given operator (composite or otherwise) is supported by an external compiler.

That would lose some of the control that could have been gained by writing a custom annotation pass, but I'm not sure what use cases would leverage that.

[comaniac]

That's exactly what I was thinking -- we can debate if we should merge annotation to this pass, or just get rid of the annotation and make this pass general enough. Anyway can be discussed in the follow-up steps.

[masahi]

do we still need partitioning if we get rid of the annotation pass? At least for simple patterns like conv + bias + relu, I can get the same composite (or partitioning, whatever) without writing a clunky custom annotator and doing the partitioning pass.

[mbaret]

It's important to emphasise the difference between partitioned subgraphs (which correspond to 'External' functions), and the merged patterns here (which correspond to 'Composite' functions).

Composite functions are a way to group together relay patterns which correspond to a particular external codegen function (eg. qnn_conv2d + bias_add + requantize maps to a single convolution function in Arm Compute Library). External functions are a way to represent a subgraph that should not go via TVM codegen but should instead go through an external compiler (eg. DNNL).

For libraries like ACL or DNNL, it happens to be the case that you can just put every supported operator/composite function in it's own external function (partitioned subgraph) and because the library calls are standalone that works well enough.

A more complex external compiler though may want to do some advanced fusion or manipulation which means it needs to see the full supported subgraph, not just a single operator/composite function at a time. It will produce some binary artifact that computes the entire subgraph rather than a particular operator within it. All that complexity should be handled by the external compiler's codegen pass, the responsibility of the partitioner is to ensure that pass receives a valid subgraph that is as large as possible while still being fully supported and introducing no dataflow problems.

In summary, I don't think this pass replaces partitioning which is essential for the 'complex compiler' case, but I think it may replace the annotation pass.

[masahi]

thanks, such "complex compiler" use cases sounds quite interesting and the difference in Composite/External there makes a lot of sense. I also understand your earlier comments better now.

[comaniac]

While @mbarrett97 has pointed out the difference between composite/annotation and partition, I just provide one more clarification: The main task for the partition pass is transforming compiler annotations in a Relay graph to external functions and it is compiler agnostic. It means we may even consider the possibility of partitioning a Relay graph to more than one external compilers.

[mbaret]

I've updated the PR to account for the priority ordering of patterns. You no longer specify a pattern map, but a list of tuples (pattern_name, pattern) where the order or the tuples corresponds to the order of the matching. To get this to work, I had to pass strings between Python/C++ requiring this change: #4806.

[comaniac]

I think the priority mechanism for multiple patterns is sufficient for now. Overall LGTM.
Just one minor question left in comments.

[mbaret]

I've done a big clean up + simplified the code a bit and added some comments.

[zhiics]

@mbarrett97 Considering that this is a relatively standalone pass and it helps fusion with external functions, I think it is okay to take it in.

For the pattern matching cases, one possible way I am thinking of is automatically generating some possible pattern given some metadata.

@comaniac and I will start an RFC to discuss how we can add the whitelist based annotation back. It probably will leverage the pattern matching here.

@masahi @anijain2305 PTAL and share your comments if there is any. Thanks.

[anijain2305]

My only request is to add more test cases. In my experience, things get very ugly as the networks get bigger and the things that seems corner cases become very common. But, I am ok with the scope of this PR and delaying aggressive testing to later PR. I think it helps in quick flush of e2e pileline.

[anijain2305]

Oh just saw. The tests are already added. Thanks, its good from my side :)

[masahi]

@zhiics @comaniac It is worth discussing if we can use composite and partitioning passes to remove the annotation pass, as mentioned by @mbarrett97

[masahi]

Thanks @mbarrett97 @comaniac @zhiics @anijain2305 this is merged.

[mbaret]

@masahi @comaniac @zhiics Thanks for the reviews. An RFC on alternative annotation mechanisms would be great.