Fold compares irrespective of whether allocation can be elided

[anna]

Summary

When a non-escaping pointer is compared to a global value, the comparison can be folded even if the corresponding malloc/allocation call cannot be elided.
We need to make sure the global value is not null, since comparisons to null cannot be folded.

In future, we should also handle cases when the the comparison instruction dominates the pointer escape.

[sanjoy]

I think we cannot fold comparisons with values loaded from globals, since the global we're loading from could contain null. Maybe we need to use !nonnull metadata?

lib/Analysis/CaptureTracking.cpp

312 ↗

(On Diff #55049)

Nit: I'd rephrase this as "Comparison against value stored in global variable"

lib/Analysis/InstructionSimplify.cpp

2091 ↗

(On Diff #55049)

Nit: use nullptr

[anna]

fixes issue Sanjoy pointed out: avoiding compare folding when null value loaded in global.

[majnemer]

lib/Analysis/CaptureTracking.cpp

313 ↗

(On Diff #55084)

Please format this appropriately.

lib/Analysis/InstructionSimplify.cpp

2087–2092 ↗

(On Diff #55084)

Why not use isKnownNonNull or isKnownNonNullAt ?

[sanjoy]

lib/Analysis/CaptureTracking.cpp

312 ↗

(On Diff #55093)

Please add a sentence or two here on why this is correct i.e. you're
optimistically assuming that the pointer does not escape, and given
that it does not escape, its value could could not have been "guessed
into existence" and separately stored in a global variable.

I.e. we're avoiding situations like this:

store <some value> to <global>
%val = load <global>
if (%val == %ptr)
  escape(%val)

since <some value> cannot be equal to %ptr if %ptr does not escape.

lib/Analysis/InstructionSimplify.cpp

2087 ↗

(On Diff #55093)

s/Comparisons/comparisons and s/NULL/null. Can you also make the description one continuous paragraph?

2100 ↗

(On Diff #55093)

I think you need to pass true for both ReturnCaptures and StoreCaptures. Can you also add a test case that would've caught this?

[anna]

lib/Analysis/InstructionSimplify.cpp

2100 ↗

(On Diff #55093)

I had ReturnCaptures as false because of this case:
compare dominates the ret statement which returns the pointer. We can fold the compare since pointer has not escaped at the time of compare.

In case compare does not dominate the ret statement, if we fold the cmp and some optimizations later on reordere the statements such that the ret no longer returns the right value, we have incorrect functionality in the IR.
However, the test case should be such that the ret value is different, but the only statement causing the pointer to be captured should be the ret. If there were other statements causing pointer to be captured, we would have returned false for PointerMayBeCaptured()

TC I wrote for the case where cmp for %m does not dominate the ret

define i8* @compare_ret_escape(i8* %c) {

  %m = call i8* @malloc(i64 4)
  %n = call i8* @malloc(i64 4)
  %cmp = icmp eq i8* %n, %c
  br i1 %cmp, label %retst, label %chk

retst:
  ret i8* %m

chk:
  %bc = bitcast i8* %m to i32*
  %lgp = load i32*, i32** @gp, align 8, !nonnull !0
  %cmp2 = icmp eq i32* %bc, %lgp
  br i1 %cmp2, label %retst,  label %chk2

chk2:
  ret i8* %n
}

Right now the %cmp2 folds to false, and it looks right. Is there a case where we really need ReturnCaptures to be true?

[anna]

lib/Analysis/InstructionSimplify.cpp

2100 ↗

(On Diff #55093)

Also, just to add to this, the third case is cmp post dominates the ret. In all cases, the important point is that ret is a terminator instruction that exits the function, so there is no way to go from the ret to cmp.

[sanjoy]

lib/Analysis/InstructionSimplify.cpp

2100 ↗

(On Diff #55093)

You're right in that the optimizer needs to preserve single threaded
data dependencies. But I think there can be problems with
cross-thread data dependencies:

define i8* @f() {
entry:
  %ptr = call i8* @malloc()
  %global = load atomic i8*, i8** @global_var
  %cmp = icmp ne i8* %global, %ptr
  if (%cmp)
    *%global = 42;
  else
    *%global = 50;
  ret i8* %ptr
}

define void @g() {
  %val2 = call i8* @f()
  store atomic i8* %val2, i8** @global_var_2  ;; noalias with @global_var
  ret void
}

opt folds %cmp in @f to true, and then inlines @f into @g:

define void @g() {
entry:
  %ptr = call i8* @malloc()
  %global = load atomic i8*, i8** @global_var
  *%global = 42;
  store atomic i8* %ptr, i8** @global_var_2  ;; noalias with @global_var
  ret void
}

and re-orders:

define void @g() {
entry:
  %ptr = call i8* @malloc()
  store atomic i8* %ptr, i8** @global_var_2  ;; noalias with @global_var
  %global = load atomic i8*, i8** @global_var
  *%global = 42;
  ret void
}

The atomic markers on the loads and stores mean that these memory
accesses are visible on other threads (without atomic, sharing
memory across threads is undefined behavior), but it does not imply
any specific ordering.

Let's say we also have a thread doing this:

define void @another_thread() {
  %val = load atomic i8*, i8** @global_var_2
  store atomic i8* %val, i8** @global_var_2
  ret void
}

The original program would only ever write 50 to the value stored in
@global_var if the value stored in @global_var was from the
@malloc, so you could later do this safely:

%val = load i8*, i8** @global_var
if (*%val != 50)
  // no need to free %val

If another thread reads (atomically) from @global_var_2 and writes
that value to @global_var then in the transformed program the above
invariant is no longer true, since you could have this trace:

define void @g() {
entry:
  %ptr = call i8* @malloc()
  store atomic i8* %ptr, i8** @global_var_2  ;; noalias with @global_var

              %val = load atomic i8*, i8** @global_var_2
              store atomic i8* %val, i8** @global_var_2

  %global = load atomic i8*, i8** @global_var
  *%global = 42;
  ret void
}

and write 42 to a malloc'ed pointer stored in @global_var (illegal
state according to the original program).

Now I think we can get around this by restricting the load instruction
to be non-atomic, but for now I'd say it is safer to just bail out if
we see stores or returns.

[sanjoy]

lib/Analysis/CaptureTracking.cpp

312 ↗

(On Diff #55584)

Is this comment block wrapped to 80 chars?

313 ↗

(On Diff #55584)

Nit: "its"

314 ↗

(On Diff #55584)

Nit: I'd remove the "If the pointer escapes ... bit". I think the first sentence is enough, and the "captured = true" bit can be misleading since there isn't a bool capture in this function.

lib/Analysis/InstructionSimplify.cpp

2100 ↗

(On Diff #55584)

You can use CmpInst::isFalseWhenEqual here, but did you intend to put this under the previous Pred == CmpInst::ICMP_EQ || Pred == CmpInst::ICMP_NE check? Right now you're handling inequalities too, I think.

Inequalities are tricky, since it is difficult to justify folding %newPtr ULE -1 to false, for instance.

[anna]

lib/Analysis/InstructionSimplify.cpp

2100 ↗

(On Diff #55584)

I see your point. Although we can never know the %newPtr value when it is dynamically allocated, there *might be* some inequalities which can not automatically be folded to false, such as making sure the value is ULE or UGE a 'boundary' value.

[sanjoy]

This LGTM, but can you please add a test for ne specifically and one or two negative tests for inequality?

lib/Analysis/InstructionSimplify.cpp

3046 ↗

(On Diff #55865)

This line is too long. I think just changing Constant to auto will fix it, otherwise please wrap.

[anna]

added more tests for compare as ne and scenarios when we cannot fold the comparison.

llvm/trunk/lib/Analysis/CaptureTracking.cpp

case Instruction::AddrSpaceCast:
         if (Count++ >= Threshold)
           return Tracker->tooManyUses();
 
         if (Visited.insert(&UU).second)
           if (Tracker->shouldExplore(&UU))
             Worklist.push_back(&UU);
       }
       break;
-    case Instruction::ICmp:
+    case Instruction::ICmp: {
       // Don't count comparisons of a no-alias return value against null as
       // captures. This allows us to ignore comparisons of malloc results
       // with null, for example.
       if (ConstantPointerNull *CPN =
           dyn_cast<ConstantPointerNull>(I->getOperand(1)))
         if (CPN->getType()->getAddressSpace() == 0)
           if (isNoAliasCall(V->stripPointerCasts()))
             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.
+      unsigned OtherIndex = (I->getOperand(0) == V) ? 1 : 0;
+      auto *LI = dyn_cast<LoadInst>(I->getOperand(OtherIndex));
+      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.
       if (Tracker->captured(U))
         return;
       break;
     }
   }
 
   // All uses examined.
 }

llvm/trunk/lib/Analysis/InstructionSimplify.cpp

 // they have not been simplified then results are correct but maybe suboptimal).
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AliasAnalysis.h"
+#include "llvm/Analysis/CaptureTracking.h"
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/Analysis/VectorUtils.h"
 #include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/GetElementPtrTypeIterator.h"

 //
 // C11's version is more restrictive, however there's no reason why an argument
 // couldn't be a one-past-the-end value for a stack object in the caller and be
 // equal to the beginning of a stack object in the callee.
 //
 // 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(const DataLayout &DL,
-                                    const TargetLibraryInfo *TLI,
-                                    CmpInst::Predicate Pred, Value *LHS,
-                                    Value *RHS) {
+static Constant *
+computePointerICmp(const DataLayout &DL, const TargetLibraryInfo *TLI,
+                   const DominatorTree *DT, CmpInst::Predicate Pred,
+                   const Instruction *CxtI, Value *LHS, Value *RHS) {
   // 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 (llvm::isKnownNonNull(LHS, TLI) && isa<ConstantPointerNull>(RHS) &&
       (Pred == CmpInst::ICMP_EQ || Pred == CmpInst::ICMP_NE))
     return ConstantInt::get(GetCompareTy(LHS),
auto IsAllocDisjoint = [](SmallVectorImpl<Value *> &Objects) {
         return false;
       });
     };
 
     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.
+    Value *MI = nullptr;
+    if (isAllocLikeFn(LHS, TLI) && llvm::isKnownNonNullAt(RHS, CxtI, DT, TLI))
+      MI = LHS;
+    else if (isAllocLikeFn(RHS, TLI) &&
+             llvm::isKnownNonNullAt(LHS, CxtI, DT, TLI))
+      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;
 }
 
 /// Given operands for an ICmpInst, see if we can fold the result.
 /// If not, this returns null.
if (match(LHS, m_SMax(m_Value(A), m_Value(B))) &&
     if (Pred == CmpInst::ICMP_UGT)
       // Always false.
       return getFalse(ITy);
   }
 
   // Simplify comparisons of related pointers using a powerful, recursive
   // GEP-walk when we have target data available..
   if (LHS->getType()->isPointerTy())
-    if (Constant *C = computePointerICmp(Q.DL, Q.TLI, Pred, LHS, RHS))
+    if (auto *C = computePointerICmp(Q.DL, Q.TLI, Q.DT, Pred, Q.CxtI, LHS, RHS))
       return C;
 
   if (GetElementPtrInst *GLHS = dyn_cast<GetElementPtrInst>(LHS)) {
     if (GEPOperator *GRHS = dyn_cast<GEPOperator>(RHS)) {
       if (GLHS->getPointerOperand() == GRHS->getPointerOperand() &&
           GLHS->hasAllConstantIndices() && GRHS->hasAllConstantIndices() &&
           (ICmpInst::isEquality(Pred) ||
            (GLHS->isInBounds() && GRHS->isInBounds() &&

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

 
 define i1 @compare_global_trivialne() {
   %m = call i8* @malloc(i64 4)
   %bc = bitcast i8* %m to i32*
   %lgp = load i32*, i32** @gp, align 8
   %cmp = icmp ne i32* %bc, %lgp
   ret i1 %cmp
 ; CHECK-LABEL: compare_global_trivialne
-; CHECK: ret i1 true 
+; CHECK: ret i1 true
 }
 
 
 ; Although the %m is marked nocapture in the deopt operand in call to function f,
 ; we cannot remove the alloc site: call to malloc
-; FIXME: The comparison should fold to false irrespective of whether the call to malloc can be elided or not
+; The comparison should fold to false irrespective of whether the call to malloc can be elided or not
 declare void @f()
-define i32 @compare_and_call_with_deopt() {
+define i1 @compare_and_call_with_deopt() {
 ; CHECK-LABEL: compare_and_call_with_deopt
   %m = call i8* @malloc(i64 24)
   %bc = bitcast i8* %m to i32*
-  %lgp = load i32*, i32** @gp, align 8
-  %cmp = icmp eq i32* %bc, %lgp
-  %rt = zext i1 %cmp to i32
+  %lgp = load i32*, i32** @gp, align 8, !nonnull !0
+  %cmp = icmp eq i32* %lgp, %bc
+  tail call void @f() [ "deopt"(i8* %m) ]
+  ret i1 %cmp
+; CHECK: ret i1 false
+}
+
+; Same functon as above with deopt operand in function f, but comparison is NE
+define i1 @compare_ne_and_call_with_deopt() {
+; CHECK-LABEL: compare_ne_and_call_with_deopt
+  %m = call i8* @malloc(i64 24)
+  %bc = bitcast i8* %m to i32*
+  %lgp = load i32*, i32** @gp, align 8, !nonnull !0
+  %cmp = icmp ne i32* %lgp, %bc
   tail call void @f() [ "deopt"(i8* %m) ]
-  ret i32 %rt 
-; CHECK: ret i32 %rt
+  ret i1 %cmp
+; CHECK: ret i1 true
+}
+
+; Same function as above, but global not marked nonnull, and we cannot fold the comparison
+define i1 @compare_ne_global_maybe_null() {
+; CHECK-LABEL: compare_ne_global_maybe_null
+  %m = call i8* @malloc(i64 24)
+  %bc = bitcast i8* %m to i32*
+  %lgp = load i32*, i32** @gp
+  %cmp = icmp ne i32* %lgp, %bc
+  tail call void @f() [ "deopt"(i8* %m) ]
+  ret i1 %cmp
+; CHECK: ret i1 %cmp
+}
+
+; FIXME: The comparison should fold to false since %m escapes (call to function escape)
+; after the comparison.
+declare void @escape(i8*)
+define i1 @compare_and_call_after() {
+; CHECK-LABEL: compare_and_call_after
+  %m = call i8* @malloc(i64 24)
+  %bc = bitcast i8* %m to i32*
+  %lgp = load i32*, i32** @gp, align 8, !nonnull !0
+  %cmp = icmp eq i32* %bc, %lgp
+  br i1 %cmp, label %escape_call, label %just_return
+
+escape_call:
+ call void @escape(i8* %m)
+ ret i1 true
+
+just_return:
+ ret i1 %cmp
 }
 
 define i1 @compare_distinct_mallocs() {
   %m = call i8* @malloc(i64 4)
   %n = call i8* @malloc(i64 4)
   %cmp = icmp eq i8* %m, %n
   ret i1 %cmp
   ; CHECK-LABEL: compare_distinct_mallocs
   ; CHECK: ret i1 false
 }
 
 ; the compare is folded to true since the folding compare looks through bitcasts. 
 ; call to malloc and the bitcast instructions are elided after that since there are no uses of the malloc 
 define i1 @compare_samepointer_under_bitcast() {
   %m = call i8* @malloc(i64 4)
   %bc = bitcast i8* %m to i32*
   %bcback = bitcast i32* %bc to i8*
   %cmp = icmp eq i8* %m, %bcback
   ret i1 %cmp
 ; CHECK-LABEL: compare_samepointer_under_bitcast
 ; CHECK: ret i1 true 
 }
 
 ; the compare is folded to true since the folding compare looks through bitcasts. 
-; call to malloc and the bitcast instructions are elided after that since there are no uses of the malloc 
+; The malloc call for %m cannot be elided since it is used in the call to function f.
 define i1 @compare_samepointer_escaped() {
   %m = call i8* @malloc(i64 4)
   %bc = bitcast i8* %m to i32*
   %bcback = bitcast i32* %bc to i8*
   %cmp = icmp eq i8* %m, %bcback
   call void @f() [ "deopt"(i8* %m) ]
   ret i1 %cmp
 ; CHECK-LABEL: compare_samepointer_escaped
 ; CHECK-NEXT: %m = call i8* @malloc(i64 4)
 ; CHECK-NEXT: call void @f() [ "deopt"(i8* %m) ]
 ; CHECK: ret i1 true 
 }
 
+; Technically, we can fold the %cmp2 comparison, even though %m escapes through
+; the ret statement since `ret` terminates the function and we cannot reach from
+; the ret to cmp. 
+; FIXME: Folding this %cmp2 when %m escapes through ret could be an issue with
+; cross-threading data dependencies since we do not make the distinction between
+; atomic and non-atomic loads in capture tracking.
+define i8* @compare_ret_escape(i8* %c) {
+  %m = call i8* @malloc(i64 4)
+  %n = call i8* @malloc(i64 4)
+  %cmp = icmp eq i8* %n, %c
+  br i1 %cmp, label %retst, label %chk
+
+retst:
+  ret i8* %m
+
+chk:
+  %bc = bitcast i8* %m to i32*
+  %lgp = load i32*, i32** @gp, align 8, !nonnull !0
+  %cmp2 = icmp eq i32* %bc, %lgp
+  br i1 %cmp2, label %retst,  label %chk2
+
+chk2:
+  ret i8* %n
+; CHECK-LABEL: compare_ret_escape
+; CHECK: %cmp = icmp eq i8* %n, %c
+; CHECK: %cmp2 = icmp eq i32* %bc, %lgp
+}
+
 ; The malloc call for %m cannot be elided since it is used in the call to function f.
 ; However, the cmp can be folded to true as %n doesnt escape and %m, %n are distinct allocations
 define i1 @compare_distinct_pointer_escape() {
   %m = call i8* @malloc(i64 4)
   %n = call i8* @malloc(i64 4)
   tail call void @f() [ "deopt"(i8* %m) ]
   %cmp = icmp ne i8* %m, %n
   ret i1 %cmp
 ; CHECK-LABEL: compare_distinct_pointer_escape
 ; CHECK-NEXT: %m = call i8* @malloc(i64 4)
 ; CHECK-NEXT: tail call void @f() [ "deopt"(i8* %m) ]
 ; CHECK-NEXT: ret i1 true
 }
+
+!0 = !{}