Predicated Rewrite Patterns?

Hi folks,

Recently, I found myself reaching for ‘ad-hoc’ predicated rewrite patterns, i.e., patterns that take some function object supplied by the pass and use it to decide whether to fail to match early or not. I saw two existing mechanisms to achieve this:

  1. Dynamic: Adding an extra constructor argument with a ‘filter function’. An example would be patterns in VectorRewritePatterns.h.
  2. Static: Adding a templated wrapper pattern to inherit from a base pattern and add extra checks. An example would be common type checks in MathToSPIRV.cpp. Note that this requires the base pattern to be exposed via its class declaration and not be marked as final, which doesn’t make it useful with majority of patterns that are exposed through populate* functions.

I wondering if it would be useful to have the ability to specify dynamic and/or static rewrite pattern predicates in a more generic and universally available way. I’m thinking about something like this:

  1. For dynamic predicates, add a new pattern rewriter class OpRewritePatternWithFilter<typename SourceOp> with a constructor taking an optional filter function of signature LogicalResult(SourceOp).
  2. For static predicates, add a new pattern rewriter class PredicatedOpRewritePattern<typename SourceOp, auto FilterFn>. Derived classes could either specify a function pointer directly or take a template argument on their own and ‘forward’ it to the second template argument of PredicatedOpPatternRewriter. I’d expect the main use to be conversion patterns with a common set of preconditions.

(Naming suggestions welcome.)

Does this sound like something that would be useful to folks? I’m not sure if this was suggested or discussed before.

-Jakub


cc: @River707, @matthias-springer, @ftynse, @antiagainst, @Mogball

Do you see a way for the driver to check the condition before calling the pattern? I would be interested to see if we can check the function pointer for example and call it once for multiple patterns and save runtime.

You’ll notice in DRR there is predication using function call. It happens as part of match, so in match and rewrite, the predication check is in match and treated uniformly with other matches.

To Mehdi’s question: is this convenience wrapper for predicating in match?

That would be cool in the case of conversion passes that use patterns that mostly share the same precondition, but that would require this function to be ‘pure’. I think we could add a FunctionPtr -> LogicalResult map in PatternApplicator::matchAndRewrite and cache predicate results over the same op until the first rewrite. I can look into this if this sounds like something useful.

In the dynamic case, I was thinking of of a pre-check in matchAndRewrite (or internally before).

That would be exactly what DRR generates today if given a dynamic predicate.

I’m not sure how much we could cache without knowing contents. E.g., if predicate is based on input operand, then we’d need to know what feature of it is being checked in whether to invalidate the cached state and so if anything that a predicate may query could be changed we’d need to invalidate, which seems like tricky bookkeeping. A convenience wrapper that effectively calls predicate function at start of matchAndRewrite seems simple/intuitive.

I wasn’t thinking of caching across modifications of the IR actually. When visiting an operation, the driver will try all pattern attached to this operation and call matchAndRewrite on them. If multiple patterns shared the match, the driver could call it once at this point and for this traversal of the pattern. No caching beyond this :slight_smile:

By the way, PDL was designed to allow this kind of transformations automatically while merging multiple patterns matching an operation!

This reminds me that we still have separate match and rewrite functions that can be overridden on RewritePattern. (The default implementation of matchAndRewrite is to call one than the other. We added it because the data sharing mechanism between match and rewrite was awkward.) Specifically, a predicate sounds like a common match or a part of it for multiple patterns. Maybe we can have a sort of a list that contains multiple matchers and compare their addresses.