[InstCombine] Always rewrite multi-use GEP for pointer difference

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

@@ -2101,29 +2101,34 @@ Value *InstCombinerImpl::OptimizePointerDifference(Value *LHS, Value *RHS,
   if (!GEP1)
     return nullptr;
 
+  // Emit the offset of the GEP as an intptr_t.
   // To avoid duplicating the offset arithmetic, rewrite the GEP to use the
   // computed offset. This may erase the original GEP, so be sure to cache the
   // nowrap flags before emitting the offset.
-  // TODO: We should probably do this even if there is only one GEP.
-  bool RewriteGEPs = GEP2 != nullptr;
-
-  // Emit the offset of the GEP and an intptr_t.
   GEPNoWrapFlags GEP1NW = GEP1->getNoWrapFlags();
-  Value *Result = EmitGEPOffset(GEP1, RewriteGEPs);
+  Value *Result = EmitGEPOffset(GEP1, /*RewriteGEP=*/true);
 
   // If this is a single inbounds GEP and the original sub was nuw,
   // then the final multiplication is also nuw.
   if (auto *I = dyn_cast<Instruction>(Result))
     if (IsNUW && !GEP2 && !Swapped && GEP1NW.isInBounds() &&
-        I->getOpcode() == Instruction::Mul)
+        I->getOpcode() == Instruction::Mul && !I->hasNoUnsignedWrap()) {
+      if (!I->use_empty()) {
+        // If the offset calculation is reused by the GEP, add the nuw flag to
+        // a separate clone. This may improve folds and will get CSEd if not
+        // useful.
+        Result = I = I->clone();
+        Builder.Insert(I);
+      }
       I->setHasNoUnsignedWrap();
+    }
 
   // If we have a 2nd GEP of the same base pointer, subtract the offsets.
   // If both GEPs are inbounds, then the subtract does not have signed overflow.
   // If both GEPs are nuw and the original sub is nuw, the new sub is also nuw.
   if (GEP2) {
     GEPNoWrapFlags GEP2NW = GEP2->getNoWrapFlags();
-    Value *Offset = EmitGEPOffset(GEP2, RewriteGEPs);
+    Value *Offset = EmitGEPOffset(GEP2, /*RewriteGEP=*/true);
     Result = Builder.CreateSub(Result, Offset, "gepdiff",
                                IsNUW && GEP1NW.hasNoUnsignedWrap() &&
                                    GEP2NW.hasNoUnsignedWrap(),

llvm/test/Transforms/InstCombine/sub-gep.ll

@@ -760,6 +760,39 @@ entry:
   ret i64 %ret
 }
 
+declare void @use(ptr)
+
+define i64 @sub_multi_use(ptr %base, i64 %idx) {
+; CHECK-LABEL: @sub_multi_use(
+; CHECK-NEXT:    [[P2_IDX:%.*]] = shl nsw i64 [[IDX:%.*]], 2
+; CHECK-NEXT:    [[P2:%.*]] = getelementptr inbounds i8, ptr [[BASE:%.*]], i64 [[P2_IDX]]
+; CHECK-NEXT:    call void @use(ptr [[P2]])
+; CHECK-NEXT:    ret i64 [[P2_IDX]]
+;
+  %p2 = getelementptr inbounds [0 x i32], ptr %base, i64 0, i64 %idx
+  call void @use(ptr %p2)
+  %i1 = ptrtoint ptr %base to i64
+  %i2 = ptrtoint ptr %p2 to i64
+  %d = sub i64 %i2, %i1
+  ret i64 %d
+}
+
+define i64 @sub_multi_use_nuw(ptr %base, i64 %idx) {
+; CHECK-LABEL: @sub_multi_use_nuw(
+; CHECK-NEXT:    [[P2_IDX:%.*]] = shl nsw i64 [[IDX:%.*]], 2
+; CHECK-NEXT:    [[P2:%.*]] = getelementptr inbounds i8, ptr [[BASE:%.*]], i64 [[P2_IDX]]
+; CHECK-NEXT:    call void @use(ptr [[P2]])
+; CHECK-NEXT:    [[D:%.*]] = shl nuw nsw i64 [[IDX]], 2
+; CHECK-NEXT:    ret i64 [[D]]
+;
+  %p2 = getelementptr inbounds [0 x i32], ptr %base, i64 0, i64 %idx
+  call void @use(ptr %p2)
+  %i1 = ptrtoint ptr %base to i64
+  %i2 = ptrtoint ptr %p2 to i64
+  %d = sub nuw i64 %i2, %i1
+  ret i64 %d
+}
+
 define i1 @_gep_phi1(ptr %str1) {
 ; CHECK-LABEL: @_gep_phi1(
 ; CHECK-NEXT:  entry: