[InstCombine] Expand usub_sat patterns to handle constants

[dmgreen]

Summary

The constants come through as add %x, -C, not a sub as would be expected. They need some extra matchers to canonicalise them towards usub_sat.

----------------------------------------
Optimization: Forwards
Precondition: true
  %cmp = icmp ugt i32 %a, 10
  %sub = add i32 %a, -10
  %sel = select i1 %cmp, i32 %sub, i32 0
=>
  %sel = usub_sat %a, 10

Done: 1
Optimization is correct!

----------------------------------------
Optimization: Backwards
Precondition: true
  %cmp = icmp ult i32 %a, 2
  %sub = add i32 %a, -2
  %sel = select i1 %cmp, i32 %sub, i32 0
=>
  %x = usub_sat 2, %a
  %sel = sub 0, %x

Done: 1
Optimization is correct!

[dmgreen]

Looks like there should be some extra patterns for "one-off" constants too.

[lebedev.ri]

Thank you for looking into this.

llvm/test/Transforms/InstCombine/unsigned_saturated_sub.ll

357–358

In this pattern either one of these two instructions must be one-use and go away.

[dmgreen]

Sorry for the delay. I was trying to look into the "off by one" patterns, the last two of these:

Name: Forwards
Pre: C1 == -C2
  %cmp = icmp ugt i8 %a, C1
  %sub = add i8 %a, C2
  %sel = select i1 %cmp, i8 %sub, i8 0
=>
  %sel = usub_sat %a, C1

Name: Backwards
Pre: C1 == -C2
  %cmp = icmp ult i8 %a, C1
  %sub = add i8 %a, C2
  %sel = select i1 %cmp, i8 %sub, i8 0
=>
  %x = usub_sat C1, %a
  %sel = sub 0, %x

Name: Forwards_off
Pre: C1 + 1 == -C2 && C2 != 0
  %cmp = icmp ugt i8 %a, C1
  %sub = add i8 %a, C2
  %sel = select i1 %cmp, i8 %sub, i8 0
=>
  %sel = usub_sat %a, C1 + 1

Name: Backwards_off
Pre: C1 - 1 == -C2 && C1 != 0
  %cmp = icmp ult i8 %a, C1
  %sub = add i8 %a, C2
  %sel = select i1 %cmp, i8 %sub, i8 0
=>
  %x = usub_sat C1 - 1, %a
  %sel = sub 0, %x

They are a bit fiddly though and deserve to be their own commit.

[lebedev.ri]

(marking as reviewed)

Will you also look into the edge-case patterns, like:

unsigned t0(unsigned num) {
    return num > 0 ? num-1 : num;
}

https://godbolt.org/z/DsBxd8

Name: unsigned minus 1
  %2 = icmp ugt i32 %0, 1
  %3 = add i32 %0, -1
  %4 = select i1 %2, i32 %3, i32 0
  ret i32 %4
=>
  %r = usub_sat %0, 1
  ret %r
  %3 = add i32 %0, -1
  %2 = icmp ugt i32 %0, 1
  %4 = select i1 %2, i32 %3, i32 0

Done: 1
Optimization is correct!

(not sure what other ones are there, this is the one that i noticed in real code..)

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

722–724

Err, this isn't quite right. I should have been more specific.

We produce a single instruction @llvm.usub.sat.,
so there shouldn't be a one-use check in general.

But we also may need to negate the result.
So *if* we need to negate the result,
*then* we need ensure that a *single* instruction
goes away - either icmp or sub.

[dmgreen]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

722–724

OK. I see. I was thinking you were referring to assembly instructions being longer if the icmp had multiple uses.

There may be some benefit to making sure the output is actually smaller.

I've updated the code. Let me know if this was or wasn't what you were thinking.

[dmgreen]

Sorry for the delay. Adjust this one use check.

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

707–710

I'm having doubt about the comment.
If b < a, then b - a will wrap, and we'll return wrapped result.
And if b >= a, then b - a will not wrap, but we will saturate to 0.
What am i missing?

722–724

Yep, this seems what i meant.

[dmgreen]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

707–710

From the comment below, this looks like it is referring to
(a > b) ? b - a : 0 -> -usub.sat(a, b)
I'll try and update it to be clearer.

[lebedev.ri]

LGTM, thank you.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

707–710

Aha, okay, thank you.

722–725

I would recommend commit this as a separate commit first.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

if (Pred == ICmpInst::ICMP_ULE || Pred == ICmpInst::ICMP_ULT) {
     // (b < a) ? a - b : 0 -> (a > b) ? a - b : 0
     std::swap(A, B);
     Pred = ICmpInst::getSwappedPredicate(Pred);
   }
 
   assert((Pred == ICmpInst::ICMP_UGE || Pred == ICmpInst::ICMP_UGT) &&
          "Unexpected isUnsigned predicate!");
 
-  // Account for swapped form of subtraction: ((a > b) ? b - a : 0).
+  // Ensure the sub is of the form:
+  //  (a > b) ? a - b : 0 -> usub.sat(a, b)
+  //  (a > b) ? b - a : 0 -> -usub.sat(a, b)
+  // Checking for both a-b and a+(-b) as a constant.
   bool IsNegative = false;
-  if (match(TrueVal, m_Sub(m_Specific(B), m_Specific(A))))
+  const APInt *C;
+  if (match(TrueVal, m_Sub(m_Specific(B), m_Specific(A))) ||
+      (match(A, m_APInt(C)) &&
+       match(TrueVal, m_Add(m_Specific(B), m_SpecificInt(-*C)))))
     IsNegative = true;
-  else if (!match(TrueVal, m_Sub(m_Specific(A), m_Specific(B))))
+  else if (!match(TrueVal, m_Sub(m_Specific(A), m_Specific(B))) &&
+           !(match(B, m_APInt(C)) &&
+             match(TrueVal, m_Add(m_Specific(A), m_SpecificInt(-*C)))))
     return nullptr;
 
   // If we are adding a negate and the sub and icmp are used anywhere else, we
   // would end up with more instructions.
   if (IsNegative && !TrueVal->hasOneUse() && !ICI->hasOneUse())
     return nullptr;
 
   // (a > b) ? a - b : 0 -> usub.sat(a, b)
   // (a > b) ? b - a : 0 -> -usub.sat(a, b)
   Value *Result = Builder.CreateBinaryIntrinsic(Intrinsic::usub_sat, A, B);
   if (IsNegative)
     Result = Builder.CreateNeg(Result);
   return Result;
 }

llvm/test/Transforms/InstCombine/builtin-dynamic-object-size.ll

entry:
   %ptr = call i8* @malloc(i64 %sz)
   %ptr2 = getelementptr inbounds i8, i8* %ptr, i32 2
   %calc_size = call i64 @llvm.objectsize.i64.p0i8(i8* %ptr2, i1 false, i1 true, i1 true)
   ret i64 %calc_size
 }
 
 ; CHECK:      define i64 @internal_pointer(i64 %sz)
 ; CHECK-NEXT: entry:
-; CHECK-NEXT:   %0 = add i64 %sz, -2
-; CHECK-NEXT:   %1 = icmp ult i64 %sz, 2
-; CHECK-NEXT:   %2 = select i1 %1, i64 0, i64 %0
-; CHECK-NEXT:   ret i64 %2
+; CHECK-NEXT:   %0 = call i64 @llvm.usub.sat.i64(i64 %sz, i64 2)
+; CHECK-NEXT:   ret i64 %0
 ; CHECK-NEXT: }
 
 define i64 @uses_nullptr_no_fold() {
 entry:
   %res = call i64 @llvm.objectsize.i64.p0i8(i8* null, i1 false, i1 true, i1 true)
   ret i64 %res
 }
 

llvm/test/Transforms/InstCombine/unsigned_saturated_sub.ll

