[SROA] Teach promote to register about droppable instructions

[jdoerfert]

Summary

This is the second of two patches to address PR46753. We basically allow
SROA to promote allocas that are used in doppable instructions, for
now that means llvm.assume. The (transitive) uses are replaced by
undef in the droppable instructions.

See also D83976.

[Tyker]

llvm/lib/Transforms/Scalar/SROA.cpp

3094

Same as https://reviews.llvm.org/D83976#inline-773557

[jdoerfert]

Use proper function to drop uses

[jdoerfert]

Rebase on top of D84085

[Tyker]

just one nit

llvm/lib/Transforms/Scalar/SROA.cpp

3094

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]

Use explicit API call

llvm/lib/Transforms/Scalar/SROA.cpp

void visitMemTransferInst(MemTransferInst &II) {
     assert(AS.Slices[PrevIdx].getUse()->getUser() == &II &&
            "Map index doesn't point back to a slice with this user.");
   }
 
   // Disable SRoA for any intrinsics except for lifetime invariants.
   // FIXME: What about debug intrinsics? This matches old behavior, but
   // doesn't make sense.
   void visitIntrinsicInst(IntrinsicInst &II) {
+    if (II.isDroppable())
+      return;
+
     if (!IsOffsetKnown)
       return PI.setAborted(&II);
 
     if (II.isLifetimeStartOrEnd()) {
       ConstantInt *Length = cast<ConstantInt>(II.getArgOperand(0));
       uint64_t Size = std::min(AllocSize - Offset.getLimitedValue(),
                                Length->getLimitedValue());
       insertUse(II, Offset, Size, true);
static bool isVectorPromotionViableForSlice(Partition &P, const Slice &S,
   Use *U = S.getUse();
 
   if (MemIntrinsic *MI = dyn_cast<MemIntrinsic>(U->getUser())) {
     if (MI->isVolatile())
       return false;
     if (!S.isSplittable())
       return false; // Skip any unsplittable intrinsics.
   } else if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(U->getUser())) {
-    if (!II->isLifetimeStartOrEnd())
+    if (!II->isLifetimeStartOrEnd() && !II->isDroppable())
       return false;
   } else if (U->get()->getType()->getPointerElementType()->isStructTy()) {
     // Disable vector promotion when there are loads or stores of an FCA.
     return false;
   } else if (LoadInst *LI = dyn_cast<LoadInst>(U->getUser())) {
     if (LI->isVolatile())
       return false;
     Type *LTy = LI->getType();
if (IntegerType *ITy = dyn_cast<IntegerType>(ValueTy)) {
       return false;
     }
   } else if (MemIntrinsic *MI = dyn_cast<MemIntrinsic>(U->getUser())) {
     if (MI->isVolatile() || !isa<Constant>(MI->getLength()))
       return false;
     if (!S.isSplittable())
       return false; // Skip any unsplittable intrinsics.
   } else if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(U->getUser())) {
-    if (!II->isLifetimeStartOrEnd())
+    if (!II->isLifetimeStartOrEnd() && !II->isDroppable())
       return false;
   } else {
     return false;
   }
 
   return true;
 }
 
if (!isa<Constant>(II.getLength())) {
       return false;
     }
 
     // Record this instruction for deletion.
     Pass.DeadInsts.insert(&II);
 
     Type *AllocaTy = NewAI.getAllocatedType();
     Type *ScalarTy = AllocaTy->getScalarType();
-    
     const bool CanContinue = [&]() {
       if (VecTy || IntTy)
         return true;
       if (BeginOffset > NewAllocaBeginOffset ||
           EndOffset < NewAllocaEndOffset)
         return false;
       auto *C = cast<ConstantInt>(II.getLength());
       if (C->getBitWidth() > 64)
StoreInst *Store = cast<StoreInst>(
         IRB.CreateAlignedStore(Src, DstPtr, DstAlign, II.isVolatile()));
     if (AATags)
       Store->setAAMetadata(AATags);
     LLVM_DEBUG(dbgs() << "          to: " << *Store << "\n");
     return !II.isVolatile();
   }
 
   bool visitIntrinsicInst(IntrinsicInst &II) {
-    assert(II.isLifetimeStartOrEnd());
+    assert((II.isLifetimeStartOrEnd() || II.isDroppable()) &&
+           "Unexpected intrinsic!");
     LLVM_DEBUG(dbgs() << "    original: " << II << "\n");
-    assert(II.getArgOperand(1) == OldPtr);
 
     // Record this instruction for deletion.
     Pass.DeadInsts.insert(&II);
 
+    if (II.isDroppable()) {
+      assert(II.getIntrinsicID() == Intrinsic::assume && "Expected assume");
+      // TODO For now we forget assumed information, this can be improved.
+      OldPtr->dropDroppableUses(
+          [&](const Use *U) { return U->getUser() == &II; });
+      return true;
+    }
+
+    assert(II.getArgOperand(1) == OldPtr);
     // Lifetime intrinsics are only promotable if they cover the whole alloca.
     // Therefore, we drop lifetime intrinsics which don't cover the whole
     // alloca.
     // (In theory, intrinsics which partially cover an alloca could be
     // promoted, but PromoteMemToReg doesn't handle that case.)
     // FIXME: Check whether the alloca is promotable before dropping the
     // lifetime intrinsics?
     if (NewBeginOffset != NewAllocaBeginOffset ||

llvm/test/Transforms/SROA/ignore-droppable.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -sroa -S -o - < %s | FileCheck %s
+; RUN: opt -passes=sroa -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:%.*]]), "ignore"(i32* undef, i64 2) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "ignore"(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:    [[A:%.*]] = alloca { i8, i16 }, align 8
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "align"({ i8, i16 }* [[A]], i64 8), "align"({ i8, i16 }* [[A]], i64 16) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"({ i8, i16 }* [[A]]), "align"({ i8, i16 }* [[A]], 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:    [[A:%.*]] = alloca { i8, i16 }, align 8
+; CHECK-NEXT:    [[B:%.*]] = getelementptr { i8, i16 }, { i8, i16 }* [[A]], i32 0, i32 0
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "align"(i8* [[B]], i64 8), "align"(i8* [[B]], i64 16) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "nonnull"(i8* [[B]]), "align"(i8* [[B]], 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) [ "ignore"(i8* undef), "ignore"(i8* undef, i64 8), "ignore"(i8* undef, i64 16) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "ignore"(i8* undef), "ignore"(i8* undef, i64 2), "ignore"(i8* undef) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "ignore"(i32* undef), "ignore"(i32* undef, i64 2), "ignore"(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
+}