[instcombine] Generalize one-compare rule from foldAllocaCmp for noalias calls

[reames]

Summary

This change generalizes the handling for a potentially capturing comparison of an unescaped alloca to handle noalias calls as well. In the process, it also generalizes the existing alloca handling to reuse the generic capture tracking logic.

Doing this requires reverting 43d7e1c at the same time, and thus fixing https://github.com/llvm/llvm-project/issues/54002. If we add nocapture argument support without reverting that change, we take a currently mostly silent miscompile and risk increasing it's frequency.

[reames]

Uploaded wrong patch version. Fixing...

[reames]

Corrected patch. Sorry for the noise.

[reames]

llvm/test/Transforms/InstCombine/compare-unescaped.ll

302

To help others avoid the same confusion I had: This case isn't being caught by foldAllocCmp. We have another transform in visitAllocSite which is doing escape analysis and removing the allocation in one go - and thus doing consistent folding.

I should probably add a non-capturing use to this test.

[nikic]

This looks basically right to me.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

1035

Can you please add an explicit check here that this is an icmp user? I guess you omitted the check on the premise that you have a known icmp user, so if you see a single capture it must be that icmp. That may be true, but the fact that we also have certain non-capturing icmps makes the reasoning here very non-obvious.

6053

malloc(4) + 1 == 1 would be a hidden null check. Though from a quick look, we would not recognize malloc(4) + 1 as non-zero.

malloc(4) < 4 would be another possibility, making use of malloc alignment guarantees to perform a hidden null pointer check.

Not sure what to do about these, it's rather contrived...

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

1035

Or rather, we should check that this user it the same one we're trying to fold.

And it may be worth pointing out that this critically depends on the fact that captures are per-use, not per-user, so if we do something like alloc + x == alloc + y, then we'll end up not folding because we have two capturing uses in the same icmp.

[reames]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

1035

I'd tried to clarify the later point in the comment, but I will expand.

6053

Oh, yuck. I'd completely missed this possibility.

I see two variants here, and they're a bit different.

1. If we ptrtoint the allocation (to do the comparison), then that's a capture, and we're fine.

2. If we do inttoptr of the RHS instead, that's essentially a *guess* of the address (i.e. it's null), but the logic about heap layouts doesn't apply because we can't just decide the allocation didn't fail (since some later use might fault.)

I'm going to add some test cases of this variety, then restrict the transform to the non-null allocation case OR when directly comparing a non-offset allocation.

Once that lands, I'll revisit and see if I can do something stronger.

[reames]

Closing review I'm not currently working on, and am unlikely to get back to in near future. Will reopen if priorities change.

llvm/lib/Analysis/CaptureTracking.cpp

case Instruction::ICmp: {
           auto *O = I->getOperand(Idx)->stripPointerCastsSameRepresentation();
           // Comparing a dereferenceable_or_null pointer against null cannot
           // lead to pointer escapes, because if it is not null it must be a
           // valid (in-bounds) pointer.
           if (Tracker->isDereferenceableOrNull(O, I->getModule()->getDataLayout()))
             break;
         }
       }
-      // Comparison against value stored in global variable. Given the pointer
-      // does not escape, its value cannot be guessed and stored separately in a
-      // global variable.
-      auto *LI = dyn_cast<LoadInst>(I->getOperand(OtherIdx));
-      if (LI && isa<GlobalVariable>(LI->getPointerOperand()))
-        break;
       // Otherwise, be conservative. There are crazy ways to capture pointers
       // using comparisons.
       if (Tracker->captured(U))
         return;
       break;
     }
     default:
       // Something else - be conservative and say it is captured.

