[RISCV][GlobalISel] Fix selectShiftMask when shift mask is created from G_AND

[bvlgah]

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.

[llvmbot]

@llvm/pr-subscribers-llvm-globalisel

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

Author: Hongbin Jin (bvlgah)

Changes

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.

---

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

2 Files Affected:

• (modified) llvm/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp (+15-1)
• (modified) llvm/test/CodeGen/RISCV/GlobalISel/instruction-select/shift-rv64.mir (+86)

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

[bvlgah]

Hi. I have been working on make RISC-V's global instruction selection usable and performant.
This patch is intended to fix one of issues encountered in the LLVM test suite. Before the patch,
the following gMIR

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

is lowered to something like

    %0:gpr = COPY $x10
    %1:gpr = COPY $x11
    %6:gpr = SRL %1, %0

This is incorrect because the the following condition always evaluates to true.

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

Lines 187 to 188 in f0cbdd3

	KnownBits Known = KB->getKnownBits(ShAmtReg);
	if (ShMask.isSubsetOf(AndMask | Known.Zero))

[michaelmaitland]

LGTM

[bvlgah]

@topperc Hi, could you please review this patch, thanks

[topperc]

LGTM

[bvlgah]

@michaelmaitland Hi Michael, could please help me with merging the patch (I have no write permission), thanks.
One thing to note is I have committed several fix-ups, so squash is needed

[michaelmaitland]

Committed! Thanks for this patch :)

[bvlgah]

Committed! Thanks for this patch :)

Feel grateful to work with people in this community