[mem2reg] Enhance to ignore @llvm.assume(icmp ne null) uses.

[DiamondLovesYou]

Summary

These checks should be dead after promotion, so don't block promotion because of them.

[jdoerfert]

While I see how this is generally useful, I fear the implementation goes in the wrong direction.
What I mean is that we are moving away from instruction-based encodings for assume, see for example D74209.
What we need here is basically D73832.

@Tyker @lebedev.ri Maybe lifetime markers are the second use case for the "drop users" interface (D73404).

[DiamondLovesYou]

In D75462#1902205, @jdoerfert wrote:

While I see how this is generally useful, I fear the implementation goes in the wrong direction.
What I mean is that we are moving away from instruction-based encodings for assume, see for example D74209.
What we need here is basically D73832.

@Tyker @lebedev.ri Maybe lifetime markers are the second use case for the "drop users" interface (D73404).

Interesting. I see your point, and I agree. Perhaps accept this until that lands, then revert this (or just leave inplace until -mem2reg is deleted? If that's the long term plan?)? I also have another patch similar to this which ignores alignment assumptions, FYI. These missed optimizations are very expensive for a target like AMDGPU because of unnecessary stack usage.

Sadly, my use case is for Rust, which ATM is a bit behind LLVM's release cadence (they are currently on a patched 9.0; see https://github.com/rust-lang/rust/pull/67759). So, at the very least I'll have to keep this patch local until that work lands and Rust is upgraded.

[jdoerfert]

Given that you want/have to keep this downstream for a while anyway, I would prefer that is all we do, though that is not to say I can't be convinced.

[lebedev.ri]

Rebase now that there has been changes wrt assume handling?

[lebedev.ri]

This review seems to be stuck/dead, consider abandoning if no longer relevant.

llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp

if (const LoadInst *LI = dyn_cast<LoadInst>(U)) {
         return false;
     } else if (const GetElementPtrInst *GEPI = dyn_cast<GetElementPtrInst>(U)) {
       if (GEPI->getType() != Type::getInt8PtrTy(U->getContext(), AS))
         return false;
       if (!GEPI->hasAllZeroIndices())
         return false;
       if (!onlyUsedByLifetimeMarkers(GEPI))
         return false;
+    } else if (const auto *Cmp = dyn_cast<ICmpInst>(U)) {
+      // check if this is a @llvm.assume(icmp ne null), which can be deleted
+      // while promoting.
+      if (Cmp->getPredicate() != CmpInst::ICMP_NE) {
+        return false;
+      }
+      unsigned OtherOperandIdx = 1;
+      if (Cmp->getOperand(0) != AI) {
+        OtherOperandIdx = 0;
+      }
+      auto *Lhs = Cmp->getOperand(OtherOperandIdx)->stripPointerCasts();
+      if (!isa<ConstantPointerNull>(Lhs)) {
+        return false;
+      }
+      for (const User *ICmpU : Cmp->users()) {
+        if (const auto *II = dyn_cast<IntrinsicInst>(ICmpU)) {
+          if (II->getIntrinsicID() == Intrinsic::assume) {
+            continue;
+          }
+        }
+        return false;
+      }
     } else {
       return false;
     }
   }
 
   return true;
 }
 
static void addAssumeNonNull(AssumptionCache *AC, LoadInst *LI) {
   ICmpInst *LoadNotNull = new ICmpInst(ICmpInst::ICMP_NE, LI,
                                        Constant::getNullValue(LI->getType()));
   LoadNotNull->insertAfter(LI);
   CallInst *CI = CallInst::Create(AssumeIntrinsic, {LoadNotNull});
   CI->insertAfter(LoadNotNull);
   AC->registerAssumption(CI);
 }
 
-static void removeLifetimeIntrinsicUsers(AllocaInst *AI) {
+static void removeLifetimeAndAssumeIntrinsicUsers(AllocaInst *AI) {
   // Knowing that this alloca is promotable, we know that it's safe to kill all
   // instructions except for load and store.
 
   for (auto UI = AI->user_begin(), UE = AI->user_end(); UI != UE;) {
     Instruction *I = cast<Instruction>(*UI);
     ++UI;
     if (isa<LoadInst>(I) || isa<StoreInst>(I))
       continue;
 
     if (!I->getType()->isVoidTy()) {
-      // The only users of this bitcast/GEP instruction are lifetime intrinsics.
+      // Remove lifetime/assume intrinsics.
       // Follow the use/def chain to erase them now instead of leaving it for
       // dead code elimination later.
       for (auto UUI = I->user_begin(), UUE = I->user_end(); UUI != UUE;) {
         Instruction *Inst = cast<Instruction>(*UUI);
         ++UUI;
         Inst->eraseFromParent();
       }
     }
void PromoteMem2Reg::run() {
 
   for (unsigned AllocaNum = 0; AllocaNum != Allocas.size(); ++AllocaNum) {
     AllocaInst *AI = Allocas[AllocaNum];
 
     assert(isAllocaPromotable(AI) && "Cannot promote non-promotable alloca!");
     assert(AI->getParent()->getParent() == &F &&
            "All allocas should be in the same function, which is same as DF!");
 
-    removeLifetimeIntrinsicUsers(AI);
+    removeLifetimeAndAssumeIntrinsicUsers(AI);
 
     if (AI->use_empty()) {
       // If there are no uses of the alloca, just delete it now.
       AI->eraseFromParent();
 
       // Remove the alloca from the Allocas list, since it has been processed
       RemoveFromAllocasList(AllocaNum);
       ++NumDeadAlloca;

llvm/test/Transforms/Mem2Reg/ignore-assume-nonnull.ll

+; RUN: opt -mem2reg -S -o - < %s | FileCheck %s
+
+declare void @llvm.assume(i1)
+
+define void @test1() {
+; CHECK: test1
+; CHECK-NOT: alloca
+  %A = alloca i32
+  %b = icmp ne i32* %A, null
+  call void @llvm.assume(i1 %b)
+  store i32 1, i32* %A
+  ret void
+}