fixup! [InstCombine] Relax the same-underlying-object constraint for the GEP canonicalization

Author: dtcxzyw

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

@@ -2477,13 +2477,31 @@ Instruction *InstCombinerImpl::visitGetElementPtrInst(GetElementPtrInst &GEP) {
         Value *X = GEP.getPointerOperand();
         Value *Y;
         if (match(GEP.getOperand(1),
-                  m_Sub(m_PtrToInt(m_Value(Y)), m_PtrToInt(m_Specific(X)))) &&
-            (all_of(GEP.users(),
-                    [](User *U) {
-                      return isa<ICmpInst>(U) || isa<PtrToIntInst>(U);
-                    }) ||
-             getUnderlyingObject(X) == getUnderlyingObject(Y)))
-          return CastInst::CreatePointerBitCastOrAddrSpaceCast(Y, GEPType);
+                  m_Sub(m_PtrToInt(m_Value(Y)), m_PtrToInt(m_Specific(X))))) {
+          std::optional<bool> HasSameUnderlyingObjectCache;
+          auto HasSameUnderlyingObject = [&] {
+            if (!HasSameUnderlyingObjectCache)
+              HasSameUnderlyingObjectCache =
+                  getUnderlyingObject(X) == getUnderlyingObject(Y);
+            return *HasSameUnderlyingObjectCache;
+          };
+          std::optional<Value *> CastedValueCache;
+          auto GetCastedValue = [&] {
+            if (!CastedValueCache)
+              CastedValueCache =
+                  Builder.CreatePointerBitCastOrAddrSpaceCast(Y, GEPType);
+            return *CastedValueCache;
+          };
+
+          bool Changed = false;
+          for (User *U : GEP.users())
+            if (isa<ICmpInst>(U) || isa<PtrToIntInst>(U) ||
+                HasSameUnderlyingObject()) {
+              U->replaceUsesOfWith(&GEP, GetCastedValue());
+              Changed = true;
+            }
+          return Changed ? &GEP : nullptr;
+        }
       } else {
         // Canonicalize (gep T* X, V / sizeof(T)) to (gep i8* X, V)
         Value *V;

llvm/test/Transforms/InstCombine/getelementptr.ll

@@ -1590,13 +1590,13 @@ if.else:
 
 define i64 @test_used_by_both_invalid(ptr %a, ptr %b, ptr %c) {
 ; CHECK-LABEL: @test_used_by_both_invalid(
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq ptr [[B:%.*]], [[C:%.*]]
+; CHECK-NEXT:    br i1 [[CMP]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    [[PB:%.*]] = ptrtoint ptr [[B]] to i64
 ; CHECK-NEXT:    [[PA:%.*]] = ptrtoint ptr [[A:%.*]] to i64
-; CHECK-NEXT:    [[PB:%.*]] = ptrtoint ptr [[B:%.*]] to i64
 ; CHECK-NEXT:    [[SUB:%.*]] = sub i64 [[PB]], [[PA]]
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i8, ptr [[A]], i64 [[SUB]]
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq ptr [[GEP]], [[C:%.*]]
-; CHECK-NEXT:    br i1 [[CMP]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
-; CHECK:       if.then:
 ; CHECK-NEXT:    [[VAL:%.*]] = load i64, ptr [[GEP]], align 8
 ; CHECK-NEXT:    ret i64 [[VAL]]
 ; CHECK:       if.else: