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 agetPromotableSlotsmethod returning an array ofMemorySlotstructs, 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.-
canUsesBeRemovedwhich 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. -
removeBlockingUseswhich 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!