llvm/lib/Analysis/InstructionSimplify.cpp

 // If the C and C++ standards are ever made sufficiently restrictive in this
 // area, it may be possible to update LLVM's semantics accordingly and reinstate
 // this optimization.
 static Constant *
 computePointerICmp(CmpInst::Predicate Pred, Value *LHS, Value *RHS,
                    const SimplifyQuery &Q) {
   const DataLayout &DL = Q.DL;
   const TargetLibraryInfo *TLI = Q.TLI;
-  const DominatorTree *DT = Q.DT;
-  const Instruction *CxtI = Q.CxtI;
   const InstrInfoQuery &IIQ = Q.IIQ;
 
   // First, skip past any trivial no-ops.
   LHS = LHS->stripPointerCasts();
   RHS = RHS->stripPointerCasts();
 
   // A non-null pointer is not equal to a null pointer.
   if (isa<ConstantPointerNull>(RHS) && ICmpInst::isEquality(Pred) &&
if (Pred == CmpInst::ICMP_EQ || Pred == CmpInst::ICMP_NE) {
     // noalias calls.  We assume that indexing from such disjoint storage
     // into the heap is undefined, and thus offsets can be safely ignored.
     auto IsAllocDisjoint = [](ArrayRef<const Value *> Objects) {
       return all_of(Objects, ::IsAllocDisjoint);
     };
 
     if ((IsNAC(LHSUObjs) && IsAllocDisjoint(RHSUObjs)) ||
         (IsNAC(RHSUObjs) && IsAllocDisjoint(LHSUObjs)))
         return ConstantInt::get(GetCompareTy(LHS),
                                 !CmpInst::isTrueWhenEqual(Pred));
-
-    // Fold comparisons for non-escaping pointer even if the allocation call
-    // cannot be elided. We cannot fold malloc comparison to null. Also, the
-    // dynamic allocation call could be either of the operands.  Note that
-    // the other operand can not be based on the alloc - if it were, then
-    // the cmp itself would be a capture.
-    Value *MI = nullptr;
-    if (isAllocLikeFn(LHS, TLI) &&
-        llvm::isKnownNonZero(RHS, DL, 0, nullptr, CxtI, DT))
-      MI = LHS;
-    else if (isAllocLikeFn(RHS, TLI) &&
-             llvm::isKnownNonZero(LHS, DL, 0, nullptr, CxtI, DT))
-      MI = RHS;
-    // FIXME: We should also fold the compare when the pointer escapes, but the
-    // compare dominates the pointer escape
-    if (MI && !PointerMayBeCaptured(MI, true, true))
-      return ConstantInt::get(GetCompareTy(LHS),
-                              CmpInst::isFalseWhenEqual(Pred));
   }
 
   // Otherwise, fail.
   return nullptr;
 }
 
 /// Fold an icmp when its operands have i1 scalar type.
 static Value *simplifyICmpOfBools(CmpInst::Predicate Pred, Value *LHS,

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

 //===- InstCombineCompares.cpp --------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the visitICmp and visitFCmp functions.
 //
 //===----------------------------------------------------------------------===//
 
 #include "InstCombineInternal.h"
 #include "llvm/ADT/APSInt.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/AliasAnalysis.h"
+#include "llvm/Analysis/CaptureTracking.h"
 #include "llvm/Analysis/CmpInstAnalysis.h"
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/GetElementPtrTypeIterator.h"
 #include "llvm/IR/IntrinsicInst.h"
if (GEPLHS->isInBounds() && ICmpInst::isEquality(Cond) &&
     }
   }
 
   // Try convert this to an indexed compare by looking through PHIs/casts as a
   // last resort.
   return transformToIndexedCompare(GEPLHS, RHS, Cond, DL);
 }
 
-Instruction *InstCombinerImpl::foldAllocaCmp(ICmpInst &ICI,
-                                             const AllocaInst *Alloca) {
+Instruction *InstCombinerImpl::foldAllocCmp(ICmpInst &ICI,
+                                            const Value *Allocation) {
   assert(ICI.isEquality() && "Cannot fold non-equality comparison.");
 
-  // It would be tempting to fold away comparisons between allocas and any
-  // pointer not based on that alloca (e.g. an argument). However, even
-  // though such pointers cannot alias, they can still compare equal.
+  // It would be tempting to fold away comparisons between allocas or noalias
+  // calls and any pointer not based on that allocation (e.g. an argument).
+  // However, even though such pointers cannot alias, they can still compare
+  // equal.
   //
-  // But LLVM doesn't specify where allocas get their memory, so if the alloca
-  // doesn't escape we can argue that it's impossible to guess its value, and we
-  // can therefore act as if any such guesses are wrong.
+  // But LLVM doesn't specify where allocas or noalias calls get their memory,
+  // so if the allocation doesn't escape we can argue that it's impossible to
+  // guess its value, and we can therefore act as if any such guesses are wrong.
   //
-  // The code below checks that the alloca doesn't escape, and that it's only
-  // used in a comparison once (the current instruction). The
+  // The code below checks that the allocation doesn't escape, and that it's
+  // only used in a comparison once (the current instruction). The
   // single-comparison-use condition ensures that we're trivially folding all
   // comparisons against the alloca consistently, and avoids the risk of
   // erroneously folding a comparison of the pointer with itself.
 
-  unsigned MaxIter = 32; // Break cycles and bound to constant-time.
+  struct CustomCaptureTracker : public CaptureTracker {
+    explicit CustomCaptureTracker() {}
 
-  SmallVector<const Use *, 32> Worklist;
-  for (const Use &U : Alloca->uses()) {
-    if (Worklist.size() >= MaxIter)
-      return nullptr;
-    Worklist.push_back(&U);
+    void tooManyUses() override { Captured = true; }
+
+    bool captured(const Use *U) override {
+      NumCmps++;
+      if (NumCmps == 1)
+        // Exactly one compare use seen
+        return false;
+
+      Captured = true;
+      return true;
     }
 
     unsigned NumCmps = 0;
-  while (!Worklist.empty()) {
-    assert(Worklist.size() <= MaxIter);
-    const Use *U = Worklist.pop_back_val();
-    const Value *V = U->getUser();
-    --MaxIter;
-
-    if (isa<BitCastInst>(V) || isa<GetElementPtrInst>(V) || isa<PHINode>(V) ||
-        isa<SelectInst>(V)) {
-      // Track the uses.
-    } else if (isa<LoadInst>(V)) {
-      // Loading from the pointer doesn't escape it.
-      continue;
-    } else if (const auto *SI = dyn_cast<StoreInst>(V)) {
-      // Storing *to* the pointer is fine, but storing the pointer escapes it.
-      if (SI->getValueOperand() == U->get())
-        return nullptr;
-      continue;
-    } else if (isa<ICmpInst>(V)) {
-      if (NumCmps++)
-        return nullptr; // Found more than one cmp.
-      continue;
-    } else if (const auto *Intrin = dyn_cast<IntrinsicInst>(V)) {
-      switch (Intrin->getIntrinsicID()) {
-        // These intrinsics don't escape or compare the pointer. Memset is safe
-        // because we don't allow ptrtoint. Memcpy and memmove are safe because
-        // we don't allow stores, so src cannot point to V.
-        case Intrinsic::lifetime_start: case Intrinsic::lifetime_end:
-        case Intrinsic::memcpy: case Intrinsic::memmove: case Intrinsic::memset:
-          continue;
-        default:
-          return nullptr;
-      }
-    } else {
-      return nullptr;
-    }
-    for (const Use &U : V->uses()) {
-      if (Worklist.size() >= MaxIter)
-        return nullptr;
-      Worklist.push_back(&U);
-    }
-  }
+    bool Captured = false;
+  };
+
+  CustomCaptureTracker Tracker;
+  const unsigned MaxIter = 32; // Break cycles and bound to constant-time.
+  PointerMayBeCaptured(Allocation, &Tracker, MaxIter); 
+
+  if (Tracker.Captured)
 
   auto *Res = ConstantInt::get(ICI.getType(),
                                !CmpInst::isTrueWhenEqual(ICI.getPredicate()));
   return replaceInstUsesWith(ICI, Res);
 }
 
 /// Fold "icmp pred (X+C), X".
 Instruction *InstCombinerImpl::foldICmpAddOpConst(Value *X, const APInt &C,
Instruction *InstCombinerImpl::visitICmpInst(ICmpInst &I) {
   // Try to optimize 'icmp GEP, P' or 'icmp P, GEP'.
   if (auto *GEP = dyn_cast<GEPOperator>(Op0))
     if (Instruction *NI = foldGEPICmp(GEP, Op1, I.getPredicate(), I))
       return NI;
   if (auto *GEP = dyn_cast<GEPOperator>(Op1))
     if (Instruction *NI = foldGEPICmp(GEP, Op0, I.getSwappedPredicate(), I))
       return NI;
 
-  // Try to optimize equality comparisons against alloca-based pointers.
+  // Try to optimize equality comparisons against memory allocations
+  // (e.g. alloca or noalias calls)
   if (Op0->getType()->isPointerTy() && I.isEquality()) {
     assert(Op1->getType()->isPointerTy() && "Comparing pointer with non-pointer?");
-    if (auto *Alloca = dyn_cast<AllocaInst>(getUnderlyingObject(Op0)))
-      if (Instruction *New = foldAllocaCmp(I, Alloca))
+    auto *UO0 = getUnderlyingObject(Op0);
+    if (isa<AllocaInst>(UO0))
+      if (Instruction *New = foldAllocCmp(I, UO0))
+        return New;
+    auto *UO1 = getUnderlyingObject(Op1);
+    if (isa<AllocaInst>(UO1))
+      if (Instruction *New = foldAllocCmp(I, UO1))
+        return New;
+
+    if (isNoAliasCall(UO0) || isNoAliasCall(UO1)) {
+      // noalias results are allowed to be null.  We need to make sure that
+      // either a) the result isn't null at the use, or b) the result being
+      // null still results in it not being equal to the other side.
+      if (isKnownNonZero(Op0, Q.DL, /*Depth=*/0, Q.AC, Q.CxtI, Q.DT) ||
+          isKnownNonZero(Op1, Q.DL, /*Depth=*/0, Q.AC, Q.CxtI, Q.DT)) {
+        if (isNoAliasCall(UO0))
+          if (Instruction *New = foldAllocCmp(I, UO0))
             return New;
-    if (auto *Alloca = dyn_cast<AllocaInst>(getUnderlyingObject(Op1)))
-      if (Instruction *New = foldAllocaCmp(I, Alloca))
+        if (isNoAliasCall(UO1))
+          if (Instruction *New = foldAllocCmp(I, UO1))
             return New;
       }
+    }
+  }
 
   if (Instruction *Res = foldICmpBitCast(I))
     return Res;
 
   // TODO: Hoist this above the min/max bailout.
   if (Instruction *R = foldICmpWithCastOp(I))
     return R;
 

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

public:
   Instruction *foldIntegerTypedPHI(PHINode &PN);
 
   /// Helper function for FoldPHIArgXIntoPHI() to set debug location for the
   /// folded operation.
   void PHIArgMergedDebugLoc(Instruction *Inst, PHINode &PN);
 
   Instruction *foldGEPICmp(GEPOperator *GEPLHS, Value *RHS,
                            ICmpInst::Predicate Cond, Instruction &I);
-  Instruction *foldAllocaCmp(ICmpInst &ICI, const AllocaInst *Alloca);
+  Instruction *foldAllocCmp(ICmpInst &ICI, const Value *Allocation);
   Instruction *foldCmpLoadFromIndexedGlobal(LoadInst *LI,
                                             GetElementPtrInst *GEP,
                                             GlobalVariable *GV, CmpInst &ICI,
                                             ConstantInt *AndCst = nullptr);
   Instruction *foldFCmpIntToFPConst(FCmpInst &I, Instruction *LHSI,
                                     Constant *RHSC);
   Instruction *foldICmpAddOpConst(Value *X, const APInt &C,
                                   ICmpInst::Predicate Pred);

llvm/test/Transforms/InstCombine/compare-alloca.ll

 }
 
 ; A nocapture call can't cause a consistent result issue as it is (by
 ; assumption) not able to contain a comparison which might capture the
 ; address.
 
 declare void @unknown(i8*)
 
-; TODO: Missing optimization
 define i1 @consistent_nocapture_inttoptr() {
 ; CHECK-LABEL: @consistent_nocapture_inttoptr(
 ; CHECK-NEXT:    [[M1:%.*]] = alloca [4 x i8], align 1
 ; CHECK-NEXT:    [[M1_SUB:%.*]] = getelementptr inbounds [4 x i8], [4 x i8]* [[M1]], i32 0, i32 0
 ; CHECK-NEXT:    call void @unknown(i8* nocapture nonnull [[M1_SUB]])
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8* [[M1_SUB]], inttoptr (i64 2048 to i8*)
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
   %m = alloca i8, i32 4
   call void @unknown(i8* nocapture %m)
   %rhs = inttoptr i64 2048 to i8*
   %cmp = icmp eq i8* %m, %rhs
   ret i1 %cmp
 }
 
 define i1 @consistent_nocapture_offset() {
 ; CHECK-LABEL: @consistent_nocapture_offset(
 ; CHECK-NEXT:    [[M1:%.*]] = alloca [4 x i8], align 1
 ; CHECK-NEXT:    [[M1_SUB:%.*]] = getelementptr inbounds [4 x i8], [4 x i8]* [[M1]], i32 0, i32 0
 ; CHECK-NEXT:    call void @unknown(i8* nocapture nonnull [[M1_SUB]])
 ; CHECK-NEXT:    ret i1 false
 ;
   %m = alloca i8, i32 4
   call void @unknown(i8* nocapture %m)
   %n = alloca i8, i32 4
   %rhs = getelementptr i8, i8* %n, i32 4
   %cmp = icmp eq i8* %m, %rhs
   ret i1 %cmp
 }
 
 @gp = global i32* null, align 8
-; TODO: Missing optimization
 define i1 @consistent_nocapture_through_global() {
 ; CHECK-LABEL: @consistent_nocapture_through_global(
 ; CHECK-NEXT:    [[M1:%.*]] = alloca i32, align 1
 ; CHECK-NEXT:    [[M1_SUB:%.*]] = bitcast i32* [[M1]] to i8*
 ; CHECK-NEXT:    call void @unknown(i8* nocapture nonnull [[M1_SUB]])
-; CHECK-NEXT:    [[LGP:%.*]] = load i32*, i32** @gp, align 8, !nonnull !0
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32* [[M1]], [[LGP]]
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
   %m = alloca i8, i32 4
   call void @unknown(i8* nocapture %m)
   %bc = bitcast i8* %m to i32*
   %lgp = load i32*, i32** @gp, align 8, !nonnull !{}
   %cmp = icmp eq i32* %bc, %lgp
   ret i1 %cmp
 }

llvm/test/Transforms/InstCombine/compare-unescaped.ll

 ; the same point that applies to allocas, applied to noaiias/malloc.
 
 ; These two functions represents either a) forging a pointer via inttoptr or
 ; b) indexing off an adjacent allocation.  In either case, the operation is
 ; obscured by an uninlined helper and not visible to instcombine.
 declare i8* @hidden_inttoptr()
 declare i8* @hidden_offset(i8* %other)
 
-; FIXME: Missed oppurtunity
 define i1 @ptrtoint_single_cmp() {
 ; CHECK-LABEL: @ptrtoint_single_cmp(
-; CHECK-NEXT:    [[M:%.*]] = call dereferenceable_or_null(4) i8* @malloc(i64 4)
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8* [[M]], inttoptr (i64 2048 to i8*)
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
   %m = call i8* @malloc(i64 4)
   %rhs = inttoptr i64 2048 to i8*
   %cmp = icmp eq i8* %m, %rhs
   ret i1 %cmp
 }
 
 define i1 @offset_single_cmp() {
;
   %lgp = load i32*, i32** @gp, align 8
   %rhs2 = call i8* @hidden_inttoptr()
   %cmp1 = icmp eq i32* %bc, %lgp
   %cmp2 = icmp eq i8* %m, %rhs2
   call void @witness(i1 %cmp1, i1 %cmp2)
   ret void
 }
 
-; FIXME: This appears correct, but the current implementation relies
-; on visiting both cmps in the same pass.  We may have an simplification order
-; under which one is missed, and that would be a bug.
 define void @neg_consistent_fold4() {
 ; CHECK-LABEL: @neg_consistent_fold4(
 ; CHECK-NEXT:    call void @witness(i1 false, i1 false)
 ; CHECK-NEXT:    ret void
 ;
   %m = call i8* @malloc(i64 4)
   %bc = bitcast i8* %m to i32*
   %lgp = load i32*, i32** @gp, align 8
   %cmp1 = icmp eq i32* %bc, %lgp
   %cmp2 = icmp eq i32* %bc, %lgp
   call void @witness(i1 %cmp1, i1 %cmp2)
   ret void
 }
 
 declare void @unknown(i8*)
 
 ; Points out that a nocapture call can't cause a consistent result issue
 ; as it is (by assumption) not able to contain a comparison which might
 ; capture the address.
 
 define i1 @consistent_nocapture_inttoptr() {
 ; CHECK-LABEL: @consistent_nocapture_inttoptr(
 ; CHECK-NEXT:    [[M:%.*]] = call dereferenceable_or_null(4) i8* @malloc(i64 4)
 ; CHECK-NEXT:    call void @unknown(i8* nocapture [[M]])
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8* [[M]], inttoptr (i64 2048 to i8*)
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
   %m = call i8* @malloc(i64 4)
   call void @unknown(i8* nocapture %m)
   %rhs = inttoptr i64 2048 to i8*
   %cmp = icmp eq i8* %m, %rhs
   ret i1 %cmp
 }