[RFC] Support region control-flow in Mem2Reg

Moxinilian · March 6, 2026, 4:42pm

Hey!

We have been interested at NVIDIA to have support for region-based control-flow in the upstream Mem2Reg pass. We have been working on an implementation and we would like to make it available upstream. The corresponding PR is already up for review under the assumptions of the current state of the proposal, and is available to consult implementation details.

There are two aspects we would like to get feedback on before proceeding: 1. whether the proposed new interface for region control-flow fits the use cases you can imagine for your dialects, 2. the new implement introduces a new “edge case” in support, which we would like to get confirmation is acceptable.

Quick context / lexicon

Mem2Reg, also known formally as SSA construction, is the process of turning variables described in terms of load/store semantics into direct usages of SSA values. Historically, this pass has been useful in LLVM to turn stack-based allocations into SSA values (for ease of optimization and for convenience for front-ends that do not want to generate SSA form).

In the generic Mem2Reg implementation in MLIR, the abstraction is centered around the notion of a memory slot. This represents an isolated slot of memory created by an allocator. The MemorySlot C++ struct contains a slot pointer, an SSA value used to refer to the slot (it is typically of a pointer-like type), and the type of the value contained in the slot. The goal of the Mem2Reg transformation is to render the slot pointer unused, a process called slot promotion.

Which operations are allocators, and which operations can access the content of the memory is governed by a set of interfaces that cover it. They live upstream, in MemorySlotInterfaces.td. The interfaces allow operations accessing the memory slot to drop their uses of the slot by receiving a reaching definition of the slot. The reaching definition at a particular operation is an SSA value that contains what the slot would have contained at that point. For example, to drop its uses of the slot pointer, a load-like operation would typically replace its result value with the reaching definition, and delete itself.

The new interface

In order to communicate the control-flow semantics of region-bearing operations to the pass, we introduce a new interface. The role of the interface is to specify:

From the reaching definitions entering the region operation, what the reaching definition will be at the entry of its inner regions.
From the reaching definitions at the exit of all the blocks in its inner regions, what the reaching definition will be at the exit of the region operations.

As an easy example, scf.if will forward the entry reaching definition as-is to both its inner regions. In the case where there is a store in one of the regions (information provided to the interface), the interface will then add a result value to the if (and the region’s scf.yields) that will serve as the replacement reaching definition. In practice:

store %ptr, %old_reaching_definition
if %cond {
  store %ptr, %new_value
  yield
} else {
  yield
}

becomes

%res = if %cond {
  yield %new_value
} else {
  yield %old_reaching_definition
}

and %res becomes the new reaching definition.

The same principle can be extended to loops (by adding new block arguments in case of stores), and other operations. We were able to implement this interface for all the SCF operations, which you can find in the accompanying PR. They are, we believe, not very complex.

The interface definition is the following:

def PromotableRegionOpInterface
    : OpInterface<"PromotableRegionOpInterface"> {
  let description = [{
    Describes an operation for which memory slots can be promoted to SSA values
    within the operation's regions.
  }];
  let cppNamespace = "::mlir";

  let methods = [
    InterfaceMethod<[{
        Returns true when the provided region of the operation can be analyzed
        for promotion. The provided region must be a child of the operation's
        region.
        The `hasValueStores` flag indicates whether the region contains
        store-like operations that write to the memory slot.
      }], "bool", "isRegionPromotable",
      (ins
        "const ::mlir::MemorySlot &":$slot,
        "::mlir::Region *":$region,
        "bool":$hasValueStores
      )
    >,
    InterfaceMethod<[{
        Called before processing the nested regions in the operation.

        Based on the `reachingDef` value representing the value in the memory
        slot at the entry into the operation, `setupPromotion` fills in the
        `regionsToProcess` with the reaching definition at the entry of
        all its promotable regions.

        `setupPromotion` is allowed to mutate
        the operation in place, including its nested regions, but cannot
        delete existing operations or modify successor-bearing terminators.
        Other mutations are not allowed.

        The `hasValueStores` flag indicates whether the regions contain
        `store`-like operations that write to the memory slot. This field can be
        used to reduce the amount of book-keeping required to track the reaching
        definitions.
      }], "void", "setupPromotion",
      (ins
        "const ::mlir::MemorySlot &":$slot,
        "::mlir::Value":$reachingDef,
        "bool":$hasValueStores,
        "::llvm::SmallMapVector<::mlir::Region *, ::mlir::Value, 2> &":$regionsToProcess
      )
    >,
    InterfaceMethod<[{
        Called after promotion has been completed in all the relevant regions.

        Returns the new reaching definition at the exit of the operation. For
        this purpose, it is allowed to mutate the operation using the provided
        `builder`, along with its region contents. However, all blocks within
        the existing regions must remain valid and no new blocks may be added.
        As a result, the operation is allowed to be cloned and replaced only
        if its region content is moved from the original operation and not
        copied. Operations with an effect on the value of the slot must not
        change said effect (for example, new control flow that could change
        reaching definitions for a block is not allowed).

        The `entryReachingDef` is the reaching definition at the entry of the
        region operation.

        The `reachingAtBlockEnd` map contains the reaching definitions after all
        the terminators within the regions of the operation. If a block of the
        region is not present in the map, it is either dead code or within a
        region that does not interact with the value of the slot.

        The `hasValueStores` flag indicates whether the regions contain
        `store`-like operations that write to the memory slot. This field can be
        used to reduce the amount of book-keeping required to track the reaching
        definitions.
      }],
      "::mlir::Value", "finalizePromotion",
      (ins
        "const ::mlir::MemorySlot &":$slot,
        "::mlir::Value":$entryReachingDef,
        "bool":$hasValueStores,
        "::llvm::DenseMap<::mlir::Block *, ::mlir::Value> &":$reachingAtBlockEnd,
        "::mlir::OpBuilder &":$builder
      )
    >,
  ];
}

Opinions on the flexibility of this interface for your operations would be appreciated. Minor comments are also appreciated, and probably best suited on the accompanying PR.

Alternatives

In a different post, the option of not adding a new interface in favor of using RegionBranchOpInterface was explored. We did not pursue this option for two reasons:

Implementing RegionBranchOpInterface might be a bigger commitment with regards to semantics that what we request with this interface; and
Most importantly, RegionBranchOpInterface does not cover the mutation of control-flow edges necessary to add the reaching definition values (and adding such a feature would be difficult in a generic way as discussed in the post mentioned above). If we need an extra interface to do the rewiring, we might as well offer more flexibility, especially considering that RegionBranchOpInterface is not very useful to implement the interface above anyway^[1].

The new edge case

TLDR: there is a support regression in Mem2Reg for a super niche edge-case. Not supporting it improves performance and simplicity of both the algorithm and the interfaces.

On control-flow graphs, the Mem2Reg analysis must compute on which blocks two predecessing reaching definitions would disagree, which are called merge points. Indeed, at those places, there is more than a single SSA value that one could use to describe the content of the slot. The merge point analysis thus allows inserting a block argument at each merge point that will act as the replacement reaching definition. The predecessors will then forward their reaching definition as block operands, resolving the conflict.

This conflict resolution is only possible if all the predecessors implement the BranchOpInterface, as it is the standard mechanism the Mem2Reg implementation uses to add operands simply forwarding the reaching definitions. Not all terminators in the region need to be BranchOpInterface, only the ones preceding a merge point.

In practice, this means that the more merge points you have, the more constraint you have on terminators implementing BranchOpInterface. It is not a big constraint, but it does disallow promotion if it is hit.

Not all merge points can be considered live with respect to the value of the slot. Consider the following example pseudo-code:

func @dead_merge_point(%a, %b) {
  %ptr = alloca
  cond_branch ^store1, ^store2
^store1:
  store %ptr, %a
  branch ^merge
^store2:
  store %ptr, %b
  branch ^merge
^merge:
  return
}

^merge is technically a merge point, as two predecessors have different reaching definitions. However, with the current upstream implementation, promotion will still proceed if branch does not implement BranchOpInterface. This is because in the Mem2Reg pass, a liveness analysis designed to track the slot value prunes dead merge points before assessing the suitability of their predecessors: indeed, here, there is no need to add the block argument as we can see it would not be used.

This liveness analysis is a dataflow analysis that is efficiently and simply implemented by exploiting that basic blocks see no “interruptions”, and thus liveness can be reasoned about at basic block granularity. When introducing region control-flow, this is no longer the case. To be correct, the liveness analysis needs to traverse operations one by one, traversing through the nested regions. This is not only a performance burden but also complexifies the region interface significantly as it will now need to reason about propagating liveness information. All of this for this edge case of an edge case where non-branch ops are used to branch after a store that is not used.

