[InstCombine] Canonicalize uadd.with.overflow to uadd.sat

[dmgreen]

Summary

This adds some patterns to transform uadd.with.overflow to uadd.sat (with usub.with.overflow to usub.sat too). The patterns selects from UINTMAX (or 0 for subs) depending on whether the operation overflowed.

Signed patterns are a little more involved (they can wrap in two directions), but can be added here in a followup patch too.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

1744

We might want to add an m_ExtractValue<0>(Matcher) matcher...

[dmgreen]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

1744

Righteo.

I've only made it single index, as opposed to try and get fancy with some recursive templates. That is likely the most common case. I've also added a m_BinOpIntrinsic to more easily grab the BinaryOpIntrinsic.

[dmgreen]

Added a m_ExtractValue matcher.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

1744

I think it would make more sense to make this a matcher for WithOverflowInst intrinsics -- that's what's needed here, and matching both saturating and with.overflow in one matcher doesn't seem super useful in general.

[dmgreen]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

1744

Oh yeah. That sounds like what I need. I could also add a generic "match one of these things" operator. Something like

diff --git a/llvm/include/llvm/IR/PatternMatch.h b/llvm/include/llvm/IR/PatternMatch.h
index 0cf8e62..649b047 100644
--- a/llvm/include/llvm/IR/PatternMatch.h
+++ b/llvm/include/llvm/IR/PatternMatch.h
@@ -559,8 +559,9 @@ inline bind_ty<const Value> m_Value(const Value *&V) { return V; }
 inline bind_ty<Instruction> m_Instruction(Instruction *&I) { return I; }
 /// Match a binary operator, capturing it if we match.
 inline bind_ty<BinaryOperator> m_BinOp(BinaryOperator *&I) { return I; }