;
   %sel = select i1 %cmp, i64 0 ,i64 %sub
   %extrasub = sub i64 %a, %b
   call void @use(i64 %extrasub)
   ret i64 %sel
 }
 
 define i32 @max_sub_ugt_c1(i32 %a) {
 ; CHECK-LABEL: @max_sub_ugt_c1(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i32 [[A:%.*]], 1
-; CHECK-NEXT:    [[SUB:%.*]] = add i32 [[A]], -1
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i32 [[SUB]], i32 0
-; CHECK-NEXT:    ret i32 [[SEL]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @llvm.usub.sat.i32(i32 [[A:%.*]], i32 1)
+; CHECK-NEXT:    ret i32 [[TMP1]]
 ;
   %cmp = icmp ugt i32 %a, 1
   %sub = add i32 %a, -1
   %sel = select i1 %cmp, i32 %sub ,i32 0
   ret i32 %sel
 }
 
 define i32 @max_sub_ugt_c01(i32 %a) {
 ; CHECK-LABEL: @max_sub_ugt_c01(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[A:%.*]], 0
 ; CHECK-NEXT:    [[SUB:%.*]] = add i32 [[A]], -1
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i32 0, i32 [[SUB]]
 ; CHECK-NEXT:    ret i32 [[SEL]]
 ;
   %cmp = icmp ugt i32 %a, 0
   %sub = add i32 %a, -1
   %sel = select i1 %cmp, i32 %sub ,i32 0
   ret i32 %sel
 }
 
 define i32 @max_sub_ugt_c10(i32 %a) {
 ; CHECK-LABEL: @max_sub_ugt_c10(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i32 [[A:%.*]], 10
-; CHECK-NEXT:    [[SUB:%.*]] = add i32 [[A]], -10
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i32 [[SUB]], i32 0
-; CHECK-NEXT:    ret i32 [[SEL]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @llvm.usub.sat.i32(i32 [[A:%.*]], i32 10)
+; CHECK-NEXT:    ret i32 [[TMP1]]
 ;
   %cmp = icmp ugt i32 %a, 10
   %sub = add i32 %a, -10
   %sel = select i1 %cmp, i32 %sub, i32 0
   ret i32 %sel
 }
 
 define i32 @max_sub_ugt_c910(i32 %a) {
;
   %cmp = icmp ult i32 %a, 1
   %sub = add i32 %a, -1
   %sel = select i1 %cmp, i32 %sub, i32 0
   ret i32 %sel
 }
 
 define i32 @max_sub_ult_c2(i32 %a) {
 ; CHECK-LABEL: @max_sub_ult_c2(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[A:%.*]], 2
-; CHECK-NEXT:    [[SUB:%.*]] = add i32 [[A]], -2
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i32 [[SUB]], i32 0
-; CHECK-NEXT:    ret i32 [[SEL]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @llvm.usub.sat.i32(i32 2, i32 [[A:%.*]])
+; CHECK-NEXT:    [[TMP2:%.*]] = sub nsw i32 0, [[TMP1]]
+; CHECK-NEXT:    ret i32 [[TMP2]]
 ;
   %cmp = icmp ult i32 %a, 2
   %sub = add i32 %a, -2
   %sel = select i1 %cmp, i32 %sub, i32 0
   ret i32 %sel
 }
 
 define i32 @max_sub_ult_c2_oneuseicmp(i32 %a) {
 ; CHECK-LABEL: @max_sub_ult_c2_oneuseicmp(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[A:%.*]], 2
-; CHECK-NEXT:    [[SUB:%.*]] = add i32 [[A]], -2
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i32 [[SUB]], i32 0
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @llvm.usub.sat.i32(i32 2, i32 [[A]])
+; CHECK-NEXT:    [[TMP2:%.*]] = sub nsw i32 0, [[TMP1]]
 ; CHECK-NEXT:    call void @usei1(i1 [[CMP]])
-; CHECK-NEXT:    ret i32 [[SEL]]
+; CHECK-NEXT:    ret i32 [[TMP2]]
 ;
   %cmp = icmp ult i32 %a, 2
   %sub = add i32 %a, -2
   %sel = select i1 %cmp, i32 %sub, i32 0
   call void @usei1(i1 %cmp)
   ret i32 %sel
 }
 
 define i32 @max_sub_ult_c2_oneusesub(i32 %a) {
 ; CHECK-LABEL: @max_sub_ult_c2_oneusesub(
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[A:%.*]], 2
-; CHECK-NEXT:    [[SUB:%.*]] = add i32 [[A]], -2
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[CMP]], i32 [[SUB]], i32 0
+; CHECK-NEXT:    [[SUB:%.*]] = add i32 [[A:%.*]], -2
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @llvm.usub.sat.i32(i32 2, i32 [[A]])
+; CHECK-NEXT:    [[TMP2:%.*]] = sub nsw i32 0, [[TMP1]]
 ; CHECK-NEXT:    call void @usei32(i32 [[SUB]])
-; CHECK-NEXT:    ret i32 [[SEL]]
+; CHECK-NEXT:    ret i32 [[TMP2]]
 ;
   %cmp = icmp ult i32 %a, 2
   %sub = add i32 %a, -2
   %sel = select i1 %cmp, i32 %sub, i32 0
   call void @usei32(i32 %sub)
   ret i32 %sel
 }