[InstCombine] Support multiple comparisons in foldAllocaCmp()

[nikic]

Summary

foldAllocaCmp() needs to fold all comparisons of an alloca at the same time, to ensure that there is a consistent view of the alloca address. Currently, it folds "all" comparisons by limiting to the case where there is only one. This patch switches the algorithm to instead actually collect and fold all comparisons.

Something we need to be careful about here is that there may be comparisons where both sides of the icmp are based on the alloca. Such comparisons are comparing offsets of the alloca, and as such can be ignored here, but shouldn't be folded to false!

[nikic]

ping

[nikic]

ping :)

[nikic]

ping ;)

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

918

Instead of just deleting this, can you replace it with an equally useful comment about what the new logic is doing?

932

I don't understand this TODO, isn't getUnderlyingObject(*U) == Alloca doing just that?

934

Can you add TODO to canonicalize ule/uge/sle/sge -> ult/ugt/slt/sgt?

984

!CmpInst::isTrueWhenEqual(ICmp->getPredicate()) -> ICmp->getPredicate() == ICmpInst::ICMP_NE is clearer imo.

987

Is there a reason this isn't in InstructionSimplify?

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

990

You could add a todo here for some simplifications i.e

define i1 @foo(i64 %x) {
  %alloc = alloca i64, i64 8
  %p = getelementptr i64, ptr %alloc, i64 %x
  %q = getelementptr i64, ptr %alloc, i64 3
  %cmp = icmp eq ptr %p, %q
  ret i1 %cmp
}

is really icmp eq i64 %x, 3.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

932

getUnderlyingObject() does not look through phi nodes.

987

This can't be in InstSimplify because it is a correctness requirement that *all* comparisons of the alloca are replaced at the same time, something that InstSimplify cannot guarantee.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

990

I went ahead and implemented this in https://github.com/llvm/llvm-project/commit/e8ec42b80b5dc75186dca543572416f2f2e21475 and https://github.com/llvm/llvm-project/commit/289542b1e72dc4fe17093952dfb1b04cce259183.

[nikic]

Improve comments.

[nikic]

Upload correct patch.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

934

I don't think that would be legal, because it effectively retains a capture of the alloca, through which the address may be determined or constrained, leading to a contradiction with another fold.

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

938

Why does something like:

target datalayout = "p:32:32"
@gp = global ptr null, align 8
declare ptr @hidden_inttoptr()
declare void @witness(i1, i1)
define void @src() {
  %m = alloca i8, i32 4
  %lgp = load ptr, ptr @gp, align 8
  %rhs2 = call ptr @hidden_inttoptr()
  %cmp1 = icmp eq ptr %m, %rhs2
  %cmp2 = icmp eq ptr %lgp, %m
  call void @witness(i1 %cmp1, i1 %cmp2)
  ret void
}

Work? Isn't %m on both sides? (Although it is fine to reduce this snippet), just dont understand what U->getOperandNo() is actually checking for.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

938

What we care about is whether the alloca is on both sides of the same icmp. icmp based_on(%alloca), %other can be assumed to be false, icmp %other, based_on(%alloca) can be assumed to be false, but icmp based_on(%alloca), based_on(%alloca) may be true or false (but conversely, is not an escape, so we can ignore it entirely in this fold).

[goldstein.w.n]

LGTM.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

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) {
-  assert(ICI.isEquality() && "Cannot fold non-equality comparison.");
-
+bool InstCombinerImpl::foldAllocaCmp(AllocaInst *Alloca) {
   // 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.
   //
   // 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.
   //
-  // The code below checks that the alloca 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.
+  // However, we need to ensure that this folding is consistent: We can't fold
+  // one comparison to false, and then leave a different comparison against the
+  // same value alone (as it might evaluate to true at runtime, leading to a
+  // contradiction). As such, this code ensures that all comparisons are folded
+  // at the same time, and there are no other escapes.
 
   struct CmpCaptureTracker : public CaptureTracker {
+    AllocaInst *Alloca;
     bool Captured = false;
-    unsigned NumCmps = 0;
+    /// The value of the map is a bit mask of which icmp operands the alloca is
+    /// used in.
+    SmallMapVector<ICmpInst *, unsigned, 4> ICmps;
+
+    CmpCaptureTracker(AllocaInst *Alloca) : Alloca(Alloca) {}
 
     void tooManyUses() override { Captured = true; }
 
     bool captured(const Use *U) override {
-      if (isa<ICmpInst>(U->getUser()) && ++NumCmps == 1) {
-        // Ignore one icmp capture.
+      auto *ICmp = dyn_cast<ICmpInst>(U->getUser());
+      // We need to check that U is based *only* on the alloca, and doesn't
+      // have other contributions from a select/phi operand.
+      // TODO: We could check whether getUnderlyingObjects() reduces to one
+      // object, which would allow looking through phi nodes.
+      if (ICmp && ICmp->isEquality() && getUnderlyingObject(*U) == Alloca) {
+        // Collect equality icmps of the alloca, and don't treat them as
+        // captures.
+        auto Res = ICmps.insert({ICmp, 0});
+        Res.first->second |= 1u << U->getOperandNo();
         return false;
       }
 
       Captured = true;
       return true;
     }
   };
 
-  CmpCaptureTracker Tracker;
+  CmpCaptureTracker Tracker(Alloca);
   PointerMayBeCaptured(Alloca, &Tracker);
   if (Tracker.Captured)
-    return nullptr;
+    return false;
 
-  auto *Res = ConstantInt::get(ICI.getType(),
-                               !CmpInst::isTrueWhenEqual(ICI.getPredicate()));
-  return replaceInstUsesWith(ICI, Res);
+  bool Changed = false;
+  for (auto [ICmp, Operands] : Tracker.ICmps) {
+    switch (Operands) {
+    case 1:
+    case 2: {
+      // The alloca is only used in one icmp operand. Assume that the
+      // equality is false.
+      auto *Res = ConstantInt::get(
+          ICmp->getType(), ICmp->getPredicate() == ICmpInst::ICMP_NE);
+      replaceInstUsesWith(*ICmp, Res);
+      eraseInstFromFunction(*ICmp);
+      Changed = true;
+      break;
+    }
+    case 3:
+      // Both icmp operands are based on the alloca, so this is comparing
+      // pointer offsets, without leaking any information about the address
+      // of the alloca. Ignore such comparisons.
+      break;
+    default:
+      llvm_unreachable("Cannot happen");
+    }
+  }
+
+  return Changed;
 }
 
 /// Fold "icmp pred (X+C), X".
 Instruction *InstCombinerImpl::foldICmpAddOpConst(Value *X, const APInt &C,
                                                   ICmpInst::Predicate Pred) {
   // From this point on, we know that (X+C <= X) --> (X+C < X) because C != 0,
   // so the values can never be equal.  Similarly for all other "or equals"
   // operators.
   assert(!!C && "C should not be zero!");
 
   // (X+1) <u X        --> X >u (MAXUINT-1)        --> X == 255
   // (X+2) <u X        --> X >u (MAXUINT-2)        --> X > 253
   // (X+MAXUINT) <u X  --> X >u (MAXUINT-MAXUINT)  --> X != 0
   if (Pred == ICmpInst::ICMP_ULT || Pred == ICmpInst::ICMP_ULE) {
     Constant *R = ConstantInt::get(X->getType(),
                                    APInt::getMaxValue(C.getBitWidth()) - C);
     return new ICmpInst(ICmpInst::ICMP_UGT, X, R);
   }
 
   // (X+1) >u X        --> X <u (0-1)        --> X != 255
   // (X+2) >u X        --> X <u (0-2)        --> X <u 254
   // (X+MAXUINT) >u X  --> X <u (0-MAXUINT)  --> X <u 1  --> X == 0
Instruction *InstCombinerImpl::visitICmpInst(ICmpInst &I) {
   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.
   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))
-        return New;
+      if (foldAllocaCmp(Alloca))
+        return nullptr;
     if (auto *Alloca = dyn_cast<AllocaInst>(getUnderlyingObject(Op1)))
-      if (Instruction *New = foldAllocaCmp(I, Alloca))
-        return New;
+      if (foldAllocaCmp(Alloca))
+        return nullptr;
   }
 
   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:
   /// 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);
+  bool foldAllocaCmp(AllocaInst *Alloca);
   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

;
   call void @escape(ptr %alloc)
   %cmp = icmp eq ptr %alloc, %arg
   ret i1 %cmp
 }
 
 declare void @check_compares(i1, i1)
 define void @alloca_argument_compare_two_compares(ptr %p) {
 ; CHECK-LABEL: @alloca_argument_compare_two_compares(
-; CHECK-NEXT:    [[Q1:%.*]] = alloca [8 x i64], align 8
-; CHECK-NEXT:    [[R:%.*]] = getelementptr i64, ptr [[P:%.*]], i32 1
-; CHECK-NEXT:    [[S:%.*]] = getelementptr inbounds i64, ptr [[Q1]], i32 2
-; CHECK-NEXT:    [[CMP1:%.*]] = icmp eq ptr [[Q1]], [[P]]
-; CHECK-NEXT:    [[CMP2:%.*]] = icmp eq ptr [[R]], [[S]]
-; CHECK-NEXT:    call void @check_compares(i1 [[CMP1]], i1 [[CMP2]])
+; CHECK-NEXT:    call void @check_compares(i1 false, i1 false)
 ; CHECK-NEXT:    ret void
 ;
   %q = alloca i64, i64 8
   %r = getelementptr i64, ptr %p, i64 1
   %s = getelementptr i64, ptr %q, i64 2
   %cmp1 = icmp eq ptr %p, %q
   %cmp2 = icmp eq ptr %r, %s
   call void @check_compares(i1 %cmp1, i1 %cmp2)
;
   %n = alloca i8, i32 4
   %rhs = getelementptr i8, ptr %n, i32 4
   %cmp = icmp eq ptr %m, %rhs
   ret i1 %cmp
 }
 
 declare void @witness(i1, i1)
 
-define void @neg_consistent_fold1() {
-; CHECK-LABEL: @neg_consistent_fold1(
-; CHECK-NEXT:    [[M1:%.*]] = alloca [4 x i8], align 1
+define void @consistent_fold1() {
+; CHECK-LABEL: @consistent_fold1(
 ; CHECK-NEXT:    [[RHS2:%.*]] = call ptr @hidden_inttoptr()
-; CHECK-NEXT:    [[CMP1:%.*]] = icmp eq ptr [[M1]], inttoptr (i64 2048 to ptr)
-; CHECK-NEXT:    [[CMP2:%.*]] = icmp eq ptr [[M1]], [[RHS2]]
-; CHECK-NEXT:    call void @witness(i1 [[CMP1]], i1 [[CMP2]])
+; CHECK-NEXT:    call void @witness(i1 false, i1 false)
 ; CHECK-NEXT:    ret void
 ;
   %m = alloca i8, i32 4
   %rhs = inttoptr i64 2048 to ptr
   %rhs2 = call ptr @hidden_inttoptr()
   %cmp1 = icmp eq ptr %m, %rhs
   %cmp2 = icmp eq ptr %m, %rhs2
   call void @witness(i1 %cmp1, i1 %cmp2)
   ret void
 }
 
-define void @neg_consistent_fold2() {
-; CHECK-LABEL: @neg_consistent_fold2(
-; CHECK-NEXT:    [[M1:%.*]] = alloca [4 x i8], align 1
+define void @consistent_fold2() {
+; CHECK-LABEL: @consistent_fold2(
 ; CHECK-NEXT:    [[N2:%.*]] = alloca [4 x i8], align 1
-; CHECK-NEXT:    [[RHS:%.*]] = getelementptr inbounds i8, ptr [[N2]], i32 4
 ; CHECK-NEXT:    [[RHS2:%.*]] = call ptr @hidden_offset(ptr nonnull [[N2]])
-; CHECK-NEXT:    [[CMP1:%.*]] = icmp eq ptr [[M1]], [[RHS]]
-; CHECK-NEXT:    [[CMP2:%.*]] = icmp eq ptr [[M1]], [[RHS2]]
-; CHECK-NEXT:    call void @witness(i1 [[CMP1]], i1 [[CMP2]])
+; CHECK-NEXT:    call void @witness(i1 false, i1 false)
 ; CHECK-NEXT:    ret void
 ;
   %m = alloca i8, i32 4
   %n = alloca i8, i32 4
   %rhs = getelementptr i8, ptr %n, i32 4
   %rhs2 = call ptr @hidden_offset(ptr %n)
   %cmp1 = icmp eq ptr %m, %rhs
   %cmp2 = icmp eq ptr %m, %rhs2
   call void @witness(i1 %cmp1, i1 %cmp2)
   ret void
 }
 
-define void @neg_consistent_fold3() {
-; CHECK-LABEL: @neg_consistent_fold3(
-; CHECK-NEXT:    [[M1:%.*]] = alloca [4 x i8], align 1
-; CHECK-NEXT:    [[LGP:%.*]] = load ptr, ptr @gp, align 8
+define void @consistent_fold3() {
+; CHECK-LABEL: @consistent_fold3(
 ; CHECK-NEXT:    [[RHS2:%.*]] = call ptr @hidden_inttoptr()
-; CHECK-NEXT:    [[CMP1:%.*]] = icmp eq ptr [[M1]], [[LGP]]
-; CHECK-NEXT:    [[CMP2:%.*]] = icmp eq ptr [[M1]], [[RHS2]]
-; CHECK-NEXT:    call void @witness(i1 [[CMP1]], i1 [[CMP2]])
+; CHECK-NEXT:    call void @witness(i1 false, i1 false)
 ; CHECK-NEXT:    ret void
 ;
   %m = alloca i8, i32 4
   %lgp = load ptr, ptr @gp, align 8
   %rhs2 = call ptr @hidden_inttoptr()
   %cmp1 = icmp eq ptr %m, %lgp
   %cmp2 = icmp eq ptr %m, %rhs2
   call void @witness(i1 %cmp1, i1 %cmp2)