[InstCombine] icmp eq/ne (gep inbounds P, Idx..), null -> icmp eq/ne P, null for vectors

[reames]

Summary

Extend the transform introduced in https://reviews.llvm.org/D66608 to work for vector geps as well.

[grandinj]

test/Transforms/InstCombine/gep-inbounds-null.ll

78

TODO can be removed

[jdoerfert]

The TODO nit was noted, otherwise LGTM.

lib/Transforms/InstCombine/InstCombineCompares.cpp

if (PtrBase == RHS && GEPLHS->isInBounds()) {
     Value *Offset = evaluateGEPOffsetExpression(GEPLHS, *this, DL);
 
     // If not, synthesize the offset the hard way.
     if (!Offset)
       Offset = EmitGEPOffset(GEPLHS);
     return new ICmpInst(ICmpInst::getSignedPredicate(Cond), Offset,
                         Constant::getNullValue(Offset->getType()));
   } else if (GEPLHS->isInBounds() && ICmpInst::isEquality(Cond) &&
-             GEPLHS->getType()->isPointerTy() && // TODO: extend to vector geps
              isa<Constant>(RHS) && cast<Constant>(RHS)->isNullValue() &&
              !NullPointerIsDefined(I.getFunction(),
                                    RHS->getType()->getPointerAddressSpace())) {
     // For most address spaces, an allocation can't be placed at null, but null
     // itself is treated as a 0 size allocation in the in bounds rules.  Thus,
     // the only valid inbounds address derived from null, is null itself.
     // Thus, we have four cases to consider:
     // 1) Base == nullptr, Offset == 0 -> inbounds, null
     // 2) Base == nullptr, Offset != 0 -> poison as the result is out of bounds
     // 3) Base != nullptr, Offset == (-base) -> poison (crossing allocations)
     // 4) Base != nullptr, Offset != (-base) -> nonnull (and possibly poison)
     //
     // (Note if we're indexing a type of size 0, that simply collapses into one
     //  of the buckets above.)
     //
     // In general, we're allowed to make values less poison (i.e. remove
     //   sources of full UB), so in this case, we just select between the two
     //   non-poison cases (1 and 4 above).
+    //
+    // For vectors, we apply the same reasoning on a per-lane basis.
     auto *Base = GEPLHS->getPointerOperand();
+    if (GEPLHS->getType()->isVectorTy() && Base->getType()->isPointerTy()) {
+      int NumElts = GEPLHS->getType()->getVectorNumElements();
+      Base = Builder.CreateVectorSplat(NumElts, Base);
+    }
     return new ICmpInst(Cond, Base,
                         ConstantExpr::getBitCast(cast<Constant>(RHS),
                                                  Base->getType()));
   } else if (GEPOperator *GEPRHS = dyn_cast<GEPOperator>(RHS)) {
     // If the base pointers are different, but the indices are the same, just
     // compare the base pointer.
     if (PtrBase != GEPRHS->getOperand(0)) {
       bool IndicesTheSame = GEPLHS->getNumOperands()==GEPRHS->getNumOperands();

test/Transforms/InstCombine/gep-inbounds-null.ll

 ; CHECK-NEXT:    ret i1 [[CND]]
 ;
 entry:
   %gep = getelementptr inbounds i8, i8* %base, i64 %idx
   %cnd = icmp eq i8* %gep, null
   ret i1 %cnd
 }
 
 ;; TODO: vectors not yet handled
 define <2 x i1> @test_vector_base(<2 x i8*> %base, i64 %idx) {
 ; CHECK-LABEL: @test_vector_base(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i8, <2 x i8*> [[BASE:%.*]], i64 [[IDX:%.*]]
-; CHECK-NEXT:    [[CND:%.*]] = icmp eq <2 x i8*> [[GEP]], zeroinitializer
+; CHECK-NEXT:    [[CND:%.*]] = icmp eq <2 x i8*> [[BASE:%.*]], zeroinitializer
 ; CHECK-NEXT:    ret <2 x i1> [[CND]]
 ;
 entry:
   %gep = getelementptr inbounds i8, <2 x i8*> %base, i64 %idx
   %cnd = icmp eq <2 x i8*> %gep, zeroinitializer
   ret <2 x i1> %cnd
 }
 
 define <2 x i1> @test_vector_index(i8* %base, <2 x i64> %idx) {
 ; CHECK-LABEL: @test_vector_index(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i8, i8* [[BASE:%.*]], <2 x i64> [[IDX:%.*]]
-; CHECK-NEXT:    [[CND:%.*]] = icmp eq <2 x i8*> [[GEP]], zeroinitializer
+; CHECK-NEXT:    [[DOTSPLATINSERT:%.*]] = insertelement <2 x i8*> undef, i8* [[BASE:%.*]], i32 0
+; CHECK-NEXT:    [[DOTSPLAT:%.*]] = shufflevector <2 x i8*> [[DOTSPLATINSERT]], <2 x i8*> undef, <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[CND:%.*]] = icmp eq <2 x i8*> [[DOTSPLAT]], zeroinitializer
 ; CHECK-NEXT:    ret <2 x i1> [[CND]]
 ;
 entry:
   %gep = getelementptr inbounds i8, i8* %base, <2 x i64> %idx
   %cnd = icmp eq <2 x i8*> %gep, zeroinitializer
   ret <2 x i1> %cnd
 }
 
 define <2 x i1> @test_vector_both(<2 x i8*> %base, <2 x i64> %idx) {
 ; CHECK-LABEL: @test_vector_both(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i8, <2 x i8*> [[BASE:%.*]], <2 x i64> [[IDX:%.*]]
-; CHECK-NEXT:    [[CND:%.*]] = icmp eq <2 x i8*> [[GEP]], zeroinitializer
+; CHECK-NEXT:    [[CND:%.*]] = icmp eq <2 x i8*> [[BASE:%.*]], zeroinitializer
 ; CHECK-NEXT:    ret <2 x i1> [[CND]]
 ;
 entry:
   %gep = getelementptr inbounds i8, <2 x i8*> %base, <2 x i64> %idx
   %cnd = icmp eq <2 x i8*> %gep, zeroinitializer
   ret <2 x i1> %cnd
 }