From 8e0d4bbd2459ce3db1c9e39545342615bdcb32f5 Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper@sifive.com>
Date: Sun, 28 Jul 2024 20:27:35 -0700
Subject: [PATCH] [RISCV] Add isel special case for (and (srl X, c2), c1) ->
 (slli_uw (srli x, c2+c3), c3).

Where c1 is a shifted mask with 32 set bits and c3 trailing zeros.

Fixes #100936.
---
 llvm/lib/Target/RISCV/RISCVISelDAGToDAG.cpp | 12 ++++++++++++
 llvm/test/CodeGen/RISCV/rv64zba.ll          |  6 ++----
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVISelDAGToDAG.cpp b/llvm/lib/Target/RISCV/RISCVISelDAGToDAG.cpp
index eef6ae677ac85..01b2bc08d3ba0 100644
--- a/llvm/lib/Target/RISCV/RISCVISelDAGToDAG.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelDAGToDAG.cpp
@@ -1393,6 +1393,18 @@ void RISCVDAGToDAGISel::Select(SDNode *Node) {
           ReplaceNode(Node, SLLI);
           return;
         }
+        // If we have 32 bits in the mask, we can use SLLI_UW instead of SLLI.
+        if (Trailing > 0 && Leading + Trailing == 32 && C2 + Trailing < XLen &&
+            OneUseOrZExtW && Subtarget->hasStdExtZba()) {
+          SDNode *SRLI = CurDAG->getMachineNode(
+              RISCV::SRLI, DL, VT, X,
+              CurDAG->getTargetConstant(C2 + Trailing, DL, VT));
+          SDNode *SLLI_UW = CurDAG->getMachineNode(
+              RISCV::SLLI_UW, DL, VT, SDValue(SRLI, 0),
+              CurDAG->getTargetConstant(Trailing, DL, VT));
+          ReplaceNode(Node, SLLI_UW);
+          return;
+        }
       }
 
       // Turn (and (shl x, c2), c1) -> (slli (srli x, c3-c2), c3) if c1 is a
diff --git a/llvm/test/CodeGen/RISCV/rv64zba.ll b/llvm/test/CodeGen/RISCV/rv64zba.ll
index 61be5ee458e9d..20a0484464018 100644
--- a/llvm/test/CodeGen/RISCV/rv64zba.ll
+++ b/llvm/test/CodeGen/RISCV/rv64zba.ll
@@ -2962,8 +2962,7 @@ define i64 @srli_slliuw_2(i64 %1) {
 ;
 ; RV64ZBA-LABEL: srli_slliuw_2:
 ; RV64ZBA:       # %bb.0: # %entry
-; RV64ZBA-NEXT:    srli a0, a0, 15
-; RV64ZBA-NEXT:    srli a0, a0, 3
+; RV64ZBA-NEXT:    srli a0, a0, 18
 ; RV64ZBA-NEXT:    slli.uw a0, a0, 3
 ; RV64ZBA-NEXT:    ret
 entry:
@@ -2985,8 +2984,7 @@ define i64 @srli_slliuw_canonical_2(i64 %0) {
 ;
 ; RV64ZBA-LABEL: srli_slliuw_canonical_2:
 ; RV64ZBA:       # %bb.0: # %entry
-; RV64ZBA-NEXT:    srli a0, a0, 15
-; RV64ZBA-NEXT:    srli a0, a0, 3
+; RV64ZBA-NEXT:    srli a0, a0, 18
 ; RV64ZBA-NEXT:    slli.uw a0, a0, 3
 ; RV64ZBA-NEXT:    ret
 entry: