[InstCombine] Fold icmp(constants[x]) when the range of x is given

[XChy]

Summary

This patch extends foldCmpLoadFromIndexedGlobal to fold IR below:

define i1 @cmp_load_constant_array0(i64 %x){
entry:
  %cond = icmp ult i64 %x, 2
  br i1 %cond, label %case1, label %case2

case2:
  ret i1 0

case1:
  %isOK_ptr = getelementptr inbounds i32, ptr @CG, i64 %x
  %isOK = load i32, ptr %isOK_ptr
  %cond_inferred = icmp ult i32 %isOK, 3
  ret i1 %cond_inferred
}

Problem still:
Not globally optimized best.
Here I think if D152838 is accepted, that would be better.

Proof:
alive2

Related issue:
issue

[nikic]

This should be handled by the foldCmpLoadFromIndexedGlobal() fold, by first converting the icmp(load(gep(g, idx))) into icmp(idx) and then folding that icmp against the dominating condition.

That fold currently requires a very specific GEP source element type, but should be easy to generalize.

[XChy]

In D156657#4546067, @nikic wrote:

This should be handled by the foldCmpLoadFromIndexedGlobal() fold, by first converting the icmp(load(gep(g, idx))) into icmp(idx) and then folding that icmp against the dominating condition.

That fold currently requires a very specific GEP source element type, but should be easy to generalize.

Sure, I'll have a check.

[XChy]

Extends foldCmpLoadFromIndexedGlobal to fold.

[XChy]

Update tests.

[nikic]

It would be better to base this on collectOffset() and work in offset representation. Especially if you're interested in handling Rust, the type of the global and the GEP / loads will typically not line up (constant globals in Rust are usually i8 arrays, regardless of "real" type.)

[XChy]

In D156657#4547304, @nikic wrote:

It would be better to base this on collectOffset() and work in offset representation. Especially if you're interested in handling Rust, the type of the global and the GEP / loads will typically not line up (constant globals in Rust are usually i8 arrays, regardless of "real" type.)

Yes, such generlization in units of byte is necessary. It seems to be related to ptradd RFC?
Currently, this fold is still in units of array element . I'll try to refactor and generlize it.

llvm/test/Transforms/InstCombine/load-cmp.ll

351

Optimized well locally, but not globally.

[XChy]

Complete basically correct fold.

TODOs:

• Multi-variable testcases, Bitwidth-changed testcases
• Comments
• Performance/optimization to improve. When index is ax + by + .... + C, the longest distance we can step by is gcd(a,b,c,d, ....), not just one. Based on this we fold more precisely.

[nikic]

I think to start with, you should limit this to just the case where there is exactly one variable index (i.e. C*x or C*x+C2). Handling complex expressions is unlikely to be useful.

[XChy]

Already limit it to one variable. But I don't know why it is unlikely to be useful. From my perspective For multi-dimensional arrays, it's possible to have more than one index variable.

Besides, there exists miscompilation, I found just now. Current state-machine doesn't emit the CmpInst/Poison to guarrantee the index/offset is always not negative. Lack of such range-check would lead to something triggering UB not to doing so. I need to check it later.

[XChy]

In D156657#4550996, @XChy wrote:

Already limit it to one variable. But I don't know why it is unlikely to be useful. From my perspective For multi-dimensional arrays, it's possible to have more than one index variable.

Besides, there exists miscompilation, I found just now. Current state-machine doesn't emit the CmpInst/Poison to guarrantee the index/offset is always not negative. Lack of such range-check would lead to something triggering UB not to doing so. I need to check it later.

After trying, I don't think there is any way to do boundary-check (let things out of boundary lead to load poison) without modifying CFG. But InstCombine doesn't modify CFG. I haven't better idea here.

[XChy]

Update tests and FIXME: Boundary-check

[nikic]

Besides, there exists miscompilation, I found just now. Current state-machine doesn't emit the CmpInst/Poison to guarrantee the index/offset is always not negative. Lack of such range-check would lead to something triggering UB not to doing so. I need to check it later.

It's okay for transforms to remove UB (this is called "refinement" and done by many transforms). You just can't introduce new UB / poison that did not exist before.

[XChy]

Eliminate boundary-check.

[XChy]

Simplify LazyGetIndex function, and ready to be reviewed.

[XChy]

ping.

[goldstein.w.n]

In D156657#4639543, @XChy wrote:

ping.

nikic is on vacation for the next week, since he requested changes plan to let him get to it when he comes back.
Might have time to review later in the week, however.

[XChy]

In D156657#4639867, @goldstein.w.n wrote:

In D156657#4639543, @XChy wrote:

ping.

nikic is on vacation for the next week, since he requested changes plan to let him get to it when he comes back.
Might have time to review later in the week, however.

Thanks. That's OK. I hope this patch could be finished before the Phabricator shutdowns.

[XChy]

Polish comments.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

 /// we can optimize "icmp eq (load (gep "foo", 0, i)), 0" into "icmp eq i, 3".
 ///
 /// If AndCst is non-null, then the loaded value is masked with that constant
 /// before doing the comparison. This handles cases like "A[i]&4 == 0".
 Instruction *InstCombinerImpl::foldCmpLoadFromIndexedGlobal(
     LoadInst *LI, GetElementPtrInst *GEP, GlobalVariable *GV, CmpInst &ICI,
     ConstantInt *AndCst) {
   if (LI->isVolatile() || LI->getType() != GEP->getResultElementType() ||
-      GV->getValueType() != GEP->getSourceElementType() ||
       !GV->isConstant() || !GV->hasDefinitiveInitializer())
     return nullptr;
 
   Constant *Init = GV->getInitializer();
   if (!isa<ConstantArray>(Init) && !isa<ConstantDataArray>(Init))
     return nullptr;
 
   uint64_t ArrayElementCount = Init->getType()->getArrayNumElements();
   // Don't blow up on huge arrays.
   if (ArrayElementCount > MaxArraySizeForCombine)
     return nullptr;
 
+  uint64_t IndexPosInOperands;
   // There are many forms of this optimization we can handle, for now, just do
   // the simple index into a single-dimensional array.
   //
-  // Require: GEP GV, 0, i {{, constant indices}}
-  if (GEP->getNumOperands() < 3 ||
-      !isa<ConstantInt>(GEP->getOperand(1)) ||
-      !cast<ConstantInt>(GEP->getOperand(1))->isZero() ||
-      isa<Constant>(GEP->getOperand(2)))
+  // Require: GEP ArrayTy GV, 0, i {{, constant indices}}
+  // Or: GEP ElementTy GV, i {{. constant indices}}
+  if (GV->getValueType() &&
+      GV->getValueType()->getArrayElementType() ==
+          GEP->getSourceElementType() &&
+      GEP->getNumOperands() >= 2) {
+    IndexPosInOperands = 1;
+  }
+
+  else if (GV->getValueType() == GEP->getSourceElementType() &&
+           GEP->getNumOperands() >= 3 && isa<ConstantInt>(GEP->getOperand(1)) &&
+           cast<ConstantInt>(GEP->getOperand(1))->isZero() &&
+           !isa<Constant>(GEP->getOperand(2))) {
+    IndexPosInOperands = 2;
+  } else {
     return nullptr;
+  }
 
   // Check that indices after the variable are constants and in-range for the
   // type they index.  Collect the indices.  This is typically for arrays of
   // structs.
   SmallVector<unsigned, 4> LaterIndices;
-
   Type *EltTy = Init->getType()->getArrayElementType();
-  for (unsigned i = 3, e = GEP->getNumOperands(); i != e; ++i) {
+
+  for (unsigned i = IndexPosInOperands + 1, e = GEP->getNumOperands(); i != e;
+       ++i) {
     ConstantInt *Idx = dyn_cast<ConstantInt>(GEP->getOperand(i));
     if (!Idx) return nullptr;  // Variable index.
 
     uint64_t IdxVal = Idx->getZExtValue();
     if ((unsigned)IdxVal != IdxVal) return nullptr; // Too large array index.
 
     if (StructType *STy = dyn_cast<StructType>(EltTy))
       EltTy = STy->getElementType(IdxVal);
for (unsigned i = 0, e = ArrayElementCount; i != e; ++i) {
     if ((i & 8) == 0 && i >= 64 && SecondTrueElement == Overdefined &&
         SecondFalseElement == Overdefined && TrueRangeEnd == Overdefined &&
         FalseRangeEnd == Overdefined)
       return nullptr;
   }
 
   // Now that we've scanned the entire array, emit our new comparison(s).  We
   // order the state machines in complexity of the generated code.
-  Value *Idx = GEP->getOperand(2);
+  Value *Idx = GEP->getOperand(IndexPosInOperands);
 
   // If the index is larger than the pointer offset size of the target, truncate
   // the index down like the GEP would do implicitly.  We don't have to do this
   // for an inbounds GEP because the index can't be out of range.
   if (!GEP->isInBounds()) {
     Type *PtrIdxTy = DL.getIndexType(GEP->getType());
     unsigned OffsetSize = PtrIdxTy->getIntegerBitWidth();
     if (Idx->getType()->getPrimitiveSizeInBits().getFixedValue() > OffsetSize)

llvm/test/Transforms/InstCombine/load-cmp.ll

;
   %q = load i32, ptr %p
   %r = icmp eq i32 %q, 9
   ret i1 %r
 }
 
 
 @CG = constant [4 x i32] [i32 1, i32 2, i32 3, i32 4]
 
+; TODO: Fold it globally.
 define i1 @cmp_load_constant_array0(i64 %x){
 ; CHECK-LABEL: @cmp_load_constant_array0(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[COND:%.*]] = icmp ult i64 [[X:%.*]], 2
 ; CHECK-NEXT:    br i1 [[COND]], label [[CASE1:%.*]], label [[CASE2:%.*]]
 ; CHECK:       case2:
 ; CHECK-NEXT:    ret i1 false
 ; CHECK:       case1:
-; CHECK-NEXT:    [[TMP0:%.*]] = trunc i64 [[X]] to i32
-; CHECK-NEXT:    [[ISOK_PTR:%.*]] = getelementptr inbounds i32, ptr @CG, i32 [[TMP0]]
-; CHECK-NEXT:    [[ISOK:%.*]] = load i32, ptr [[ISOK_PTR]], align 4
-; CHECK-NEXT:    [[COND_INFERRED:%.*]] = icmp ult i32 [[ISOK]], 3
+; CHECK-NEXT:    [[TMP0:%.*]] = and i64 [[X]], 4294967294
+; CHECK-NEXT:    [[COND_INFERRED:%.*]] = icmp eq i64 [[TMP0]], 0
 ; CHECK-NEXT:    ret i1 [[COND_INFERRED]]
 ;
 entry:
   %cond = icmp ult i64 %x, 2
   br i1 %cond, label %case1, label %case2
 
 case2:
   ret i1 0
 
 case1:
   %isOK_ptr = getelementptr inbounds i32, ptr @CG, i64 %x
   %isOK = load i32, ptr %isOK_ptr
   %cond_inferred = icmp ult i32 %isOK, 3
   ret i1 %cond_inferred
 }
 
 define i1 @cmp_load_constant_array1(i64 %x){
 ; CHECK-LABEL: @cmp_load_constant_array1(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[COND:%.*]] = icmp ult i64 [[X:%.*]], 2
 ; CHECK-NEXT:    br i1 [[COND]], label [[CASE1:%.*]], label [[CASE2:%.*]]
 ; CHECK:       case2:
 ; CHECK-NEXT:    ret i1 false
 ; CHECK:       case1:
-; CHECK-NEXT:    [[TMP0:%.*]] = trunc i64 [[X]] to i32
-; CHECK-NEXT:    [[ISOK_PTR:%.*]] = getelementptr inbounds i32, ptr @CG, i32 [[TMP0]]
-; CHECK-NEXT:    [[ISOK:%.*]] = load i32, ptr [[ISOK_PTR]], align 4
-; CHECK-NEXT:    [[COND_INFERRED:%.*]] = icmp ugt i32 [[ISOK]], 10
-; CHECK-NEXT:    ret i1 [[COND_INFERRED]]
+; CHECK-NEXT:    ret i1 false
 ;
 entry:
   %cond = icmp ult i64 %x, 2
   br i1 %cond, label %case1, label %case2
 
 case2:
   ret i1 0
 

 ; CHECK-LABEL: @cmp_load_constant_array_messy(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[COND:%.*]] = icmp slt i64 [[X:%.*]], 6
 ; CHECK-NEXT:    br i1 [[COND]], label [[CASE1:%.*]], label [[CASE2:%.*]]
 ; CHECK:       case2:
 ; CHECK-NEXT:    ret i1 false
 ; CHECK:       case1:
 ; CHECK-NEXT:    [[TMP0:%.*]] = trunc i64 [[X]] to i32
-; CHECK-NEXT:    [[ISOK_PTR:%.*]] = getelementptr i32, ptr @CG_MESSY, i32 [[TMP0]]
-; CHECK-NEXT:    [[ISOK:%.*]] = load i32, ptr [[ISOK_PTR]], align 4
-; CHECK-NEXT:    [[COND_INFERRED:%.*]] = icmp slt i32 [[ISOK]], 5
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[TMP0]], 1073741823
+; CHECK-NEXT:    [[TMP2:%.*]] = lshr i32 373, [[TMP1]]
+; CHECK-NEXT:    [[TMP3:%.*]] = and i32 [[TMP2]], 1
+; CHECK-NEXT:    [[COND_INFERRED:%.*]] = icmp ne i32 [[TMP3]], 0
 ; CHECK-NEXT:    ret i1 [[COND_INFERRED]]
 ;
 entry:
   %cond = icmp slt i64 %x, 6
   br i1 %cond, label %case1, label %case2
 
 case2:
   ret i1 0
case2:
   ret i1 0
 
 case1:
   %isOK_ptr = getelementptr inbounds i32, ptr @CG, i64 %x
   %isOK = load i32, ptr %isOK_ptr
   %cond_inferred = icmp ult i32 %isOK, %y
   ret i1 %cond_inferred
 }
-