LLVM Discussion Forums

Pass Infrastructure: Define post/pre-processing hooks on specific passes

Hello !

I am new to MLIR and my question is quit basic so I’m sorry if I have missed something obvious.

I would like to implement a pass that along with some transformations it also collects information.
I would then like to use this info but only when that specific pass finishes. Alternatively, I would like to pass information from one pass to the next. Either would do.

Here is what I know:

  1. The Pass class provides a hook called runOnOperation which - from what I understand - runs multiple times for each matching Op. However, I would like to collect information across all these calls.
  2. The PassInstrumentation class looked promising at first but it seems that it runs before/after every pass which seems a bit too much for me (I would like to run something after only one specific pass).

I just couldn’t find anything related in the documentation, examples, or the code yet.

Thank you in advance !

Would an analysis work for your needs? https://mlir.llvm.org/docs/PassManagement/#analysis-management

Thank you!

Perhaps analysis could work, but I haven’t figured out how yet.
Here are a few more details on my situation that might help:

  1. On step one I would like to find specific ops and remove them
  2. I would also like to call signalPassFailure if I find inconsistencies in these ops
  3. Before removing these ops I would like to collect info about them
  4. Finally I want to use this info to modify some other op but only once I have collected all the info from all the ops matched by (1)

Please correct me if I’m wrong but the analysis framework will maintain
information for the duration of a specific pass, right ?
Is there a way to communicate analysis results to subsequent passes ?

Alternatively, I tried an OperationPass<ModuleOp> and then used an OpRewritePattern to collect the information.
This can work if I perform step (4) above within runOnOperation and after the rewrite-pattern was applied to all Ops from (1).
But the problem I ran into was that I can no longer call signalPassFailure from OpRewritePattern. This is a smaller issue though for another thread perhaps.

In the meantime, I will look more closely on Analysis Management see if this can help.

Thank you !

Analyses can be preserved between passes, although the API is still a bit clunky. Here’s one of our analyses that we just mark as always preserved :smiley:

https://sourcegraph.com/github.com/google/iree/-/blob/iree/compiler/Dialect/Flow/Analysis/Dispatchability.h

The tricky thing for me is that analyses provide information about an operation, and you want information about an operation after you’ve already deleted it. That seems a bit odd. Can you give an example of what you’re actually trying to do with your pass?

But I think you could also do this with a mondo pass that does basic manipulations and doesn’t do anything with patterns and such. Just make it a pass over the op with the region you’re trying to modify and do whatever you want :smiley:

Here’s an example of a pass that operates on the IR at a more raw level, without rewrite patterns: https://sourcegraph.com/github.com/google/iree/-/blob/iree/compiler/Dialect/Flow/Transforms/OutlineDispatchRegions.cpp

Thank You, that helped!

I looked more carefully into the analysis framework and I was able to put something together that works. I’ll come back to that after I answer your question though:

In our dialect we have an Op (say foo::FooOp) that is close to a source-level construct. This construct cannot be lowered to the LLVM dialect so it needs to be removed. The plan is to retrieve the information that it holds and save them as attributes in the enclosing ModuleOp. Why we want to do that ? That’s a different discussion :laughing:

In the meantime, here is what I did:

First I defined a module pass as follows:

struct PreProcessPass : 
       public PassWrapper<PreProcessPass, OperationPass<ModuleOp>> {

  void runOnOperation() {
    getAnalysis<foo::Info>(); // foo::Info is some class holding "info"
  }
};

I then defined a FooOp pass as follows:

struct MainPass :
       public PassWrapper<MainPass, OperationPass<foo::FooOp>> {
    
  void runOnFunction() {
    auto info = getCachedParentAnalysis<foo::Info>();
      assert(info); // This assertion is successful
      // Modify the info as I please here
  }
};

Finally, I define another module pass to consume the information as follows:

struct PostProcessPass :
       public PassWrapper<PostProcessPass, OperationPass<ModuleOp>> {

  void runOnFunction() {
    auto info = getCachedAnalysis<foo::Info>();
    // This works !
  }
};

The main program would then look like this:

mlir::PassManager pm(&context, true);
pm.addPass(std::make_unique<PreProcessPass>());
pm.addPass(std::make_unique<MainPass>());
pm.addPass(std::make_unique<PostProcessPass>());

What do you think ? Is there a better - more idiomatic - way to do this ?
I don’t think this will work if FooOp is nested deeper into the module (e.g., inside a FunctionOp). So in such a scenario this will start becoming a bit ugly with various special cases. But perhaps that’s OK.

If you want to lookup information and attach them on the module, you likely should make this a ModuleOp pass.
Inside the pass you can then do:

void runOnOperation() {
  ModuleOp m = getOperation();
  m.walk([&] (FooOp fooOp) {
    // will visit all the FooOp nested inside the module.
    ...
  }); 
}

The reason is that the pass manager is multi-threaded: a pass isn’t allowed to modify a parent or sibling op (so a FunctionPass shouldn’t modify the module for example).

Also when you create a pipeline like:

mlir::PassManager pm(&context, true);
pm.addPass(std::make_unique<PreProcessPass>());
pm.addPass(std::make_unique<MainPass>());
pm.addPass(std::make_unique<PostProcessPass>());

The MainPass will be scheduled on FooOp immediately nested inside the module, not the one inside functions and other level of nesting.

Aha, I totally missed the walk method. Thanks !
That seems to work. So here is what I did:

void runOnOperation() {
  ModuleOp m = getOperation();
  Info info; // The info I want to gather
  m.walk([&] (FooOp fooOp) {
    if (has_property(fooOp)) { // I am looking for specific FooOps
      // Do some processing and save to the Info object
      // ...
      // Finally, remove the Op
      OwningRewritePatternList patterns;
      patterns.insert<FooRemover>(&getContext());
      applyOpPatternsAndFold(fooOp, patterns);
    }
  });
  // Now I can use info as I please and modify the module
} // This Works !

Notice how I am using applyOpPatternsAndFold at the end just so I can call eraseOp for the fooOp. Is that the right way to do this ? (final question :slight_smile: )

Thanks Again !!

Using a pattern to remove the op seems heavy to me, you can just do fooOp.erase(); (assuming you replaced all the uses).

1 Like