Demanded elements support for vector GEPs

[reames]

Summary

GEPs can produce either scalar or vector results. If we're extracting only a subset of the vector lanes, simplifying the operands is helpful in eliminating redundant computation, and (eventually) allowing further optimizations.

This is the first concrete patch split out of https://reviews.llvm.org/D57140. If you're curious where this is going, that combined patch may be useful for context.

[reames]

ping

[spatel]

I haven't looked at vector geps very much, but the code change looks obviously ok, and the tests are thorough. LGTM.

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

auto simplifyAndSetOp = [&](Instruction *Inst, unsigned OpNum,
     }
   };
 
   APInt UndefElts2(VWidth, 0);
   APInt UndefElts3(VWidth, 0);
   switch (I->getOpcode()) {
   default: break;
 
+  case Instruction::GetElementPtr: {
+    // Conservatively track the demanded elements back through any vector
+    // operands we may have.  We know there must be at least one, or we
+    // wouldn't have a vector result to get here. Note that we intentionally
+    // merge the undef bits here since gepping with either an undef base or
+    // index results in undef. 
+    for (unsigned i = 0; i < I->getNumOperands(); i++)
+      if (I->getOperand(i)->getType()->isVectorTy())
+        simplifyAndSetOp(I, i, DemandedElts, UndefElts);
+
+    break;
+  }
   case Instruction::InsertElement: {
     // If this is a variable index, we don't know which element it overwrites.
     // demand exactly the same input as we produce.
     ConstantInt *Idx = dyn_cast<ConstantInt>(I->getOperand(2));
     if (!Idx) {
       // Note that we can't propagate undef elt info, because we don't know
       // which elt is getting updated.
       simplifyAndSetOp(I, 0, DemandedElts, UndefElts2);

test/Transforms/InstCombine/vec_demanded_elts.ll

 ;
   %gep = getelementptr i32, <2 x i32*> %base, i64 1
   %ee = extractelement <2 x i32*> %gep, i32 1
   ret i32* %ee
 }
 
 define i32* @gep_splat_base_w_s_idx(i32* %base) {
 ; CHECK-LABEL: @gep_splat_base_w_s_idx(
-; CHECK-NEXT:    [[BASEVEC1:%.*]] = insertelement <2 x i32*> undef, i32* [[BASE:%.*]], i32 0
-; CHECK-NEXT:    [[BASEVEC2:%.*]] = shufflevector <2 x i32*> [[BASEVEC1]], <2 x i32*> undef, <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[BASEVEC2:%.*]] = insertelement <2 x i32*> undef, i32* [[BASE:%.*]], i32 1
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i32, <2 x i32*> [[BASEVEC2]], i64 1
 ; CHECK-NEXT:    [[EE:%.*]] = extractelement <2 x i32*> [[GEP]], i32 1
 ; CHECK-NEXT:    ret i32* [[EE]]
 ;
   %basevec1 = insertelement <2 x i32*> undef, i32* %base, i32 0
   %basevec2 = shufflevector <2 x i32*> %basevec1, <2 x i32*> undef, <2 x i32> zeroinitializer
   %gep = getelementptr i32, <2 x i32*> %basevec2, i64 1
   %ee = extractelement <2 x i32*> %gep, i32 1
   ret i32* %ee
 }
 
 
 define i32* @gep_splat_base_w_cv_idx(i32* %base) {
 ; CHECK-LABEL: @gep_splat_base_w_cv_idx(
-; CHECK-NEXT:    [[BASEVEC1:%.*]] = insertelement <2 x i32*> undef, i32* [[BASE:%.*]], i32 0
-; CHECK-NEXT:    [[BASEVEC2:%.*]] = shufflevector <2 x i32*> [[BASEVEC1]], <2 x i32*> undef, <2 x i32> zeroinitializer
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i32, <2 x i32*> [[BASEVEC2]], <2 x i64> <i64 0, i64 1>
+; CHECK-NEXT:    [[BASEVEC2:%.*]] = insertelement <2 x i32*> undef, i32* [[BASE:%.*]], i32 1
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i32, <2 x i32*> [[BASEVEC2]], <2 x i64> <i64 undef, i64 1>
 ; CHECK-NEXT:    [[EE:%.*]] = extractelement <2 x i32*> [[GEP]], i32 1
 ; CHECK-NEXT:    ret i32* [[EE]]
 ;
   %basevec1 = insertelement <2 x i32*> undef, i32* %base, i32 0
   %basevec2 = shufflevector <2 x i32*> %basevec1, <2 x i32*> undef, <2 x i32> zeroinitializer
   %gep = getelementptr i32, <2 x i32*> %basevec2, <2 x i64> <i64 0, i64 1>
   %ee = extractelement <2 x i32*> %gep, i32 1
   ret i32* %ee
 }
 
 define i32* @gep_splat_base_w_vidx(i32* %base, <2 x i64> %idxvec) {
 ; CHECK-LABEL: @gep_splat_base_w_vidx(
-; CHECK-NEXT:    [[BASEVEC1:%.*]] = insertelement <2 x i32*> undef, i32* [[BASE:%.*]], i32 0
-; CHECK-NEXT:    [[BASEVEC2:%.*]] = shufflevector <2 x i32*> [[BASEVEC1]], <2 x i32*> undef, <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[BASEVEC2:%.*]] = insertelement <2 x i32*> undef, i32* [[BASE:%.*]], i32 1
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i32, <2 x i32*> [[BASEVEC2]], <2 x i64> [[IDXVEC:%.*]]
 ; CHECK-NEXT:    [[EE:%.*]] = extractelement <2 x i32*> [[GEP]], i32 1
 ; CHECK-NEXT:    ret i32* [[EE]]
 ;
   %basevec1 = insertelement <2 x i32*> undef, i32* %base, i32 0
   %basevec2 = shufflevector <2 x i32*> %basevec1, <2 x i32*> undef, <2 x i32> zeroinitializer
   %gep = getelementptr i32, <2 x i32*> %basevec2, <2 x i64> %idxvec
   %ee = extractelement <2 x i32*> %gep, i32 1
   ret i32* %ee
 }
 
 
 @GLOBAL = internal global i32 zeroinitializer
 
 define i32* @gep_cvbase_w_s_idx(<2 x i32*> %base, i64 %raw_addr) {
 ; CHECK-LABEL: @gep_cvbase_w_s_idx(
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i32, <2 x i32*> <i32* @GLOBAL, i32* @GLOBAL>, i64 [[RAW_ADDR:%.*]]
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i32, <2 x i32*> <i32* undef, i32* @GLOBAL>, i64 [[RAW_ADDR:%.*]]
 ; CHECK-NEXT:    [[EE:%.*]] = extractelement <2 x i32*> [[GEP]], i32 1
 ; CHECK-NEXT:    ret i32* [[EE]]
 ;
   %gep = getelementptr i32, <2 x i32*> <i32* @GLOBAL, i32* @GLOBAL>, i64 %raw_addr
   %ee = extractelement <2 x i32*> %gep, i32 1
   ret i32* %ee
 }
 
 define i32* @gep_cvbase_w_cv_idx(<2 x i32*> %base, i64 %raw_addr) {
 ; CHECK-LABEL: @gep_cvbase_w_cv_idx(
 ; CHECK-NEXT:    ret i32* extractelement (<2 x i32*> getelementptr (i32, <2 x i32*> <i32* @GLOBAL, i32* @GLOBAL>, <2 x i64> <i64 0, i64 1>), i32 1)
 ;
   %gep = getelementptr i32, <2 x i32*> <i32* @GLOBAL, i32* @GLOBAL>, <2 x i64> <i64 0, i64 1>
   %ee = extractelement <2 x i32*> %gep, i32 1
   ret i32* %ee
 }
 
 
 define i32* @gep_sbase_w_cv_idx(i32* %base) {
 ; CHECK-LABEL: @gep_sbase_w_cv_idx(
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i32, i32* [[BASE:%.*]], <2 x i64> <i64 0, i64 1>
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i32, i32* [[BASE:%.*]], <2 x i64> <i64 undef, i64 1>
 ; CHECK-NEXT:    [[EE:%.*]] = extractelement <2 x i32*> [[GEP]], i32 1
 ; CHECK-NEXT:    ret i32* [[EE]]
 ;
   %gep = getelementptr i32, i32* %base, <2 x i64> <i64 0, i64 1>
   %ee = extractelement <2 x i32*> %gep, i32 1
   ret i32* %ee
 }
 
 define i32* @gep_sbase_w_splat_idx(i32* %base, i64 %idx) {
 ; CHECK-LABEL: @gep_sbase_w_splat_idx(
-; CHECK-NEXT:    [[IDXVEC1:%.*]] = insertelement <2 x i64> undef, i64 [[IDX:%.*]], i32 0
-; CHECK-NEXT:    [[IDXVEC2:%.*]] = shufflevector <2 x i64> [[IDXVEC1]], <2 x i64> undef, <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[IDXVEC2:%.*]] = insertelement <2 x i64> undef, i64 [[IDX:%.*]], i32 1
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i32, i32* [[BASE:%.*]], <2 x i64> [[IDXVEC2]]
 ; CHECK-NEXT:    [[EE:%.*]] = extractelement <2 x i32*> [[GEP]], i32 1
 ; CHECK-NEXT:    ret i32* [[EE]]
 ;
   %idxvec1 = insertelement <2 x i64> undef, i64 %idx, i32 0
   %idxvec2 = shufflevector <2 x i64> %idxvec1, <2 x i64> undef, <2 x i32> zeroinitializer
   %gep = getelementptr i32, i32* %base, <2 x i64> %idxvec2
   %ee = extractelement <2 x i32*> %gep, i32 1
   ret i32* %ee
 }
 define i32* @gep_splat_both(i32* %base, i64 %idx) {
 ; CHECK-LABEL: @gep_splat_both(
-; CHECK-NEXT:    [[BASEVEC1:%.*]] = insertelement <2 x i32*> undef, i32* [[BASE:%.*]], i32 0
-; CHECK-NEXT:    [[BASEVEC2:%.*]] = shufflevector <2 x i32*> [[BASEVEC1]], <2 x i32*> undef, <2 x i32> zeroinitializer
-; CHECK-NEXT:    [[IDXVEC1:%.*]] = insertelement <2 x i64> undef, i64 [[IDX:%.*]], i32 0
-; CHECK-NEXT:    [[IDXVEC2:%.*]] = shufflevector <2 x i64> [[IDXVEC1]], <2 x i64> undef, <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[BASEVEC2:%.*]] = insertelement <2 x i32*> undef, i32* [[BASE:%.*]], i32 1
+; CHECK-NEXT:    [[IDXVEC2:%.*]] = insertelement <2 x i64> undef, i64 [[IDX:%.*]], i32 1
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i32, <2 x i32*> [[BASEVEC2]], <2 x i64> [[IDXVEC2]]
 ; CHECK-NEXT:    [[EE:%.*]] = extractelement <2 x i32*> [[GEP]], i32 1
 ; CHECK-NEXT:    ret i32* [[EE]]
 ;
   %basevec1 = insertelement <2 x i32*> undef, i32* %base, i32 0
   %basevec2 = shufflevector <2 x i32*> %basevec1, <2 x i32*> undef, <2 x i32> zeroinitializer
   %idxvec1 = insertelement <2 x i64> undef, i64 %idx, i32 0
   %idxvec2 = shufflevector <2 x i64> %idxvec1, <2 x i64> undef, <2 x i32> zeroinitializer
   %gep = getelementptr i32, <2 x i32*> %basevec2, <2 x i64> %idxvec2
   %ee = extractelement <2 x i32*> %gep, i32 1
   ret i32* %ee
 }