RFC: Generic mem2reg in MLIR

Right now there is no mem2reg implementation in MLIR that would operate on arbitrary dialects. LLVM has of course its own mem2reg implementation, and Polygeist has an implementation of mem2reg but it is highly specialized for the memref dialect. We propose a design of a generic pass for MLIR that lets user specify mem2reg-related semantics for their dialects so they can easily promote memory slot patterns into direct uses of values. An implementation of this RFC is available for review.

Users can define allocating operations (for example alloca in LLVMIR or in memref) that define memory slots. A memory slot is a slot in memory containing a single value of a specific type, for which the content is accessed using memory operations (for example load and store in LLVMIR or in memref) via a unique non-aliased pointer value (for example llvm.ptr in LLVMIR or memref in memref).

Memory slots are modeled abstractly as a simple struct

struct MemorySlot {
  /// Pointer to the memory slot, used by operations to refer to it.
  Value ptr;
  /// Type of the value contained in the slot.
  Type elemType;
  /// Location of the slot in the source code. This may for example point to a
  /// local variable representing a memory slot on the stack.
  Location location;

such that memory slots are not bound to the results of allocating operations, allowing support for traditional result-style allocas, or region-style allocas (for which the slot pointer is passed as a region argument).

Sometimes the memory slot pointer is not used only by memory operations. For example, in LLVM IR, lifetime intrinsics can use the pointer, sometimes indirectly though bitcasts or similar operations. While it would be incorrect to allow those uses by default to ensure the slot pointer does not get aliased, such uses should not block promotion either. Indeed, there may exist situations in which the operation knows how to safely stop using the slot pointer (in the lifetime intrinsic example, by simply deleting itself).

Therefore, operations can be promotable operations, describing operations that have a procedure to stop using a set of uses, in the context of slot promotion. Promotable operations sometimes need to resolve further blocking uses to resolve their own blocking uses. This is notably the case when an operation can only resolve its uses by deleting itself, and must therefore request the users of its results to stop using them.

Note that memory operations can also be promotable operations, and in fact typically should considering they need to get rid of their use of the slot pointer.

The design is based around three interfaces. Their exact definitions can be found in the revision I submitted containing the implementation of this RFC. The following is a brief overall of their definitions.

  • The PromotableAllocationOpInterface, which represents allocating operations. This interface has a getPromotableSlots method returning an array of MemorySlot structs, describing which slots should be promoted provided they are used correctly (this is checked by the mem2reg pass, not the interface itself). The interface also has hooks to run logic for all new block arguments inserted and at the end of a successful promotion, typically for cleanup opportunities.
  • The PromotableMemOpInterface, which represents memory operations. This interface simply allows accessing which SSA value is stored or loaded from a given slot.
  • The PromotableOpInterface, which represents promotable operations. This interface is divided in two methods.
    • canUsesBeRemoved which decides given the set of uses it will have to drop if it can in fact drop them. This is used to abort promotion before mutating if promotion is impossible.
    • removeBlockingUses which performs the rewiring without deleting any operation. This method can return if the current promotable operation should be deleted. The deletion only happens at the very end of promotion to not delete operations the removal of other operations may depend on. This method also gets the reaching definition of the memory slot at this point (allowing for example load operations to replace their own uses with the reaching definition).

This design allows limiting dependencies between dialects. Indeed, with it, using a memory operation from dialect B on a slot allocated by an operation from dialect A can still lead to promotion. As long as dialect B implements the memory operation interface, the mem2reg pass will attempt to resolve the use no matter where the memory slot comes from.

The following is an overall view of how the pass operates:

  • The pass itself first looks for all operations that must be promoted in some way (failing if they cannot do the promotion).
  • It then finds all blocks for which two stores will compete for the definition of the slot and where a block argument should be created (failing if the block argument cannot be populated by terminators).
  • At this point, promotion is guaranteed to succeed so we can start mutating IR. The pass computes the definitions reaching each operation that must be promoted, inserting block arguments where necessary.
  • Finally, the promotable operations get to remove their blocking uses.

Feel free to leave any comment on this design for changes, and to drop a review on the code, that would be super appreciated!


Thanks a lot for your work. This looks like a well-thought-through interface!

Is the type in the MemorySlot structure really needed? It would seem to me that typing the memory location could create issues akin to the ones that lead to the introduction of opaque pointers in LLVM. A more detailed discussion on this topic is available here: Opaque Pointers — LLVM 17.0.0git documentation

Thank you!

The reason slots are typed is because sometimes memory operations may cause implicit casting, which would make promotion invalid if attempted. The type on the memory slot allows checking that types remain consistent across the uses of the slot (see the implementation of canUsesBeRemoved for LLVM’s StoreOp in the patch).

For example, the following LLVM dialect program is valid (albeit undefined behavior):

%size = llvm.mlir.constant(1 : i32) : i32
%alloca = llvm.alloca %size x i32 {alignment = 4 : i64} : (i32) -> !llvm.ptr
%some_i16 = llvm.mlir.constant(0 : i16) : i16
llvm.store %some_i16, %alloca {alignment = 4 : i64} : i16, !llvm.ptr
%res = llvm.load %alloca {alignment = 4 : i64} : !llvm.ptr -> i32
llvm.return %res : i32

If promotion of this was attempted, this would lead to invalid IR:

%some_i16 = llvm.mlir.constant(0 : i16) : i16
llvm.return %some_i16 : i32

because the store/load hides an implicit cast, which would be difficult to accomodate for. So instead, promotion is only performed in LLVM IR when types are consistent (the LLVM implementation does the same thing). Typing the slot allows making sure this is enforced.

When looking at the link many of the points that speak against typed pointers are actually not a disadvantage in our case. We are not worried about code size and also not introducing any bitcasts. Considering the memory slot structure is ephemeral, I think the problems with typed pointers do not really apply here. But do tell me if you can think of something.

Right, thank you for the detailed explanation! Let’s see what the code review leads to :smiley:

I have made a few changes to the details of the implementation of the interfaces.

I was reaching a case where I wanted to remove the promotable op infrastructure for trivially discardable operations outside of mem2reg. However, because providing a reaching definition was required by the interface, this was not possible. However, only memory operations require a reaching definition (because requiring a reaching definition to eliminate uses means the opereration accesses the slot in some way). So I made specialized methods for use elimination in the memory operation interface, and dropped the reaching definition requirement for promotable ops. Memory ops should thus no longer be promotable ops, the doubling of interfaces was weird in the first place.

This has been updated in the revision.