[InstCombine] Support noalias calls in foldAllocaCmp()

[nikic]

Summary

This is a reboot of D120371, which adds support for noalias calls in foldAllocaCmp(). The only difference to alloca handling is that allocators might return null, in which case we must make sure that we don't fold away comparisons to null.

Once this lands, I'll commit the removal of the old incorrect InstSimplify/CaptureTracking code in a separate commit, to reduce risk.

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

6458

Maybe cache these in an optional so the second call to canFold doesn't need to recompute?

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

6458

These can't be easily cached because they may work on different values (underlying object vs plain operand).

[nikic]

I've thought about this a bit more, and I think there's another problem here: Noalias calls are really the wrong criterion for this. Ideally, noalias would only be about provenance, especially now that we have allockind, and this optimization is not provenance-related.

However, even if we switch this to use isAllocationFn() and require that allocators have no heap layout guarantees, then we may run into problems with partial inlining. Consider this cute custom allocator:

// Let's ignore all the missing safety checks here...
void *cached_alloc;
void *my_malloc(size_t size) {
    if (!cached_alloc)
        return cached_alloc;
    return malloc(size);
}
void free(void *ptr) {
    if (ptr == cached_alloc) {
        cached_alloc = NULL;
        return;
    }
    free(ptr);
}

Now, if this gets partially inlined, and we see my_malloc as an allocator, but free is inlined, then we're going to optimize away that ptr == cached_alloc comparison, resulting in a miscompile.

Incidentally, this specific example could already be a problem with the existing InstSimplify code (because it uses the icmp of load of global pattern, just missing the nonnull bit), but this patch has the potential to make things worse, so we probably shouldn't land it in this form.

Probably the proper way to support this would be to add an additional allockind like any_heap_layout and then only annotate non-replacable allocators with it (e.g. malloc but not new) -- though there are still LTO hazards.

In the meantime I've landed https://github.com/llvm/llvm-project/commit/4d192f2d2a545c8cd51ef86f9e83209eccbe229e to move the CaptureTracking special case into the InstSimplify fold, so it at least doesn't affect other code.

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/VectorUtils.h"
 #include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/GetElementPtrTypeIterator.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 *Alloc) {
   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 alloca 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 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.
   //
-  // 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.
 
   struct CmpCaptureTracker : public CaptureTracker {
     bool Captured = false;
     unsigned NumCmps = 0;
 
     void tooManyUses() override { Captured = true; }
 
     bool captured(const Use *U) override {
       if (isa<ICmpInst>(U->getUser()) && ++NumCmps == 1) {
         // Ignore one icmp capture.
         return false;
       }
 
       Captured = true;
       return true;
     }
   };
 
   CmpCaptureTracker Tracker;
-  PointerMayBeCaptured(Alloca, &Tracker);
+  PointerMayBeCaptured(Alloc, &Tracker);
   if (Tracker.Captured)
     return nullptr;
 
   auto *Res = ConstantInt::get(ICI.getType(),
                                !CmpInst::isTrueWhenEqual(ICI.getPredicate()));
   return replaceInstUsesWith(ICI, Res);
 }
 
Instruction *InstCombinerImpl::visitICmpInst(ICmpInst &I) {
 
   if (auto *SI = dyn_cast<SelectInst>(Op0))
     if (Instruction *NI = foldSelectICmp(I.getPredicate(), SI, Op1, I))
       return NI;
   if (auto *SI = dyn_cast<SelectInst>(Op1))
     if (Instruction *NI = foldSelectICmp(I.getSwappedPredicate(), SI, Op0, I))
       return NI;
 
-  // Try to optimize equality comparisons against alloca-based pointers.
+  // Try to optimize equality comparisons against memory allocations.
   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))
+    assert(Op1->getType()->isPointerTy() &&
+           "Comparing pointer with non-pointer?");
+    auto CanFold = [&](const Value *UO, const Value *Op, const Value *Other) {
+      if (isa<AllocaInst>(UO))
+        return true;
+      // Allocators might return null values. We can only fold comparisons to
+      // false if either the allocator can't return null, or we have a direct
+      // comparison against a non-null value. (We don't want to reason about
+      // hidden null checks like (gep malloc, X) == inttoptr(X) here.)
+      if (isNoAliasCall(UO))
+        return isKnownNonZero(UO, Q.DL, /*Depth*/ 0, Q.AC, &I, Q.DT) ||
+               (UO == Op &&
+                isKnownNonZero(Other, Q.DL, /*Depth*/ 0, Q.AC, &I, Q.DT));
+      return false;
+    };
+
+    Value *UO0 = getUnderlyingObject(Op0);
+    if (CanFold(UO0, Op0, Op1))
+      if (Instruction *New = foldAllocCmp(I, UO0))
         return New;
-    if (auto *Alloca = dyn_cast<AllocaInst>(getUnderlyingObject(Op1)))
-      if (Instruction *New = foldAllocaCmp(I, Alloca))
+    Value *UO1 = getUnderlyingObject(Op1);
+    if (CanFold(UO1, Op1, Op0))
+      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))

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

public:
   /// 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 *foldSelectICmp(ICmpInst::Predicate Pred, SelectInst *SI,
                               Value *RHS, const ICmpInst &I);
-  Instruction *foldAllocaCmp(ICmpInst &ICI, const AllocaInst *Alloca);
+  Instruction *foldAllocCmp(ICmpInst &ICI, const Value *Alloc);
   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-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 ptr @hidden_inttoptr()
 declare ptr @hidden_offset(ptr %other)
 
-; FIXME: Missed oppurtunity
 define i1 @ptrtoint_single_cmp() {
 ; CHECK-LABEL: @ptrtoint_single_cmp(
-; CHECK-NEXT:    [[M:%.*]] = call dereferenceable_or_null(4) ptr @malloc(i64 4)
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq ptr [[M]], inttoptr (i64 2048 to ptr)
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
   %m = call ptr @malloc(i64 4)
   %rhs = inttoptr i64 2048 to ptr
   %cmp = icmp eq ptr %m, %rhs
   ret i1 %cmp
 }
 
 define i1 @offset_single_cmp() {

 ; 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) ptr @malloc(i64 4)
 ; CHECK-NEXT:    call void @unknown(ptr nocapture [[M]])
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq ptr [[M]], inttoptr (i64 2048 to ptr)
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
   %m = call ptr @malloc(i64 4)
   call void @unknown(ptr nocapture %m)
   %rhs = inttoptr i64 2048 to ptr
   %cmp = icmp eq ptr %m, %rhs
   ret i1 %cmp
 }
 

 ;
   %m = call nonnull ptr @malloc(i64 4)
   %n = call nonnull ptr @malloc(i64 4)
   call void @unknown(ptr %n)
   %cmp = icmp eq ptr %m, %n
   ret i1 %cmp
 }
 
-; TODO: We can fold this, but don't with the current scheme.
 define i1 @two_nonnull_mallocs_hidden() {
 ; CHECK-LABEL: @two_nonnull_mallocs_hidden(
-; CHECK-NEXT:    [[M:%.*]] = call nonnull dereferenceable(4) ptr @malloc(i64 4)
-; CHECK-NEXT:    [[N:%.*]] = call nonnull dereferenceable(4) ptr @malloc(i64 4)
-; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr i8, ptr [[M]], i64 1
-; CHECK-NEXT:    [[GEP2:%.*]] = getelementptr i8, ptr [[N]], i64 2
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq ptr [[GEP1]], [[GEP2]]
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
   %m = call nonnull ptr @malloc(i64 4)
   %n = call nonnull ptr @malloc(i64 4)
   %gep1 = getelementptr i8, ptr %m, i32 1
   %gep2 = getelementptr i8, ptr %n, i32 2
   %cmp = icmp eq ptr %gep1, %gep2
   ret i1 %cmp
 }
 
 
 !0 = !{}