[llvm] [AArch64] Combine SEXT_INREG(CSET) to CSETM. (PR #156429)

[Ricardo Jesus via llvm-commits]

https://github.com/rj-jesus created https://github.com/llvm/llvm-project/pull/156429

Add the following patterns to performSignExtendInRegCombine:
* SIGN_EXTEND_INREG (CSEL 0, 1, cc), i1 --> CSEL 0, -1, cc
* SIGN_EXTEND_INREG (CSEL 1, 0, cc), i1 --> CSEL -1, 0, cc

The combined forms can be matched to a CSETM.

I'd like to eventually propose enabling `convertSelectOfConstantsToMath`, assuming others think that would be a good idea, primarily to improve a few cases such as [this](https://godbolt.org/z/hz744KzW1). However, doing so currently introduces a few code gen regressions. This patch fixes some of them, and hopefully is a worthwhile addition on its own (please let me know otherwise).

>From d84b0c3cd9effbf371b3449a4a82780063a69d39 Mon Sep 17 00:00:00 2001
From: Ricardo Jesus <rjj at nvidia.com>
Date: Mon, 1 Sep 2025 01:32:29 -0700
Subject: [PATCH] [AArch64] Combine SEXT_INREG(CSET) to CSETM.

Add the following patterns to performSignExtendInRegCombine:
  SIGN_EXTEND_INREG (CSEL 0, 1, cc), i1 --> CSEL 0, -1, cc
  SIGN_EXTEND_INREG (CSEL 1, 0, cc), i1 --> CSEL -1, 0, cc

The combined forms can be matched to a CSETM.
---
 .../Target/AArch64/AArch64ISelLowering.cpp    | 20 +++++++++++++++++++
 llvm/test/CodeGen/AArch64/arm64-ccmp.ll       |  6 +++---
 .../CodeGen/AArch64/extract-vector-cmp.ll     |  4 +---
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index b7011e0ea1669..bd79b43c8e484 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -26585,6 +26585,26 @@ performSignExtendInRegCombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI,
     return DAG.getNode(SOpc, DL, N->getValueType(0), Ext);
   }
 
+  // Sign extend of CSET -> CSETM.
+  if (Opc == AArch64ISD::CSEL &&
+      cast<VTSDNode>(N->getOperand(1))->getVT() == MVT::i1) {
+    EVT VT = N->getValueType(0);
+    SDValue TVal = Src.getOperand(0);
+    SDValue FVal = Src.getOperand(1);
+
+    // SIGN_EXTEND_INREG (CSEL 0, 1, cc), i1 --> CSEL 0, -1, cc
+    if (isNullConstant(TVal) && isOneConstant(FVal))
+      return DAG.getNode(AArch64ISD::CSEL, DL, VT, TVal,
+                         DAG.getAllOnesConstant(DL, VT), Src.getOperand(2),
+                         Src.getOperand(3));
+
+    // SIGN_EXTEND_INREG (CSEL 1, 0, cc), i1 --> CSEL -1, 0, cc
+    if (isOneConstant(TVal) && isNullConstant(FVal))
+      return DAG.getNode(AArch64ISD::CSEL, DL, VT,
+                         DAG.getAllOnesConstant(DL, VT), FVal,
+                         Src.getOperand(2), Src.getOperand(3));
+  }
+
   if (DCI.isBeforeLegalizeOps())
     return SDValue();
 
diff --git a/llvm/test/CodeGen/AArch64/arm64-ccmp.ll b/llvm/test/CodeGen/AArch64/arm64-ccmp.ll
index 4fe01e838771d..cad5df0d9655e 100644
--- a/llvm/test/CodeGen/AArch64/arm64-ccmp.ll
+++ b/llvm/test/CodeGen/AArch64/arm64-ccmp.ll
@@ -632,11 +632,11 @@ define i64 @select_noccmp2(i64 %v1, i64 %v2, i64 %v3, i64 %r) {
 ; CHECK-SD-NEXT:    cmp x0, #0
 ; CHECK-SD-NEXT:    ccmp x0, #13, #0, pl
 ; CHECK-SD-NEXT:    cset w8, gt
+; CHECK-SD-NEXT:    csetm w9, gt
 ; CHECK-SD-NEXT:    cmp w8, #0
 ; CHECK-SD-NEXT:    csel x0, xzr, x3, ne
-; CHECK-SD-NEXT:    sbfx w8, w8, #0, #1
-; CHECK-SD-NEXT:    adrp x9, _g at PAGE
-; CHECK-SD-NEXT:    str w8, [x9, _g at PAGEOFF]
+; CHECK-SD-NEXT:    adrp x8, _g at PAGE
+; CHECK-SD-NEXT:    str w9, [x8, _g at PAGEOFF]
 ; CHECK-SD-NEXT:    ret
 ;
 ; CHECK-GI-LABEL: select_noccmp2:
diff --git a/llvm/test/CodeGen/AArch64/extract-vector-cmp.ll b/llvm/test/CodeGen/AArch64/extract-vector-cmp.ll
index f076ee12427d8..832e34b664fbe 100644
--- a/llvm/test/CodeGen/AArch64/extract-vector-cmp.ll
+++ b/llvm/test/CodeGen/AArch64/extract-vector-cmp.ll
@@ -142,14 +142,12 @@ for.cond.cleanup:
 }
 
 
-; TODO: Combine the sbfx(cset) into a csetm
 define i32 @issue_121372(<4 x i32> %v) {
 ; CHECK-LABEL: issue_121372:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    fmov w8, s0
 ; CHECK-NEXT:    cmp w8, #0
-; CHECK-NEXT:    cset w8, eq
-; CHECK-NEXT:    sbfx w8, w8, #0, #1
+; CHECK-NEXT:    csetm w8, eq
 ; CHECK-NEXT:    cmp w8, #1
 ; CHECK-NEXT:    csetm w0, lt
 ; CHECK-NEXT:    ret