[NaryReassociate] allow candidate to have a different type

[jingyue]

Summary

If Candiadte may have a different type from GEP, we should bitcast or
pointer cast it to GEP's type so that the later RAUW doesn't complain.

Added a test in nary-gep.ll

[meheff]

LGTM

[mcrosier]

lib/Transforms/Scalar/NaryReassociate.cpp

432 ↗

(On Diff #43159)

This might be a silly question, but am I correct in assuming the IR builder will just return the Candidate if the types already match? Hence, the reason we don't need a check prior to the call to CreateBitOrPointerCast. Of course if we did the check up front we wouldn't need to create the builder... ..and if the common case is that they do match, then we might consider predicating all this logic.

[jingyue]

lib/Transforms/Scalar/NaryReassociate.cpp

432 ↗

(On Diff #43159)

Right, IRBuilder will just return Candidate if the types already match, so we don't need predicating.

llvm/trunk/lib/Transforms/Scalar/NaryReassociate.cpp

if (isKnownNonNegative(LHS, *DL, 0, AC, GEP, DT) &&
     // @reassociate_gep_assume for an example of this canonicalization.
     IndexExprs[I] =
         SE->getZeroExtendExpr(IndexExprs[I], GEP->getOperand(I)->getType());
   }
   const SCEV *CandidateExpr = SE->getGEPExpr(
       GEP->getSourceElementType(), SE->getSCEV(GEP->getPointerOperand()),
       IndexExprs, GEP->isInBounds());
 
-  auto *Candidate = findClosestMatchingDominator(CandidateExpr, GEP);
+  Value *Candidate = findClosestMatchingDominator(CandidateExpr, GEP);
   if (Candidate == nullptr)
     return nullptr;
 
-  PointerType *TypeOfCandidate = dyn_cast<PointerType>(Candidate->getType());
-  // Pretty rare but theoretically possible when a numeric value happens to
-  // share CandidateExpr.
-  if (TypeOfCandidate == nullptr)
-    return nullptr;
+  IRBuilder<> Builder(GEP);
+  // Candidate does not necessarily have the same pointer type as GEP. Use
+  // bitcast or pointer cast to make sure they have the same type, so that the
+  // later RAUW doesn't complain.
+  Candidate = Builder.CreateBitOrPointerCast(Candidate, GEP->getType());
+  assert(Candidate->getType() == GEP->getType());
 
   // NewGEP = (char *)Candidate + RHS * sizeof(IndexedType)
   uint64_t IndexedSize = DL->getTypeAllocSize(IndexedType);
-  Type *ElementType = TypeOfCandidate->getElementType();
+  Type *ElementType = GEP->getType()->getElementType();
   uint64_t ElementSize = DL->getTypeAllocSize(ElementType);
   // Another less rare case: because I is not necessarily the last index of the
   // GEP, the size of the type at the I-th index (IndexedSize) is not
   // necessarily divisible by ElementSize. For example,
   //
   // #pragma pack(1)
   // struct S {
   //   int a[3];
   //   int64 b[8];
   // };
   // #pragma pack()
   //
   // sizeof(S) = 100 is indivisible by sizeof(int64) = 8.
   //
   // TODO: bail out on this case for now. We could emit uglygep.
   if (IndexedSize % ElementSize != 0)
     return nullptr;
 
   // NewGEP = &Candidate[RHS * (sizeof(IndexedType) / sizeof(Candidate[0])));
-  IRBuilder<> Builder(GEP);
-  Type *IntPtrTy = DL->getIntPtrType(TypeOfCandidate);
+  Type *IntPtrTy = DL->getIntPtrType(GEP->getType());
   if (RHS->getType() != IntPtrTy)
     RHS = Builder.CreateSExtOrTrunc(RHS, IntPtrTy);
   if (IndexedSize != ElementSize) {
     RHS = Builder.CreateMul(
         RHS, ConstantInt::get(IntPtrTy, IndexedSize / ElementSize));
   }
   GetElementPtrInst *NewGEP =
       cast<GetElementPtrInst>(Builder.CreateGEP(Candidate, RHS));

llvm/trunk/test/Transforms/NaryReassociate/NVPTX/nary-gep.ll

; CHECK: call void @foo(float* [[t1]])
   %3 = getelementptr float, float* %a, i128 %1
 ; CHECK: [[truncj:[^ ]+]] = trunc i128 %j to i64
 ; CHECK: [[t2:[^ ]+]] = getelementptr float, float* [[t1]], i64 [[truncj]]
   call void @foo(float* %3)
 ; CHECK: call void @foo(float* [[t2]])
   ret void
 }
 
+%struct.complex = type { float, float }
+
+declare void @bar(%struct.complex*)
+
+define void @different_types(%struct.complex* %input, i64 %i) {
+; CHECK-LABEL: @different_types(
+  %t1 = getelementptr %struct.complex, %struct.complex* %input, i64 %i
+  call void @bar(%struct.complex* %t1)
+  %j = add i64 %i, 5
+  %t2 = getelementptr %struct.complex, %struct.complex* %input, i64 %j, i32 0
+; CHECK: [[cast:[^ ]+]] = bitcast %struct.complex* %t1 to float*
+; CHECK-NEXT: %t2 = getelementptr float, float* [[cast]], i64 10
+; CHECK-NEXT: call void @foo(float* %t2)
+  call void @foo(float* %t2)
+  ret void
+}
+
 declare void @llvm.assume(i1)