[InstCombine] Fix incongruous GEP type addrspace

[EwanCrawford]

Summary

Currently running the @insertelem_after_gep function below through the InstCombine pass with opt produces invalid IR.

Input:

define void @insertelem_after_gep(<16 x i32>* %t0) {
   %t1 = bitcast <16 x i32>* %t0 to [16 x i32]*
   %t2 = addrspacecast [16 x i32]* %t1 to [16 x i32] addrspace(3)*
   %t3 = getelementptr inbounds [16 x i32], [16 x i32] addrspace(3)* %t2, i64 0, i64 0
   %t4 = insertelement <16 x i32 addrspace(3)*> undef, i32 addrspace(3)* %t3, i32 0
   call void @extern_vec_pointers_func(<16 x i32 addrspace(3)*> %t4)
   ret void
}

Output:

define void @insertelem_after_gep(<16 x i32>* %t0) {
  %t3 = getelementptr inbounds <16 x i32>, <16 x i32>* %t0, i64 0, i64 0
  %t4 = insertelement <16 x i32 addrspace(3)*> undef, i32 addrspace(3)* %t3, i32 0
  call void @my_extern_func(<16 x i32 addrspace(3)*> %t4)
  ret void
}

Which although causes no complaints when produced, isn't valid IR as the insertelement use of the %t3 GEP expects an address space.

opt: /tmp/bad.ll:52:73: error: '%t3' defined with type 'i32*' but expected 'i32 addrspace(3)*'
  %t4 = insertelement <16 x i32 addrspace(3)*> undef, i32 addrspace(3)* %t3, i32 0

I've fixed this by adding an addrspacecast after the GEP in the InstCombine pass, and including a check for this type mismatch to the verifier.

[lebedev.ri]

lib/Transforms/InstCombine/InstructionCombining.cpp

2056–2058

I do not understand, why do we need to create a new GEP here, as compared to the old code?
Why can't we just only create AddrSpaceCastInst if needed?

[EwanCrawford]

lib/Transforms/InstCombine/InstructionCombining.cpp

2056–2058

Thanks for taking a look @lebedev.ri!

Doing setOperand then setSourceElementType in the old code doesn't actually change the type of the GEP Value for this test case. I get i32 addrspace(3)* for GEP.getType() both before and after these operation, despite the fact that the GEP now points to a i32* afterwards.

Therefore if we call new addrSpaceCastInst(GEP, GEPType) we hit an assert since the types of GEP and GEPType are identical. However if we create a new GEP, then it is correctly instantiated with the i32* type and the addrSpaceCastInst can be created successfully.

That was my thinking anyway.

[EwanCrawford]

ping

[spatel]

Looking at the existing code in this file, this seems like the way to handle it. Reviewing D44833 (where I added this transform), can we minimize the test to something like:

define i32 addrspace(3)* @bitcast_vec_to_array_gep(<7 x i32>* %x, i64 %y, i64 %z) {
  %arr_ptr = bitcast <7 x i32>* %x to [7 x i32]*
  %asc = addrspacecast [7 x i32]* %arr_ptr to [7 x i32] addrspace(3)*
  %gep = getelementptr [7 x i32], [7 x i32] addrspace(3)* %asc, i64 %y, i64 %z
  ret i32 addrspace(3)* %gep
}

There should be an 'inbounds' test too to cover the other code path. Given that this doesn't trigger the IR verifier (should we fix that too?), you could add the tests showing the bogus IR as a preliminary commit.

[EwanCrawford]

Thanks for the feedback @spatel.
I've modified the test according to your suggestions and added a check to the verifier which fires without the instcombine change.

[spatel]

LGTM, but let's see if @lebedev.ri has any other feedback before marking as approved.

lib/Transforms/InstCombine/InstructionCombining.cpp

2059

Nit: since we know there are only 3 operands here, it would be less ambiguous to use something like { Op[1], Op[2] } here?

[lebedev.ri]

In D52294#1255824, @spatel wrote:

LGTM, but let's see if @lebedev.ri has any other feedback before marking as approved.

No, no real feedback, that is why i'm not commenting.

lib/IR/Verifier.cpp

3142

if (auto *PTy = dyn_cast<PointerType>(GEP.getType())) {

lib/Transforms/InstCombine/InstructionCombining.cpp

2055–2070

nit: Maybe add a comment explaining why we are creating a new GEP + AddrSpaceCast?

[EwanCrawford]

done, thanks again

[spatel]

LGTM - thanks!

lib/IR/Verifier.cpp

for (Value *Idx : Idxs) {
       if (IndexTy->isVectorTy()) {
         unsigned IndexWidth = IndexTy->getVectorNumElements();
         Assert(IndexWidth == GEPWidth, "Invalid GEP index vector width", &GEP);
       }
       Assert(IndexTy->isIntOrIntVectorTy(),
              "All GEP indices should be of integer type");
     }
   }
