[LoongArch] Combine ((x & shifted_mask) >> shamt) to bstrpick.{w/d}

[SixWeining]

Summary

This is an improvement to LoongArch codegen. In D127206 we combined
((x >> shamt) & shifted_mask) to bstrpick and here we do a similar
combination when certain conditions are met.

Thanks to @xen0n for reminding me.

[xen0n]

Nice overall improvements, thanks!

Only one minor suggestion. Do you think an extra test case that verifies something like (x & 0xfff0) >> 17 are properly optimized to 0? Or is such an optimization already covered by all the generic mechanisms so it's not needed here?

[SixWeining]

In D128191#3596399, @xen0n wrote:

Nice overall improvements, thanks!

Only one minor suggestion. Do you think an extra test case that verifies something like (x & 0xfff0) >> 17 are properly optimized to 0? Or is such an optimization already covered by all the generic mechanisms so it's not needed here?

Yes, I have meant to write such a negative test but I find it have already been optimized to 0. So I think we don’t need it here.

[xen0n]

LGTM, thanks!

[xen0n]

In D128191#3596468, @SixWeining wrote:

In D128191#3596399, @xen0n wrote:

Nice overall improvements, thanks!

Only one minor suggestion. Do you think an extra test case that verifies something like (x & 0xfff0) >> 17 are properly optimized to 0? Or is such an optimization already covered by all the generic mechanisms so it's not needed here?

Yes, I have meant to write such a negative test but I find it have already been optimized to 0. So I think we don’t need it here.

I guessed so. So this is fine, and I've approved. We may wait for a couple of days for others to take a look at this though.

llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp

LoongArchTargetLowering::LoongArchTargetLowering(const TargetMachine &TM,
 
   setBooleanContents(ZeroOrOneBooleanContent);
 
   // Function alignments.
   const Align FunctionAlignment(4);
   setMinFunctionAlignment(FunctionAlignment);
 
   setTargetDAGCombine(ISD::AND);
+  setTargetDAGCombine(ISD::SRL);
 }
 
 SDValue LoongArchTargetLowering::LowerOperation(SDValue Op,
                                                 SelectionDAG &DAG) const {
   switch (Op.getOpcode()) {
   default:
     report_fatal_error("unimplemented operand");
   case ISD::SHL_PARTS:
if (FirstOperandOpc == ISD::SRA || FirstOperandOpc == ISD::SRL) {
     NewOperand = FirstOperand;
   }
   msb = lsb + SMLen - 1;
   return DAG.getNode(LoongArchISD::BSTRPICK, DL, ValTy, NewOperand,
                      DAG.getConstant(msb, DL, GRLenVT),
                      DAG.getConstant(lsb, DL, GRLenVT));
 }
 
+static SDValue performSRLCombine(SDNode *N, SelectionDAG &DAG,
+                                 TargetLowering::DAGCombinerInfo &DCI,
+                                 const LoongArchSubtarget &Subtarget) {
+  if (DCI.isBeforeLegalizeOps())
+    return SDValue();
+
+  // $dst = srl (and $src, Mask), Shamt
+  // =>
+  // BSTRPICK $dst, $src, MaskIdx+MaskLen-1, Shamt
+  // when Mask is a shifted mask, and MaskIdx <= Shamt <= MaskIdx+MaskLen-1
+  //
+
+  SDValue FirstOperand = N->getOperand(0);
+  ConstantSDNode *CN;
+  EVT ValTy = N->getValueType(0);
+  SDLoc DL(N);
+  MVT GRLenVT = Subtarget.getGRLenVT();
+  unsigned MaskIdx, MaskLen;
+  uint64_t Shamt;
+
+  // The first operand must be an AND and the second operand of the AND must be
+  // a shifted mask.
+  if (FirstOperand.getOpcode() != ISD::AND ||
+      !(CN = dyn_cast<ConstantSDNode>(FirstOperand.getOperand(1))) ||
+      !isShiftedMask_64(CN->getZExtValue(), MaskIdx, MaskLen))
+    return SDValue();
+
+  // The second operand (shift amount) must be an immediate.
+  if (!(CN = dyn_cast<ConstantSDNode>(N->getOperand(1))))
+    return SDValue();
+
+  Shamt = CN->getZExtValue();
+  if (MaskIdx <= Shamt && Shamt <= MaskIdx + MaskLen - 1)
+    return DAG.getNode(LoongArchISD::BSTRPICK, DL, ValTy,
+                       FirstOperand->getOperand(0),
+                       DAG.getConstant(MaskIdx + MaskLen - 1, DL, GRLenVT),
+                       DAG.getConstant(Shamt, DL, GRLenVT));
+
+  return SDValue();
+}
+
 SDValue LoongArchTargetLowering::PerformDAGCombine(SDNode *N,
                                                    DAGCombinerInfo &DCI) const {
   SelectionDAG &DAG = DCI.DAG;
   switch (N->getOpcode()) {
   default:
     break;
   case ISD::AND:
     return performANDCombine(N, DAG, DCI, Subtarget);
+  case ISD::SRL:
+    return performSRLCombine(N, DAG, DCI, Subtarget);
   }
   return SDValue();
 }
 
 const char *LoongArchTargetLowering::getTargetNodeName(unsigned Opcode) const {
   switch ((LoongArchISD::NodeType)Opcode) {
   case LoongArchISD::FIRST_NUMBER:
     break;

llvm/test/CodeGen/LoongArch/bstrpick_d.ll

 define i64 @and4095(i64 %a) {
 ; CHECK-LABEL: and4095:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    andi $a0, $a0, 4095
 ; CHECK-NEXT:    jirl $zero, $ra, 0
   %and = and i64 %a, 4095
   ret i64 %and
 }
+
+;; (srl (and a, 0xff0), 4) => (BSTRPICK a, 11, 4)
+define i64 @and0xff0_lshr4(i64 %a) {
+; CHECK-LABEL: and0xff0_lshr4:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    bstrpick.d $a0, $a0, 11, 4
+; CHECK-NEXT:    jirl $zero, $ra, 0
+  %and = and i64 %a, 4080
+  %shr = lshr i64 %and, 4
+  ret i64 %shr
+}
+
+;; (sra (and a, 0xff0), 5) can also be combined to (BSTRPICK a, 11, 5).
+;; This is because (sra (and a, 0xff0)) would be combined to (srl (and a, 0xff0), 5)
+;; firstly by DAGCombiner::SimplifyDemandedBits.
+define i64 @and4080_ashr5(i64 %a) {
+; CHECK-LABEL: and4080_ashr5:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    bstrpick.d $a0, $a0, 11, 5
+; CHECK-NEXT:    jirl $zero, $ra, 0
+  %and = and i64 %a, 4080
+  %shr = ashr i64 %and, 5
+  ret i64 %shr
+}
+
+;; Negative test: the second operand of AND is not a shifted mask
+define i64 @and0xf30_lshr4(i64 %a) {
+; CHECK-LABEL: and0xf30_lshr4:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    andi $a0, $a0, 3888
+; CHECK-NEXT:    srli.d $a0, $a0, 4
+; CHECK-NEXT:    jirl $zero, $ra, 0
+  %and = and i64 %a, 3888
+  %shr = lshr i64 %and, 4
+  ret i64 %shr
+}
+
+;; Negative test: Shamt < MaskIdx
+define i64 @and0xff0_lshr3(i64 %a) {
+; CHECK-LABEL: and0xff0_lshr3:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    andi $a0, $a0, 4080
+; CHECK-NEXT:    srli.d $a0, $a0, 3
+; CHECK-NEXT:    jirl $zero, $ra, 0
+  %and = and i64 %a, 4080
+  %shr = lshr i64 %and, 3
+  ret i64 %shr
+}

llvm/test/CodeGen/LoongArch/bstrpick_w.ll

 define i32 @and4095(i32 %a) {
 ; CHECK-LABEL: and4095:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    andi $a0, $a0, 4095
 ; CHECK-NEXT:    jirl $zero, $ra, 0
   %and = and i32 %a, 4095
   ret i32 %and
 }
+
+;; (srl (and a, 0xff0), 4) => (BSTRPICK a, 11, 4)
+define i32 @and0xff0_lshr4(i32 %a) {
+; CHECK-LABEL: and0xff0_lshr4:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    bstrpick.w $a0, $a0, 11, 4
+; CHECK-NEXT:    jirl $zero, $ra, 0
+  %and = and i32 %a, 4080
+  %shr = lshr i32 %and, 4
+  ret i32 %shr
+}
+
+;; (sra (and a, 0xff0), 5) can also be combined to (BSTRPICK a, 11, 5).
+;; This is because (sra (and a, 0xff0)) would be combined to (srl (and a, 0xff0), 5)
+;; firstly by DAGCombiner::SimplifyDemandedBits.
+define i32 @and4080_ashr5(i32 %a) {
+; CHECK-LABEL: and4080_ashr5:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    bstrpick.w $a0, $a0, 11, 5
+; CHECK-NEXT:    jirl $zero, $ra, 0
+  %and = and i32 %a, 4080
+  %shr = ashr i32 %and, 5
+  ret i32 %shr
+}
+
+;; Negative test: the second operand of AND is not a shifted mask
+define i32 @and0xf30_lshr4(i32 %a) {
+; CHECK-LABEL: and0xf30_lshr4:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    andi $a0, $a0, 3888
+; CHECK-NEXT:    srli.w $a0, $a0, 4
+; CHECK-NEXT:    jirl $zero, $ra, 0
+  %and = and i32 %a, 3888
+  %shr = lshr i32 %and, 4
+  ret i32 %shr
+}
+
+;; Negative test: Shamt < MaskIdx
+define i32 @and0xff0_lshr3(i32 %a) {
+; CHECK-LABEL: and0xff0_lshr3:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    andi $a0, $a0, 4080
+; CHECK-NEXT:    srli.w $a0, $a0, 3
+; CHECK-NEXT:    jirl $zero, $ra, 0
+  %and = and i32 %a, 4080
+  %shr = lshr i32 %and, 3
+  ret i32 %shr
+}

llvm/test/CodeGen/LoongArch/ir-instruction/lshr.ll

 ; LA64-NEXT:    jirl $zero, $ra, 0
   %lshr = lshr i1 %x, 3
   ret i1 %lshr
 }
 
 define i8 @lshr_i8_3(i8 %x) {
 ; LA32-LABEL: lshr_i8_3:
 ; LA32:       # %bb.0:
-; LA32-NEXT:    andi $a0, $a0, 248
-; LA32-NEXT:    srli.w $a0, $a0, 3
+; LA32-NEXT:    bstrpick.w $a0, $a0, 7, 3
 ; LA32-NEXT:    jirl $zero, $ra, 0
 ;
 ; LA64-LABEL: lshr_i8_3:
 ; LA64:       # %bb.0:
-; LA64-NEXT:    andi $a0, $a0, 248
-; LA64-NEXT:    srli.d $a0, $a0, 3
+; LA64-NEXT:    bstrpick.d $a0, $a0, 7, 3
 ; LA64-NEXT:    jirl $zero, $ra, 0
   %lshr = lshr i8 %x, 3
   ret i8 %lshr
 }
 
 define i16 @lshr_i16_3(i16 %x) {
 ; LA32-LABEL: lshr_i16_3:
 ; LA32:       # %bb.0:
-; LA32-NEXT:    lu12i.w $a1, 15
-; LA32-NEXT:    ori $a1, $a1, 4088
-; LA32-NEXT:    and $a0, $a0, $a1
-; LA32-NEXT:    srli.w $a0, $a0, 3
+; LA32-NEXT:    bstrpick.w $a0, $a0, 15, 3
 ; LA32-NEXT:    jirl $zero, $ra, 0
 ;
 ; LA64-LABEL: lshr_i16_3:
 ; LA64:       # %bb.0:
-; LA64-NEXT:    lu12i.w $a1, 15
-; LA64-NEXT:    ori $a1, $a1, 4088
-; LA64-NEXT:    and $a0, $a0, $a1
-; LA64-NEXT:    srli.d $a0, $a0, 3
+; LA64-NEXT:    bstrpick.d $a0, $a0, 15, 3
 ; LA64-NEXT:    jirl $zero, $ra, 0
   %lshr = lshr i16 %x, 3
   ret i16 %lshr
 }
 
 define i32 @lshr_i32_3(i32 %x) {
 ; LA32-LABEL: lshr_i32_3:
 ; LA32:       # %bb.0:
 ; LA32-NEXT:    srli.w $a0, $a0, 3
 ; LA32-NEXT:    jirl $zero, $ra, 0
 ;
 ; LA64-LABEL: lshr_i32_3:
 ; LA64:       # %bb.0:
-; LA64-NEXT:    addi.w $a1, $zero, -8
-; LA64-NEXT:    lu32i.d $a1, 0
-; LA64-NEXT:    and $a0, $a0, $a1
-; LA64-NEXT:    srli.d $a0, $a0, 3
+; LA64-NEXT:    bstrpick.d $a0, $a0, 31, 3
 ; LA64-NEXT:    jirl $zero, $ra, 0
   %lshr = lshr i32 %x, 3
   ret i32 %lshr
 }
 
 define i64 @lshr_i64_3(i64 %x) {
 ; LA32-LABEL: lshr_i64_3:
 ; LA32:       # %bb.0: