1 files changed, 39 insertions, 15 deletions

diff --git a/llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp b/llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp
index 3daff3b..e6e5272 100644
--- a/llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp
+++ b/llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp
@@ -6,8 +6,6 @@
 //
 //===----------------------------------------------------------------------===//
 
-#define DEBUG_TYPE "assume-builder"
-
 #include "llvm/Transforms/Utils/AssumeBundleBuilder.h"
 #include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/ADT/MapVector.h"
@@ -37,6 +35,8 @@ cl::opt<bool> EnableKnowledgeRetention(
     cl::desc(
         "enable preservation of attributes throughout code transformation"));
 
+#define DEBUG_TYPE "assume-builder"
+
 STATISTIC(NumAssumeBuilt, "Number of assume built by the assume builder");
 STATISTIC(NumBundlesInAssumes, "Total number of Bundles in the assume built");
 STATISTIC(NumAssumesMerged,
@@ -65,7 +65,7 @@ bool isUsefullToPreserve(Attribute::AttrKind Kind) {
 
 /// This function will try to transform the given knowledge into a more
 /// canonical one. the canonical knowledge maybe the given one.
-RetainedKnowledge canonicalizedKnowledge(RetainedKnowledge RK, Module *M) {
+RetainedKnowledge canonicalizedKnowledge(RetainedKnowledge RK, DataLayout DL) {
   switch (RK.AttrKind) {
   default:
     return RK;
@@ -76,8 +76,7 @@ RetainedKnowledge canonicalizedKnowledge(RetainedKnowledge RK, Module *M) {
     Value *V = RK.WasOn->stripInBoundsOffsets([&](const Value *Strip) {
       if (auto *GEP = dyn_cast<GEPOperator>(Strip))
         RK.ArgValue =
-            MinAlign(RK.ArgValue,
-                     GEP->getMaxPreservedAlignment(M->getDataLayout()).value());
+            MinAlign(RK.ArgValue, GEP->getMaxPreservedAlignment(DL).value());
     });
     RK.WasOn = V;
     return RK;
@@ -85,8 +84,8 @@ RetainedKnowledge canonicalizedKnowledge(RetainedKnowledge RK, Module *M) {
   case Attribute::Dereferenceable:
   case Attribute::DereferenceableOrNull: {
     int64_t Offset = 0;
-    Value *V = GetPointerBaseWithConstantOffset(
-        RK.WasOn, Offset, M->getDataLayout(), /*AllowNonInBounds*/ false);
+    Value *V = GetPointerBaseWithConstantOffset(RK.WasOn, Offset, DL,
+                                                /*AllowNonInBounds*/ false);
     if (Offset < 0)
       return RK;
     RK.ArgValue = RK.ArgValue + Offset;
@@ -103,16 +102,16 @@ struct AssumeBuilderState {
 
   using MapKey = std::pair<Value *, Attribute::AttrKind>;
   SmallMapVector<MapKey, unsigned, 8> AssumedKnowledgeMap;
-  Instruction *InstBeingRemoved = nullptr;
+  Instruction *InstBeingModified = nullptr;
   AssumptionCache* AC = nullptr;
   DominatorTree* DT = nullptr;
 
   AssumeBuilderState(Module *M, Instruction *I = nullptr,
                      AssumptionCache *AC = nullptr, DominatorTree *DT = nullptr)
-      : M(M), InstBeingRemoved(I), AC(AC), DT(DT) {}
+      : M(M), InstBeingModified(I), AC(AC), DT(DT) {}
 
   bool tryToPreserveWithoutAddingAssume(RetainedKnowledge RK) {
-    if (!InstBeingRemoved || !RK.WasOn)
+    if (!InstBeingModified || !RK.WasOn)
       return false;
     bool HasBeenPreserved = false;
     Use* ToUpdate = nullptr;
@@ -120,13 +119,12 @@ struct AssumeBuilderState {
         RK.WasOn, {RK.AttrKind}, AC,
         [&](RetainedKnowledge RKOther, Instruction *Assume,
             const CallInst::BundleOpInfo *Bundle) {
-          if (!isValidAssumeForContext(Assume, InstBeingRemoved, DT))
+          if (!isValidAssumeForContext(Assume, InstBeingModified, DT))
             return false;
           if (RKOther.ArgValue >= RK.ArgValue) {
             HasBeenPreserved = true;
             return true;
-          } else if (isValidAssumeForContext(InstBeingRemoved, Assume,
-                                             DT)) {
+          } else if (isValidAssumeForContext(InstBeingModified, Assume, DT)) {
             HasBeenPreserved = true;
             IntrinsicInst *Intr = cast<IntrinsicInst>(Assume);
             ToUpdate = &Intr->op_begin()[Bundle->Begin + ABA_Argument];
@@ -162,14 +160,14 @@ struct AssumeBuilderState {
         if (RK.WasOn->use_empty())
           return false;
         Use *SingleUse = RK.WasOn->getSingleUndroppableUse();
-        if (SingleUse && SingleUse->getUser() == InstBeingRemoved)
+        if (SingleUse && SingleUse->getUser() == InstBeingModified)
           return false;
       }
     return true;
   }
 
   void addKnowledge(RetainedKnowledge RK) {
-    RK = canonicalizedKnowledge(RK, M);
+    RK = canonicalizedKnowledge(RK, M->getDataLayout());
 
     if (!isKnowledgeWorthPreserving(RK))
       return;
@@ -299,6 +297,32 @@ void llvm::salvageKnowledge(Instruction *I, AssumptionCache *AC,
   }
 }
 
+IntrinsicInst *
+llvm::buildAssumeFromKnowledge(ArrayRef<RetainedKnowledge> Knowledge,
+                               Instruction *CtxI, AssumptionCache *AC,
+                               DominatorTree *DT) {
+  AssumeBuilderState Builder(CtxI->getModule(), CtxI, AC, DT);
+  for (const RetainedKnowledge &RK : Knowledge)
+    Builder.addKnowledge(RK);
+  return Builder.build();
+}
+
+RetainedKnowledge llvm::simplifyRetainedKnowledge(CallBase *Assume,
+                                                  RetainedKnowledge RK,
+                                                  AssumptionCache *AC,
+                                                  DominatorTree *DT) {
+  assert(Assume->getIntrinsicID() == Intrinsic::assume);
+  AssumeBuilderState Builder(Assume->getModule(), Assume, AC, DT);
+  RK = canonicalizedKnowledge(RK, Assume->getModule()->getDataLayout());
+
+  if (!Builder.isKnowledgeWorthPreserving(RK))
+    return RetainedKnowledge::none();
+
+  if (Builder.tryToPreserveWithoutAddingAssume(RK))
+    return RetainedKnowledge::none();
+  return RK;
+}
+
 namespace {
 
 struct AssumeSimplify {