[SCCP] Consider provenance when propagating constant ptrs

Similarly to what it is being already done in GVN (fb632ed),
make sure pointers equalities derived via PredicatedInfo may
be propagated, taking into account their provenance as well.

Fixes: #159565.

Author: antoniofrighetto

llvm/include/llvm/Transforms/Utils/SCCPSolver.h

@@ -218,6 +218,10 @@ class SCCPSolver {
   // old ValueLatticeElement::isOverdefined() and is intended to be used in the
   // transition to ValueLatticeElement.
   LLVM_ABI static bool isOverdefined(const ValueLatticeElement &LV);
+
+  // Helper to check whether \p V may be carrying a PredicateInfo derived copy
+  // of a pointer.
+  LLVM_ABI bool mayForwardPointerPredicatedCopy(Value *V);
 };
 } // namespace llvm
 

llvm/lib/Analysis/Loads.cpp

@@ -799,7 +799,7 @@ Value *llvm::FindAvailableLoadedValue(LoadInst *Load, BatchAAResults &AA,
 
 // Returns true if a use is either in an ICmp/PtrToInt or a Phi/Select that only
 // feeds into them.
-static bool isPointerUseReplacable(const Use &U) {
+static bool isPointerUseReplaceable(const Use &U) {
   unsigned Limit = 40;
   SmallVector<const User *> Worklist({U.getUser()});
   SmallPtrSet<const User *, 8> Visited;
@@ -847,7 +847,7 @@ bool llvm::canReplacePointersInUseIfEqual(const Use &U, const Value *To,
 
   if (isPointerAlwaysReplaceable(&*U, To, DL))
     return true;
-  return isPointerUseReplacable(U);
+  return isPointerUseReplaceable(U);
 }
 
 bool llvm::canReplacePointersIfEqual(const Value *From, const Value *To,

llvm/lib/Transforms/Utils/SCCPSolver.cpp

@@ -16,6 +16,7 @@
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/InstructionSimplify.h"
+#include "llvm/Analysis/Loads.h"
 #include "llvm/Analysis/ValueLattice.h"
 #include "llvm/Analysis/ValueLatticeUtils.h"
 #include "llvm/Analysis/ValueTracking.h"
@@ -82,6 +83,24 @@ bool SCCPSolver::tryToReplaceWithConstant(Value *V) {
     return false;
   }
 
+  // Perform constant pointer propagation as long as assuming PredicateInfo
+  // derived equality between the two holds, and their provenance is the same.
+  if (mayForwardPointerPredicatedCopy(V)) {
+    bool MadeChange = false;
+    const auto &DL = cast<Instruction>(V)->getDataLayout();
+
+    V->replaceUsesWithIf(Const, [&](Use &U) {
+      bool CanReplace = canReplacePointersInUseIfEqual(U, Const, DL);
+      if (CanReplace)
+        LLVM_DEBUG(dbgs() << "  Constant pointer: " << *Const << " = " << *V
+                          << '\n');
+
+      MadeChange |= CanReplace;
+      return CanReplace;
+    });
+    return MadeChange;
+  }
+
   LLVM_DEBUG(dbgs() << "  Constant: " << *Const << " = " << *V << '\n');
 
   // Replaces all of the uses of a variable with uses of the constant.
@@ -350,11 +369,12 @@ bool SCCPSolver::simplifyInstsInBlock(BasicBlock &BB,
     if (Inst.getType()->isVoidTy())
       continue;
     if (tryToReplaceWithConstant(&Inst)) {
-      if (wouldInstructionBeTriviallyDead(&Inst))
+      if (isInstructionTriviallyDead(&Inst)) {
         Inst.eraseFromParent();
+        ++InstRemovedStat;
+      }
 
       MadeChanges = true;
-      ++InstRemovedStat;
     } else if (replaceSignedInst(*this, InsertedValues, Inst)) {
       MadeChanges = true;
       ++InstReplacedStat;
@@ -569,6 +589,9 @@ class SCCPInstVisitor : public InstVisitor<SCCPInstVisitor> {
   // The BasicBlock work list
   SmallVector<BasicBlock *, 64> BBWorkList;
 
+  /// List of users that carry forward the predicated copy.
+  SmallPtrSet<Value *, 16> VisitedUsers;
+
   /// KnownFeasibleEdges - Entries in this set are edges which have already had
   /// PHI nodes retriggered.
   using Edge = std::pair<BasicBlock *, BasicBlock *>;
@@ -1030,6 +1053,27 @@ class SCCPInstVisitor : public InstVisitor<SCCPInstVisitor> {
     }
     Invalidated.clear();
   }
+
+  bool mayForwardPointerPredicatedCopy(Value *V) {
+    if (auto *I = dyn_cast<Instruction>(V); I && I->getType()->isPointerTy()) {
+      auto It = FnPredicateInfo.find(I->getFunction());
+      if (It == FnPredicateInfo.end())
+        return false;
+      if (It->second->getPredicateInfoFor(I))
+        return true;
+      // See whether there are intermediate copies of the predicated value.
+      // As per how the solver propagates constants, it should suffice to visit
+      // the operands of the currently examined instruction, adding it to the
+      // already visited users.
+      for (Value *Op : I->operand_values()) {
+        if (It->second->getPredicateInfoFor(Op) || VisitedUsers.count(Op)) {
+          VisitedUsers.insert(I);
+          return true;
+        }
+      }
+    }
+    return false;
+  }
 };
 
 } // namespace llvm
@@ -2379,3 +2423,7 @@ void SCCPSolver::markFunctionUnreachable(Function *F) {
 void SCCPSolver::visit(Instruction *I) { Visitor->visit(I); }
 
 void SCCPSolver::visitCall(CallInst &I) { Visitor->visitCall(I); }
+
+bool SCCPSolver::mayForwardPointerPredicatedCopy(Value *V) {
+  return Visitor->mayForwardPointerPredicatedCopy(V);
+}

llvm/test/Transforms/SCCP/assume-equality-pointers.ll

@@ -0,0 +1,86 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt < %s -passes=ipsccp -S | FileCheck %s
+
+define ptr @assume_pointers_equality_maybe_different_provenance_1(ptr %x) {
+; CHECK-LABEL: define ptr @assume_pointers_equality_maybe_different_provenance_1(
+; CHECK-SAME: ptr [[X:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq ptr [[X]], inttoptr (i64 12345678 to ptr)
+; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
+; CHECK-NEXT:    ret ptr [[X]]
+;
+entry:
+  %cmp = icmp eq ptr %x, inttoptr (i64 12345678 to ptr)
+  call void @llvm.assume(i1 %cmp)
+  ret ptr %x
+}
+
+define ptr @assume_pointers_equality_maybe_different_provenance_2(ptr %x, i1 %cond) {
+; CHECK-LABEL: define ptr @assume_pointers_equality_maybe_different_provenance_2(
+; CHECK-SAME: ptr [[X:%.*]], i1 [[COND:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq ptr [[X]], inttoptr (i64 12345678 to ptr)
+; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
+; CHECK-NEXT:    br i1 [[COND]], label %[[NEXT:.*]], label %[[EXIT:.*]]
+; CHECK:       [[NEXT]]:
+; CHECK-NEXT:    br label %[[EXIT]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    [[PHI:%.*]] = phi ptr [ [[X]], %[[ENTRY]] ], [ inttoptr (i64 12345678 to ptr), %[[NEXT]] ]
+; CHECK-NEXT:    ret ptr [[PHI]]
+;
+entry:
+  %cmp = icmp eq ptr %x, inttoptr (i64 12345678 to ptr)
+  call void @llvm.assume(i1 %cmp)
+  br i1 %cond, label %next, label %exit
+
+next:
+  br label %exit
+
+exit:
+  %phi = phi ptr [ %x, %entry ], [ inttoptr (i64 12345678 to ptr), %next ]
+  ret ptr %phi
+}
+
+define internal ptr @callee(ptr returned %p) memory(none) {
+; CHECK-LABEL: define internal ptr @callee(
+; CHECK-SAME: ptr [[P:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    ret ptr poison
+;
+entry:
+  ret ptr %p
+}
+
+define ptr @assume_pointers_equality_maybe_different_provenance_3(ptr %x, i1 %cond) {
+; CHECK-LABEL: define ptr @assume_pointers_equality_maybe_different_provenance_3(
+; CHECK-SAME: ptr [[X:%.*]], i1 [[COND:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq ptr [[X]], inttoptr (i64 12345678 to ptr)
+; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
+; CHECK-NEXT:    br i1 [[COND]], label %[[NEXT:.*]], label %[[EXIT:.*]]
+; CHECK:       [[NEXT]]:
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 true, ptr [[X]], ptr inttoptr (i64 12345678 to ptr)
+; CHECK-NEXT:    br label %[[EXIT]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    [[PHI:%.*]] = phi ptr [ [[X]], %[[ENTRY]] ], [ [[SEL]], %[[NEXT]] ]
+; CHECK-NEXT:    [[Q:%.*]] = call ptr @callee(ptr [[PHI]])
+; CHECK-NEXT:    [[SEL2:%.*]] = select i1 [[COND]], ptr [[Q]], ptr [[X]]
+; CHECK-NEXT:    ret ptr [[SEL2]]
+;
+entry:
+  %cmp = icmp eq ptr %x, inttoptr (i64 12345678 to ptr)
+  call void @llvm.assume(i1 %cmp)
+  br i1 %cond, label %next, label %exit
+
+next:
+  %sel = select i1 %cond, ptr %x, ptr inttoptr (i64 12345678 to ptr)
+  br label %exit
+
+exit:
+  %phi = phi ptr [ %x, %entry ], [ %sel, %next ]
+  %q = call ptr @callee(ptr %phi)
+  %sel2 = select i1 %cond, ptr %q, ptr %x
+  ret ptr %sel2
+}
+
+declare void @llvm.assume(i1)

llvm/test/Transforms/SCCP/replace-dereferenceable-ptr-with-undereferenceable.ll

@@ -11,7 +11,7 @@ define i32 @eq_undereferenceable(ptr %p) {
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq ptr [[P:%.*]], getelementptr inbounds (i32, ptr @x, i64 1)
 ; CHECK-NEXT:    br i1 [[CMP]], label [[IF_THEN:%.*]], label [[IF_END:%.*]]
 ; CHECK:       if.then:
-; CHECK-NEXT:    store i32 2, ptr getelementptr inbounds (i32, ptr @x, i64 1), align 4
+; CHECK-NEXT:    store i32 2, ptr [[P]], align 4
 ; CHECK-NEXT:    br label [[IF_END]]
 ; CHECK:       if.end:
 ; CHECK-NEXT:    [[TMP0:%.*]] = load i32, ptr @y, align 4
@@ -65,7 +65,7 @@ define i1 @eq_undereferenceable_cmp_simp(ptr %p) {
 ; CHECK-NEXT:    [[CMP_0:%.*]] = icmp eq ptr [[P:%.*]], getelementptr inbounds (i32, ptr @x, i64 1)
 ; CHECK-NEXT:    br i1 [[CMP_0]], label [[IF_THEN:%.*]], label [[IF_END:%.*]]
 ; CHECK:       if.then:
-; CHECK-NEXT:    store i32 2, ptr getelementptr inbounds (i32, ptr @x, i64 1), align 4
+; CHECK-NEXT:    store i32 2, ptr [[P]], align 4
 ; CHECK-NEXT:    ret i1 true
 ; CHECK:       if.end:
 ; CHECK-NEXT:    ret i1 false