[llvm] [WIP][GISel][RISCV] Implement selectShiftMask. (PR #77572)

[Mikhail Gudim via llvm-commits]

https://github.com/mgudim updated https://github.com/llvm/llvm-project/pull/77572

>From 2b5a25d4e158f83eb9e068f8fb4ff251deb6fb24 Mon Sep 17 00:00:00 2001
From: Mikhail Gudim <mgudim at gmail.com>
Date: Wed, 10 Jan 2024 04:47:46 -0500
Subject: [PATCH] [WIP][GISel][RISCV] Implement selectShiftMask.

---
 .../RISCV/GISel/RISCVInstructionSelector.cpp  | 52 ++++++++++++++++--
 .../instruction-select/shift-rv64.mir         | 54 +++++++++++++++++++
 2 files changed, 103 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/CodeGen/RISCV/GlobalISel/instruction-select/shift-rv64.mir

diff --git a/llvm/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp b/llvm/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp
index 61bdbfc47d947f..82860eb4b9e134 100644
--- a/llvm/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp
+++ b/llvm/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp
@@ -159,9 +159,55 @@ RISCVInstructionSelector::RISCVInstructionSelector(
 
 InstructionSelector::ComplexRendererFns
 RISCVInstructionSelector::selectShiftMask(MachineOperand &Root) const {
-  // TODO: Also check if we are seeing the result of an AND operation which
-  // could be bypassed since we only check the lower log2(xlen) bits.
-  return {{[=](MachineInstrBuilder &MIB) { MIB.add(Root); }}};
+  // TODO: ShiftWidth can be 64.
+  unsigned ShiftWidth = 32;
+
+  using namespace llvm::MIPatternMatch;
+  MachineFunction &MF = *Root.getParent()->getParent()->getParent();
+  MachineRegisterInfo &MRI = MF.getRegInfo();
+
+  if (!Root.isReg())
+    return std::nullopt;
+  Register RootReg = Root.getReg();
+  Register ShAmtReg = RootReg;
+  // Peek through zext.
+  Register ZExtSrcReg;
+  if (mi_match(ShAmtReg, MRI, m_GZExt(m_Reg(ZExtSrcReg)))) {
+    ShAmtReg = ZExtSrcReg;
+  }
+
+  APInt AndMask;
+  Register AndSrcReg;
+  if (mi_match(ShAmtReg, MRI, m_GAnd(m_Reg(AndSrcReg), m_ICst(AndMask)))) {
+    APInt ShMask(AndMask.getBitWidth(), ShiftWidth - 1);
+    if (ShMask.isSubsetOf(AndMask)) {
+      ShAmtReg = AndSrcReg;
+    } else {
+      // TODO:
+      // SimplifyDemandedBits may have optimized the mask so try restoring any
+      // bits that are known zero.
+    }
+  }
+
+  APInt Imm;
+  Register Reg;
+  if (mi_match(ShAmtReg, MRI, m_GAdd(m_Reg(Reg), m_ICst(Imm)))) {
+    if (Imm != 0 && Imm.urem(ShiftWidth) == 0)
+      // If we are shifting by X+N where N == 0 mod Size, then just shift by X
+      // to avoid the ADD.
+      ShAmtReg = Reg;
+  } else if (mi_match(ShAmtReg, MRI, m_GSub(m_ICst(Imm), m_Reg(Reg)))) {
+    if (Imm != 0 && Imm.urem(ShiftWidth) == 0) {
+      // If we are shifting by N-X where N == 0 mod Size, then just shift by -X
+      // to generate a NEG instead of a SUB of a constant.
+      ShAmtReg = Reg;
+    }
+  }
+
+  if (ShAmtReg == RootReg)
+    return std::nullopt;
+
+  return {{[=](MachineInstrBuilder &MIB) { MIB.addReg(ShAmtReg); }}};
 }
 
 InstructionSelector::ComplexRendererFns
diff --git a/llvm/test/CodeGen/RISCV/GlobalISel/instruction-select/shift-rv64.mir b/llvm/test/CodeGen/RISCV/GlobalISel/instruction-select/shift-rv64.mir
new file mode 100644
index 00000000000000..d92f414e8d9fac
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/GlobalISel/instruction-select/shift-rv64.mir
@@ -0,0 +1,54 @@
+# RUN: llc -mtriple=riscv64 -run-pass=instruction-select \
+# RUN:   -simplify-mir -verify-machineinstrs %s -o -
+
+---
+name:            shl_zext
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x10
+
+    %0:gprb(s64) = COPY $x10
+    %1:gprb(s32) = G_CONSTANT i32 1
+    %2:gprb(s64) = G_ZEXT %1
+    %3:gprb(s64) = G_SHL %0, %2(s64)
+    $x10 = COPY %3(s64)
+    PseudoRET implicit $x10
+
+---
+name:            shl_and
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x10, $x11
+
+    %0:gprb(s64) = COPY $x10
+    %1:gprb(s64) = COPY $x11
+    %2:gprb(s64) = G_CONSTANT i64 31
+    %3:gprb(s64) = G_AND %1, %2
+    %4:gprb(s64) = G_SHL %0, %3(s64)
+    $x10 = COPY %4(s64)
+    PseudoRET implicit $x10
+
+---
+name:            shl_and_zext
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x10, $x11
+
+    %0:gprb(s64) = COPY $x10
+    %addr:gprb(p0) = COPY $x11
+    %1:gprb(s32) = G_LOAD %addr(p0) :: (load (s8))
+    %2:gprb(s32) = G_CONSTANT i32 31
+    %3:gprb(s32) = G_AND %1, %2
+    %4:gprb(s64) = G_ZEXT %3
+    %5:gprb(s64) = G_SHL %0, %4(s64)
+    $x10 = COPY %5(s64)
+    PseudoRET implicit $x10