Instead, to simplify the implementation, simplify the interface and improve performance on already supported cases, we decided to remove the currently O(basic block amount) liveness analysis entirely. Instead, a quick O(merge point amount) analysis at the end removes block arguments that were added and turned out unused (before linking them to the predecessors). This produces the exact same result as before on IR that saw successful promotion, except for the edge case above where it does block promotion, as the early-exit before mutation decision does not know yet that the merge point argument will not be used.

While technically a regression in capabilities, we believe it is acceptable that Mem2Reg does not handle it, for the benefit it brings. In particular, this is similar to the case where two branches would happen to have the same reaching definition value obtained from a different store (the merge point would also technically not be necessary yet the early exit decision cannot know this). Considering how rare this case is, it seems acceptable. To re-emphasize, this requires:

That non-BranchOpInterface terminators are used in a CFG context, and
That a store happens as the last use, with no load or other escape to observe it, and
That said store is laid out in the CFG such that it conflicts with another such store.

We cannot imagine this situation being frequent, and can be mitigated by the user by implementing BranchOpInterface. However, since it remains a regression, we would like to know if anyone depends on this being supported.

Are we good to go?

Thank you for reading and for your feedback!

I originally wanted to make the SCF implementation use generic helpers that exploited the RegionBranchOpInterface definitions, but it turned out much more cumbersome and less efficient to write that the quasi-one-liners needed without the interface. ↩︎

rengolin · March 6, 2026, 7:18pm

Moxinilian:

In practice:

store %ptr, %old_reaching_definition
if %cond {
  store %ptr, %new_value
  yield
} else {
  yield
}

becomes

%res = if %cond {
  yield %new_value
} else {
  yield %old_reaching_definition
}

and %res becomes the new reaching definition.

Does that mean that, for every op with regions, you’ll collect all reaching definitions and yield a union of all possible values on all regions?

In your example, there’s only one %ptr being tracked, but any store in either regions will have to be returned by the if operation (like a multi-phi node) and then use those SSA values in following code.

I think this is what you mean here:

Also…

Does that mean that liveness analysis now relies on merge points? Or is it just a different analysis when you do have them? Or rather, you’re counting any branch-to block a merge point and just changing the current analysis to match?

Does this benefit require Mem2Reg to have run? If there are no benefits without it, then it will just be a regression on the analysis side, and the trade-off will be different.

Moxinilian · March 6, 2026, 7:34pm

In the case of an if statement, each region will yield its own reaching definition. But of course the behavior is different for loops, this is all implemented using the interface. I think I may be misunderstanding the question.

I’m not sure I understand what the question is.

Merge points are only relevant to CFG control-flow, not region control flow. The reason for this is that CFG control flow is generic enough that the “merge resolution” is standard and incorporated in the Mem2Reg pass itself. In contrast, in region control-flow it is resolved on a per-operation basis, by the interface. This is tractable because the rigidity of region-based control-flow operations allows for the resolution to be simple. If you look at the implementation for SCF, for IfOp it’s doing the resolution in the finalizePromotion step where it collects all the reaching definition at exit points and sets the new if op result as replacement reaching definition.

The liveness analysis is completely gone, I have deleted it. It was only there to prune unused merge points, in a sense it was “optional” for the purpose of transforming the IR.

I’m sorry I don’t understand the question, could you rephrase? The benefits are in the execution time of the pass, in the simplicity of pass code, and in the simplicity of the interface.

gysit · March 8, 2026, 7:44pm

This is an exciting proposal and a long-awaited extension of the Mem2Reg pass!

After reviewing the PR, I’m wondering if it would be possible to generalize the interface beyond the scope of Mem2Reg. We have quite a few transformations downstream that use a homegrown SSA construction interface to introduce new SSA values into loops and conditionals.

It seems like only a few minor changes would be needed to make the interface more versatile:

Passing the SSA value type to the interface methods instead of the MemorySlot struct.
Renaming hasValueStores to something more generic like hasRedefinition.
Giving the interface a more generic name.

This would allow the logic to be reused for other transformations that require SSA construction but don’t strictly involve memory slots. Or am I missing some roadblock there?

rengolin · March 9, 2026, 9:44am

Turns out I took some statements in a much more general scope than they should have. Ignore this question.

For the same value. Take this code:

store %p0, %v0
...
store %pN, %vN
...
// 0 <= i, j <= N
if %cond {
  store %pi, %new_vi
  yield
} else {
  store %pj, %new_vj
  yield
}

In this case the different regions store to the different addresses, so the merge point isn’t the same, right? IE, will control flow operations with multiple regions need to return multiple merge points (one for each address written) and keep track of the unmodified addresses, too?

This would be the conversion for the example above:

%res0, %res1 = if %cond {
  yield %new_vi, %old_vj
} else {
  yield %old_vi, %new_vj
}

Otherwise, with a runtime %cond, you wouldn’t know what value to bind %res.

What about this conditional:

if %cond {
  // Runtime offset
  %new_address = add %address_of_some_p, %runtime_value
  // Pointer of pointer
  load %p, %new_address
  // No way to know what happens here
  store %p, %new_value
  yield
}

Here I imagine you just bail your analysis and don’t try to convert, right?

Moxinilian · March 9, 2026, 2:53pm

I’m very happy with generalizing the interface if it is useful, especially if it does not change it much! My general perspective on interface specificity is that it is good to start with specialized interfaces and from there define more general ones, as we can then offer users the option of implementing the specific one if the general one is too general, and provide opt-in implementations of the specialized one for all operations implementing the general one for free. But here it does not seem to be this much of a generalization.

I originally thought knowing the slot pointer could be useful but I struggle to come up with a case, so your described option sounds reasonable to me. I believe you suggest we do this for general load-store propagation, is that the case?

Moxinilian · March 9, 2026, 3:01pm

In the sense of this RFC and the Mem2Reg implementation, a merge point is a block that has two CFG predecessor blocks in which the exit reaching definition is not the same. It is not relevant to region control flow in itself.

I believe you are concerned about cases with multiple allocas. Each alloca is treated independently, and their order of processing does not matter. If an if statement contains stores about two different allocas, processing one will add a block argument, then processing the second will add a second argument (doing so with fewer arguments is not possible in the general case anyway). Those are effectively two separate invocations of the mem2reg pass.

When the interface says hasValueStores, it implicitly only covers stores to the provided MemorySlot.

Moxinilian · March 9, 2026, 3:03pm

I’m not sure I understand exactly what the example covers, but if you are concerned about pointer aliasing of the MemorySlot, yes, this is currently not supported and would bail out (unless the aliasing implements the interfaces in an incorrect way, of course). I have a draft design to have something for that, but it is out of the scope of the current proposal.

gysit · March 9, 2026, 3:04pm

I would say it doesn’t even need to be load-store propagation. We may just want to construct a new SSA value. An example could be that we want to add a boolean variable to a loop nest that is set to true if a specific conditional region inside the loop nest did execute at least once. Your interface could be used for this directly - rather allocating a stack variable and running Mem2Reg to get rid of it again.

Moxinilian · March 16, 2026, 10:45am

As you pointed out in the PR review, the main challenge with generalising this interface turns out to be less in the signature than in the implicit contract of the implementation. Mem2Reg has expectations about the behaviour of the interface, notably with regards to mutation, that are easy to fulfil but rich to specify. It is to a point that I’m worried this part of the contract is very Mem2Reg specific.

If you don’t mind, I would like to take on your offer to postpone generalisation. If we find out a more general interface, we can either deprecate this one or offer a generic implementation of this one from any op that implements the more generic interface. I think neither of these options are a concern.

gysit · March 16, 2026, 1:46pm

Postponing the generalization sound reasonable to me.

Moxinilian · March 23, 2026, 9:59pm

The PR has been merged! Thanks everyone for the comments.

Topic		Replies	Views
Current status of mem2reg for scf/region-based IR MLIR	7	507	March 6, 2026
[RFC] Structural and Reachable Control Flow Edges in `RegionBranchOpInterface` MLIR	20	537	April 22, 2026
Structured Control Flow is Not Necessarily Regions MLIR	18	1938	April 16, 2020
[RFC] Region-based control-flow with early exits in MLIR MLIR	39	3885	March 24, 2026
[PSA] RegionBranchOpInterface revamping MLIR	4	455	October 18, 2025

[RFC] Support region control-flow in Mem2Reg

Quick context / lexicon

The new interface

Alternatives

The new edge case

Related topics