[InstCombine] Extend (add (sext x), cst) --> (sext (add x, cst')) and (add (zext x), cst) --> (zext (add x, cst')) to work for vectors

[craig.topper]

Summary

This patch extends this transform to work for vectors, including non-splat vectors that contain all ConstantInts. I've added a new matcher to PatternMatch to detect vectors of all ConstantInts. We can probably also skip undefs in it, but I've left that for the future.

[spatel]

lib/Transforms/InstCombine/InstCombineAddSub.cpp

1199

I don't understand what benefit the new pattern matcher offers. We're already guaranteed that the constant is (vector) integer type because this is an 'add'. Do we need to disallow constant expressions? If not, use the existing m_Constant() matcher?

[craig.topper]

lib/Transforms/InstCombine/InstCombineAddSub.cpp

1199

I did it to disallow constant expressions and be better than just using m_APInt. I don't think willNotOverflowSignedAdd will give an answer for constant expressions in general. So we'd just be creating new constant expressions with the call get to ConstantExpr::getTrunc for no reason. And we'd need tests for it.

[lebedev.ri]

include/llvm/IR/PatternMatch.h

211–212

1. This should ideally describe how this is different from the normal m_Constant() matcher (i.e. that this ignores constant expressions)
2. Can we go the other way around, and add a matcher for the constant expressions?

[craig.topper]

include/llvm/IR/PatternMatch.h

211–212

I'm not sure I understand what you're asking for 2.

[craig.topper]

Add more to comment. Fix typo in comment

[spatel]

I appreciate the concern about having a test for constant expressions if we just use m_Constant() here, but I wouldn't hold this patch up for that. We have dozens of uses of m_Constant() and barely any tests for constant expressions (presumably because those are rarely the motivation for adding a transform).

The rarity of constant expressions would also suggest that we don't really need to be worried about efficiency in handling those.

So my preference would be:

1. Add the tests with baseline checks.
2. Switch this transform over to m_Constant(), so it can handle arbitrary vector constants.
3. Add m_ConstantIntOrConstantIntVec and apply it wherever we think it would make the code clearer/faster. I'm not convinced yet that it's worth the effort to exclude constant expressions, but if it is, then let's do it throughout instcombine?

[craig.topper]

Switch to m_Constant

[spatel]

LGTM

include/llvm/IR/PatternMatch.h

 /// Match a ConstantInt or splatted ConstantVector, binding the
 /// specified pointer to the contained APInt.
 inline apint_match m_APInt(const APInt *&Res) { return Res; }
 
 /// Match a ConstantFP or splatted ConstantVector, binding the
 /// specified pointer to the contained APFloat.
 inline apfloat_match m_APFloat(const APFloat *&Res) { return Res; }
 
+struct constantint_or_constantint_vec_match {
+  Constant *&Res;
+  constantint_or_constantint_vec_match(Constant *&R) : Res(R) {}
+  template <typename ITy> bool match(ITy *V) {
+    if (auto *CI = dyn_cast<ConstantInt>(V)) {
+      Res = CI;
+      return true;
+    }
+
+    if (V->getType()->isVectorTy()) {
+      if (auto *C = dyn_cast<Constant>(V)) {
+        unsigned NumElts = V->getType()->getVectorNumElements();
+        assert(NumElts != 0 && "Constant vector with no elements?");
+        for (unsigned i = 0; i != NumElts; ++i) {
+          Constant *Elt = C->getAggregateElement(i);
+          if (!Elt || !isa<ConstantInt>(Elt))
+            return false;
+        }
+
+        Res = C;
+        return true;
+      }
+    }
+
+    return false;
+  }
+};
+
+/// Match a ConstantInt or a vector of ConstantInts, binding the specified
+/// pointer to the Constant. This is more specific than m_Constant in that it
+/// explicitly ensures no ConstantExpr or Undef elements.
+inline constantint_or_constantint_vec_match
+m_ConstantIntOrConstantIntVec(Constant *&Res) { return Res; }
+
 template <int64_t Val> struct constantint_match {
   template <typename ITy> bool match(ITy *V) {
     if (const auto *CI = dyn_cast<ConstantInt>(V)) {
       const APInt &CIV = CI->getValue();
       if (Val >= 0)
         return CIV == static_cast<uint64_t>(Val);
       // If Val is negative, and CI is shorter than it, truncate to the right
       // number of bits.  If it is larger, then we have to sign extend.  Just

lib/Transforms/InstCombine/InstCombineAddSub.cpp

if (SI && SI->hasOneUse()) {
         return SelectInst::Create(SI->getCondition(), A, N);
     }
   }
 
   // Check for (add (sext x), y), see if we can merge this into an
   // integer add followed by a sext.
   if (SExtInst *LHSConv = dyn_cast<SExtInst>(LHS)) {
     // (add (sext x), cst) --> (sext (add x, cst'))
-    if (ConstantInt *RHSC = dyn_cast<ConstantInt>(RHS)) {
+    Constant *RHSC;
+    if (match(RHS, m_ConstantIntOrConstantIntVec(RHSC))) {
       if (LHSConv->hasOneUse()) {
         Constant *CI =
             ConstantExpr::getTrunc(RHSC, LHSConv->getOperand(0)->getType());
         if (ConstantExpr::getSExt(CI, Ty) == RHSC &&
             willNotOverflowSignedAdd(LHSConv->getOperand(0), CI, I)) {
           // Insert the new, smaller add.
           Value *NewAdd =
               Builder.CreateNSWAdd(LHSConv->getOperand(0), CI, "addconv");
if (SExtInst *RHSConv = dyn_cast<SExtInst>(RHS)) {
       }
     }
   }
 
   // Check for (add (zext x), y), see if we can merge this into an
   // integer add followed by a zext.
   if (auto *LHSConv = dyn_cast<ZExtInst>(LHS)) {
     // (add (zext x), cst) --> (zext (add x, cst'))
-    if (ConstantInt *RHSC = dyn_cast<ConstantInt>(RHS)) {
+    Constant *RHSC;
+    if (match(RHS, m_ConstantIntOrConstantIntVec(RHSC))) {
       if (LHSConv->hasOneUse()) {
         Constant *CI =
             ConstantExpr::getTrunc(RHSC, LHSConv->getOperand(0)->getType());
         if (ConstantExpr::getZExt(CI, Ty) == RHSC &&
             willNotOverflowUnsignedAdd(LHSConv->getOperand(0), CI, I)) {
           // Insert the new, smaller add.
           Value *NewAdd =
               Builder.CreateNUWAdd(LHSConv->getOperand(0), CI, "addconv");

test/Transforms/InstCombine/sink-zext.ll

 ;
   %call1 = call i32 @callee(), !range !0
   %call2 = call i32 @callee(), !range !0
   %add = mul i32 %call1, %call2
   %zext = sext i32 %add to i64
   ret i64 %zext
 }
 
+define i64 @test5(i32 %V) {
+; CHECK-LABEL: @test5(
+; CHECK-NEXT:    [[ASHR:%.*]] = ashr i32 [[V:%.*]], 1
+; CHECK-NEXT:    [[ADDCONV:%.*]] = add nsw i32 [[ASHR]], 1073741823
+; CHECK-NEXT:    [[ADD:%.*]] = sext i32 [[ADDCONV]] to i64
+; CHECK-NEXT:    ret i64 [[ADD]]
+;
+  %ashr = ashr i32 %V, 1
+  %sext = sext i32 %ashr to i64
+  %add = add i64 %sext, 1073741823
+  ret i64 %add
+}
+
+define <2 x i64> @test5_splat(<2 x i32> %V) {
+; CHECK-LABEL: @test5_splat(
+; CHECK-NEXT:    [[ASHR:%.*]] = ashr <2 x i32> [[V:%.*]], <i32 1, i32 1>
+; CHECK-NEXT:    [[ADDCONV:%.*]] = add nsw <2 x i32> [[ASHR]], <i32 1073741823, i32 1073741823>
+; CHECK-NEXT:    [[ADD:%.*]] = sext <2 x i32> [[ADDCONV]] to <2 x i64>
+; CHECK-NEXT:    ret <2 x i64> [[ADD]]
+;
+  %ashr = ashr <2 x i32> %V, <i32 1, i32 1>
+  %sext = sext <2 x i32> %ashr to <2 x i64>
+  %add = add <2 x i64> %sext, <i64 1073741823, i64 1073741823>
+  ret <2 x i64> %add
+}
+
+define <2 x i64> @test5_vec(<2 x i32> %V) {
+; CHECK-LABEL: @test5_vec(
+; CHECK-NEXT:    [[ASHR:%.*]] = ashr <2 x i32> [[V:%.*]], <i32 1, i32 1>
+; CHECK-NEXT:    [[ADDCONV:%.*]] = add nsw <2 x i32> [[ASHR]], <i32 1, i32 2>
+; CHECK-NEXT:    [[ADD:%.*]] = sext <2 x i32> [[ADDCONV]] to <2 x i64>
+; CHECK-NEXT:    ret <2 x i64> [[ADD]]
+;
+  %ashr = ashr <2 x i32> %V, <i32 1, i32 1>
+  %sext = sext <2 x i32> %ashr to <2 x i64>
+  %add = add <2 x i64> %sext, <i64 1, i64 2>
+  ret <2 x i64> %add
+}
+
+define i64 @test6(i32 %V) {
+; CHECK-LABEL: @test6(
+; CHECK-NEXT:    [[ASHR:%.*]] = ashr i32 [[V:%.*]], 1
+; CHECK-NEXT:    [[ADDCONV:%.*]] = add nsw i32 [[ASHR]], -1073741824
+; CHECK-NEXT:    [[ADD:%.*]] = sext i32 [[ADDCONV]] to i64
+; CHECK-NEXT:    ret i64 [[ADD]]
+;
+  %ashr = ashr i32 %V, 1
+  %sext = sext i32 %ashr to i64
+  %add = add i64 %sext, -1073741824
+  ret i64 %add
+}
+
+define <2 x i64> @test6_splat(<2 x i32> %V) {
+; CHECK-LABEL: @test6_splat(
+; CHECK-NEXT:    [[ASHR:%.*]] = ashr <2 x i32> [[V:%.*]], <i32 1, i32 1>
+; CHECK-NEXT:    [[ADDCONV:%.*]] = add nsw <2 x i32> [[ASHR]], <i32 -1073741824, i32 -1073741824>
+; CHECK-NEXT:    [[ADD:%.*]] = sext <2 x i32> [[ADDCONV]] to <2 x i64>
+; CHECK-NEXT:    ret <2 x i64> [[ADD]]
+;
+  %ashr = ashr <2 x i32> %V, <i32 1, i32 1>
+  %sext = sext <2 x i32> %ashr to <2 x i64>
+  %add = add <2 x i64> %sext, <i64 -1073741824, i64 -1073741824>
+  ret <2 x i64> %add
+}
+
+define <2 x i64> @test6_vec(<2 x i32> %V) {
+; CHECK-LABEL: @test6_vec(
+; CHECK-NEXT:    [[ASHR:%.*]] = ashr <2 x i32> [[V:%.*]], <i32 1, i32 1>
+; CHECK-NEXT:    [[ADDCONV:%.*]] = add nsw <2 x i32> [[ASHR]], <i32 -1, i32 -2>
+; CHECK-NEXT:    [[ADD:%.*]] = sext <2 x i32> [[ADDCONV]] to <2 x i64>
+; CHECK-NEXT:    ret <2 x i64> [[ADD]]
+;
+  %ashr = ashr <2 x i32> %V, <i32 1, i32 1>
+  %sext = sext <2 x i32> %ashr to <2 x i64>
+  %add = add <2 x i64> %sext, <i64 -1, i64 -2>
+  ret <2 x i64> %add
+}
+
+define <2 x i64> @test6_vec2(<2 x i32> %V) {
+; CHECK-LABEL: @test6_vec2(
+; CHECK-NEXT:    [[ASHR:%.*]] = ashr <2 x i32> [[V:%.*]], <i32 1, i32 1>
+; CHECK-NEXT:    [[ADDCONV:%.*]] = add nsw <2 x i32> [[ASHR]], <i32 -1, i32 1>
+; CHECK-NEXT:    [[ADD:%.*]] = sext <2 x i32> [[ADDCONV]] to <2 x i64>
+; CHECK-NEXT:    ret <2 x i64> [[ADD]]
+;
+  %ashr = ashr <2 x i32> %V, <i32 1, i32 1>
+  %sext = sext <2 x i32> %ashr to <2 x i64>
+  %add = add <2 x i64> %sext, <i64 -1, i64 1>
+  ret <2 x i64> %add
+}
+
+define i64 @test7(i32 %V) {
+; CHECK-LABEL: @test7(
+; CHECK-NEXT:    [[LSHR:%.*]] = lshr i32 [[V:%.*]], 1
+; CHECK-NEXT:    [[ADDCONV:%.*]] = add nuw i32 [[LSHR]], 2147483647
+; CHECK-NEXT:    [[ADD:%.*]] = zext i32 [[ADDCONV]] to i64
+; CHECK-NEXT:    ret i64 [[ADD]]
+;
+  %lshr = lshr i32 %V, 1
+  %zext = zext i32 %lshr to i64
+  %add = add i64 %zext, 2147483647
+  ret i64 %add
+}
+
+define <2 x i64> @test7_splat(<2 x i32> %V) {
+; CHECK-LABEL: @test7_splat(
+; CHECK-NEXT:    [[LSHR:%.*]] = lshr <2 x i32> [[V:%.*]], <i32 1, i32 1>
+; CHECK-NEXT:    [[ADDCONV:%.*]] = add nuw <2 x i32> [[LSHR]], <i32 2147483647, i32 2147483647>
+; CHECK-NEXT:    [[ADD:%.*]] = zext <2 x i32> [[ADDCONV]] to <2 x i64>
+; CHECK-NEXT:    ret <2 x i64> [[ADD]]
+;
+  %lshr = lshr <2 x i32> %V, <i32 1, i32 1>
+  %zext = zext <2 x i32> %lshr to <2 x i64>
+  %add = add <2 x i64> %zext, <i64 2147483647, i64 2147483647>
+  ret <2 x i64> %add
+}
+
+define <2 x i64> @test7_vec(<2 x i32> %V) {
+; CHECK-LABEL: @test7_vec(
+; CHECK-NEXT:    [[LSHR:%.*]] = lshr <2 x i32> [[V:%.*]], <i32 1, i32 1>
+; CHECK-NEXT:    [[ADDCONV:%.*]] = add nuw <2 x i32> [[LSHR]], <i32 1, i32 2>
+; CHECK-NEXT:    [[ADD:%.*]] = zext <2 x i32> [[ADDCONV]] to <2 x i64>
+; CHECK-NEXT:    ret <2 x i64> [[ADD]]
+;
+  %lshr = lshr <2 x i32> %V, <i32 1, i32 1>
+  %zext = zext <2 x i32> %lshr to <2 x i64>
+  %add = add <2 x i64> %zext, <i64 1, i64 2>
+  ret <2 x i64> %add
+}
+
 !0 = !{ i32 0, i32 2000 }