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:
A pattern.
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
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:
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
If you need more control, use the Partition check if there are more complicated/external constraints
if you need even more control, use the Rewriter pass for more complicated graph rewriting.
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.
