Restore shift mask from known bits.

Author: mgudim

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

@@ -16,6 +16,7 @@
 #include "RISCVSubtarget.h"
 #include "RISCVTargetMachine.h"
 #include "llvm/CodeGen/GlobalISel/GIMatchTableExecutorImpl.h"
+#include "llvm/CodeGen/GlobalISel/GISelKnownBits.h"
 #include "llvm/CodeGen/GlobalISel/GenericMachineInstrs.h"
 #include "llvm/CodeGen/GlobalISel/InstructionSelector.h"
 #include "llvm/CodeGen/GlobalISel/MIPatternMatch.h"
@@ -183,9 +184,11 @@ RISCVInstructionSelector::selectShiftMask(MachineOperand &Root) const {
     if (ShMask.isSubsetOf(AndMask)) {
       ShAmtReg = AndSrcReg;
     } else {
-      // TODO:
       // SimplifyDemandedBits may have optimized the mask so try restoring any
       // bits that are known zero.
+      KnownBits Known = KB->getKnownBits(ShAmtReg);
+      if (ShMask.isSubsetOf(AndMask | Known.Zero))
+        ShAmtReg = AndSrcReg;
     }
   }
 
@@ -200,18 +203,18 @@ RISCVInstructionSelector::selectShiftMask(MachineOperand &Root) const {
     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);
+      ShAmtReg = MRI.createVirtualRegister(&RISCV::GPRRegClass);
       unsigned NegOpc = Subtarget->is64Bit() ? RISCV::SUBW : RISCV::SUB;
       return {{[=](MachineInstrBuilder &MIB) {
         MachineIRBuilder(*MIB.getInstr())
             .buildInstr(NegOpc, {ShAmtReg}, {Register(RISCV::X0), Reg});
         MIB.addReg(ShAmtReg);
       }}};
     }
-    if ((Imm.urem(ShiftWidth) & (ShiftWidth - 1)) == ShiftWidth - 1) {
+    if (Imm.urem(ShiftWidth) == ShiftWidth - 1) {
       // If we are shifting by N-X where N == -1 mod Size, then just shift by ~X
       // to generate a NOT instead of a SUB of a constant.
-      ShAmtReg = MRI.createGenericVirtualRegister(ShiftLLT);
+      ShAmtReg = MRI.createVirtualRegister(&RISCV::GPRRegClass);
       return {{[=](MachineInstrBuilder &MIB) {
         MachineIRBuilder(*MIB.getInstr())
             .buildInstr(RISCV::XORI, {ShAmtReg}, {Reg})

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

@@ -53,6 +53,35 @@ body:             |
     PseudoRET implicit $x10
 ...
 
+---
+name:            shl_and_with_simplified_mask
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x10, $x11
+
+    ; CHECK-LABEL: name: shl_and_with_simplified_mask
+    ; 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 [[COPY1]], 31
+    ; CHECK-NEXT: [[SLL:%[0-9]+]]:gpr = SLL [[COPY]], [[ANDI]]
+    ; CHECK-NEXT: $x10 = COPY [[SLL]]
+    ; CHECK-NEXT: PseudoRET implicit $x10
+    %0:gprb(s32) = COPY $x10
+    %1:gprb(s32) = COPY $x11
+    %2:gprb(s32) = G_CONSTANT i32 31
+    %3:gprb(s32) = G_AND %1, %2
+    %4:gprb(s32) = G_CONSTANT i32 31
+    %5:gprb(s32) = G_AND %3, %4
+    %6:gprb(s32) = G_SHL %0, %5(s32)
+    $x10 = COPY %6(s32)
+    PseudoRET implicit $x10
+...
+
 ---
 name:            shl_add
 legalized:       true

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

@@ -77,6 +77,35 @@ body:             |
     PseudoRET implicit $x10
 ...
 
+---
+name:            shl_and_with_simplified_mask
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x10, $x11
+
+    ; CHECK-LABEL: name: shl_and_with_simplified_mask
+    ; 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 [[COPY1]], 62
+    ; CHECK-NEXT: [[SLL:%[0-9]+]]:gpr = SLL [[COPY]], [[ANDI]]
+    ; CHECK-NEXT: $x10 = COPY [[SLL]]
+    ; CHECK-NEXT: PseudoRET implicit $x10
+    %0:gprb(s64) = COPY $x10
+    %1:gprb(s64) = COPY $x11
+    %2:gprb(s64) = G_CONSTANT i64 62
+    %3:gprb(s64) = G_AND %1, %2
+    %4:gprb(s64) = G_CONSTANT i64 62
+    %5:gprb(s64) = G_AND %3, %4
+    %6:gprb(s64) = G_SHL %0, %5(s64)
+    $x10 = COPY %6(s64)
+    PseudoRET implicit $x10
+...
+
 ---
 name:            shl_add
 legalized:       true
@@ -197,18 +226,18 @@ body:             |
     ; CHECK: liveins: $x10, $x11
     ; CHECK-NEXT: {{  $}}
     ; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr = COPY $x10
-    ; CHECK-NEXT: %addr:gpr = COPY $x11
-    ; CHECK-NEXT: [[LBU:%[0-9]+]]:gpr = LBU %addr, 0 :: (load (s8))
-    ; CHECK-NEXT: [[SLL:%[0-9]+]]:gpr = SLL [[COPY]], [[LBU]]
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:gpr = COPY $x11
+    ; CHECK-NEXT: [[LW:%[0-9]+]]:gpr = LW [[COPY1]], 0 :: (load (s32))
+    ; CHECK-NEXT: [[SLL:%[0-9]+]]:gpr = SLL [[COPY]], [[LW]]
     ; CHECK-NEXT: $x10 = COPY [[SLL]]
     ; CHECK-NEXT: PseudoRET implicit $x10
     %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 63
-    %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)
+    %1:gprb(p0) = COPY $x11
+    %2:gprb(s32) = G_LOAD %1(p0) :: (load (s32))
+    %3:gprb(s32) = G_CONSTANT i32 63
+    %4:gprb(s32) = G_AND %2, %3
+    %5:gprb(s64) = G_ZEXT %4
+    %6:gprb(s64) = G_SHL %0, %5(s64)
+    $x10 = COPY %6(s64)
     PseudoRET implicit $x10
 ...