[BYOC] Use pattern language to create composite functions

Thanks for the comments. Here are a short summary of supporting check using pattern language based on my POC (https://gist.github.com/comaniac/1f399dfdfee05a2f7a087c65c21f550c):

A. Callback

The check logic is used in callbacks to determine if this subgraph really matches the pattern. If matching, the callback needs to create a composite function as the POC I provided.

  • User API would be similar as now. Specifically, user provides the following inputs to PatternCallback:
    1. A pattern.
    2. A map of check functions for this pattern (e.g., {'nn.conv2d': lambda ...: true}).
  • A redundant visit to the matched subgraph in order to create the composite function seems inevitable.

B. A More Powerful but Complex Pattern Logic

The check logic is implemented along with the pattern using AttrPattern.

  • User API would be just a pattern.
  • We can use pattern.partition() to easily create composite functions.
  • The logic of implementing a pattern might be more complex, implying a more steep learning curve.

Here are my two cents. Since we already have pattern.partition(), it seems not practical to improve the rewriter API to make the partition efficient. As a result, I prefer solution B. On the other hand, I am not sure how complex this logic would be and if it’s reasonable for 3rd-party backend developers to learn and implement. It would be great if @mbrookhart you could provide an illustration or even an RFC to give us a more concrete idea :slight_smile:

@comaniac @mbaret

I’ve extended the AttrPattern to handle attributes of CallNodes and FunctionNodes here: https://github.com/apache/incubator-tvm/pull/5637

Would you mind taking a look and letting me know if you think that would be sufficiently powerful for solution B?

1 Like

I was mainly referring to this case https://github.com/apache/incubator-tvm/pull/5637/files#diff-f9920485e5e341787129348ce1985db9R247-R251 and it does achieve the same functionality as the current check functions in MergeComposite pass in my opinion.

In terms of ease-of-use, to me this programming model is even more straightforward than the check functions.

@mbaret @masahi what do you guys think?

For our codegen, the check functions are actually quite complex. They’re not just attribute checks, there’s also things relating to tensor sizes, number of dimensions and quantisation parameters. Additionally, our codegen library exposes an API to check whether a given operator is supported for a particular set of parameters and we query that function in the check. Therefore, in my view the check function needs to be able to support an arbitrary packed function.

Hmm. Okay. I’ll add a check callback to the partition pass.

I think that leaves us with three levels of possible user control in level of priority:

  1. Use the Pattern Language for any constraints possible. We might be able to add shape checking to the language, we can’t use an external codegen type checker :slight_smile:
  2. If you need more control, use the Partition check if there are more complicated/external constraints
  3. if you need even more control, use the Rewriter pass for more complicated graph rewriting.
1 Like

@mbrookhart this sounds great!

Then I can use the second method to implement MergeComposite to allow optional check functions as it currently does. Meanwhile, if a user’s pattern only needs to check attributes and shapes, the user can use the first method to specify the pattern.

Updated the gist with merged features (configurable partition function attributes; pattern with more powerful attributes).

Once the last feature (partition callback) is in, we should be fine to rewrite MergeComposite pass without any loss.

@comaniac @mbaret @masahi

Sorry for the delay! The Meetup this morning ate more of my time than I expected.

2 Likes

New MergeComposite pass PR filed.