[PSA] RegionBranchOpInterface revamping

Just some PSA about breaking changes on-going to the RegionBranchOpInterface. In preparation of the introduction of early-exits in structured control-flow I’ve been looking at this interface, and I’m making already some NFC breaking changes addressing some technical debt with the existing design: the concepts of RegionBranchPoint and RegionSuccessor have been quite incorrectly set up in the past, and this PR will align it better: [MLIR] Revamp RegionBranchOpInterface by joker-eph · Pull Request #161575 · llvm/llvm-project · GitHub
Not everything is completely solved though, we still have some issues with properly tracking the dataflow when trying to infer it in a backward way using this interface and more updates may be needed: so you might expected more breaking changes to this interface in the future. Hopefully this PR will provide better foundations to get there though.

I am not able to follow the details of what you are trying to fix. Could you clarify what you mean by “early-exits in structured control flow”. At the outset that feels like something structured control flow (like scf.for ) does not support. Making this happen could break analysis that might not be accounting for it/assuming that scf.for does not allow this.

Sorry, I didn’t mean to confuse anyone with this, I meant it as context of what I’m doing (I’ve been on and off trying to get a concrete patch and proposal to land the early-exit feature we’ve been discussing for the last 2 years or so). This patch does not include anything related to early-exit though: I just encountered some APIs issue and went down a rabbit hole with this interface.

Let me paste the description of the patch here:

Now let’s look at these two classes, they already exist in the codebase, I didn’t make these up.
Here is the RegionSuccessor doc and here is the RegionBranchPoint one.

I didn’t change much these class. You can mostly note a change to the RegionBranchPoint: it was modeling the “branch point” as a “region” instead of the terminator, now it captures the actual terminator that is the branch point.
There still are some issues (I left some TODOs to not bloat the patch too much): a RegionSuccessor is implicitly constructible from a Region because the second argument of the constructor has a default value:

  RegionSuccessor(Region *region, Block::BlockArgListType regionInputs = {})

This is a problem because the RegionSuccessor is supposed to define the block args that are used set when the control flow is transferred to this region. The fact that there is implicit conversion means that we’re sometimes calling some API with just a Region and the successor is incorrectly constructed. I found a few bugs, but not all,

The important changes are on the interface itself. Because RegionBranchPoint and RegionSuccessor could both be the parent operation or a region, that means they could easily be mixed up, and the methods of the interface were just incorrect with respect to this, e.g. getEntrySuccessorOperands is documented as

Returns the operands of this operation that are forwarded to the region successor’s block arguments or this operation’s results when branching to point"

You may notice the issue: the RegionBranchPoint point is not something you branch to but something you branch from. This API was meant to take a RegionBranchSuccessor here (which is where I started my update). The same thing for getMutableSuccessorOperands.

Finally, as part of this, I also introduced some helper methods on the interface with a default implementation: getPredecessors/getPredecessorValues which were helpers in LocalAliasAnalysis.cpp. The interface is designed so that you can get successors for any given branch point, however the opposite (finding the predecessors) isn’t easy with the interface: for a given successor, you need to iterate all the branch point for the RegionBranchOpInterface operation, query all the successors for each of these branch point, and filter to find the target successor. The default implementation should be functional, but not efficient, the method can be overridden to provide a straight implementation though.
I think we should rethink the API further to make it more efficient to walk the control-flow in both direction without implementing separately both for each op. This could be done by asking the op to populate a “CFG” structure instead of just “successor edges” for a given node.
This would be a much larger change though, I tried to keep it minimal here.

Hope this helps.

I would add: I’m interested to brainstorm this further if anyone is motivated. I’ll be at the LLVM conference, maybe we can have an impromptu whiteboarding session for thinking about this interface and the API?

Another topic about this interface is its compatibility with parallelism: mapping a regular CFG on our structured control-flow has been the underlying principle of this API, however it falls apart somehow when considering things like a scf.forall operation: the interface does not convey the fork/join happening as part of transferring the control to the region and back to the parent. This may or may not be that much of a problem, but can be interesting to discuss as well.

Progress on early-exits is exciting! I have a talk at the upcoming MLIR workshop which discusses the tricks we use to support early exits in JSIR. It would be much better if everything is directly supported by upstream MLIR so we don’t have to maintain our own workarounds.

TLDR: We have a wrapper on top of the MLIR dataflow analysis API, in which we define a GeneralCfgEdge that represents a jump between a pair of arbitrary mlir::ProgramPoints. This unifies cf.{br,cond_br} and region branches including early exits, which makes the dataflow analysis API easier to understand.

I’m hoping that this idea of a general CFG edge can be used to simplify the upstream MLIR dataflow analysis API a bit. Currently, cf.{br,cond_br} and Region{BranchPoint,Successor} are handled differently, which seems to lead to code duplication.