+
+  if (auto *PTy = dyn_cast<PointerType>(GEP.getType())) {
+    Assert(GEP.getAddressSpace() == PTy->getAddressSpace(),
+           "GEP address space doesn't match type", &GEP);
+  }
+
   visitInstruction(GEP);
 }
 
 static bool isContiguous(const ConstantRange &A, const ConstantRange &B) {
   return A.getUpper() == B.getLower() || A.getLower() == B.getUpper();
 }
 
 void Verifier::visitRangeMetadata(Instruction &I, MDNode *Range, Type *Ty) {

lib/Transforms/InstCombine/InstructionCombining.cpp

auto areMatchingArrayAndVecTypes = [](Type *ArrTy, Type *VecTy) {
       return ArrTy->getArrayElementType() == VecTy->getVectorElementType() &&
              ArrTy->getArrayNumElements() == VecTy->getVectorNumElements();
     };
     if (GEP.getNumOperands() == 3 &&
         ((GEPEltType->isArrayTy() && SrcEltType->isVectorTy() &&
           areMatchingArrayAndVecTypes(GEPEltType, SrcEltType)) ||
          (GEPEltType->isVectorTy() && SrcEltType->isArrayTy() &&
           areMatchingArrayAndVecTypes(SrcEltType, GEPEltType)))) {
-      GEP.setOperand(0, SrcOp);
-      GEP.setSourceElementType(SrcEltType);
-      return &GEP;
+
+      // Create a new GEP here, as using `setOperand()` followed by
+      // `setSourceElementType()` won't actually update the type of the
+      // existing GEP Value. Causing issues if this Value is accessed when
+      // constructing an AddrSpaceCastInst
+      Value *NGEP =
+          GEP.isInBounds()
+              ? Builder.CreateInBoundsGEP(nullptr, SrcOp, {Ops[1], Ops[2]})
+              : Builder.CreateGEP(nullptr, SrcOp, {Ops[1], Ops[2]});
+      NGEP->takeName(&GEP);
+
+      // Preserve GEP address space to satisfy users
+      if (NGEP->getType()->getPointerAddressSpace() != GEP.getAddressSpace())
+        return new AddrSpaceCastInst(NGEP, GEPType);
+
+      return replaceInstUsesWith(GEP, NGEP);
     }
 
     // See if we can simplify:
     //   X = bitcast A* to B*
     //   Y = gep X, <...constant indices...>
     // into a gep of the original struct. This is important for SROA and alias
     // analysis of unions. If "A" is also a bitcast, wait for A/X to be merged.
     unsigned OffsetBits = DL.getIndexTypeSizeInBits(GEPType);

test/Transforms/InstCombine/gep-vector.ll

 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds [3 x i32], [3 x i32]* [[X:%.*]], i64 [[Y:%.*]], i64 [[Z:%.*]]
 ; CHECK-NEXT:    ret i32* [[GEP]]
 ;
   %vec_ptr = bitcast [3 x i32]* %x to <3 x i32>*
   %gep = getelementptr inbounds <3 x i32>, <3 x i32>* %vec_ptr, i64 %y, i64 %z
   ret i32* %gep
 }
 
+define i32 addrspace(3)* @bitcast_vec_to_array_addrspace(<7 x i32>* %x, i64 %y, i64 %z) {
+; CHECK-LABEL: @bitcast_vec_to_array_addrspace(
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr <7 x i32>, <7 x i32>* [[X:%.*]], i64 [[Y:%.*]], i64 [[Z:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast i32* [[GEP]] to i32 addrspace(3)*
+; CHECK-NEXT:    ret i32 addrspace(3)* [[TMP1]]
+;
+  %arr_ptr = bitcast <7 x i32>* %x to [7 x i32]*
+  %asc = addrspacecast [7 x i32]* %arr_ptr to [7 x i32] addrspace(3)*
+  %gep = getelementptr [7 x i32], [7 x i32] addrspace(3)* %asc, i64 %y, i64 %z
+  ret i32 addrspace(3)* %gep
+}
+
+define i32 addrspace(3)* @inbounds_bitcast_vec_to_array_addrspace(<7 x i32>* %x, i64 %y, i64 %z) {
+; CHECK-LABEL: @inbounds_bitcast_vec_to_array_addrspace(
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds <7 x i32>, <7 x i32>* [[X:%.*]], i64 [[Y:%.*]], i64 [[Z:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast i32* [[GEP]] to i32 addrspace(3)*
+; CHECK-NEXT:    ret i32 addrspace(3)* [[TMP1]]
+;
+  %arr_ptr = bitcast <7 x i32>* %x to [7 x i32]*
+  %asc = addrspacecast [7 x i32]* %arr_ptr to [7 x i32] addrspace(3)*
+  %gep = getelementptr inbounds [7 x i32], [7 x i32] addrspace(3)* %asc, i64 %y, i64 %z
+  ret i32 addrspace(3)* %gep
+}