[DAG] Fold (umin (sub a b) a) -> (usubo a b); (select usubo.1 a usubo.0)

[ckoparkar]

Fixes #161036.

[ckoparkar]

/cc @RKSimon

[llvmbot]

@llvm/pr-subscribers-backend-aarch64
@llvm/pr-subscribers-llvm-selectiondag

@llvm/pr-subscribers-backend-x86

Author: Chaitanya Koparkar (ckoparkar)

Changes

Fixes #161036.

---

Full diff: https://github.com/llvm/llvm-project/pull/161651.diff

2 Files Affected:

• (modified) llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp (+22)
• (added) llvm/test/CodeGen/X86/underflow-compare-fold.ll (+16)

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 558c5a0390228..99d7000c3b62e 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -6199,6 +6199,28 @@ SDValue DAGCombiner::visitIMINMAX(SDNode *N) {
                                         SDLoc(N), VT, N0, N1))
     return SD;
 
+  // (umin (sub a, b) a) -> (usubo a, b); (select usubo.1, a, usubo.0)
+  //
+  // IR:
+  //   %sub  = sub %a, %b
+  //   %cond = umin %sub, %a
+  //   ->
+  //   %usubo    = usubo %a, %b
+  //   %overflow = extractvalue %usubo, 1
+  //   %sub      = extractvalue %usubo, 0
+  //   %cond     = select %overflow, %a, %sub
+  if (N0.getOpcode() == ISD::SUB) {
+    SDValue A, B, C;
+    if (sd_match(N, m_UMin(m_Sub(m_Value(A), m_Value(B)), m_Value(C)))) {
+      EVT AVT = A.getValueType();
+      if (A == C && TLI.isOperationLegalOrCustom(ISD::USUBO, AVT)) {
+        SDVTList VTs = DAG.getVTList(AVT, MVT::i1);
+        SDValue USO = DAG.getNode(ISD::USUBO, DL, VTs, A, B);
+        return DAG.getSelect(DL, VT, USO.getValue(1), A, USO.getValue(0));
+      }
+    }
+  }
+
   // Simplify the operands using demanded-bits information.
   if (SimplifyDemandedBits(SDValue(N, 0)))
     return SDValue(N, 0);
diff --git a/llvm/test/CodeGen/X86/underflow-compare-fold.ll b/llvm/test/CodeGen/X86/underflow-compare-fold.ll
new file mode 100644
index 0000000000000..2416bcb909485
--- /dev/null
+++ b/llvm/test/CodeGen/X86/underflow-compare-fold.ll
@@ -0,0 +1,16 @@
+; RUN: llc < %s -mtriple=x86_64 | FileCheck %s
+
+; GitHub issue #161036
+
+define i64 @subIfNoUnderflow_umin(i64 %a, i64 %b) {
+; CHECK-LABEL: subIfNoUnderflow_umin
+; CHECK-LABEL: %bb.0
+; CHECK-NEXT: movq    %rdi, %rax
+; CHECK-NEXT: subq    %rsi, %rax
+; CHECK-NEXT: cmovbq  %rdi, %rax
+; CHECK-NEXT: retq
+entry:
+  %sub = sub i64 %a, %b
+  %cond = tail call i64 @llvm.umin.i64(i64 %sub, i64 %a)
+  ret i64 %cond
+}

[RKSimon]

MVT::i1 ?

[ckoparkar]

It'll always be MVT::i1 unless umin is applied to a vector? And the vector case is already handled here, so we shouldn't see a vector type after that line.

[RKSimon]

please add aarch64 test coverage

[ckoparkar]

@RKSimon this is ready for another review.

[arsenm]

Worried this might break after legalize; probably should use getSetCCResultType instead of i1

[arsenm]

Broken for vectors?

[arsenm]

You're matching the opcode here and in the pattern, so pick one way or the other

[arsenm]

Test same pattern with vectors. Also add multiple use test, and illegal type

[ckoparkar]

Apologies, I was traveling for a few days. I'll update this today.