[llvm] [RISCV][GlobalISel] Fix selectShiftMask when shift mask is created from G_AND (PR #89602)

[Hongbin Jin via llvm-commits]

https://github.com/bvlgah created https://github.com/llvm/llvm-project/pull/89602

This patch fixes cases where G_AND creating the shift mask is eliminated if one
of its source operands is a constant, resulting from an incorrect predicate.

>From e24983778c619dd8c1383b253dc30ab2fdc8d0bd Mon Sep 17 00:00:00 2001
From: bvlgah <octopus.busts_0w at icloud.com>
Date: Mon, 22 Apr 2024 18:35:34 +0800
Subject: [PATCH] [RISCV][GlobalISel] Fix selectShiftMask when shift mask is
 created from G_AND

---
 .../RISCV/GISel/RISCVInstructionSelector.cpp  | 16 +++-
 .../instruction-select/shift-rv64.mir         | 86 +++++++++++++++++++
 2 files changed, 101 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp b/llvm/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp
index 3103992a86c099..e31d907a2d9bd7 100644
--- a/llvm/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp
+++ b/llvm/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp
@@ -177,6 +177,20 @@ RISCVInstructionSelector::selectShiftMask(MachineOperand &Root) const {
 
   APInt AndMask;
   Register AndSrcReg;
+  // Try to combine the following pattern (applicable to other shift
+  // instructions as well as 32-bit ones):
+  //
+  //   %4:gprb(s64) = G_AND %3, %2
+  //   %5:gprb(s64) = G_LSHR %1, %4(s64)
+  //
+  // According to RISC-V's ISA manual, SLL, SRL, and SRA ignore other bits than
+  // the lowest log2(XLEN) bits of register rs2. As for the above pattern, if
+  // the lowest bits of register rd and rs2 of G_AND are the same, then it can
+  // be eliminated. Given register rs1 or rs2 holding a constant (the mask),
+  // there are two cases G_ADD can be erased:
+  //
+  // 1. the lowest log2(XLEN) bits of the and mask are all set
+  // 2. the bits of the register being masked are already unset (zero set)
   if (mi_match(ShAmtReg, MRI, m_GAnd(m_Reg(AndSrcReg), m_ICst(AndMask)))) {
     APInt ShMask(AndMask.getBitWidth(), ShiftWidth - 1);
     if (ShMask.isSubsetOf(AndMask)) {
@@ -184,7 +198,7 @@ RISCVInstructionSelector::selectShiftMask(MachineOperand &Root) const {
     } else {
       // SimplifyDemandedBits may have optimized the mask so try restoring any
       // bits that are known zero.
-      KnownBits Known = KB->getKnownBits(ShAmtReg);
+      KnownBits Known = KB->getKnownBits(AndSrcReg);
       if (ShMask.isSubsetOf(AndMask | Known.Zero))
         ShAmtReg = AndSrcReg;
     }
diff --git a/llvm/test/CodeGen/RISCV/GlobalISel/instruction-select/shift-rv64.mir b/llvm/test/CodeGen/RISCV/GlobalISel/instruction-select/shift-rv64.mir
index 1e6890098498ea..dbb1a55531f997 100644
--- a/llvm/test/CodeGen/RISCV/GlobalISel/instruction-select/shift-rv64.mir
+++ b/llvm/test/CodeGen/RISCV/GlobalISel/instruction-select/shift-rv64.mir
@@ -241,3 +241,89 @@ body:             |
     $x10 = COPY %6(s64)
     PseudoRET implicit $x10
 ...
+
+---
+name:            srl_and_needed
+alignment:       2
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: gprb }
+  - { id: 1, class: gprb }
+  - { id: 2, class: gprb }
+  - { id: 3, class: gprb }
+  - { id: 4, class: gprb }
+  - { id: 5, class: gprb }
+  - { id: 6, class: gprb }
+liveins:
+  - { reg: '$x10' }
+  - { reg: '$x11' }
+body:             |
+  bb.1.entry:
+    liveins: $x10, $x11
+
+    ; CHECK-LABEL: name: srl_and_needed
+    ; CHECK: liveins: $x10, $x11
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr = COPY $x10
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:gpr = COPY $x11
+    ; CHECK-NEXT: [[ANDI:%[0-9]+]]:gpr = ANDI [[COPY]], 15
+    ; CHECK-NEXT: [[SRL:%[0-9]+]]:gpr = SRL [[COPY1]], [[ANDI]]
+    ; CHECK-NEXT: $x10 = COPY [[SRL]]
+    ; CHECK-NEXT: PseudoRET implicit $x10
+    %0:gprb(s64) = COPY $x10
+    %1:gprb(s64) = COPY $x11
+    %2:gprb(s32) = G_CONSTANT i32 15
+    %3:gprb(s32) = G_TRUNC %0(s64)
+    %4:gprb(s32) = G_AND %3, %2
+    %5:gprb(s64) = nneg G_ZEXT %4(s32)
+    %6:gprb(s64) = G_LSHR %1, %5(s64)
+    $x10 = COPY %6(s64)
+    PseudoRET implicit $x10
+...
+
+---
+name:            srl_and_eliminated
+alignment:       2
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: gprb }
+  - { id: 1, class: gprb }
+  - { id: 2, class: gprb }
+  - { id: 3, class: gprb }
+  - { id: 4, class: gprb }
+  - { id: 5, class: gprb }
+  - { id: 6, class: gprb }
+  - { id: 7, class: gprb }
+  - { id: 8, class: gprb }
+liveins:
+  - { reg: '$x10' }
+  - { reg: '$x11' }
+body:             |
+  bb.1.entry:
+    liveins: $x10, $x11
+
+    ; CHECK-LABEL: name: srl_and_eliminated
+    ; CHECK: liveins: $x10, $x11
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr = COPY $x10
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:gpr = COPY $x11
+    ; CHECK-NEXT: [[ANDI:%[0-9]+]]:gpr = ANDI [[COPY]], 79
+    ; CHECK-NEXT: [[SRL:%[0-9]+]]:gpr = SRL [[COPY1]], [[ANDI]]
+    ; CHECK-NEXT: $x10 = COPY [[SRL]]
+    ; CHECK-NEXT: PseudoRET implicit $x10
+    %0:gprb(s64) = COPY $x10
+    %1:gprb(s64) = COPY $x11
+    %2:gprb(s32) = G_CONSTANT i32 15
+    %3:gprb(s32) = G_TRUNC %0(s64)
+    %7:gprb(s32) = G_CONSTANT i32 79
+    %8:gprb(s32) = G_AND %3, %7
+    %4:gprb(s32) = G_AND %8, %2
+    %5:gprb(s64) = nneg G_ZEXT %4(s32)
+    %6:gprb(s64) = G_LSHR %1, %5(s64)
+    $x10 = COPY %6(s64)
+    PseudoRET implicit $x10
+...