[InstCombine] Combine lshr of add that intends to get the carry as llvm.uadd.with.overflow

[abinavpp]

Summary

Inst-combine version of https://reviews.llvm.org/D106139.

[dmgreen]

I was trying out some examples, to see how this performs on ARM architectures. This pattern was relatively rare and didn't come up a lot in the benchmarks I tried (or optimized to be the same performance/codesize).

Testing the specific pattern, the first one I tried was i8's, a type not legal on those archs. It doesn't do great with this (although aarch64 does do OK given the instructions it can use):
https://godbolt.org/z/hbP445xsT
i32 looks OK for the same thing. AArch64 even looks better, I think because ARM already uses addcarry and splits things in a way that turns out to be equivalent.
https://godbolt.org/z/3zbhKPs1e

I was trying to look at across basic block too, but didn't do very well at getting something that wasn't just optimized away:
i8: https://godbolt.org/z/n8sxv3rqE
i32: https://godbolt.org/z/1sEa7Wd68
And then there was also this, but it crashed putting instructions in the wrong places:
https://godbolt.org/z/48abYPq8M

I have no strong objections to creating uadd_with_overflow in instcombine, but we have not done that in the past and I don't know if I see a huge reason to start now. It appears that we should at least be limiting it to legal types.

[abinavpp]

Rebased; Addressed review comments.

[abinavpp]

In D107552#2930663, @dmgreen wrote:

It appears that we should at least be limiting it to legal types.

I agree. The i8 cases doesn't look good on thumbv6m and thumbv7m. I've added DL.isLegalInteger() checks now.

[lebedev.ri]

Most of non-correctness bailouts here don't really make sense to me.
Presumably, if it is profitable for a particular backend, said backend should be expanding whatever "bad" intrinsics into it's original form.

[abinavpp]

In D107552#2933906, @lebedev.ri wrote:

Most of non-correctness bailouts here don't really make sense to me.
Presumably, if it is profitable for a particular backend, said backend should be expanding whatever "bad" intrinsics into it's original form.

I agree that we shouldn't ideally add non-correctness bailouts, but I would like to do whatever I can in this patch to avoid regressing other backends.

[abinavpp]

Fixed the case in which the add of the uaddo won't dominate the original trunc use. For e.g.:

target datalayout = "n32"

declare void @true()
declare void @false()

define i32 @lshr_32_add_zext_trunc(i32 %a, i32 %b) {

%zext.a = zext i32 %a to i64
%zext.b = zext i32 %b to i64
%add = add i64 %zext.a, %zext.b
%trunc.add = trunc i64 %add to i32
%cmp = icmp ult i32 %a, %b
br i1 %cmp, label %brTrue, label %brFalse

brTrue:

%shr = lshr i64 %add, 32
%trunc.shr = trunc i64 %shr to i32
%retTrue = add i32 %trunc.add, %trunc.shr
call void @true()
br label %brMerge

brFalse:

%retFalse = add i32 %trunc.add, 12
call void @false()
br label %brMerge

brMerge:

%ret = phi i32 [ %retTrue, %brTrue ], [ %retFalse, %brFalse ]
ret i32 %ret

}

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1060

I really don't think we should be doing this..

1069

Early return

1083–1084

Presumably you should just set the expansion point to be right after the old add?

[craig.topper]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1105

We probably need to do something here to schedule the dead truncates to revisited to remove them. But I'm not sure.

[craig.topper]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1105

I think this should use InstCombinerImpl::replaceInstUsesWith rather than Instruction::replaceAllUsesWith. Then I think it should call InstCombinerImpl::eraseInstFromFunction.

@spatel or @lebedev.ri does that sound right?

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1105

Right, the pass-specific call should be used.
I don't think we need to call eraseInstFromFunction() explcitly,
that should be handled automatically once you return non-null from this function.
Generally, i would think not having an intermediate vector would be better, but not sure.

[craig.topper]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1105

The truncates aren’t directly part of the chain of instructions we’re changing so I’m not sure if the get queued properly for DCE once we return.

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1105

InstCombinerImpl::replaceInstUsesWith() adds the instruction uses of which we replaced into worklist,
so the next time it is revisited, it will be dropped.

[abinavpp]

Rebased; Addressed review comments.

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1060

I think I'm being too cautious here. I just realized that InstCombinerImpl::foldICmpWithConstant() can produce sadd_with_overflow, if I read correctly it doesn't seem to worry about cross basic-block use. I have removed this check now.

[abinavpp]

Ping.

[arsenm]

In D107552#2930663, @dmgreen wrote:

I was trying out some examples, to see how this performs on ARM architectures. This pattern was relatively rare and didn't come up a lot in the benchmarks I tried (or optimized to be the same performance/codesize).

Testing the specific pattern, the first one I tried was i8's, a type not legal on those archs. It doesn't do great with this (although aarch64 does do OK given the instructions it can use):
https://godbolt.org/z/hbP445xsT
i32 looks OK for the same thing. AArch64 even looks better, I think because ARM already uses addcarry and splits things in a way that turns out to be equivalent.
https://godbolt.org/z/3zbhKPs1e

I was trying to look at across basic block too, but didn't do very well at getting something that wasn't just optimized away:
i8: https://godbolt.org/z/n8sxv3rqE
i32: https://godbolt.org/z/1sEa7Wd68
And then there was also this, but it crashed putting instructions in the wrong places:
https://godbolt.org/z/48abYPq8M

I have no strong objections to creating uadd_with_overflow in instcombine, but we have not done that in the past and I don't know if I see a huge reason to start now. It appears that we should at least be limiting it to legal types.

It's arguably a worse canonical form since fewer optimizations will understand this (e.g. ScalarEvolution, though not sure it really will matter here).

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1056

Demorgan this

1059–1060

I also think a legality check doesn't belong here. I also don't think the legal types in the datalayout is particularly well defined

1063

demorgan

[abinavpp]

D138814

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

public:
   Instruction *visitShl(BinaryOperator &I);
   Value *reassociateShiftAmtsOfTwoSameDirectionShifts(
       BinaryOperator *Sh0, const SimplifyQuery &SQ,
       bool AnalyzeForSignBitExtraction = false);
   Instruction *canonicalizeCondSignextOfHighBitExtractToSignextHighBitExtract(
       BinaryOperator &I);
   Instruction *foldVariableSignZeroExtensionOfVariableHighBitExtract(
       BinaryOperator &OldAShr);
+  Instruction *foldXShrToOverflow(BinaryOperator &I);
   Instruction *visitAShr(BinaryOperator &I);
   Instruction *visitLShr(BinaryOperator &I);
   Instruction *commonShiftTransforms(BinaryOperator &I);
   Instruction *visitFCmpInst(FCmpInst &I);
   CmpInst *canonicalizeICmpPredicate(CmpInst &I);
   Instruction *visitICmpInst(ICmpInst &I);
   Instruction *FoldShiftByConstant(Value *Op0, Constant *Op1,
                                    BinaryOperator &I);

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

Instruction *InstCombinerImpl::visitShl(BinaryOperator &I) {
   if (match(Op0, m_One()) &&
       match(Op1, m_Sub(m_SpecificInt(BitWidth - 1), m_Value(X))))
     return BinaryOperator::CreateLShr(
         ConstantInt::get(Ty, APInt::getSignMask(BitWidth)), X);
 
   return nullptr;
 }
 
+// Tries to perform (Xshr (add (zext a, i2^n), (zext b, i2^n)), 2^(n-1)) ->
+// (llvm.uadd.with.overflow a, b).overflow where Xshr can be ashr or lshr; a and
+// b has type i2^(n-1).
+Instruction *InstCombinerImpl::foldXShrToOverflow(BinaryOperator &I) {
+  assert(I.getOpcode() == Instruction::AShr ||
+         I.getOpcode() == Instruction::LShr);
+
+  Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
+  Type *Ty = I.getType();
+  const APInt *ShAmtAPInt;
+  if (!match(Op1, m_APInt(ShAmtAPInt)))
+    return nullptr;
+
+  unsigned ShAmt = ShAmtAPInt->getZExtValue();
+  unsigned BitWidth = Ty->getScalarSizeInBits();
+  if (!(isPowerOf2_32(BitWidth) && Ty->isIntegerTy(ShAmt * 2)))
+    return nullptr;
+
+  // Be conservative if the newly created overflow would be a live-out of it's
+  // basic-block.
+  for (const User *Usr : I.users()) {
+    const Instruction *UsrInst = dyn_cast<Instruction>(Usr);
+    if (UsrInst && UsrInst->getParent() != I.getParent())
+      return nullptr;
+  }
+
+  Value *X = nullptr, *Y = nullptr;
+  if (match(Op0, m_Add(m_ZExt(m_Value(X)), m_ZExt(m_Value(Y)))) &&
+      X->getType()->getScalarSizeInBits() == ShAmt &&
+      Y->getType()->getScalarSizeInBits() == ShAmt) {
+
+    // Make sure that `Op0` is only used by `I` and `ShAmt`-truncates.
+    SmallVector<TruncInst *, 2> Truncs;
+    for (User *Usr : Op0->users()) {
+      if (Usr == &I)
+        continue;
+
+      TruncInst *Trunc = dyn_cast<TruncInst>(Usr);
+      if (!Trunc)
+        return nullptr;
+
+      if (Trunc->getType()->getScalarSizeInBits() != ShAmt)
+        return nullptr;
+
+      Truncs.push_back(Trunc);
+    }
+
+    // If we get here, we can be sure that `Op0` is only used by `Truncs` and
+    // `I`.
+
+    Value *UAddOverflow = Builder.CreateBinaryIntrinsic(
+        Intrinsic::uadd_with_overflow, X, Y, /* FMFSource */ nullptr, "uaddo");
+    Value *UAdd = Builder.CreateExtractValue(UAddOverflow, 0,
+                                             UAddOverflow->getName() + ".add");
+    Value *Overflow = Builder.CreateExtractValue(
+        UAddOverflow, 1, UAddOverflow->getName() + ".overflow");
+
+    // Replace the uses of truncated `Op0` with `UAdd` since `UAddOverflow`
+    // performs the truncated version of the addition performed by `Op0`.
+    for (TruncInst *Trunc : Truncs) {
+      Trunc->replaceAllUsesWith(UAdd);
+    }
+
+    // Replace the use of `Op0` by `I` with `Overflow`.
+    return new ZExtInst(Overflow, Ty);
+  }
+
+  return nullptr;
+}
+
 Instruction *InstCombinerImpl::visitLShr(BinaryOperator &I) {
   if (Value *V = SimplifyLShrInst(I.getOperand(0), I.getOperand(1), I.isExact(),
                                   SQ.getWithInstruction(&I)))
     return replaceInstUsesWith(I, V);
 
   if (Instruction *X = foldVectorBinop(I))
     return X;
 
if (match(Op1, m_APInt(ShAmtAPInt))) {
     // If the shifted-out value is known-zero, then this is an exact shift.
     if (!I.isExact() &&
         MaskedValueIsZero(Op0, APInt::getLowBitsSet(BitWidth, ShAmt), 0, &I)) {
       I.setIsExact();
       return &I;
     }
   }
 
+  if (Instruction *Overflow = foldXShrToOverflow(I))
+    return Overflow;
+
   // Transform  (x << y) >> y  to  x & (-1 >> y)
   Value *X;
   if (match(Op0, m_OneUse(m_Shl(m_Value(X), m_Specific(Op1))))) {
     Constant *AllOnes = ConstantInt::getAllOnesValue(Ty);
     Value *Mask = Builder.CreateLShr(AllOnes, Op1);
     return BinaryOperator::CreateAnd(Mask, X);
   }
 

llvm/test/Transforms/InstCombine/shift-add.ll

 ;
   %A = zext i16 %I to i32
   %B = insertelement <4 x i32> undef, i32 %A, i32 0
   %C = shufflevector <4 x i32> %B, <4 x i32> undef, <4 x i32> zeroinitializer
   %D = add <4 x i32> %C, <i32 0, i32 1, i32 50, i32 16>
   %E = lshr <4 x i32> <i32 6, i32 2, i32 1, i32 -7>, %D
   ret <4 x i32> %E
 }
