[RISCV][GlobalISel] Fix selectShiftMask when shift mask is created from G_AND (#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.

Author: bvlgah

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

@@ -177,14 +177,28 @@ 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 log2(XLEN) 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 and mask), there are two cases G_AND 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)) {
       ShAmtReg = AndSrcReg;
     } 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;
     }

llvm/test/CodeGen/RISCV/GlobalISel/instruction-select/shift-rv32.mir

@@ -188,3 +188,59 @@ body:             |
     $x10 = COPY %4(s32)
     PseudoRET implicit $x10
 ...
+
+---
+name:            srl_and_needed
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+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(s32) = COPY $x10
+    %1:gprb(s32) = COPY $x11
+    %2:gprb(s32) = G_CONSTANT i32 15
+    %3:gprb(s32) = G_AND %0, %2
+    %4:gprb(s32) = G_LSHR %1, %3(s32)
+    $x10 = COPY %4(s32)
+    PseudoRET implicit $x10
+...
+
+---
+name:            srl_and_eliminated
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+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]], 47
+    ; CHECK-NEXT: [[SRL:%[0-9]+]]:gpr = SRL [[COPY1]], [[ANDI]]
+    ; CHECK-NEXT: $x10 = COPY [[SRL]]
+    ; CHECK-NEXT: PseudoRET implicit $x10
+    %0:gprb(s32) = COPY $x10
+    %1:gprb(s32) = COPY $x11
+    %2:gprb(s32) = G_CONSTANT i32 15
+    %3:gprb(s32) = G_CONSTANT i32 47
+    %4:gprb(s32) = G_AND %0, %3
+    %5:gprb(s32) = G_AND %4, %2
+    %6:gprb(s32) = G_LSHR %1, %5(s32)
+    $x10 = COPY %6(s32)
+    PseudoRET implicit $x10
+...

llvm/test/CodeGen/RISCV/GlobalISel/instruction-select/shift-rv64.mir

@@ -241,3 +241,63 @@ body:             |
     $x10 = COPY %6(s64)
     PseudoRET implicit $x10
 ...
+
+---
+name:            srl_and_needed
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+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
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+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
+...