[SROA] Improve handleling of assumes bundles by SROA

[Tyker]

Summary

This patch fixes this crash https://gcc.godbolt.org/z/Ps8d1e
And gives SROA the ability to remove assumes if it allows promoting an alloca to register
Without removing assumes when it can't promote to register.

[jdoerfert]

LGTM.

llvm/include/llvm/IR/Value.h

#endif
   /// uses.
   void dropDroppableUses(llvm::function_ref<bool(const Use *)> ShouldDrop =
                              [](const Use *) { return true; });
 
   /// Remove every use of this value in \p User that can safely be removed.
   void dropDroppableUsesIn(User &Usr);
 
   /// Remove the droppable use \p U.
-  void dropDroppableUse(Use &U);
+  static void dropDroppableUse(Use &U);
 
   /// Check if this value is used in the specified basic block.
   bool isUsedInBasicBlock(const BasicBlock *BB) const;
 
   /// This method computes the number of uses of this Value.
   ///
   /// This is a linear time operation.  Use hasOneUse, hasNUses, or
   /// hasNUsesOrMore to check for specific values.

llvm/lib/IR/Value.cpp

void Value::dropDroppableUse(Use &U) {
   if (auto *Assume = dyn_cast<IntrinsicInst>(U.getUser())) {
     assert(Assume->getIntrinsicID() == Intrinsic::assume);
     unsigned OpNo = U.getOperandNo();
     if (OpNo == 0)
       U.set(ConstantInt::getTrue(Assume->getContext()));
     else {
       U.set(UndefValue::get(U.get()->getType()));
       CallInst::BundleOpInfo &BOI = Assume->getBundleOpInfoForOperand(OpNo);
-      BOI.Tag = getContext().pImpl->getOrInsertBundleTag("ignore");
+      BOI.Tag = Assume->getContext().pImpl->getOrInsertBundleTag("ignore");
     }
     return;
   }
 
   llvm_unreachable("unkown droppable use");
 }
 
 bool Value::isUsedInBasicBlock(const BasicBlock *BB) const {

llvm/lib/Transforms/Scalar/SROA.cpp

public:
   // Forward declare the iterator and range accessor for walking the
   // partitions.
   class partition_iterator;
   iterator_range<partition_iterator> partitions();
 
   /// Access the dead users for this alloca.
   ArrayRef<Instruction *> getDeadUsers() const { return DeadUsers; }
 
+  /// Access Uses that should be dropped if the alloca is promotable.
+  ArrayRef<Use *> getDeadUsesIfPromotable() const {
+    return DeadUseIfPromotable;
+  }
+
   /// Access the dead operands referring to this alloca.
   ///
   /// These are operands which have cannot actually be used to refer to the
   /// alloca as they are outside its range and the user doesn't correct for
   /// that. These mostly consist of PHI node inputs and the like which we just
   /// need to replace with undef.
   ArrayRef<Use *> getDeadOperands() const { return DeadOperands; }
 
#endif
   /// Instructions which will become dead if we rewrite the alloca.
   ///
   /// Note that these are not separated by slice. This is because we expect an
   /// alloca to be completely rewritten or not rewritten at all. If rewritten,
   /// all these instructions can simply be removed and replaced with undef as
   /// they come from outside of the allocated space.
   SmallVector<Instruction *, 8> DeadUsers;
 
+  /// Uses which will become dead if can promote the alloca.
+  SmallVector<Use *, 8> DeadUseIfPromotable;
+
   /// Operands which will become dead if we rewrite the alloca.
   ///
   /// These are operands that in their particular use can be replaced with
   /// undef when we rewrite the alloca. These show up in out-of-bounds inputs
   /// to PHI nodes and the like. They aren't entirely dead (there might be
   /// a GEP back into the bounds using it elsewhere) and nor is the PHI, but we
   /// want to swap this particular input for undef to simplify the use lists of
   /// the alloca.
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())
+    if (II.isDroppable()) {
+      AS.DeadUseIfPromotable.push_back(U);
       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());
for (SelectInst *Sel : SelectUsers)
     if (!isSafeSelectToSpeculate(*Sel)) {
       Promotable = false;
       PHIUsers.clear();
       SelectUsers.clear();
       break;
     }
 
   if (Promotable) {
+    for (Use *U : AS.getDeadUsesIfPromotable()) {
+      auto *OldInst = dyn_cast<Instruction>(U->get());
+      Value::dropDroppableUse(*U);
+      if (OldInst)
+        if (isInstructionTriviallyDead(OldInst))
+          DeadInsts.insert(OldInst);
+    }
     if (PHIUsers.empty() && SelectUsers.empty()) {
       // Promote the alloca.
       PromotableAllocas.push_back(NewAI);
     } else {
       // If we have either PHIs or Selects to speculate, add them to those
       // worklists and re-queue the new alloca so that we promote in on the
       // next iteration.
       for (PHINode *PHIUser : PHIUsers)

llvm/test/Transforms/SROA/assume.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -sroa -S | FileCheck %s
+
+source_filename = "tmp.cpp"
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+%struct.n.2 = type { i8 }
+%struct.bi = type { %class.as }
+%class.as = type { %class.ao }
+%class.ao = type { %class.ai }
+%class.ai = type { i32, i32 }
+
+define linkonce_odr dso_local void @_ZN2bg2baIiEC2ES_() unnamed_addr align 2 {
+; CHECK-LABEL: @_ZN2bg2baIiEC2ES_(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "ignore"(%struct.bi* undef) ]
+; CHECK-NEXT:    ret void
+;
+entry:
+  %y = alloca i64, align 8
+  store i64 undef, i64* %y, align 8
+  %0 = bitcast i64* %y to %struct.n.2*
+  %1 = bitcast %struct.n.2* %0 to %struct.bi*
+  call void @llvm.assume(i1 true) [ "nonnull"(%struct.bi* %1) ]
+  ret void
+}
+
+; Function Attrs: nofree norecurse nounwind willreturn
+declare void @llvm.assume(i1) #0
+
+attributes #0 = { nofree norecurse nounwind willreturn }

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

;
   %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:    call void @llvm.assume(i1 true) [ "ignore"({ i8, i16 }* undef, i64 8), "ignore"({ i8, i16 }* undef, i64 16) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "ignore"({ i8, i16 }* undef), "ignore"({ 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:    [[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:    call void @llvm.assume(i1 true) [ "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) ]
 ; 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)