InstCombine: Fold comparisons between unguessable allocas and other pointers

[hans]

Summary

This will allow us to optimize code such as:

int f(int *p) {
  int x;
  return p == &x;
}

as well as:

int *allocate(void);
int f() {
  int x;
  int *p = allocate();
  return p == &x;
}

The folding can only be done under certain circumstances. Even though p and &x cannot alias, the comparison must still return true if the pointer representations are equal. If a user successfully generates a p that's a correct guess for &x, comparison should return true even though p is an invalid pointer.

This patch argues that if the address of the alloca isn't observable outside the function, the function can act as-if the address is impossible to guess from the outside. The tricky part is keeping the act consistent: if we fold p == &x to false in one place, we must make sure to fold any other comparisons based on those pointers similarly. To ensure that, we only fold when &x is involved exactly once in comparison instructions.

Let me know what you think.

[sanjoy]

lib/Analysis/InstructionSimplify.cpp

2190

I have a question about this second bit -- do we have to ensure that we only *fold* such comparisons consistently, or do we have to ensure that they also evaluate consistently? IOW, say I have

void f(i8* %arg) {
 i8* %val = alloca
 %t0 = %val == %arg
 ...
 %t1 = complex_xform(%val) == %arg

such that complex_xform(x) == x always, but the compiler cannot prove that. Consequently we fold %t0 to false, but don't fold complex_xform(x) == x at all, and it ends up evaluating to true at runtime.

Are you relying on the fact that a complex_xform cannot be formed out of the operators you look through in CanTrackAlluses? I'd be more comfortable if you folded *all* comparisons based off of %alloca in one go.

2197

Can this be rewritten to use PointerMayBeCaptured?

2198

Minor: I'd call this MaxIter.

2223

Why not just use APInt's constructor?

2228

patchpoint, stackmap and gc.statepoint all can escape values -- they're basically a call to an unknown function (I don't know if they'll show up as an IntrinsicInst though).

[hans]

lib/Analysis/InstructionSimplify.cpp

2190

Yes, exactly: the comparisons have to evaluate consistently, so we have to fold them all, and I'm relying on CanTrackAllUses (maybe not a great name actually) to ensure that we can do that.

I don't think I could fold all comparisons based on an %alloca in one go from InstSimplify as it's actually not changing the instruction, just returning a new value for it. Maybe there's a better place to do this?

On the other hand, I don't think multiple comparisons is very common. Maybe we should just bail out if we see another comparison? I'll do that for now.

2197

My check is much more conservative. For example, it won't follow the use through a phi node, and certainly won't allow it in a function call - even if the function is nocapture, it could make comparisons with the pointer.

David pointed out that I should probably break this out into a separate function though, so I'll do that.

2228

Thanks. David also pointed out he had some intrinsic that would escape a pointer. I'll build a whitelist here.

[hans]

Addressing Sanjoy's comments.

[hans]

Actually, requiring that the alloca is only used in a single comparison makes everything much simpler. I also verified that it didn't make the optimization fire any less over Chromium (this patch reduced binary size by about 10k).

Also, this means we don't need to trace the Argument value. It's fine if it escapes, because we know it can't be compared against any values based on the alloca since only allow one cmp.

And now that we don't need to be able to trace the non-Argument side of the comparison, it means we don't have to restrains that to just Arguments, in fact under these conditions we can act as if the alloca doesn't alias pointers based on any other value, including e.g. pointers from malloc. (This is similar to the "an alloca which is only used in a single cmp can be trivially folded" argument in the patch Philip sent to the list.)

I think this makes the optimization much more powerful. (Unfortunately it only knocked 100 more bytes of the Chromium build.) I'll get an updated patch out tomorrow.

[hans]

Moving this to instcombine.

The previous patch ran into a terrifying bug where instsimplify would get run on IR that wasn't fully constructed yet, which meant it couldn't see all uses of an alloca and folded a bit too aggressively.

This new patch successfully bootstraps Clang. The optimization fires 67 times during bootstrap and reduces Clang binary size with 8519 bytes.

[majnemer]

lib/Transforms/InstCombine/InstCombineCompares.cpp

780 ↗

(On Diff #36417)

Is this sufficient? What if getValueOperand is not U->get() but, say, a bitcast of U->get() ?

792–793 ↗

(On Diff #36417)

Can't memcpy escape the value just like a StoreInst would?

[hans]

lib/Transforms/InstCombine/InstCombineCompares.cpp

780 ↗

(On Diff #36417)

I think so. Such a bitcast would show up as a use of our value, so we'd track it and see if it ends up getting stored.

792–793 ↗

(On Diff #36417)

It's not possible to memcpy the value directly. The value would have to be stored in an object, and then that could be memcpy'd. But we'd notice the value escaping through the store.

[sanjoy]

Right now GetUnderlyingObject does not look through PHI nodes, but if it starts to at some point, this change will fold %cmp in

ph:
 %k = alloca i32
 br label %loop

loop:
 %t = phi i32* [ %k, %ph ], [ gep(%k, 1), %loop ]
 %cmp = icmp eq %k, %arg
 br i1 %cmp, label %loop, label %leave

I don't know if that ^ is something you want, but I wanted to point it out.

lib/Transforms/InstCombine/InstCombineCompares.cpp

786 ↗

(On Diff #36417)

How about assert(V == &ICI) here?

792 ↗

(On Diff #36417)

Might be helpful to note that memset is safe only because you don't allow ptrtoint.

[hans]

Right now GetUnderlyingObject does not look through PHI nodes, but if it starts to at some point, this change will fold %cmp [...]

That would be correct, though. Or are you saying there's a problem there?

lib/Transforms/InstCombine/InstCombineCompares.cpp

786 ↗

(On Diff #36417)

It might find a different cmp first though, and then bail out later when it finds ICI.

792 ↗

(On Diff #36417)

Thanks, I'll put a comment here.

[hans]

Add comment about the memory intrinsics.

[sanjoy]

In D13358#259819, @hans wrote:

Right now GetUnderlyingObject does not look through PHI nodes, but if it starts to at some point, this change will fold %cmp [...]

That would be correct, though. Or are you saying there's a problem there?

There's a potential problem here in justifying the transform -- in
the loop, the programmer is no longer doing a "one-off guess" anymore,
but systematically searching for the offset of a passed in argument
(which points to some slot in the caller's stack, say) from the
alloca. For instance, such a loop can allow you to implement
pointer subtraction without ptrtoint -- while (&ptra[diff++] != ptrb).

In other words, I understand the justification for the current change
to be that: an alloca can be allocated to any stack slot in the
current frame, so the compiler is allowed to fold any equality
comparisons to false. This seems fine. But a loop with an
incrementing induction variable can use an equality check to
"implement" an inequality check like ult or slt, and by folding
the != to true in while (&ptra[diff++] != ptrb), you're really
semantically folding ptra < ptrb to false. I personally think
this is not a problem, but I hope I was able to communicate what the
difference is.

lib/Transforms/InstCombine/InstCombineCompares.cpp

786 ↗

(On Diff #36524)

Ah, ok. (This is completely optional) I'd be tempted to get rid of NumCmps and write the check as

if (U != AllocaUse)
  return nullptr;

where AllocaUse is the Use * of the alloca in ICI which you either pass in or compute.

[hans]

There's a potential problem here in justifying the transform -- in
the loop, the programmer is no longer doing a "one-off guess" anymore,
but systematically searching for the offset of a passed in argument
(which points to some slot in the caller's stack, say) from the
alloca.

Thanks for clarifying. It's an interesting example :-)

I think we're still in the clear here. The way I see it, alloca has a lot of freedom in how it allocates memory, and we can act as if it allocated the memory in a place that the programmer isn't finding with their guesses.

[hans]

Sanjoy, David: thanks for your comments so far. OK to commit?

[hans]

Rebase. No change.

[hans]

Bail early if we hit an instruction with many uses. In particular, don't cause WorkList to get heap allocated by adding elements to it that we're not going to process anyway.

I can has lgtm?

[sanjoy]

FWIW, this LGTM, but given this is in territory I'm not wholly familiar with, I'll wait for @majnemer to take a look.

[majnemer]

LGTM with nits addressed.

lib/Transforms/InstCombine/InstCombineCompares.cpp

798 ↗

(On Diff #36685)

Can you make this >=, would make it a little easier to reason about.

lib/Analysis/InstructionSimplify.cpp

static Value *ExtractEquivalentCondition(Value *V, CmpInst::Predicate Pred,
   if (Pred == Cmp->getPredicate() && LHS == CmpLHS && RHS == CmpRHS)
     return Cmp;
   if (Pred == CmpInst::getSwappedPredicate(Cmp->getPredicate()) &&
       LHS == CmpRHS && RHS == CmpLHS)
     return Cmp;
   return nullptr;
 }
 
+static bool IsSafeForArgAllocaFolding(Value *V, const DataLayout &DL) {
+  unsigned NumCmps = 0;
+  SmallVector<Use *, 8> Worklist;
+  for (Use &U : V->uses())
+    Worklist.push_back(&U);
+
+  // To break any cycles and bound the search to constant-time.
+  unsigned MaxIter = 16;
+
+  while (!Worklist.empty()) {
+    if (!MaxIter--)
+      return false;
+
+    Use *U = Worklist.pop_back_val();
+    Instruction *I = cast<Instruction>(U->getUser());
+
+    if (auto *SI = dyn_cast<StoreInst>(I)) {
+      if (SI->getValueOperand() == U->get())
+        return false;
+      continue;
+    } else if (isa<LoadInst>(I)) {
+      continue;
+    } else if (isa<ICmpInst>(I)) {
+      // Bail out if we find more than one ICmp (the one we're currently
+      // investigating). This guarantees folding all or none of the cmps
+      // based on the pointer.
+      if (NumCmps++)
+        return false;
+      continue;
+    } else if (isa<BitCastInst>(I)) {
+      // Keep tracking.
+    } else if (auto *GEP = dyn_cast<GetElementPtrInst>(I)) {
+      Type *IntPtrTy =
+          DL.getIntPtrType(U->get()->getType())->getScalarType();
+      APInt Offset(IntPtrTy->getIntegerBitWidth(), 0);
+      if (!GEP->accumulateConstantOffset(DL, Offset))
+        return false;
+      // Keep tracking.
+    } else if (auto *Intrin = dyn_cast<IntrinsicInst>(I)) {
+      switch (Intrin->getIntrinsicID()) {
+        default:
+          return false;
+        case Intrinsic::lifetime_start:
+        case Intrinsic::lifetime_end:
+        case Intrinsic::dbg_declare:
+        case Intrinsic::dbg_value:
+        case Intrinsic::memcpy:
+        case Intrinsic::memmove:
+        case Intrinsic::memset:
+          continue;
+      }
+    } else {
+      return false;
+    }
+    for (Use &U : I->uses())
+      Worklist.push_back(&U);
+  }
+  return true;
+}
+
 // A significant optimization not implemented here is assuming that alloca
 // addresses are not equal to incoming argument values. They don't *alias*,
 // as we say, but that doesn't mean they aren't equal, so we take a
 // conservative approach.
 //
 // This is inspired in part by C++11 5.10p1:
 //   "Two pointers of the same type compare equal if and only if they are both
 //    null, both point to the same function, or both represent the same
auto IsAllocDisjoint = [](SmallVectorImpl<Value *> &Objects) {
                            return false;
                          });
     };
 
     if ((IsNAC(LHSUObjs) && IsAllocDisjoint(RHSUObjs)) ||
         (IsNAC(RHSUObjs) && IsAllocDisjoint(LHSUObjs)))
         return ConstantInt::get(GetCompareTy(LHS),
                                 !CmpInst::isTrueWhenEqual(Pred));
+
+    // We cannot fold comparisons between allocas and arguments in general (see
+    // comment above this function). However in some cases we can argue that
+    // it's impossible to guess the value of the alloca, and we can act as if
+    // the alloca and argument compare unequal. To do this, the alloca mustn't
+    // be observed, and *all* comparisons between pointers based on the alloca
+    // and argument must evaluate consistently, which means we should fold them
+    // all.
+    if ((isa<AllocaInst>(LHS) && isa<Argument>(RHS)) ||
+        (isa<AllocaInst>(RHS) && isa<Argument>(LHS))) {
+      AllocaInst *Alloca = cast<AllocaInst>(isa<AllocaInst>(LHS) ? LHS : RHS);
+      Argument *Arg = cast<Argument>(isa<Argument>(LHS) ? LHS : RHS);
+
+      if (IsSafeForArgAllocaFolding(Alloca, DL) &&
+          IsSafeForArgAllocaFolding(Arg, DL))
+        return ConstantInt::get(GetCompareTy(LHS),
+                                !CmpInst::isTrueWhenEqual(Pred));
+    }
   }
 
   // Otherwise, fail.
   return nullptr;
 }
 
 /// Return true if B is known to be implied by A.  A & B must be i1 (boolean)
 /// values. Note that the truth table for implication is the same as <=u on i1
 /// values (but not <=s!).  The truth table for both is:
 ///    | T | F (B)
 ///  T | T | F
 ///  F | T | T
 /// (A)
 static bool implies(Value *A, Value *B) {
   // TODO: Consider extending this to vector of i1?
   assert(A->getType()->isIntegerTy(1) && B->getType()->isIntegerTy(1));
   
   // A ==> A by definition
   if (A == B) return true;
 
   ICmpInst::Predicate APred, BPred;
   Value *I;
   Value *L;
   ConstantInt *CI;
   // i +_{nsw} C_{>0} <s L ==> i <s L
   if (match(A, m_ICmp(APred,
                       m_NSWAdd(m_Value(I), m_ConstantInt(CI)),
                       m_Value(L))) &&
       APred == ICmpInst::ICMP_SLT &&
       !CI->isNegative() &&
       match(B, m_ICmp(BPred, m_Specific(I), m_Specific(L))) &&
       BPred == ICmpInst::ICMP_SLT)
     return true;
 

test/Transforms/InstSimplify/compare.ll

define i1 @infinite_gep() {
   ret i1 1
 
 unreachableblock:
   %X = getelementptr i32, i32 *%X, i32 1
   %Y = icmp eq i32* %X, null
   ret i1 %Y
 }
 
-; It's not valid to fold a comparison of an argument with an alloca, even though
-; that's tempting. An argument can't *alias* an alloca, however the aliasing rule
-; relies on restrictions against guessing an object's address and dereferencing.
-; There are no restrictions against guessing an object's address and comparing.
+; It's not valid to fold a comparison of an argument with an alloca in the
+; general case, even though that's tempting. An argument can't *alias* an
+; alloca; however, the aliasing rule relies on restrictions against guessing
+; an object's address and dereferencing. There are no restrictions against
+; guessing an object's address and comparing.
+;
+; However, in cases where we can act as if the alloca yields an unguessable
+; address, we *can* fold comparisons against it. This works if the alloca
+; doesn't escape, as that would allow it to be observed outside the function.
+; Also, we must guarantee that if we fold one comparison of an argument and
+; alloca, we must fold *all* comparisons based on those pointers to maintain
+; consistency, so we limit ourselves to cases where there's only one cmp based
+; on these pointers.
 
 define i1 @alloca_argument_compare(i64* %arg) {
   %alloc = alloca i64
   %cmp = icmp eq i64* %arg, %alloc
   ret i1 %cmp
-  ; CHECK: alloca_argument_compare
-  ; CHECK: ret i1 %cmp
+  ; CHECK-LABEL: alloca_argument_compare
+  ; CHECK: ret i1 false
 }
 
-; As above, but with the operands reversed.
-
 define i1 @alloca_argument_compare_swapped(i64* %arg) {
   %alloc = alloca i64
   %cmp = icmp eq i64* %alloc, %arg
   ret i1 %cmp
-  ; CHECK: alloca_argument_compare_swapped
+  ; CHECK-LABEL: alloca_argument_compare_swapped
+  ; CHECK: ret i1 false
+}
+
+define i1 @alloca_argument_compare_ne(i64* %arg) {
+  %alloc = alloca i64
+  %cmp = icmp ne i64* %arg, %alloc
+  ret i1 %cmp
+  ; CHECK-LABEL: alloca_argument_compare_ne
+  ; CHECK: ret i1 true
+}
+
+define i1 @alloca_argument_compare_derived_ptrs(i64* %arg) {
+  %alloc = alloca i64, i64 8
+  %p = getelementptr i64, i64* %arg, i64 3
+  %q = getelementptr i64, i64* %alloc, i64 42
+  %cmp = icmp eq i64* %p, %q
+  ret i1 %cmp
+  ; CHECK-LABEL: alloca_argument_compare_derived_ptrs
+  ; CHECK: ret i1 false
+}
+
+declare void @escape(i64*)
+define i1 @alloca_argument_compare_escaped_alloca(i64* %arg) {
+  %alloc = alloca i64
+  call void @escape(i64* %alloc)
+  %cmp = icmp eq i64* %arg, %alloc
+  ret i1 %cmp
+  ; CHECK-LABEL: alloca_argument_compare_escaped_alloca
+  ; CHECK: %cmp = icmp eq i64* %arg, %alloc
+  ; CHECK: ret i1 %cmp
+}
+
+declare void @check_compares(i1, i1)
+define void @alloca_argument_compare_inconsistency(i64* %p, i64 %x) {
+  %q = alloca i64, i64 8
+  ; Instsimplify will not be able to track %r back to %p because the GEP has a
+  ; non-const offset. Since we're unable to fold the r == s comparison, we also
+  ; cannot fold p == q.
+  %r = getelementptr i64, i64* %p, i64 %x
+  %s = getelementptr i64, i64* %q, i64 2
+  %cmp1 = icmp eq i64* %p, %q
+  %cmp2 = icmp eq i64* %r, %s
+  call void @check_compares(i1 %cmp1, i1 %cmp2)
+  ret void
+  ; CHECK-LABEL: alloca_argument_compare_inconsistency
+  ; CHECK: call void @check_compares(i1 %cmp1, i1 %cmp2)
+}
+
+define void @alloca_argument_compare_two_compares(i64* %p) {
+  %q = alloca i64, i64 8
+  %r = getelementptr i64, i64* %p, i64 1
+  %s = getelementptr i64, i64* %q, i64 2
+  %cmp1 = icmp eq i64* %p, %q
+  %cmp2 = icmp eq i64* %r, %s
+  call void @check_compares(i1 %cmp1, i1 %cmp2)
+  ret void
+  ; Here we could fold r == s and therefore also p == q, but the optimization
+  ; is currently limited to only allow one cmp.
+  ; CHECK-LABEL: alloca_argument_compare_two_compares
+  ; CHECK: call void @check_compares(i1 %cmp1, i1 %cmp2)
+}
+
+define i1 @alloca_argument_compare_escaped_through_store(i64* %arg, i64** %ptr) {
+  %alloc = alloca i64
+  %cmp = icmp eq i64* %arg, %alloc
+  %p = getelementptr i64, i64* %alloc, i64 1
+  store i64* %p, i64** %ptr
+  ret i1 %cmp
+  ; CHECK-LABEL: alloca_argument_compare_escaped_through_store
+  ; CHECK: %cmp = icmp eq i64* %arg, %alloc
   ; CHECK: ret i1 %cmp
 }
 
+
+declare void @llvm.lifetime.start(i64, i8* nocapture)
+declare void @llvm.lifetime.end(i64, i8* nocapture)
+define i1 @alloca_argument_compare_benign_instrs(i8* %arg) {
+  %alloc = alloca i8
+  call void @llvm.lifetime.start(i64 1, i8* %alloc)
+  %cmp = icmp eq i8* %arg, %alloc
+  %x = load i8, i8* %arg
+  store i8 %x, i8* %alloc
+  call void @llvm.lifetime.end(i64 1, i8* %alloc)
+  ret i1 %cmp
+  ; CHECK-LABEL: alloca_argument_compare_benign_instrs
+  ; CHECK: ret i1 false
+}
+
+
 ; Don't assume that a noalias argument isn't equal to a global variable's
 ; address. This is an example where AliasAnalysis' NoAlias concept is
 ; different from actual pointer inequality.
 
 @y = external global i32
 define zeroext i1 @external_compare(i32* noalias %x) {
   %cmp = icmp eq i32* %x, @y
   ret i1 %cmp