How to apply a set of patterns *once*

MLIR
1

Say I want to add an argument to all functions in a module, and function calls to other functions in the module pass on the value of that additional argument. For example:

func.func public @f(%arg0: f32) -> f32 {
  %0 = func.call @times_two(%arg0) : (f32) -> f32
  func.return %0 : f32
}
func.func private @times_two(%arg0: f32) -> f32 {
  %0 = arith.addf %arg0, %arg0 : f32
  func.return %0 : f32
}

becomes:

func.func public @f(%arg0: i32, %arg1: f32) -> f32 {
  %0 = func.call @times_two(%arg0, %arg1) : (i32, f32) -> f32
  func.return %0 : f32
}
func.func private @times_two(%arg0: i32, %arg1: f32) -> f32 {
  %0 = arith.addf %arg1, %arg1 : f32
  func.return %0 : f32
}

What is the best way to implement this?

I have explored two options, neither of which really satisfy me. One is to walk the IR manually and transform the relevant ops. This would look like this:

TrivialPatternRewriter rewriter(context);
for (auto &op : module.getBodyRegion().front()) {
  if (auto funcOp = llvm::dyn_cast<FunctionOpInterface>(&op)) {
    if (!funcOp.isExternal())
      addGridArguments(funcOp, rewriter);
  }
}

TrivialPatternRewriter is copied from here; addGridArguments adds arguments to the funcOp. (I don’t have any problems implementing the latter function.)

What is still missing here is the modification of the func.call ops. I guess I can do this in a similar way but with a recursive walk.

What I don’t like is that (1) it uses TrivialPatternRewriter, which it “discouraged” according to the place I copied it from and (2) it doesn’t express the transformation as patterns and therefore, seems to be non-standard.

The second possibility is using patterns; however, I don’t know how to express this logic as patterns that are applied exactly once. I can write a pattern that adds arguments to a function but that pattern will continue to match after it has been applied, so the two pattern drivers that I know, DialectConversion and GreedyPatternRewriteDriver, will apply the pattern infinitely often.

Maybe I can get the behavior I want with GreedyRewriteConfig::maxIterations = 1?

If so, I am still not completely happy: ideally, I would like to only apply the patterns and not also possible folders. This (or something similar) has been discussed here and no consensus seems to have emerged at the time due to lack of motivation. My hesitation towards using something that also folds is mainly for testing: I would like to understand test the effect of the patterns and nothing else. Also, I don’t see the implications of applying arbitrary folders to IR that doesn’t verify: individual patterns will leave the IR in an invalid state since the transformation of the func.func op is done by a different pattern that the transformation of the func.call ops to that function, and the module will only verify after all patterns have been applied.

What’s the best decision here? Is manual IR walking adequate? Is my hesitation towards folding unjustified? Or would it make sense to have a “apply greedily but don’t fold” driver (option) upstream? (If so, I can propose an implementation.)

2

I don’t believe there is an expectation that all transformations should be implemented as “pattern”. There are many examples in the codebase (loop fusion/tiling for example). At some point we shouldn’t try to fit round pegs in a square hole.

Patterns are useful when:

  1. You want to apply greedily until fixed point.
  2. You want composability: patterns can be reused, injected, and mixed with other patterns.
  3. You’re using the dialect conversion framework.

Otherwise writing a pass that just does what you need is the way to go, you don’t need a “rewriter” either.

Also a nit, but you should be able to write your loop as:

3

Thanks a lot for the quick reply! Yeah, I think none of the three points you mention applies to my case, so there probably is no reason to use patterns.

I probably don’t need a rewriter either, but I am currently using it for RewriterBase::mergeBlocks (or RewriterBase::inlineBlockBefore), which I haven’t seen outside of a rewriter class. Is this correct? Is this something that could be added?

Thanks for pointing me to the .getOps<T> function – great to know!

4

The methods of the Rewriter (like mergeBlocks) exists so that the rewriter listeners can “listen” on these modification and the drivers (greedy or conversion framework) can take this into consideration while driving the work (including keeping the ability to rollback for the conversion framework).

They are not doing too much on top of the native APIs, often they are 1-1. If you look at the implementation of inlineBlockBefore, it is pretty straightforward:

  // Replace all of the successor arguments with the provided values.
  for (auto it : llvm::zip(source->getArguments(), argValues))
    replaceAllUsesWith(std::get<0>(it), std::get<1>(it));

  // Move operations from the source block to the dest block and erase the
  // source block.
  dest->getOperations().splice(before, source->getOperations());
  source->erase();

That said exposing an API as taking a rewriter can be a good practice when writing a utility: it allows for someone else to use the API from a pattern later (if it is something that can make sense conceptually).

5

OK, thanks, make sense! I guess I was hoping not to have to write the four lines you quote but get away with one :wink: (Wouldn’t it make sense to expose mergeBlocks or similar as a “native API”?)

Good point with the rewriter. Maybe I can actually leave it this way…

6

It makes sense to me: feel free to add it on the Block class when you have a use-case for it.

7

Btw, you can also create an IRRewriter if you just want to use a helper function from RewriterBase.

8

Aha, I can! That feel less like a hack :slight_smile: Thanks!