[InstCombine] Support (sub (sext x), (sext y)) --> (sext (sub x, y)) and (sub (zext x), (zext y)) --> (zext (sub x, y))

[craig.topper]

Summary

If the sub doesn't overflow in the original type we can move it above the sext/zext.

This is similar to what we do for add. The overflow checking for sub is currently weaker than add, so the test cases are constructed for what is supported.

[spatel]

Logic seems right, but we already had a big pile of repeated code.
I think we can handle sub the same way as add/mul after:
rL342292
The only extra work will be swapping the constant operand to LHS for sub?

[craig.topper]

Rebase to use the code from r342292.

[spatel]

LGTM - but please add at least one negative test where the math actually could overflow, so we verify that the value tracking is in place.

lib/Transforms/InstCombine/InstCombineAddSub.cpp

if (match(Op1, m_AShr(m_Value(A), m_APInt(ShAmt))) &&
     // --> (A < 0) ? -A : A
     Value *Cmp = Builder.CreateICmpSLT(A, ConstantInt::getNullValue(Ty));
     // Copy the nuw/nsw flags from the sub to the negate.
     Value *Neg = Builder.CreateNeg(A, "", I.hasNoUnsignedWrap(),
                                    I.hasNoSignedWrap());
     return SelectInst::Create(Cmp, Neg, A);
   }
 
+  // Check for (sub x, (sext y)), see if we can merge this into an
+  // integer add followed by a sext.
+  if (SExtInst *RHSConv = dyn_cast<SExtInst>(Op1)) {
+    // (sub cst, (sext y)) --> (sext (sub cst', y))
+    if (auto *LHSC = dyn_cast<Constant>(Op0)) {
+      if (RHSConv->hasOneUse()) {
+        Constant *CI =
+            ConstantExpr::getTrunc(LHSC, RHSConv->getOperand(0)->getType());
+        if (ConstantExpr::getSExt(CI, Ty) == LHSC &&
+            willNotOverflowSignedSub(CI, RHSConv->getOperand(0), I)) {
+          // Insert the new, smaller sub.
+          Value *NewSub =
+              Builder.CreateNSWSub(CI, RHSConv->getOperand(0), "subconv");
+          return new SExtInst(NewSub, Ty);
+        }
+      }
+    }
+
+    // (sub (sext x), (sext y)) --> (sext (sub int x, y))
+    if (SExtInst *LHSConv = dyn_cast<SExtInst>(Op0)) {
+      // Only do this if x/y have the same type, if at least one of them has a
+      // single use (so we don't increase the number of sexts), and if the
+      // integer add will not overflow.
+      if (LHSConv->getOperand(0)->getType() ==
+              RHSConv->getOperand(0)->getType() &&
+          (LHSConv->hasOneUse() || RHSConv->hasOneUse()) &&
+          willNotOverflowSignedSub(LHSConv->getOperand(0),
+                                   RHSConv->getOperand(0), I)) {
+        // Insert the new integer sub.
+        Value *NewSub = Builder.CreateNSWSub(LHSConv->getOperand(0),
+                                             RHSConv->getOperand(0), "subconv");
+        return new SExtInst(NewSub, Ty);
+      }
+    }
+  }
+
+  // Check for (sub x, (zext y)), see if we can merge this into an
+  // integer sub followed by a zext.
+  if (auto *RHSConv = dyn_cast<ZExtInst>(Op1)) {
+    // (sub cst, (zext y)) --> (zext (sub cst', y))
+    if (auto *LHSC = dyn_cast<Constant>(Op0)) {
+      if (RHSConv->hasOneUse()) {
+        Constant *CI =
+            ConstantExpr::getTrunc(LHSC, RHSConv->getOperand(0)->getType());
+        if (ConstantExpr::getZExt(CI, Ty) == LHSC &&
+            willNotOverflowUnsignedSub(CI, RHSConv->getOperand(0), I)) {
+          // Insert the new, smaller sub.
+          Value *NewSub =
+              Builder.CreateNUWSub(CI, RHSConv->getOperand(0), "subconv");
+          return new ZExtInst(NewSub, Ty);
+        }
+      }
+    }
+
+    // (sub (zext x), (zext y)) --> (zext (sub int x, y))
+    if (auto *LHSConv = dyn_cast<ZExtInst>(Op0)) {
+      // Only do this if x/y have the same type, if at least one of them has a
+      // single use (so we don't increase the number of zexts), and if the
+      // integer sub will not overflow.
+      if (LHSConv->getOperand(0)->getType() ==
+              RHSConv->getOperand(0)->getType() &&
+          (LHSConv->hasOneUse() || RHSConv->hasOneUse()) &&
+          willNotOverflowUnsignedSub(LHSConv->getOperand(0),
+                                     RHSConv->getOperand(0), I)) {
+        // Insert the new integer sub.
+        Value *NewSub = Builder.CreateNUWSub(
+            LHSConv->getOperand(0), RHSConv->getOperand(0), "subconv");
+        return new ZExtInst(NewSub, Ty);
+      }
+    }
+  }
   bool Changed = false;
   if (!I.hasNoSignedWrap() && willNotOverflowSignedSub(Op0, Op1, I)) {
     Changed = true;
     I.setHasNoSignedWrap(true);
   }
   if (!I.hasNoUnsignedWrap() && willNotOverflowUnsignedSub(Op0, Op1, I)) {
     Changed = true;
     I.setHasNoUnsignedWrap(true);

test/Transforms/InstCombine/sink-zext.ll

;
   %call1 = call i32 @callee(), !range !1
   %call2 = call i32 @callee(), !range !1
   %sext1 = sext i32 %call1 to i64
   %sext2 = sext i32 %call2 to i64
   %add = mul i64 %sext1, %sext2
   ret i64 %add
 }
 
+define i64 @test13(i32 %V) {
+; CHECK-LABEL: @test13(
+; CHECK-NEXT:    [[CALL1:%.*]] = call i32 @callee(), !range !2
+; CHECK-NEXT:    [[CALL2:%.*]] = call i32 @callee(), !range !3
+; CHECK-NEXT:    [[SUBCONV:%.*]] = sub nsw i32 [[CALL1]], [[CALL2]]
+; CHECK-NEXT:    [[SUB:%.*]] = sext i32 [[SUBCONV]] to i64
+; CHECK-NEXT:    ret i64 [[SUB]]
+;
+  %call1 = call i32 @callee(), !range !2
+  %call2 = call i32 @callee(), !range !3
+  %sext1 = sext i32 %call1 to i64
+  %sext2 = sext i32 %call2 to i64
+  %sub = sub i64 %sext1, %sext2
+  ret i64 %sub
+}
+
+define i64 @test14(i32 %V) {
+; CHECK-LABEL: @test14(
+; CHECK-NEXT:    [[CALL1:%.*]] = call i32 @callee(), !range !2
+; CHECK-NEXT:    [[CALL2:%.*]] = call i32 @callee(), !range !0
+; CHECK-NEXT:    [[SUBCONV:%.*]] = sub nuw nsw i32 [[CALL1]], [[CALL2]]
+; CHECK-NEXT:    [[SUB:%.*]] = zext i32 [[SUBCONV]] to i64
+; CHECK-NEXT:    ret i64 [[SUB]]
+;
+  %call1 = call i32 @callee(), !range !2
+  %call2 = call i32 @callee(), !range !0
+  %zext1 = zext i32 %call1 to i64
+  %zext2 = zext i32 %call2 to i64
+  %sub = sub i64 %zext1, %zext2
+  ret i64 %sub
+}
+
+define i64 @test15(i32 %V) {
+; CHECK-LABEL: @test15(
+; CHECK-NEXT:    [[ASHR:%.*]] = ashr i32 [[V:%.*]], 1
+; CHECK-NEXT:    [[SUBCONV:%.*]] = sub nsw i32 8, [[ASHR]]
+; CHECK-NEXT:    [[SUB:%.*]] = sext i32 [[SUBCONV]] to i64
+; CHECK-NEXT:    ret i64 [[SUB]]
+;
+  %ashr = ashr i32 %V, 1
+  %sext = sext i32 %ashr to i64
+  %sub = sub i64 8, %sext
+  ret i64 %sub
+}
+
+define <2 x i64> @test15vec(<2 x i32> %V) {
+; CHECK-LABEL: @test15vec(
+; CHECK-NEXT:    [[ASHR:%.*]] = ashr <2 x i32> [[V:%.*]], <i32 1, i32 1>
+; CHECK-NEXT:    [[SUBCONV:%.*]] = sub nsw <2 x i32> <i32 8, i32 8>, [[ASHR]]
+; CHECK-NEXT:    [[SUB:%.*]] = sext <2 x i32> [[SUBCONV]] to <2 x i64>
+; CHECK-NEXT:    ret <2 x i64> [[SUB]]
+;
+  %ashr = ashr <2 x i32> %V, <i32 1, i32 1>
+  %sext = sext <2 x i32> %ashr to <2 x i64>
+  %sub = sub <2 x i64> <i64 8, i64 8>, %sext
+  ret <2 x i64> %sub
+}
+
+define i64 @test16(i32 %V) {
+; CHECK-LABEL: @test16(
+; CHECK-NEXT:    [[LSHR:%.*]] = lshr i32 [[V:%.*]], 1
+; CHECK-NEXT:    [[SUBCONV:%.*]] = sub nuw i32 -2, [[LSHR]]
+; CHECK-NEXT:    [[SUB:%.*]] = zext i32 [[SUBCONV]] to i64
+; CHECK-NEXT:    ret i64 [[SUB]]
+;
+  %lshr = lshr i32 %V, 1
+  %zext = zext i32 %lshr to i64
+  %sub = sub i64 4294967294, %zext
+  ret i64 %sub
+}
+
+define <2 x i64> @test16vec(<2 x i32> %V) {
+; CHECK-LABEL: @test16vec(
+; CHECK-NEXT:    [[LSHR:%.*]] = lshr <2 x i32> [[V:%.*]], <i32 1, i32 1>
+; CHECK-NEXT:    [[SUBCONV:%.*]] = sub nuw <2 x i32> <i32 -2, i32 -2>, [[LSHR]]
+; CHECK-NEXT:    [[SUB:%.*]] = zext <2 x i32> [[SUBCONV]] to <2 x i64>
+; CHECK-NEXT:    ret <2 x i64> [[SUB]]
+;
+  %lshr = lshr <2 x i32> %V, <i32 1, i32 1>
+  %zext = zext <2 x i32> %lshr to <2 x i64>
+  %sub = sub <2 x i64> <i64 4294967294, i64 4294967294>, %zext
+  ret <2 x i64> %sub
+}
+
 !0 = !{ i32 0, i32 2000 }
 !1 = !{ i32 -2000, i32 0 }
+!2 = !{ i32 -512, i32 -255 }
+!3 = !{ i32 -128, i32 0 }