[AssumeBundles] Prevent generation of some redundant assumes

Summary: with this patch the assume salvageKnowledge will not generate assume if all knowledge is already available in an assume with valid context. assume bulider can also in some cases update an existing assume with better information.

Reviewers: jdoerfert

Reviewed By: jdoerfert

Subscribers: hiraditya, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D78014

llvm/include/llvm/Analysis/AssumeBundleQueries.h

@@ -142,8 +142,9 @@ RetainedKnowledge getKnowledgeFromUse(const Use *U,
 RetainedKnowledge getKnowledgeForValue(
     const Value *V, ArrayRef<Attribute::AttrKind> AttrKinds,
     AssumptionCache *AC = nullptr,
-    function_ref<bool(RetainedKnowledge, Instruction *)> Filter =
-        [](RetainedKnowledge, Instruction *) { return true; });
+    function_ref<bool(RetainedKnowledge, Instruction *,
+                            const CallBase::BundleOpInfo *)>
+        Filter = [](auto...) { return true; });
 
 /// Return a valid Knowledge associated to the Value V if its Attribute kind is
 /// in AttrKinds and the knowledge is suitable to be used in the context of

llvm/include/llvm/Transforms/Utils/AssumeBundleBuilder.h

@@ -23,6 +23,7 @@
 namespace llvm {
 class IntrinsicInst;
 class AssumptionCache;
+class DominatorTree;
 
 /// Build a call to llvm.assume to preserve informations that can be derived
 /// from the given instruction.
@@ -33,7 +34,12 @@ IntrinsicInst *buildAssumeFromInst(Instruction *I);
 /// Calls BuildAssumeFromInst and if the resulting llvm.assume is valid insert
 /// if before I. This is usually what need to be done to salvage the knowledge
 /// contained in the instruction I.
-void salvageKnowledge(Instruction *I, AssumptionCache *AC = nullptr);
+/// The AssumptionCache must be provided if it is available or the cache may
+/// become silently be invalid.
+/// The DominatorTree can optionally be provided to enable cross-block
+/// reasoning.
+void salvageKnowledge(Instruction *I, AssumptionCache *AC = nullptr,
+                      DominatorTree *DT = nullptr);
 
 /// This pass will try to build an llvm.assume for every instruction in the
 /// function. Its main purpose is testing.

llvm/lib/Analysis/AssumeBundleQueries.cpp

@@ -120,25 +120,35 @@ bool llvm::isAssumeWithEmptyBundle(CallInst &CI) {
                  });
 }
 
+static CallInst::BundleOpInfo *getBundleFromUse(const Use *U) {
+  auto *Intr = dyn_cast<IntrinsicInst>(U->getUser());
+  if (!match(U->getUser(),
+             m_Intrinsic<Intrinsic::assume>(m_Unless(m_Specific(U->get())))))
+    return nullptr;
+  return &Intr->getBundleOpInfoForOperand(U->getOperandNo());
+}
+
 RetainedKnowledge
 llvm::getKnowledgeFromUse(const Use *U,
                           ArrayRef<Attribute::AttrKind> AttrKinds) {
-  if (!match(U->getUser(),
-             m_Intrinsic<Intrinsic::assume>(m_Unless(m_Specific(U->get())))))
+  CallInst::BundleOpInfo* Bundle = getBundleFromUse(U);
+  if (!Bundle)
     return RetainedKnowledge::none();
-  auto *Intr = cast<IntrinsicInst>(U->getUser());
   RetainedKnowledge RK =
-      getKnowledgeFromOperandInAssume(*Intr, U->getOperandNo());
+      getKnowledgeFromBundle(*cast<CallInst>(U->getUser()), *Bundle);
   for (auto Attr : AttrKinds)
     if (Attr == RK.AttrKind)
       return RK;
   return RetainedKnowledge::none();
 }
 
-RetainedKnowledge llvm::getKnowledgeForValue(
-    const Value *V, ArrayRef<Attribute::AttrKind> AttrKinds,
-    AssumptionCache *AC,
-    function_ref<bool(RetainedKnowledge, Instruction *)> Filter) {
+RetainedKnowledge
+llvm::getKnowledgeForValue(const Value *V,
+                           ArrayRef<Attribute::AttrKind> AttrKinds,
+                           AssumptionCache *AC,
+                           function_ref<bool(RetainedKnowledge, Instruction *,
+                                             const CallBase::BundleOpInfo *)>
+                               Filter) {
   if (AC) {
 #ifndef NDEBUG
     RetainedKnowledge RKCheck =
@@ -150,7 +160,8 @@ RetainedKnowledge llvm::getKnowledgeForValue(
         continue;
       if (RetainedKnowledge RK = getKnowledgeFromBundle(
               *II, II->bundle_op_info_begin()[Elem.Index]))
-        if (is_contained(AttrKinds, RK.AttrKind) && Filter(RK, II)) {
+        if (is_contained(AttrKinds, RK.AttrKind) &&
+            Filter(RK, II, &II->bundle_op_info_begin()[Elem.Index])) {
           assert(!!RKCheck && "invalid Assumption cache");
           return RK;
         }
@@ -159,8 +170,13 @@ RetainedKnowledge llvm::getKnowledgeForValue(
     return RetainedKnowledge::none();
   }
   for (auto &U : V->uses()) {
-    if (RetainedKnowledge RK = getKnowledgeFromUse(&U, AttrKinds))
-      if (Filter(RK, cast<Instruction>(U.getUser())))
+    CallInst::BundleOpInfo* Bundle = getBundleFromUse(&U);
+    if (!Bundle)
+      continue;
+    if (RetainedKnowledge RK =
+            getKnowledgeFromBundle(*cast<CallInst>(U.getUser()), *Bundle))
+      if (is_contained(AttrKinds, RK.AttrKind) &&
+          Filter(RK, cast<Instruction>(U.getUser()), Bundle))
         return RK;
   }
   return RetainedKnowledge::none();
@@ -170,7 +186,7 @@ RetainedKnowledge llvm::getKnowledgeValidInContext(
     const Value *V, ArrayRef<Attribute::AttrKind> AttrKinds,
     const Instruction *CtxI, const DominatorTree *DT, AssumptionCache *AC) {
   return getKnowledgeForValue(V, AttrKinds, AC,
-                              [&](RetainedKnowledge, Instruction *I) {
+                              [&](auto, Instruction *I, auto) {
                                 return isValidAssumeForContext(I, CtxI, DT);
                               });
 }

llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp

@@ -9,7 +9,8 @@
 #include "llvm/Transforms/Utils/AssumeBundleBuilder.h"
 #include "llvm/Analysis/AssumeBundleQueries.h"
 #include "llvm/Analysis/AssumptionCache.h"
-#include "llvm/ADT/DenseSet.h"
+#include "llvm/Analysis/ValueTracking.h"
+#include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/IntrinsicInst.h"
@@ -74,10 +75,45 @@ struct AssumeBuilderState {
   using MapKey = std::pair<Value *, Attribute::AttrKind>;
   SmallDenseMap<MapKey, unsigned, 8> AssumedKnowledgeMap;
   Instruction *InsertBeforeInstruction = nullptr;
+  AssumptionCache* AC = nullptr;
+  DominatorTree* DT = nullptr;
 
-  AssumeBuilderState(Module *M) : M(M) {}
+  AssumeBuilderState(Module *M, Instruction *I = nullptr,
+                     AssumptionCache *AC = nullptr, DominatorTree *DT = nullptr)
+      : M(M), InsertBeforeInstruction(I), AC(AC), DT(DT) {}
+
+  bool tryToPreserveWithoutAddingAssume(RetainedKnowledge RK) {
+    if (!InsertBeforeInstruction || !AC || !RK.WasOn)
+      return false;
+    bool HasBeenPreserved = false;
+    Use* ToUpdate = nullptr;
+    getKnowledgeForValue(
+        RK.WasOn, {RK.AttrKind}, AC,
+        [&](RetainedKnowledge RKOther, Instruction *Assume,
+            const CallInst::BundleOpInfo *Bundle) {
+          if (!isValidAssumeForContext(Assume, InsertBeforeInstruction, DT))
+            return false;
+          if (RKOther.ArgValue >= RK.ArgValue) {
+            HasBeenPreserved = true;
+            return true;
+          } else if (isValidAssumeForContext(InsertBeforeInstruction, Assume,
+                                             DT)) {
+            HasBeenPreserved = true;
+            IntrinsicInst *Intr = cast<IntrinsicInst>(Assume);
+            ToUpdate = &Intr->op_begin()[Bundle->Begin + ABA_Argument];
+            return true;
+          }
+          return false;
+        });
+    if (ToUpdate)
+      ToUpdate->set(
+          ConstantInt::get(Type::getInt64Ty(M->getContext()), RK.ArgValue));
+    return HasBeenPreserved;
+  }
 
   void addKnowledge(RetainedKnowledge RK) {
+    if (tryToPreserveWithoutAddingAssume(RK))
+      return;
     MapKey Key{RK.WasOn, RK.AttrKind};
     auto Lookup = AssumedKnowledgeMap.find(Key);
     if (Lookup == AssumedKnowledgeMap.end()) {
@@ -185,11 +221,10 @@ IntrinsicInst *llvm::buildAssumeFromInst(Instruction *I) {
   return Builder.build();
 }
 
-void llvm::salvageKnowledge(Instruction *I, AssumptionCache *AC) {
+void llvm::salvageKnowledge(Instruction *I, AssumptionCache *AC, DominatorTree* DT) {
   if (!EnableKnowledgeRetention)
     return;
-  AssumeBuilderState Builder(I->getModule());
-  Builder.InsertBeforeInstruction = I;
+  AssumeBuilderState Builder(I->getModule(), I, AC, DT);
   Builder.addInstruction(I);
   if (IntrinsicInst *Intr = Builder.build()) {
     Intr->insertBefore(I);
@@ -200,8 +235,9 @@ void llvm::salvageKnowledge(Instruction *I, AssumptionCache *AC) {
 
 PreservedAnalyses AssumeBuilderPass::run(Function &F,
                                          FunctionAnalysisManager &AM) {
+  AssumptionCache* AC = AM.getCachedResult<AssumptionAnalysis>(F);
+  DominatorTree* DT = AM.getCachedResult<DominatorTreeAnalysis>(F);
   for (Instruction &I : instructions(F))
-    if (Instruction *Assume = buildAssumeFromInst(&I))
-      Assume->insertBefore(&I);
+    salvageKnowledge(&I, AC, DT);
   return PreservedAnalyses::all();
 }

llvm/test/Analysis/BasicAA/featuretest.ll

@@ -36,7 +36,6 @@ define i32 @different_array_test(i64 %A, i64 %B) {
 ; USE_ASSUME-NEXT:    call void @external(i32* nonnull [[ARRAY11_SUB]])
 ; USE_ASSUME-NEXT:    call void @external(i32* nonnull [[ARRAY22_SUB]])
 ; USE_ASSUME-NEXT:    [[POINTER:%.*]] = getelementptr [100 x i32], [100 x i32]* [[ARRAY11]], i64 0, i64 [[A:%.*]]
-; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[POINTER]], i64 4), "nonnull"(i32* [[POINTER]]) ]
 ; USE_ASSUME-NEXT:    [[POINTER2:%.*]] = getelementptr [200 x i32], [200 x i32]* [[ARRAY22]], i64 0, i64 [[B:%.*]]
 ; USE_ASSUME-NEXT:    store i32 7, i32* [[POINTER2]], align 4
 ; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[POINTER]], i64 4), "nonnull"(i32* [[POINTER]]) ]
@@ -78,7 +77,6 @@ define i32 @constant_array_index_test() {
 ; USE_ASSUME-NEXT:    call void @external(i32* nonnull [[ARRAY1_SUB]])
 ; USE_ASSUME-NEXT:    [[P1:%.*]] = getelementptr inbounds [100 x i32], [100 x i32]* [[ARRAY1]], i64 0, i64 7
 ; USE_ASSUME-NEXT:    [[P2:%.*]] = getelementptr inbounds [100 x i32], [100 x i32]* [[ARRAY1]], i64 0, i64 6
-; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[P1]], i64 4), "nonnull"(i32* [[P1]]) ]
 ; USE_ASSUME-NEXT:    store i32 1, i32* [[P2]], align 4
 ; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[P1]], i64 4), "nonnull"(i32* [[P1]]) ]
 ; USE_ASSUME-NEXT:    ret i32 0
@@ -105,8 +103,7 @@ define i32 @gep_distance_test(i32* %A) {
 ; NO_ASSUME-NEXT:    ret i32 0
 ;
 ; USE_ASSUME-LABEL: @gep_distance_test(
-; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[A:%.*]], i64 4), "nonnull"(i32* [[A]]) ]
-; USE_ASSUME-NEXT:    [[B:%.*]] = getelementptr i32, i32* [[A]], i64 2
+; USE_ASSUME-NEXT:    [[B:%.*]] = getelementptr i32, i32* [[A:%.*]], i64 2
 ; USE_ASSUME-NEXT:    store i32 7, i32* [[B]], align 4
 ; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[A]], i64 4), "nonnull"(i32* [[A]]) ]
 ; USE_ASSUME-NEXT:    ret i32 0
@@ -129,7 +126,6 @@ define i32 @gep_distance_test2({i32,i32}* %A, i64 %distance) {
 ;
 ; USE_ASSUME-LABEL: @gep_distance_test2(
 ; USE_ASSUME-NEXT:    [[A1:%.*]] = getelementptr { i32, i32 }, { i32, i32 }* [[A:%.*]], i64 0, i32 0
-; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[A1]], i64 4), "nonnull"(i32* [[A1]]) ]
 ; USE_ASSUME-NEXT:    [[B:%.*]] = getelementptr { i32, i32 }, { i32, i32 }* [[A]], i64 [[DISTANCE:%.*]], i32 1
 ; USE_ASSUME-NEXT:    store i32 7, i32* [[B]], align 4
 ; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[A1]], i64 4), "nonnull"(i32* [[A1]]) ]
@@ -154,8 +150,7 @@ define i32 @gep_distance_test3(i32 * %A) {
 ; NO_ASSUME-NEXT:    ret i32 0
 ;
 ; USE_ASSUME-LABEL: @gep_distance_test3(
-; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[A:%.*]], i64 4), "nonnull"(i32* [[A]]) ]
-; USE_ASSUME-NEXT:    [[C1:%.*]] = getelementptr i32, i32* [[A]], i64 1
+; USE_ASSUME-NEXT:    [[C1:%.*]] = getelementptr i32, i32* [[A:%.*]], i64 1
 ; USE_ASSUME-NEXT:    [[C:%.*]] = bitcast i32* [[C1]] to i8*
 ; USE_ASSUME-NEXT:    store i8 42, i8* [[C]], align 1
 ; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[A]], i64 4), "nonnull"(i32* [[A]]) ]
@@ -181,7 +176,6 @@ define i32 @constexpr_test() {
 ; USE_ASSUME-LABEL: @constexpr_test(
 ; USE_ASSUME-NEXT:    [[X:%.*]] = alloca i32, align 4
 ; USE_ASSUME-NEXT:    call void @external(i32* nonnull [[X]])
-; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[X]], i64 4), "nonnull"(i32* [[X]]) ]
 ; USE_ASSUME-NEXT:    store i32 5, i32* getelementptr inbounds ({ i32 }, { i32 }* @Global, i64 0, i32 0), align 4
 ; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[X]], i64 4), "nonnull"(i32* [[X]]) ]
 ; USE_ASSUME-NEXT:    ret i32 0

llvm/test/Transforms/EarlyCSE/guards.ll

@@ -229,7 +229,6 @@ define void @test08(i32 %a, i32 %b, i32* %ptr) {
 ; USE_ASSUME-NEXT:    store i32 100, i32* [[PTR:%.*]]
 ; USE_ASSUME-NEXT:    call void (i1, ...) @llvm.experimental.guard(i1 [[CMP]]) [ "deopt"() ]
 ; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[PTR]], i64 4), "nonnull"(i32* [[PTR]]) ]
-; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[PTR]], i64 4), "nonnull"(i32* [[PTR]]) ]
 ; USE_ASSUME-NEXT:    store i32 400, i32* [[PTR]]
 ; USE_ASSUME-NEXT:    ret void
 ;
@@ -271,7 +270,6 @@ define void @test09(i32 %a, i32 %b, i1 %c, i32* %ptr) {
 ; USE_ASSUME-NEXT:    store i32 100, i32* [[PTR:%.*]]
 ; USE_ASSUME-NEXT:    call void (i1, ...) @llvm.experimental.guard(i1 [[CMP]]) [ "deopt"() ]
 ; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[PTR]], i64 4), "nonnull"(i32* [[PTR]]) ]
-; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[PTR]], i64 4), "nonnull"(i32* [[PTR]]) ]
 ; USE_ASSUME-NEXT:    store i32 400, i32* [[PTR]]
 ; USE_ASSUME-NEXT:    br i1 [[C:%.*]], label [[IF_TRUE:%.*]], label [[IF_FALSE:%.*]]
 ; USE_ASSUME:       if.true:
@@ -410,8 +408,6 @@ define void @test13(i32 %a, i32 %b, i32* %ptr) {
 ; USE_ASSUME-NEXT:    [[CMP:%.*]] = icmp eq i32 [[A:%.*]], [[B:%.*]]
 ; USE_ASSUME-NEXT:    call void @llvm.assume(i1 [[CMP]])
 ; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[PTR:%.*]], i64 4), "nonnull"(i32* [[PTR]]) ]
-; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[PTR]], i64 4), "nonnull"(i32* [[PTR]]) ]
-; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[PTR]], i64 4), "nonnull"(i32* [[PTR]]) ]
 ; USE_ASSUME-NEXT:    store i32 400, i32* [[PTR]]
 ; USE_ASSUME-NEXT:    ret void
 ;
@@ -452,8 +448,6 @@ define void @test14(i32 %a, i32 %b, i1 %c, i32* %ptr) {
 ; USE_ASSUME-NEXT:    [[CMP:%.*]] = icmp eq i32 [[A:%.*]], [[B:%.*]]
 ; USE_ASSUME-NEXT:    call void @llvm.assume(i1 [[CMP]])
 ; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[PTR:%.*]], i64 4), "nonnull"(i32* [[PTR]]) ]
-; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[PTR]], i64 4), "nonnull"(i32* [[PTR]]) ]
-; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[PTR]], i64 4), "nonnull"(i32* [[PTR]]) ]
 ; USE_ASSUME-NEXT:    store i32 400, i32* [[PTR]]
 ; USE_ASSUME-NEXT:    br i1 [[C:%.*]], label [[IF_TRUE:%.*]], label [[IF_FALSE:%.*]]
 ; USE_ASSUME:       if.true:

llvm/test/Transforms/EarlyCSE/invariant-loads.ll

@@ -18,7 +18,6 @@ define void @f_0(i32* %ptr) {
 ; USE_ASSUME-NEXT:    call void @clobber_and_use(i32 [[VAL0]])
 ; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[PTR]], i64 4), "nonnull"(i32* [[PTR]]) ]
 ; USE_ASSUME-NEXT:    call void @clobber_and_use(i32 [[VAL0]])
-; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[PTR]], i64 4), "nonnull"(i32* [[PTR]]) ]
 ; USE_ASSUME-NEXT:    call void @clobber_and_use(i32 [[VAL0]])
 ; USE_ASSUME-NEXT:    ret void
 ;
@@ -194,7 +193,6 @@ define void @test_scope_start_without_load(i32* %p) {
 ; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[P]], i64 4), "nonnull"(i32* [[P]]) ]
 ; USE_ASSUME-NEXT:    [[ADD:%.*]] = add i32 [[V1]], [[V1]]
 ; USE_ASSUME-NEXT:    call void @clobber_and_use(i32 [[ADD]])
-; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[P]], i64 4), "nonnull"(i32* [[P]]) ]
 ; USE_ASSUME-NEXT:    call void @clobber_and_use(i32 [[V1]])
 ; USE_ASSUME-NEXT:    ret void
 ;
@@ -225,7 +223,6 @@ define void @test_scope_restart(i32* %p) {
 ; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[P]], i64 4), "nonnull"(i32* [[P]]) ]
 ; USE_ASSUME-NEXT:    [[ADD:%.*]] = add i32 [[V1]], [[V1]]
 ; USE_ASSUME-NEXT:    call void @clobber_and_use(i32 [[ADD]])
-; USE_ASSUME-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[P]], i64 4), "nonnull"(i32* [[P]]) ]
 ; USE_ASSUME-NEXT:    call void @clobber_and_use(i32 [[V1]])
 ; USE_ASSUME-NEXT:    ret void
 ;