[Mem2Reg] Teach promote to register about droppable instructions

[jdoerfert]

Summary

This is the first of two patches to address PR46753. We basically allow
mem2reg to promote allocas that are used in doppable instructions, for
now that means llvm.assume. The uses of the alloca (or a bitcast or
zero offset GEP from there) are replaced by undef in the droppable
instructions.

[jdoerfert]

Remove code duplication

[Tyker]

thank you for working on this

llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp

347

is there a reason you are not using dropDroppableUses here.

replacing uses in an assume by undef is usually not a safe transformation. perhaps there is UB is not exploited for operand bundles but it is very likely to be exploited for the usual boolean argument of assume.

[jdoerfert]

Use proper function to drop uses

[Tyker]

just one nit

llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp

324

using dropDroppableUses this way is a bit weird since it is already doing a traversal through the uses but it is correct.

maybe an API to drop a single use would be appropriate.

[jdoerfert]

llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp

324

I'll look into it and update this. I'll commit it if it seems straight forward, post review if anyone disagrees.

[jdoerfert]

Introduce API to drop user uses

[efriedma]

llvm/lib/IR/Value.cpp

197 ↗

(On Diff #280542)

This seems like a dubious pattern; you're iterating over the entire use list of the value each time you call dropDroppableUsesByUser, which is likely a lot more expensive than iterating over the operand list of the assume.

[fhahn]

llvm/lib/IR/Value.cpp

197 ↗

(On Diff #280542)

@jdoerfert @Tyker Any more thoughts on Eli's point?

[jdoerfert]

(I missed the comment, apologies)

llvm/lib/IR/Value.cpp

197 ↗

(On Diff #280542)

I have no particular preference. Unclear what is better (over time). I'll propose a follow up patch that uses the assume operands instead.

[efriedma]

llvm/lib/IR/Value.cpp

197 ↗

(On Diff #280542)

To be precise, my concern is that the overall algorithm is O(N^2) over the number of uses of the alloca.

[jdoerfert]

llvm/lib/IR/Value.cpp

197 ↗

(On Diff #280542)

D84804 should address this.

[Tyker]

llvm/lib/IR/Value.cpp

197 ↗

(On Diff #280542)

i agree with Eli, i tried express my concern about that in https://reviews.llvm.org/D83976?id=280542#inline-775547

llvm/include/llvm/Analysis/ValueTracking.h

class Value;
   /// ptrtoint+arithmetic+inttoptr sequences.
   bool getUnderlyingObjectsForCodeGen(const Value *V,
                             SmallVectorImpl<Value *> &Objects,
                             const DataLayout &DL);
 
   /// Return true if the only users of this pointer are lifetime markers.
   bool onlyUsedByLifetimeMarkers(const Value *V);
 
+  /// Return true if the only users of this pointer are lifetime markers or
+  /// droppable instructions.
+  bool onlyUsedByLifetimeMarkersOrDroppableInsts(const Value *V);
+
   /// Return true if speculation of the given load must be suppressed to avoid
   /// ordering or interfering with an active sanitizer.  If not suppressed,
   /// dereferenceability and alignment must be proven separately.  Note: This
   /// is only needed for raw reasoning; if you use the interface below
   /// (isSafeToSpeculativelyExecute), this is handled internally.
   bool mustSuppressSpeculation(const LoadInst &LI);
 
   /// Return true if the instruction does not have any effects besides

llvm/lib/Analysis/ValueTracking.cpp

for (const Value *V : Objs) {
         return false;
       }
       Objects.push_back(const_cast<Value *>(V));
     }
   } while (!Working.empty());
   return true;
 }
 
-/// Return true if the only users of this pointer are lifetime markers.
-bool llvm::onlyUsedByLifetimeMarkers(const Value *V) {
+static bool onlyUsedByLifetimeMarkersOrDroppableInstsHelper(
+    const Value *V, bool AllowLifetime, bool AllowDroppable) {
   for (const User *U : V->users()) {
     const IntrinsicInst *II = dyn_cast<IntrinsicInst>(U);
-    if (!II) return false;
+    if (!II)
+      return false;
+
+    if (AllowLifetime && II->isLifetimeStartOrEnd())
+      continue;
+
+    if (AllowDroppable && II->isDroppable())
+      continue;
 
-    if (!II->isLifetimeStartOrEnd())
-      return false;
   }
   return true;
 }
 
+bool llvm::onlyUsedByLifetimeMarkers(const Value *V) {
+  return onlyUsedByLifetimeMarkersOrDroppableInstsHelper(
+      V, /* AllowLifetime */ true, /* AllowDroppable */ false);
+}
+bool llvm::onlyUsedByLifetimeMarkersOrDroppableInsts(const Value *V) {
+  return onlyUsedByLifetimeMarkersOrDroppableInstsHelper(
+      V, /* AllowLifetime */ true, /* AllowDroppable */ true);
+}
+
 bool llvm::mustSuppressSpeculation(const LoadInst &LI) {
   if (!LI.isUnordered())
     return true;
   const Function &F = *LI.getFunction();
   // Speculative load may create a race that did not exist in the source.
   return F.hasFnAttribute(Attribute::SanitizeThread) ||
     // Speculative load may load data from dirty regions.
     F.hasFnAttribute(Attribute::SanitizeAddress) ||

llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp

 #define DEBUG_TYPE "mem2reg"
 
 STATISTIC(NumLocalPromoted, "Number of alloca's promoted within one block");
 STATISTIC(NumSingleStore,   "Number of alloca's promoted with a single store");
 STATISTIC(NumDeadAlloca,    "Number of dead alloca's removed");
 STATISTIC(NumPHIInsert,     "Number of PHI nodes inserted");
 
 bool llvm::isAllocaPromotable(const AllocaInst *AI) {
-  // FIXME: If the memory unit is of pointer or integer type, we can permit
-  // assignments to subsections of the memory unit.
-  unsigned AS = AI->getType()->getAddressSpace();
-
   // Only allow direct and non-volatile loads and stores...
   for (const User *U : AI->users()) {
     if (const LoadInst *LI = dyn_cast<LoadInst>(U)) {
       // Note that atomic loads can be transformed; atomic semantics do
       // not have any meaning for a local alloca.
       if (LI->isVolatile())
         return false;
     } else if (const StoreInst *SI = dyn_cast<StoreInst>(U)) {
       if (SI->getOperand(0) == AI)
         return false; // Don't allow a store OF the AI, only INTO the AI.
       // Note that atomic stores can be transformed; atomic semantics do
       // not have any meaning for a local alloca.
       if (SI->isVolatile())
         return false;
     } else if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(U)) {
-      if (!II->isLifetimeStartOrEnd())
+      if (!II->isLifetimeStartOrEnd() && !II->isDroppable())
         return false;
     } else if (const BitCastInst *BCI = dyn_cast<BitCastInst>(U)) {
-      if (BCI->getType() != Type::getInt8PtrTy(U->getContext(), AS))
-        return false;
-      if (!onlyUsedByLifetimeMarkers(BCI))
+      if (!onlyUsedByLifetimeMarkersOrDroppableInsts(BCI))
         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))
+      if (!onlyUsedByLifetimeMarkersOrDroppableInsts(GEPI))
         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 removeIntrinsicUsers(AllocaInst *AI) {
   // Knowing that this alloca is promotable, we know that it's safe to kill all
   // instructions except for load and store.
 
+  // Helper to drop the uses of \p I in \p UserI.
+  auto DropUsesIn = [](Instruction *UserI, Instruction *I,
+                       Instruction::user_iterator &UI,
+                       const Instruction::user_iterator &UE) {
+    // TODO For now we forget assumed information, this can be improved.
+    assert(isa<IntrinsicInst>(UserI) &&
+           cast<IntrinsicInst>(UserI)->getIntrinsicID() == Intrinsic::assume &&
+           "Expected assume");
+
+    // Skip ahead if I has multiple uses of AI.
+    while (UI != UE && *UI == UserI)
+      ++UI;
+
+    UserI->replaceUsesOfWith(I, UndefValue::get(I->getType()));
+  };
+
   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;
 
+    // Drop the use of AI in droppable instructions.
+    if (I->isDroppable()) {
+      DropUsesIn(I, AI, UI, UE);
+      continue;
+    }
+
     if (!I->getType()->isVoidTy()) {
       // The only users of this bitcast/GEP instruction are lifetime 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;
+
+        // Drop the use of I in droppable instructions.
+        if (Inst->isDroppable()) {
+          DropUsesIn(Inst, I, UUI, UUE);
+          continue;
+        }
         Inst->eraseFromParent();
       }
     }
     I->eraseFromParent();
   }
 }
 
 /// Rewrite as many loads as possible given a single store.
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);
+    removeIntrinsicUsers(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-droppable.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -mem2reg -S -o - < %s | FileCheck %s
+
+declare void @llvm.assume(i1)
+declare void @llvm.lifetime.start.p0i8(i64 %size, i8* nocapture %ptr)
+declare void @llvm.lifetime.end.p0i8(i64 %size, i8* nocapture %ptr)
+
+define void @positive_assume_uses(i32* %arg) {
+; CHECK-LABEL: @positive_assume_uses(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(i32* [[ARG:%.*]]), "align"(i32* undef, i64 2) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "align"(i32* undef, i64 8), "nonnull"(i32* [[ARG]]) ]
+; CHECK-NEXT:    ret void
+;
+  %A = alloca i32
+  call void @llvm.assume(i1 true) ["nonnull"(i32* %arg), "align"(i32* %A, i64 2)]
+  store i32 1, i32* %A
+  call void @llvm.assume(i1 true) ["align"(i32* %A, i64 8), "nonnull"(i32* %arg)]
+  ret void
+}
+
+define void @negative_assume_condition_use() {
+; CHECK-LABEL: @negative_assume_condition_use(
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[B:%.*]] = bitcast i32* [[A]] to i8*
+; CHECK-NEXT:    [[CND:%.*]] = icmp eq i8* [[B]], null
+; CHECK-NEXT:    call void @llvm.assume(i1 [[CND]])
+; CHECK-NEXT:    store i32 1, i32* [[A]], align 4
+; CHECK-NEXT:    ret void
+;
+  %A = alloca i32
+  %B = bitcast i32* %A to i8*
+  %cnd = icmp eq i8* %B, null
+  call void @llvm.assume(i1 %cnd)
+  store i32 1, i32* %A
+  ret void
+}
+
+define void @positive_multiple_assume_uses() {
+; CHECK-LABEL: @positive_multiple_assume_uses(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "align"({ i8, i16 }* undef, i64 8), "align"({ i8, i16 }* undef, i64 16) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"({ i8, i16 }* undef), "align"({ i8, i16 }* undef, i64 2) ]
+; CHECK-NEXT:    ret void
+;
+  %A = alloca {i8, i16}
+  call void @llvm.assume(i1 true) ["align"({i8, i16}* %A, i64 8), "align"({i8, i16}* %A, i64 16)]
+  store {i8, i16} zeroinitializer, {i8, i16}* %A
+  call void @llvm.assume(i1 true) ["nonnull"({i8, i16}* %A), "align"({i8, i16}* %A, i64 2)]
+  ret void
+}
+
+define void @positive_gep_assume_uses() {
+; CHECK-LABEL: @positive_gep_assume_uses(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "align"(i8* undef, i64 8), "align"(i8* undef, i64 16) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(i8* undef), "align"(i8* undef, i64 2) ]
+; CHECK-NEXT:    ret void
+;
+  %A = alloca {i8, i16}
+  %B = getelementptr {i8, i16}, {i8, i16}* %A, i32 0, i32 0
+  call void @llvm.lifetime.start.p0i8(i64 2, i8* %B)
+  call void @llvm.assume(i1 true) ["align"(i8* %B, i64 8), "align"(i8* %B, i64 16)]
+  store {i8, i16} zeroinitializer, {i8, i16}* %A
+  call void @llvm.lifetime.end.p0i8(i64 2, i8* %B)
+  call void @llvm.assume(i1 true) ["nonnull"(i8* %B), "align"(i8* %B, i64 2)]
+  ret void
+}
+
+define void @positive_mixed_assume_uses() {
+; CHECK-LABEL: @positive_mixed_assume_uses(
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(i8* undef), "align"(i8* undef, i64 8), "align"(i8* undef, i64 16) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(i8* undef), "align"(i8* undef, i64 2), "nonnull"(i8* undef) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(i32* undef), "align"(i32* undef, i64 2), "nonnull"(i8* undef) ]
+; CHECK-NEXT:    ret void
+;
+  %A = alloca i8
+  %B = getelementptr i8, i8* %A, i32 0
+  %C = bitcast i8* %A to i32*
+  call void @llvm.lifetime.start.p0i8(i64 2, i8* %B)
+  call void @llvm.assume(i1 true) ["nonnull"(i8* %B), "align"(i8* %A, i64 8), "align"(i8* %B, i64 16)]
+  store i8 1, i8* %A
+  call void @llvm.lifetime.end.p0i8(i64 2, i8* %B)
+  call void @llvm.assume(i1 true) ["nonnull"(i8* %B), "align"(i8* %A, i64 2), "nonnull"(i8* %A)]
+  call void @llvm.assume(i1 true) ["nonnull"(i32* %C), "align"(i32* %C, i64 2), "nonnull"(i8* %A)]
+  ret void
+}