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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-globalisel

Author: Mikhail Gudim (mgudim)

<details>
<summary>Changes</summary>

This is WIP, just posting it to let others know that I'm working on it.

---
Full diff: https://github.com/llvm/llvm-project/pull/77572.diff


2 Files Affected:

- (modified) llvm/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp (+49-3) 
- (added) llvm/test/CodeGen/RISCV/GlobalISel/instruction-select/shift-rv64.mir (+54) 


``````````diff
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..4e07819cb4b4c0
--- /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

``````````

</details>


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