[RFC] Add an interface for top level container operations

Most FIR passes only look for FIR operations inside of functions (either because they run only on func.func or they run on the module but iterate over functions internally). But there can also be FIR operations inside of fir.global, some OpenMP and OpenACC container operations.

This has worked so far for fir.global and OpenMP reductions because they only contained very simple FIR code which doesn’t need most passes to be lowered into LLVM IR. I am not sure how OpenACC works.

Some FIR passes operate on whole modules because of concurrency concerns, but in general we prefer passes to operate on specific functions, so that the MLIR pass manager can run the pass on each function in parallel. Extending this to also support container operations other than functions, requires duplicating the pass - which leads to a lot of boilerplate for each top level operation. For example, see the duplication of the Abstract Result pass: llvm-project/flang/include/flang/Optimizer/Transforms/Passes.td at d9c31ee9568277e4303715736b40925e41503596 · llvm/llvm-project · GitHub

Instead, I propose we create an interface representing top level container operations which passes should run on, and all passes should be defined to run on operations with that interface.

There are some layering concerns here, because most of these top level operations exist outside of the FIR dialect (e.g. func.func). The interface can be added programmatically inside of flang without modifying any upstream MLIR dialects. See documentation here: Interfaces - MLIR

2 Likes

Thanks for the RFC. This is indeed something we need to fix. I’m surprised this has never been discussed in MLIR core or maybe I did not see the conversation.

Would it make sense to discuss this on the MLIR discourse so a general solution can be provided?

Just to double check, would this interface be used purely as a marker for “this operation is schedulable, and passes should be run on it”? How does this differ from the IsIsolatedFromAbove trait and the use of op-agnostic pass managers?

It is already possible to make passes run on any kind of operation by simply not specifying a specific kind of op the pass to implement: llvm-project/mlir/include/mlir/Transforms/Passes.td at 2542d34522784e27850a17a8c39a1ad4b43622ab · llvm/llvm-project · GitHub
In the pass pipeline manager it is also possible to make an op-agnostic pass manager which e.g. runs over all operations directly nested underneath the top-level. As long as none of the passes added to it filter the input operations (e.g. requiring a specific op kind), it should run on all ops implementing IsolatedFromAbove directly nested underneath the module: Pass Infrastructure - MLIR

1 Like

That’s a good point. Correctly setting IsolatedFromAbove to this top level operations should work. fir.global and reduction/privatization ops in OpenACC/OpenMP already have this interface.

The OpenMP target op ('omp' Dialect - MLIR) is Isolated from above but nested in other operations.

This should not be an issue.

I was worried that “isolated from above” isn’t really the criteria for being the target of our passes. We run the passes on things that contain FIR operations. I thought mixing the two would make the code harder to understand and might bite us one day if “isolated from above” and “contains FIR operations” diverge.

But if everyone thinks it is okay, I am happy to go ahead with that solution.

You have to differentiate:

  1. The pass definition: this is when you make a requirement for a pass about the op/ops it can be scheduled on (like “only functions because I need to operate on specific aspects of a function”)
  2. The pass pipeline: this is where you actually decide on which ops the pass will run: you can be more restrictive than the pass requirements, but not the other way around (if you try to schedule a pass restricted to functions to run on another op, it’ll assert).

Thanks for taking a look Mehdi.

For 2, are you referring to static schedule filtering? Would there be some advantage to us implementing this by overriding canScheduleOn instead of just filtering by an interface? I was under the impression that this was just syntactic sugar.

I believe that canScheduleOn is actually what implements 1).
On the other hand 2) is about how you construct your pass pipeline:

PassManager pm;

pm.addNestedPass<func::FuncOp>(createFooPass());
pm.addNestedPass<fir::GlobalInitOp>(createFooPass());

Here we’d run FooPass on FuncOp and fir::GlobalInitOp (I made it up).
What happens when you use canScheduleOn is that the pass manager would crash if the FooPass would define canScheduleOn as restricting to FuncOp (this is the part about “because they run only on func.func” in the first sentence of this RFC).

1 Like