[llvm] [RISCV] Simplify (srl (and X, Mask), Const) to TH_EXTU (PR #102802)

[Wang Yaduo via llvm-commits]

https://github.com/MouseSplinter created https://github.com/llvm/llvm-project/pull/102802

None

>From 4a93fe2b6711c874fd7a746935730c03652a0e6b Mon Sep 17 00:00:00 2001
From: Wang Yaduo <wangyaduo at linux.alibaba.com>
Date: Sun, 11 Aug 2024 15:55:37 +0800
Subject: [PATCH] [RISCV] Simplify (srl (and X, Mask), Const) to TH_EXTU

---
 llvm/lib/Target/RISCV/RISCVISelDAGToDAG.cpp |  9 +++++++++
 llvm/test/CodeGen/RISCV/ctlz-cttz-ctpop.ll  | 14 ++++++--------
 llvm/test/CodeGen/RISCV/rv64xtheadbb.ll     | 16 ++++++++++++++++
 3 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVISelDAGToDAG.cpp b/llvm/lib/Target/RISCV/RISCVISelDAGToDAG.cpp
index ce3a37e194d545..63444d0dbe17a7 100644
--- a/llvm/lib/Target/RISCV/RISCVISelDAGToDAG.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelDAGToDAG.cpp
@@ -1164,6 +1164,15 @@ void RISCVDAGToDAGISel::Select(SDNode *Node) {
     }
 
     unsigned LShAmt = Subtarget->getXLen() - TrailingOnes;
+    if (Subtarget->hasVendorXTHeadBb() && TrailingOnes > ShAmt) {
+      SDNode *THEXTU = CurDAG->getMachineNode(
+          RISCV::TH_EXTU, DL, VT, N0->getOperand(0),
+          CurDAG->getTargetConstant(TrailingOnes - 1, DL, VT),
+          CurDAG->getTargetConstant(ShAmt, DL, VT));
+      ReplaceNode(Node, THEXTU);
+      return;
+    }
+
     SDNode *SLLI =
         CurDAG->getMachineNode(RISCV::SLLI, DL, VT, N0->getOperand(0),
                                CurDAG->getTargetConstant(LShAmt, DL, VT));
diff --git a/llvm/test/CodeGen/RISCV/ctlz-cttz-ctpop.ll b/llvm/test/CodeGen/RISCV/ctlz-cttz-ctpop.ll
index 8caa64c9572ce7..2c691a2de4c4de 100644
--- a/llvm/test/CodeGen/RISCV/ctlz-cttz-ctpop.ll
+++ b/llvm/test/CodeGen/RISCV/ctlz-cttz-ctpop.ll
@@ -2367,10 +2367,9 @@ define i16 @test_ctpop_i16(i16 %a) nounwind {
 ; RV32XTHEADBB-NEXT:    add a0, a2, a0
 ; RV32XTHEADBB-NEXT:    srli a1, a0, 4
 ; RV32XTHEADBB-NEXT:    add a0, a0, a1
-; RV32XTHEADBB-NEXT:    andi a1, a0, 15
-; RV32XTHEADBB-NEXT:    slli a0, a0, 20
-; RV32XTHEADBB-NEXT:    srli a0, a0, 28
-; RV32XTHEADBB-NEXT:    add a0, a1, a0
+; RV32XTHEADBB-NEXT:    th.extu a1, a0, 11, 8
+; RV32XTHEADBB-NEXT:    andi a0, a0, 15
+; RV32XTHEADBB-NEXT:    add a0, a0, a1
 ; RV32XTHEADBB-NEXT:    ret
 ;
 ; RV64XTHEADBB-LABEL: test_ctpop_i16:
@@ -2388,10 +2387,9 @@ define i16 @test_ctpop_i16(i16 %a) nounwind {
 ; RV64XTHEADBB-NEXT:    add a0, a2, a0
 ; RV64XTHEADBB-NEXT:    srli a1, a0, 4
 ; RV64XTHEADBB-NEXT:    add a0, a0, a1
-; RV64XTHEADBB-NEXT:    andi a1, a0, 15
-; RV64XTHEADBB-NEXT:    slli a0, a0, 52
-; RV64XTHEADBB-NEXT:    srli a0, a0, 60
-; RV64XTHEADBB-NEXT:    add a0, a1, a0
+; RV64XTHEADBB-NEXT:    th.extu a1, a0, 11, 8
+; RV64XTHEADBB-NEXT:    andi a0, a0, 15
+; RV64XTHEADBB-NEXT:    add a0, a0, a1
 ; RV64XTHEADBB-NEXT:    ret
   %1 = call i16 @llvm.ctpop.i16(i16 %a)
   ret i16 %1
diff --git a/llvm/test/CodeGen/RISCV/rv64xtheadbb.ll b/llvm/test/CodeGen/RISCV/rv64xtheadbb.ll
index 6cdab888ffcde7..8ce4c44d8efcc2 100644
--- a/llvm/test/CodeGen/RISCV/rv64xtheadbb.ll
+++ b/llvm/test/CodeGen/RISCV/rv64xtheadbb.ll
@@ -702,6 +702,22 @@ define i64 @zext_bf_i64(i64 %a) nounwind {
   ret i64 %and
 }
 
+define i64 @zext_bf2_i64(i64 %a) nounwind {
+; RV64I-LABEL: zext_bf2_i64:
+; RV64I:       # %bb.0:
+; RV64I-NEXT:    slli a0, a0, 48
+; RV64I-NEXT:    srli a0, a0, 49
+; RV64I-NEXT:    ret
+;
+; RV64XTHEADBB-LABEL: zext_bf2_i64:
+; RV64XTHEADBB:       # %bb.0:
+; RV64XTHEADBB-NEXT:    th.extu a0, a0, 15, 1
+; RV64XTHEADBB-NEXT:    ret
+  %t0 = and i64 %a, 65535
+  %result = lshr i64 %t0, 1
+  ret i64 %result
+}
+
 define i64 @zext_i64_srliw(i64 %a) nounwind {
 ; RV64I-LABEL: zext_i64_srliw:
 ; RV64I:       # %bb.0: