[mlir][mem2reg][RFC] Promote memory slots through transparent view operations

MLIR
1

Problem

MLIR’s mem2reg only recognizes loads and stores whose memref operand is exactly the slot pointer. Any “transparent view” op sitting between the slot and its load/store — an op that takes a slot pointer and returns another pointer aliasing the same memory, possibly at a different element type — currently blocks promotion. This pattern is quite common: memref.assume_alignment after memref.alloca, llvm.bitcast between llvm.alloca and llvm.load, fir.declare tying a Fortran variable to its storage, fir.convert casting between FIR pointer types.

Proposed API

Two new methods on PromotableOpInterface:

// Describes this op as a transparent view of a memory slot.
::std::optional<::mlir::PromotableSlotView> getPromotableSlotView();

// Bridges a value between the underlying slot's elemType and the view's
// elemType (the framework calls this in both directions for type-changing
// chains; default impl handles the identity case).
::mlir::Value convertSlotValue(::mlir::Value value, ::mlir::Type targetType,
                               ::mlir::OpBuilder &builder);

mem2reg uses these new getPromotableSlotView APIs to traverse the use-def chain, and convertSlotValue to convert the value around load/stores.

The mem2reg pass changes are minimal, the reachingDef value is still tracked with the root element type, and no new data structures are added to keep track of the slot views.

Motivation examples

memref.assume_alignment is currently blocking mem2reg in the following simple example:

%a = memref.alloca() : memref<i32>
%aa = memref.assume_alignment %a, 4 : memref<i32>     // blocks promotion today
memref.store %v, %aa[] : memref<i32>
%r = memref.load %aa[] : memref<i32>

fir.declare is used to mark what source variable the memory access are made through (for both debug and aliasing purposes). It currently has some mem2reg support which involves walking back the fir.load/store operands, but this workaround is limits loads/stores from being in different blocks. The proposal will allow removing the workaround and limitation.

%a = fir.alloca i32
%d = fir.declare %a {uniq_name = "x"} : (!fir.ref<i32>) -> !fir.ref<i32>
fir.store %v to %d : !fir.ref<i32> 
%r = fir.load %d : !fir.ref<i32>

fir.convert (typed pointer cast). The framework will chain convertSlotValue calls so loads/stores see values at the view’s element type while the reaching definition is tracked at the root. This use case is relevant for Fortran EQUIVALENCE (union):

%a = fir.alloca i32
%p = fir.convert %a : (!fir.ref<i32>) -> !fir.ref<f32>
fir.store %f to %p : !fir.ref<f32>            // mem2reg will materializes f32→i32
%r = fir.load %p : !fir.ref<f32>              // and i32→f32 around store/load

The implementation for this RFC is for review in PR196924.

2

Hey! Thanks for the proposal. This is a much needed feature.

I had already given some thought to this. This seems close to what I was thinking about already, but could you clarify a few things, like what PromotableSlotView is? The implementation strategy is not super clear to me but I’d like to understand it to compare better.

In any case, I’m a little bit wary of extending the existing PromotableOpInterface when this change is about slot pointer aliasing before anything else, in my opinion.

The way I was thinking of doing it myself was by having a PromotableAliasingOp interface, that would declare an aliaser: it would specify that a particular SSA value pointer is an alias to the memory slot, and provide a routine to resolve a read/write into the aliaser as a read/write into the original memory slot. This feels similar to your approach but I’m not sure how similar it is from the RFC alone.

cc @gysit

edit: I missed that a PR is already up. I’ll go check it out.

3

Hi @jeanPerier, I noticed that the PromotableAliaserInterface framework has been merged into MLIR.
I’m currently working on Fortran programs and very eager to try it out. Are there any upcoming PRs or active branches that implement this interface for fir.declare and fir.convert?