This is an archive of the discontinued LLVM Phabricator instance.

[mlir] Add a generic mem2reg implementation.
ClosedPublic

Authored by Moxinilian on Apr 12 2023, 3:57 AM.

Details

Summary

This patch introduces a generic implementation of mem2reg on
unstructured control-flow, along with a specialization for LLVM IR. This
is achieved by defining three new interfaces, representing 1. allocating
operations, 2. operations doing memory accesses, 3. operations that can
be rewired and/or deleted to stop using a specific use.

The file containing the core implementation of the algorithm
(Mem2Reg.cpp) contains a detailed explanation of how the algorithm
works. The contract for this pass is that given a memory slot with a
single non-aliased pointer, the pass will either remove all the uses of
the pointer or not change anything.

To help review this patch, I recommend starting by looking at the
interfaces defined in Mem2Reg.td, along with their reference
implementation for LLVM IR defined in LLVMMem2Reg.cpp. Then, the core
algorithm is implemented in Mem2Reg.cpp.

If this is all good I also have an implementation of the interfaces for
0-dimensional memref promotion that I can upstream afterwards.

Diff Detail

Event Timeline

Moxinilian created this revision.Apr 12 2023, 3:57 AM
Herald added a reviewer: dcaballe. · View Herald Transcript
Herald added a project: Restricted Project. · View Herald Transcript
Moxinilian requested review of this revision.Apr 12 2023, 3:57 AM
Moxinilian edited the summary of this revision. (Show Details)Apr 12 2023, 4:01 AM
Moxinilian edited the summary of this revision. (Show Details)Apr 12 2023, 4:01 AM
zero9178 added inline comments.Apr 12 2023, 4:44 AM
mlir/include/mlir/Transforms/Mem2Reg.td
1 ↗(On Diff #512758)

Is there any precedence for having an interface within Transforms? I think even if the interface is only used by one single transform (at the moment at least), it is more conventional to have them within the Interfaces directory.

In particular, having it within Transforms would create a few layering issues here:
MLIRLLVMDialect dialect would have to depend on the MLIRTransforms library (something this patch currently doesn't do, which would likely cause linking errors in shared library builds).

If it were to instead be within Interfaces and a library on its own there you could have both MLIRLLVMDialect and MLIRTransforms depend on the interface and be properly decoupled.

42 ↗(On Diff #512758)

ultra nit: mlir::MemorySlot is currently 24 bytes on 64 bit machines. That'd be worth passing as const mlir::MemorySlot& instead probably.
Ditto below

114 ↗(On Diff #512758)

Should this be const ::llvm::SmallPtrSetImpl<::mlir::OpOperand *> &?

Ditto below

126 ↗(On Diff #512758)

What is onlyIf here? Is this a stale description?

zero9178 added inline comments.Apr 12 2023, 4:44 AM
mlir/lib/Dialect/LLVMIR/IR/LLVMMem2Reg.cpp
35

This is a mismatch between what LLVM currently does: https://godbolt.org/z/vKTE59v7o
It creates a undef operation instead of poison. That said I think there is a transition away from undef in LLVM IR, although this has seemingly not yet been done to mem2reg. Was the choice of posion here deliberate?

44

I might be misunderstanding the logic here, but shouldn't the logic to replace DbgDeclareOp with DbgValueOp rather be done in the removeBlockingUses implementation of DbgDeclareOp?

168

Doesn't LifetimeStartOp (and end), not have any results, making this a noop?

Moxinilian added inline comments.Apr 12 2023, 5:26 AM
mlir/lib/Dialect/LLVMIR/IR/LLVMMem2Reg.cpp
35

Yeah this is I think an interesting point. In the LLVM mem2reg, some things have been replaced with poison, but for some reason not the default value for the allocas. We spent a lot of time discussing whether or not we should use undef or poison here. We settled on poison at the time because 1. there is the transition away from undef in LLVM and 2. poison would make sense semantically here, especially considering loading uninitialized variables is undefined behavior in C.

However I just realized that the langref in fact specifies that this should be undef unfortunately. I personally find that poison would be a better fit but if undef is specified I think we should stick to that then.

44

This is an unfortunate tricky point with respect to debug intrinsics. A value debug intrinsic should be generated at each store and at each merge point. If this was added to removeBlockingUses, this would only be called when removing the declare intrinsic, without knowledge of where the merge points are. So instead we have this hook on block arguments to insert it there. This is not super pretty but it is the most generic approach we could find there. If you have other suggestions that would be very welcome, but otherwise I'm not sure if there is much better to do.

Moxinilian marked 4 inline comments as done.

Address comments.

Moxinilian marked 2 inline comments as done.Apr 12 2023, 6:05 AM
Moxinilian marked an inline comment as done.

Update default for LLVM to undef

zero9178 added inline comments.Apr 12 2023, 8:02 AM
mlir/lib/Dialect/LLVMIR/IR/LLVMMem2Reg.cpp
44

I see, thank you for the explanation! I was initially mostly taken back by it being the allcoas responsibility to do so, but I don't know any better either and think it makes the most sense.

120

I think this pattern of "builder guard" & "set insertion point to *this" should be moved to the caller of removeBlockingUses, not be present in the implementation of interfaces.

Having the transformation pass do it guards against easy misuse of the API (an implementation accidently forgetting to reset the insertion point). I think it is also common expectation in MLIR (due to things like the RewritePatterns) to expect a builder to have an insertion point after the Op it is acting upon. Ops can always change their insertion point if they have different needs, but the vast majority likely needs it after the op itself.

Address builder guard comment.

Moxinilian marked an inline comment as done.Apr 13 2023, 1:53 AM

Remove default implementations for interfaces.

rriddle added inline comments.Apr 16 2023, 11:04 PM
mlir/lib/Transforms/Mem2Reg.cpp
11

Everything in the Transforms library must be dialect agnostic, is this an accidental include here?

Remove accidental include.

Adjust incorrect comments.

Relaxed interfaces a bit. The promotable op interface was a bit too restrictive for non-mem-op operations. Instead, specialized methods for memory ops where moved to the mem op interface. ALso relaxed the loaded value requirements.

gysit added inline comments.Apr 20 2023, 5:36 AM
mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
176

nit: can these be dropped?

mlir/include/mlir/Interfaces/Mem2RegInterfaces.td
2

nit: Mem2RegInterfaces.td

Note: the commit message may also need an update after renaming.

194

ultra nit: Let's drop one empty line here.

mlir/lib/Dialect/LLVMIR/IR/LLVMMem2Reg.cpp
205–206

nit: Could we use forwardToUsers here?

mlir/lib/Transforms/Mem2Reg.cpp
235

nit: drop llvm::

mlir/test/Transforms/mem2reg-llvmir.mlir
49

nit: maybe add a check promotion happened.

78

nit: I similarly would add a check if promotion happened.

gysit added inline comments.Apr 20 2023, 5:51 AM
mlir/test/Transforms/mem2reg-llvmir-dbginfo.mlir
100–101

Are these last two check lines needed? I would not expect any code there?

mlir/test/Transforms/mem2reg-llvmir.mlir
83

nit: Can you upper case the test variables like [[bb2]] or [[v3]]. Uppercase makes it easier to distinguish them from the IR.

Mogball added inline comments.Apr 20 2023, 12:26 PM
mlir/include/mlir/Interfaces/Mem2RegInterfaces.h
27

Is this location different than the location of the Value?

mlir/include/mlir/Interfaces/Mem2RegInterfaces.td
28

Couldn't you assert that the slot is a result of the operation in question? Then it is impossible for more than one operation to return the same slot.

40

can you provide examples of IR and operations and what the before/after of calling this method is? It's not clear what the operation is supposed to do when implementing this hook just from the description of the method

63

Why not have the interface provide a hook that lazily creates the default value instead? It's better to materialize IR only when necessary

mlir/lib/Dialect/LLVMIR/IR/LLVMMem2Reg.cpp
63

you could honestly provide default implementations for these methods

mlir/lib/Transforms/Mem2Reg.cpp
254

why 16? The general rule is to use the default value unless you have performance numbers to justify a particular choice.

326

drop llvm::

Mogball added inline comments.Apr 20 2023, 12:30 PM
mlir/lib/Transforms/Mem2Reg.cpp
172

This seems like something that could be added to DenseMap

Moxinilian marked an inline comment as not done.Apr 21 2023, 1:47 AM
Moxinilian added inline comments.
mlir/include/mlir/Interfaces/Mem2RegInterfaces.td
28

There are legitimate reasons for that to not be the case even though this pass supports it. For example, one could design their allocating operations using regions, if they need interesting lifetime control. Example:

dialect.region_alloca {
   ^bb0(%slot: ptr<i32>):
      // ...
}

In that case, the slot is not a result of the allocator.

63

Right. The reason I am doing it right now is because deciding the default value should be used comes in late, and lazily creating it might invalidate ongoing iterators. Let me see if I can still make it happen, the design has changed a bit since I picked this option.

mlir/lib/Transforms/Mem2Reg.cpp
254

Correct me if I am wrong but I believe SmallPtrSet simply does not have a default?

Moxinilian marked 12 inline comments as done.Apr 21 2023, 7:47 AM
Moxinilian added inline comments.
mlir/include/mlir/Interfaces/Mem2RegInterfaces.td
40

I have added more details to the documentation, although I think the LLVM IR implementation showcases a good example? This value is basically the reaching definition used if there is no reaching definition at the point of a load.

63

Okay so I have made a few changes towards this goal. Now the default value is generated lazily. The only problem is that if it is generated to provide as a reaching definition for a MemOp promotion, but the MemOp does not end up using it, it has been generated for nothing. Right now that means I have left the ability to do clean-up in the completion handler. I can picture a world where the interface method takes a lambda that would generate the default value if need be but it is becoming excessively complicated in my opinion. Tell me if you have a better idea.

mlir/lib/Dialect/LLVMIR/IR/LLVMMem2Reg.cpp
63

So we discussed this a bit internally, and the problem with default implementations here is that it is easy to forget handling a case, and that makes for hard to detect problems. Are you sure this is worth it?

Address comments.

Herald added a project: Restricted Project. · View Herald TranscriptApr 21 2023, 7:47 AM

Remove formatting changes in DenseMap.

gysit accepted this revision.Apr 24 2023, 2:33 AM

LGTM modulo the assertion comment.

mlir/lib/Transforms/Mem2Reg.cpp
189

This may trigger a warning if compiled without assertion. The standard way to avoid this is putting

(void)isResultOrNewBlockArgument;

before the assertion. Also consider starting the search from the slot.ptr and then call get defining op / cast it to a block argument, e.g:

auto isResultOrBlockArgument = [&](Value value) {
  if (value.getDefiningOp() == allocator)
    return true;
  if (BlockArgument arg = value.dyn_cast<BlockArgument>())
    if (arg.getOwner()->getParentOp() == allocator)
      return true;
  return false;
}
This revision is now accepted and ready to land.Apr 24 2023, 2:33 AM
kuhar added inline comments.Apr 24 2023, 8:55 AM
mlir/lib/Transforms/Mem2Reg.cpp
189

In this specific case, shouldn't the whole body of this constructor be guarded by #ifndef NDEBUG?

IE this block does not produce any results/side-effects used in non-debug builds.

424–425

nit: for (auto &user: llvm::make_first_range(userToBlockingUses))?

528

nit: public is the default for struct and can be omitted here

gysit added inline comments.Apr 24 2023, 9:13 AM
mlir/lib/Transforms/Mem2Reg.cpp
189

Yes that sounds right. Without NDEBUG there is probably an unused variable warning as well.

Moxinilian marked 7 inline comments as done.

Address comments.

Moxinilian marked 4 inline comments as done.Apr 25 2023, 1:42 AM
Moxinilian added inline comments.
mlir/include/mlir/Interfaces/Mem2RegInterfaces.td
28

I now assert that the block is either a result or an argument of one of its regions. Those seems like the only reasonable options anyway (if one needs something else, they should define the interface on the operation that actually generates the slot).

Rebased to latest main.

Fix formatting.

This revision was automatically updated to reflect the committed changes.
kuhar added inline comments.Apr 27 2023, 10:56 AM
mlir/lib/Transforms/Mem2Reg.cpp
189

This hasn't been fully addressed

246–247

nit: Use llvm::is_contained

It's not clear to me

mlir/include/mlir/Transforms/Passes.td
191

Is it conservative with respect to aliasing? How is this handled?

Hopefully this clarifies!

mlir/include/mlir/Transforms/Passes.td
191

Aliasing is structurally not allowed: for promotion to happen, all users of the alloca must implement one of the mem2reg interfaces, and those do not allow aliasing. I agree that this is not very clear here and in the documentation of the interfaces. I will open a subsequent revision to make this obvious.