+
+define i64 @lshr_32_add_zext_basic(i32 %a, i32 %b) {
+; CHECK-LABEL: define i64 @lshr_32_add_zext_basic(
+; CHECK-NEXT: [[uaddo:%.*]] = call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 %a, i32 %b)
+; CHECK-NEXT: [[overflow:%.*]] = extractvalue { i32, i1 } [[uaddo]], 1
+; CHECK-NEXT: [[zextOverflow:%.*]] = zext i1 %uaddo.overflow to i64
+; CHECK-NEXT: ret i64 [[zextOverflow]]
+
+  %zext.a = zext i32 %a to i64
+  %zext.b = zext i32 %b to i64
+  %add = add i64 %zext.a, %zext.b
+  %lshr = lshr i64 %add, 32
+  ret i64 %lshr
+}
+
+define i64 @ashr_32_add_zext_basic(i32 %a, i32 %b) {
+; CHECK-LABEL: define i64 @ashr_32_add_zext_basic(
+; CHECK-NEXT: [[uaddo:%.*]] = call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 %a, i32 %b)
+; CHECK-NEXT: [[overflow:%.*]] = extractvalue { i32, i1 } [[uaddo]], 1
+; CHECK-NEXT: [[zextOverflow:%.*]] = zext i1 %uaddo.overflow to i64
+; CHECK-NEXT: ret i64 [[zextOverflow]]
+
+  %zext.a = zext i32 %a to i64
+  %zext.b = zext i32 %b to i64
+  %add = add i64 %zext.a, %zext.b
+  %lshr = ashr i64 %add, 32
+  ret i64 %lshr
+}
+
+define i32 @lshr_32_add_zext_trunc(i32 %a, i32 %b) {
+; CHECK-LABEL: define i32 @lshr_32_add_zext_trunc(
+; CHECK-NEXT: [[uaddo:%.*]] = call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 %a, i32 %b)
+; CHECK-NEXT: [[add:%.*]] = extractvalue { i32, i1 } [[uaddo]], 0
+; CHECK-NEXT: [[overflow:%.*]] = extractvalue { i32, i1 } [[uaddo]], 1
+; CHECK-NEXT: [[zextOverflow:%.*]] = zext i1 [[overflow]] to i32
+; CHECK-NEXT: [[ret:%.*]] = add i32 [[add]], [[zextOverflow]]
+; CHECK-NEXT: ret i32 [[ret]]
+
+  %zext.a = zext i32 %a to i64
+  %zext.b = zext i32 %b to i64
+  %add = add i64 %zext.a, %zext.b
+  %trunc.add = trunc i64 %add to i32
+  %shr = lshr i64 %add, 32
+  %trunc.shr = trunc i64 %shr to i32
+  %ret = add i32 %trunc.add, %trunc.shr
+  ret i32 %ret
+}