Dynamically creating a function in a module

MLIR
1

Hi all,
I wrote a pass that extracts a given block into a function. However, the subsequent passes seem to not “see” the new function.

Scenario
The situation is as follows. I have two passes, say Pass1 and Pass2 that I execute like this:

    pm.addPass(createExtractBlockPass()); // Pass1
    pm.addPass(unrollLoopsInFunctionJustCreated()); // Pass2
  • The original module contained only f1
  • After Pass1 the module contains f1 and f2

Problem

  • The second pass (Pass2) is only able to see f2, but
  • If I run Pass1 → dump MLIR → Load MLIR → Pass2 then it works.

Question
Is there a way to “update” the module I am working with on the fly?

Thanks for any help,
Giuseppe

2

Is Pass1 a ModulePass?

3

Hi @jpienaar ,
Thanks for the help. No, both are function passes. Is it the case that a function pass cannot modify the module?

4

The way I’ve done this (function outlining) after looking at other examples in the MLIR and IREE sources:

Create a module pass that loops over all functions in the module. For the set of ops you want to outline, first package them into a single operation with a single region that is used in an inline manner. Perform correct closure operations to either clone in dependencies or turn dependencies into arguments of the function. Then call mlir::applyPatternsAndFooldGreedily and apply a re-writer pattern to your special single-region function.

In the rewriter, there are basically a couple steps.

  1. Create the outlined function FuncOp using FuncOp::Create, and clone the single region of your special function using region.cloneInto(new_func_op.getBody()). At this point the pointer FuncOp new_func_op is not tracked and will be leaked if your program exits
  2. Grab reference to parent ModuleOp of the rewrite target operation under consideration. Use module.push_back to push in your new function op into the module.
  3. Replace your the target operation for the rewrite with a call operation to call your outlined function.

Edit: Realize the above didn’t answer the question directly at all. But basically what I was trying to say is that you could use a single module pass that calls multiple rewriters.

5

Yes, function passes can be scheduled independently on different functions in the module and so you should only modify what is in the function they are operating on (in scope of the isolated from above op being operated on). If you want to change the module, then use a module pass and inside iterate over top level functions in module inside the pass. In your case I believe you would also see a TSAN failure in your current config - the pass pipeline is taking advantage of the requirement of independence to schedule which would result in race if requirement doesn’t hold.

Try changing to module pass and see if that resolves it

7

More, generally the restrictions for what a pass can do are at Pass Infrastructure - MLIR

1 Like