-/// Match a binary operator intrinsic, capturing it if we match.
-inline bind_ty<BinaryOpIntrinsic> m_BinOpIntrinsic(BinaryOpIntrinsic *&I) { return I; }
+/// Match a class of type T, capturing it if we match.
+template <typename T>
+inline bind_ty<T> m_Instruction(T *&I) { return I; }
 
 /// Match a ConstantInt, capturing the value if we match.
 inline bind_ty<ConstantInt> m_ConstantInt(ConstantInt *&CI) { return CI; }
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstC
index f3b0ac5..46bfb81 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -1743,8 +1743,8 @@ foldOverflowingAddSubSelect(SelectInst &SI, InstCombiner::BuilderTy &Bui
   Value *TrueVal = SI.getTrueValue();
   Value *FalseVal = SI.getFalseValue();
 
-  BinaryOpIntrinsic *II;
-  if (!match(CondVal, m_ExtractValue<1>(m_BinOpIntrinsic(II))) ||
+  WithOverflowInst *II;
+  if (!match(CondVal, m_ExtractValue<1>(m_Instruction(II))) ||
       !match(FalseVal, m_ExtractValue<0>(m_Specific(II))))
     return nullptr;

What do you think? Is there a better name than m_Instruction? (It's better than my original m_OneOfTheseThings. I kind of feel like there should be some way of doing this already that I might be missing).

[dmgreen]

Update to m_WithOverflowInst for now

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

1744

I like the idea. Overloading m_Instruction() for this purpose is probably not good though. Generally I'd prefer something that mentions the type on the matcher itself, i.e. m_Instruction<WithOverflowInst>(II), though I guess that's more of usage style question, as it probably can't be enforced in the implementation.

Some ideas...

m_a<WithOverflowInst>(II)
m_isa<WithOverflowInst>(II)
m_InstanceOf<WithOverflowInst>(II)
m_OfType<WithOverflowInst>(II)

[dmgreen]

I've gone with m_isa, which I think with the explicit WithOverflowInst looks OK.

[nikic]

LGTM, though I'd like a second opinion on the PatternMatch additions from @lebedev.ri or @spatel.

[lebedev.ri]

Looks good in general but i have reservations about matchers.

llvm/include/llvm/IR/PatternMatch.h

561–563 ↗

(On Diff #226592)

I don't like this one.
Can we change it to be a two-step process - just make m_isa<>() a predicate,
if isa<>() matches, then it should apply inner matcher, if any.
And said inner matcher can be m_Value().

[nikic]

llvm/include/llvm/IR/PatternMatch.h

561–563 ↗

(On Diff #226592)

I'm not sure I understand how that would be used. If you use something like m_isa<WithOverflowInst>(m_Value(V)), wouldn't V have to be Value * rather than WithOverflowInst *? That seems less useful to me. It's also not clear what things apart from m_Value() you would use as the inner matcher.

[lebedev.ri]

llvm/include/llvm/IR/PatternMatch.h

561–563 ↗

(On Diff #226592)

Good point. I'm mainly unhappy about the name.
m_isa<>() really looks similar to normal isa<>(), from looking at it i would not
have guessed it takes reference to the pointer of that type and binds to it if matched.

[dmgreen]

llvm/include/llvm/IR/PatternMatch.h

561–563 ↗

(On Diff #226592)

Right. I did think of that but is seemed OK as it was in the middle of a match statement.

Got a better suggestion, or a preference out of these?

m_a<WithOverflowInst>(II)
m_InstanceOf<WithOverflowInst>(II)
m_OfType<WithOverflowInst>(II)
m_OneOfTheseThings<WithOverflowInst>(II)

[lebedev.ri]

llvm/include/llvm/IR/PatternMatch.h

561–563 ↗

(On Diff #226592)

To be honest, none of those tell me that they capture..
Would m_CaptureIf_isa<>() be too ugly?
@RKSimon @spatel thoughts?

[spatel]

llvm/include/llvm/IR/PatternMatch.h

561–563 ↗

(On Diff #226592)

Catching up after travel...I might have missed something in the earlier inline comments, but we just need a WithOverflowInst sibling to:

inline bind_ty<BinaryOperator> m_BinOp(BinaryOperator *&I) { return I; }

...for this patch, right?

It's not clear to me that a more templated generic binder is needed. Either way, it would be good to split the matcher changes from the instcombine changes, so there's less risk of losing improvements if some part of this needs to be reverted.

[dmgreen]

llvm/include/llvm/IR/PatternMatch.h

561–563 ↗

(On Diff #226592)

Do you mean like this version: https://reviews.llvm.org/D69245?id=226110 :)

We can go back to that. I quite like the ability to capture a thing of type T, but you are correct that it is simpler to just add the m_WithOverflowInst case.

[dmgreen]

Back to m_WithOverflowInst

[lebedev.ri]

LGTM, thank you.

[spatel]

llvm/include/llvm/IR/PatternMatch.h

561–563 ↗

(On Diff #226592)

Yes, that's the simpler solution since we don't have a clear winner for the more generic matcher.
I'd still split this when committing (adding a couple of dedicated tests to unittests/IR/PatternMatch.cpp is even better), but either way, LGTM.

[dmgreen]

llvm/include/llvm/IR/PatternMatch.h

561–563 ↗

(On Diff #226592)

Part1: https://reviews.llvm.org/rG6cfbefbc4a7e. Let me know if that test looks dodgy.

[spatel]

llvm/include/llvm/IR/PatternMatch.h

561–563 ↗

(On Diff #226592)

Thanks! Unit tests LGTM.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

if (OtherAddOp) {
         return RI;
       } else
         return BinaryOperator::CreateAdd(SubOp->getOperand(0), NewSel);
     }
   }
   return nullptr;
 }
 
+/// Turn A, O = uadd.with.overflow X, Y; O ? -1 : A --> uadd_sat X, Y
+/// And A, O = usub.with.overflow X, Y; O ? 0 : A --> usub_sat X, Y
+static Instruction *
+foldOverflowingAddSubSelect(SelectInst &SI, InstCombiner::BuilderTy &Builder) {
+  auto *CondVal = dyn_cast<ExtractValueInst>(SI.getCondition());
+  Value *TrueVal = SI.getTrueValue();
+  auto *FalseVal = dyn_cast<ExtractValueInst>(SI.getFalseValue());
+  if (!CondVal || !FalseVal)
+    return nullptr;
+
+  // Check the extracts point to the same instruction
+  if (CondVal->getNumIndices() != 1 || CondVal->getIndices()[0] != 1)
+    return nullptr;
+  if (FalseVal->getNumIndices() != 1 || FalseVal->getIndices()[0] != 0)
+    return nullptr;
+  if (CondVal->getAggregateOperand() != FalseVal->getAggregateOperand())
+    return nullptr;
+  auto *II = dyn_cast<IntrinsicInst>(CondVal->getAggregateOperand());
+  if (!II)
+    return nullptr;
+
+  Intrinsic::ID NewIntrinsicID;
+  if (II->getIntrinsicID() == Intrinsic::uadd_with_overflow &&
+      match(TrueVal, m_AllOnes()))
+    // A, O = uadd.with.overflow X, Y; O ? -1 : A --> uadd_sat X, Y
+    NewIntrinsicID = Intrinsic::uadd_sat;
+  else if (II->getIntrinsicID() == Intrinsic::usub_with_overflow &&
+           match(TrueVal, m_Zero()))
+    // A, O = usub.with.overflow X, Y; O ? 0 : A --> usub_sat X, Y
+    NewIntrinsicID = Intrinsic::usub_sat;
+  else
+    return nullptr;
+
+  Function *F =
+      Intrinsic::getDeclaration(SI.getModule(), NewIntrinsicID, SI.getType());
+  return CallInst::Create(F, {II->getArgOperand(0), II->getArgOperand(1)});
+}
+
 Instruction *InstCombiner::foldSelectExtConst(SelectInst &Sel) {
   Constant *C;
   if (!match(Sel.getTrueValue(), m_Constant(C)) &&
       !match(Sel.getFalseValue(), m_Constant(C)))
     return nullptr;
 
   Instruction *ExtInst;
   if (!match(Sel.getTrueValue(), m_Instruction(ExtInst)) &&
Instruction *InstCombiner::visitSelectInst(SelectInst &SI) {
 
   // See if we are selecting two values based on a comparison of the two values.
   if (ICmpInst *ICI = dyn_cast<ICmpInst>(CondVal))
     if (Instruction *Result = foldSelectInstWithICmp(SI, ICI))
       return Result;
 
   if (Instruction *Add = foldAddSubSelect(SI, Builder))
     return Add;
+  if (Instruction *Add = foldOverflowingAddSubSelect(SI, Builder))
+    return Add;
 
   // Turn (select C, (op X, Y), (op X, Z)) -> (op X, (select C, Y, Z))
   auto *TI = dyn_cast<Instruction>(TrueVal);
   auto *FI = dyn_cast<Instruction>(FalseVal);
   if (TI && FI && TI->getOpcode() == FI->getOpcode())
     if (Instruction *IV = foldSelectOpOp(SI, TI, FI))
       return IV;
 

llvm/test/Transforms/InstCombine/overflow_to_sat.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -instcombine < %s | FileCheck %s
 
 define i32 @uadd(i32 %x, i32 %y) {
 ; CHECK-LABEL: @uadd(
-; CHECK-NEXT:    [[AO:%.*]] = tail call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 [[X:%.*]], i32 [[Y:%.*]])
-; CHECK-NEXT:    [[O:%.*]] = extractvalue { i32, i1 } [[AO]], 1
-; CHECK-NEXT:    [[A:%.*]] = extractvalue { i32, i1 } [[AO]], 0
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[O]], i32 -1, i32 [[A]]
+; CHECK-NEXT:    [[S:%.*]] = call i32 @llvm.uadd.sat.i32(i32 [[X:%.*]], i32 [[Y:%.*]])
 ; CHECK-NEXT:    ret i32 [[S]]
 ;
   %ao = tail call { i32, i1 } @llvm.uadd.with.overflow.i32(i32 %x, i32 %y)
   %o = extractvalue { i32, i1 } %ao, 1
   %a = extractvalue { i32, i1 } %ao, 0
   %s = select i1 %o, i32 -1, i32 %a
   ret i32 %s
 }
 
 define i32 @usub(i32 %x, i32 %y) {
 ; CHECK-LABEL: @usub(
-; CHECK-NEXT:    [[AO:%.*]] = tail call { i32, i1 } @llvm.usub.with.overflow.i32(i32 [[X:%.*]], i32 [[Y:%.*]])
-; CHECK-NEXT:    [[O:%.*]] = extractvalue { i32, i1 } [[AO]], 1
-; CHECK-NEXT:    [[A:%.*]] = extractvalue { i32, i1 } [[AO]], 0
-; CHECK-NEXT:    [[S:%.*]] = select i1 [[O]], i32 0, i32 [[A]]
+; CHECK-NEXT:    [[S:%.*]] = call i32 @llvm.usub.sat.i32(i32 [[X:%.*]], i32 [[Y:%.*]])
 ; CHECK-NEXT:    ret i32 [[S]]
 ;
   %ao = tail call { i32, i1 } @llvm.usub.with.overflow.i32(i32 %x, i32 %y)
   %o = extractvalue { i32, i1 } %ao, 1
   %a = extractvalue { i32, i1 } %ao, 0
   %s = select i1 %o, i32 0, i32 %a
   ret i32 %s
 }