[GISel][RISCV] Implement selectShiftMask. #77572

mgudim · 2024-01-10T09:53:34Z

Implement selectShiftMask in GlobalISel.

llvmbot · 2024-01-10T09:54:02Z

@llvm/pr-subscribers-backend-risc-v

Author: Mikhail Gudim (mgudim)

Changes

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 --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

llvmbot · 2024-01-10T09:54:02Z

@llvm/pr-subscribers-llvm-globalisel

Author: Mikhail Gudim (mgudim)

Changes

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 --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

arsenm · 2024-01-10T10:01:23Z

llvm/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp

+  unsigned ShiftWidth = 32;
+
+  using namespace llvm::MIPatternMatch;
+  MachineFunction &MF = *Root.getParent()->getParent()->getParent();


This should already be available in the base class

Fixed this, thanks.

github-actions · 2024-01-11T21:32:37Z

✅ With the latest revision this PR passed the C/C++ code formatter.

Implement the selectShiftMask for GlobalISel.

mgudim · 2024-01-11T21:40:03Z

llvm/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp

+        MIB.addReg(ShAmtReg);
+      }}};
+    }
+    if ((Imm.urem(ShiftWidth) & (ShiftWidth - 1)) == ShiftWidth - 1) {


I generalized this slightly from the non-global ISel version so that it works with negative numbers. This simplification works as long as lower 4 or 5 bits of the remainer are all 1. I think this is right, please correct me if I am wrong.

Why doesn't the isel version work with negative numbers?

The & (ShiftWidth - 1) seems unnecessary. Imm.urem(ShiftWidth) should return a number in the range [0, ShiftWidth -1] so the & should have no effect.

Yes, you're right. I forgot that Imm is treated as unsigned.

mgudim · 2024-01-17T04:40:32Z

This is ready for review @topperc

topperc · 2024-01-17T04:45:09Z

llvm/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp

+    if (ShMask.isSubsetOf(AndMask)) {
+      ShAmtReg = AndSrcReg;
+    } else {
+      // TODO:


What's the TODO?

It's done. Forgot to delete this line, thanks

topperc · 2024-01-17T04:55:12Z

llvm/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp

+    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 = MRI.createGenericVirtualRegister(ShiftLLT);


I wonder if this should be MRI.createVirtualRegister(&RISCV::GPRRegClass).

Hmm... Not sure, I just followed similar code (like selectSHXADDOp). The docs say about target's InstructionSelector that " It is also responsible for doing the necessary constraining of gvregs into the appropriate register classes as well as passing through COPY instructions to the register allocator." So, to me it also seems more appropriate to use createVirtualRegister. We still see gpr in the output.

You should avoid creating new generic register during selection. You can get away with it if later steps end up constraining the register, but it's just adding extra steps

Thanks for clarifying @arsenm I changed it here, I can change it in other places in separate MR

topperc

LGTM

Implement `selectShiftMask` in `GlobalISel`.

llvmbot added backend:RISC-V llvm:globalisel labels Jan 10, 2024

mgudim requested a review from topperc January 10, 2024 09:54

mgudim requested review from michaelmaitland and mshockwave January 10, 2024 09:54

mgudim force-pushed the select-shift-mask branch from 8f84ce6 to 2b5a25d Compare January 10, 2024 09:55

arsenm reviewed Jan 10, 2024

View reviewed changes

mgudim force-pushed the select-shift-mask branch from 2b5a25d to b22fe66 Compare January 11, 2024 21:30

mgudim force-pushed the select-shift-mask branch from b22fe66 to 6c90357 Compare January 11, 2024 21:34

mgudim changed the title ~~[WIP][GISel][RISCV] Implement selectShiftMask.~~ [GISel][RISCV] Implement selectShiftMask. Jan 11, 2024

[GISel][RISCV] Implement selectShiftMask.

5f51ab1

Implement the selectShiftMask for GlobalISel.

mgudim force-pushed the select-shift-mask branch from 6c90357 to 5f51ab1 Compare January 11, 2024 21:35

mgudim commented Jan 11, 2024

View reviewed changes

Update test for rotate.

3018057

topperc reviewed Jan 17, 2024

View reviewed changes

Restore shift mask from known bits.

aec7179

mgudim force-pushed the select-shift-mask branch from 6791b0f to aec7179 Compare January 17, 2024 05:51

topperc approved these changes Jan 17, 2024

View reviewed changes

mgudim merged commit c1f4338 into llvm:main Jan 17, 2024
4 checks passed

ampandey-1995 pushed a commit to ampandey-1995/llvm-project that referenced this pull request Jan 19, 2024

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

75a25c1

Implement `selectShiftMask` in `GlobalISel`.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[GISel][RISCV] Implement selectShiftMask. #77572

[GISel][RISCV] Implement selectShiftMask. #77572

mgudim commented Jan 10, 2024 •

edited

llvmbot commented Jan 10, 2024

llvmbot commented Jan 10, 2024

arsenm Jan 10, 2024

mgudim Jan 11, 2024

github-actions bot commented Jan 11, 2024 •

edited

mgudim Jan 11, 2024

topperc Jan 17, 2024

mgudim Jan 17, 2024

mgudim commented Jan 17, 2024

topperc Jan 17, 2024

mgudim Jan 17, 2024

topperc Jan 17, 2024

mgudim Jan 17, 2024

arsenm Jan 17, 2024

mgudim Jan 17, 2024

topperc left a comment

[GISel][RISCV] Implement selectShiftMask. #77572

[GISel][RISCV] Implement selectShiftMask. #77572

Conversation

mgudim commented Jan 10, 2024 • edited

llvmbot commented Jan 10, 2024

llvmbot commented Jan 10, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

github-actions bot commented Jan 11, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mgudim commented Jan 17, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

topperc left a comment

Choose a reason for hiding this comment

mgudim commented Jan 10, 2024 •

edited

github-actions bot commented Jan 11, 2024 •

edited