[LoongArch] Perform `and` combination with a shifted mask

[SixWeining]

Summary

Depends on D127203

[SixWeining]

rebase

[xen0n]

It seems only the (x >> shamt) & mask case is handled, but not (x & in_place_mask) >> shamt; is this an intentional simplification? The current behavior is good though.

[SixWeining]

In D127206#3569149, @xen0n wrote:

It seems only the (x >> shamt) & mask case is handled, but not (x & in_place_mask) >> shamt; is this an intentional simplification? The current behavior is good though.

Thanks for reminding me. (x & in_place_mask) >> shamt is not handled in this patch and it indeed can be optimized to bstrpick too. But this is for shl combination but not for and. Please let me handle it in another patch. Thanks again!

[SixWeining]

Remove the patch number in the title.

[xen0n]

In D127206#3572045, @SixWeining wrote:

In D127206#3569149, @xen0n wrote:

It seems only the (x >> shamt) & mask case is handled, but not (x & in_place_mask) >> shamt; is this an intentional simplification? The current behavior is good though.

Thanks for reminding me. (x & in_place_mask) >> shamt is not handled in this patch and it indeed can be optimized to bstrpick too. But this is for shl combination but not for and. Please let me handle it in another patch. Thanks again!

Okay. Immediates in this case need some deeper inspection for confirming it's a single chunk of 1's among 0's, so it's better to split commits.

Then this looks good to me!

[SixWeining]

rebase

[SixWeining]

Use ;; to emphasize non-RUN non-CHECK lines.

[SixWeining]

adjust the dependency

llvm/lib/Target/LoongArch/LoongArchISelDAGToDAG.cpp

 //=- LoongArchISelDAGToDAG.cpp - A dag to dag inst selector for LoongArch -===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines an instruction selector for the LoongArch target.
 //
 //===----------------------------------------------------------------------===//
 
 #include "LoongArchISelDAGToDAG.h"
+#include "LoongArchISelLowering.h"
 #include "MCTargetDesc/LoongArchMCTargetDesc.h"
 #include "MCTargetDesc/LoongArchMatInt.h"
 #include "llvm/Support/KnownBits.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "loongarch-isel"
 
if (N.getOpcode() == ISD::AND && isa<ConstantSDNode>(N.getOperand(1))) {
 
     // SimplifyDemandedBits may have optimized the mask so try restoring any
     // bits that are known zero.
     KnownBits Known = CurDAG->computeKnownBits(N->getOperand(0));
     if (ShMask.isSubsetOf(AndMask | Known.Zero)) {
       ShAmt = N.getOperand(0);
       return true;
     }
+  } else if (N.getOpcode() == LoongArchISD::BSTRPICK) {
+    // Similar to the above AND, if there is a BSTRPICK on the shift amount, we
+    // can bypass it.
+    assert(isPowerOf2_32(ShiftWidth) && "Unexpected max shift amount!");
+    assert(isa<ConstantSDNode>(N.getOperand(1)) && "Illegal msb operand!");
+    assert(isa<ConstantSDNode>(N.getOperand(2)) && "Illegal lsb operand!");
+    uint64_t msb = N.getConstantOperandVal(1), lsb = N.getConstantOperandVal(2);
+    if (lsb == 0 && Log2_32(ShiftWidth) <= msb + 1) {
+      ShAmt = N.getOperand(0);
+      return true;
+    }
   } else if (N.getOpcode() == ISD::SUB &&
              isa<ConstantSDNode>(N.getOperand(0))) {
     uint64_t Imm = N.getConstantOperandVal(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.
     if (Imm != 0 && Imm % ShiftWidth == 0) {
       SDLoc DL(N);
       EVT VT = N.getValueType();

llvm/lib/Target/LoongArch/LoongArchISelLowering.h

enum NodeType : unsigned {
   // TODO: add more LoongArchISDs
   RET,
   // 32-bit shifts, directly matching the semantics of the named LoongArch
   // instructions.
   SLL_W,
   SRA_W,
   SRL_W,
 
+  BSTRPICK,
+
 };
 } // namespace LoongArchISD
 
 class LoongArchTargetLowering : public TargetLowering {
   const LoongArchSubtarget &Subtarget;
 
 public:
   explicit LoongArchTargetLowering(const TargetMachine &TM,
                                    const LoongArchSubtarget &STI);
 
   const LoongArchSubtarget &getSubtarget() const { return Subtarget; }
 
   // Provide custom lowering hooks for some operations.
   SDValue LowerOperation(SDValue Op, SelectionDAG &DAG) const override;
   void ReplaceNodeResults(SDNode *N, SmallVectorImpl<SDValue> &Results,
                           SelectionDAG &DAG) const override;
 
+  SDValue PerformDAGCombine(SDNode *N, DAGCombinerInfo &DCI) const override;
+
   // This method returns the name of a target specific DAG node.
   const char *getTargetNodeName(unsigned Opcode) const override;
 
   // Lower incoming arguments, copy physregs into vregs.
   SDValue LowerFormalArguments(SDValue Chain, CallingConv::ID CallConv,
                                bool IsVarArg,
                                const SmallVectorImpl<ISD::InputArg> &Ins,
                                const SDLoc &DL, SelectionDAG &DAG,

llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp

LoongArchTargetLowering::LoongArchTargetLowering(const TargetMachine &TM,
 
   setStackPointerRegisterToSaveRestore(LoongArch::R3);
 
   setBooleanContents(ZeroOrOneBooleanContent);
 
   // Function alignments.
   const Align FunctionAlignment(4);
   setMinFunctionAlignment(FunctionAlignment);
+
+  setTargetDAGCombine(ISD::AND);
 }
 
 SDValue LoongArchTargetLowering::LowerOperation(SDValue Op,
                                                 SelectionDAG &DAG) const {
   switch (Op.getOpcode()) {
   default:
     report_fatal_error("unimplemented operand");
   case ISD::SHL_PARTS:
case ISD::SRL:
     if (N->getOperand(1).getOpcode() != ISD::Constant) {
       Results.push_back(customLegalizeToWOp(N, DAG));
       break;
     }
     break;
   }
 }
 
+static SDValue performANDCombine(SDNode *N, SelectionDAG &DAG,
+                                 TargetLowering::DAGCombinerInfo &DCI,
+                                 const LoongArchSubtarget &Subtarget) {
+  if (DCI.isBeforeLegalizeOps())
+    return SDValue();
+
+  SDValue FirstOperand = N->getOperand(0);
+  SDValue SecondOperand = N->getOperand(1);
+  unsigned FirstOperandOpc = FirstOperand.getOpcode();
+  EVT ValTy = N->getValueType(0);
+  SDLoc DL(N);
+  uint64_t lsb, msb;
+  unsigned SMIdx, SMLen;
+  ConstantSDNode *CN;
+  SDValue NewOperand;
+  MVT GRLenVT = Subtarget.getGRLenVT();
+
+  // Op's second operand must be a shifted mask.
+  if (!(CN = dyn_cast<ConstantSDNode>(SecondOperand)) ||
+      !isShiftedMask_64(CN->getZExtValue(), SMIdx, SMLen))
+    return SDValue();
+
+  if (FirstOperandOpc == ISD::SRA || FirstOperandOpc == ISD::SRL) {
+    // Pattern match BSTRPICK.
+    //  $dst = and ((sra or srl) $src , lsb), (2**len - 1)
+    //  => BSTRPICK $dst, $src, msb, lsb
+    //  where msb = lsb + len - 1
+
+    // The second operand of the shift must be an immediate.
+    if (!(CN = dyn_cast<ConstantSDNode>(FirstOperand.getOperand(1))))
+      return SDValue();
+
+    lsb = CN->getZExtValue();
+
+    // Return if the shifted mask does not start at bit 0 or the sum of its
+    // length and lsb exceeds the word's size.
+    if (SMIdx != 0 || lsb + SMLen > ValTy.getSizeInBits())
+      return SDValue();
+
+    NewOperand = FirstOperand.getOperand(0);
+  } else {
+    // Pattern match BSTRPICK.
+    //  $dst = and $src, (2**len- 1) , if len > 12
+    //  => BSTRPICK $dst, $src, msb, lsb
+    //  where lsb = 0 and msb = len - 1
+
+    // If the mask is <= 0xfff, andi can be used instead.
+    if (CN->getZExtValue() <= 0xfff)
+      return SDValue();
+
+    // Return if the mask doesn't start at position 0.
+    if (SMIdx)
+      return SDValue();
+
+    lsb = 0;
+    NewOperand = FirstOperand;
+  }
+  msb = lsb + SMLen - 1;
+  return DAG.getNode(LoongArchISD::BSTRPICK, DL, ValTy, NewOperand,
+                     DAG.getConstant(msb, DL, GRLenVT),
+                     DAG.getConstant(lsb, DL, GRLenVT));
+}
+
+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);
+  }
+  return SDValue();
+}
+
 const char *LoongArchTargetLowering::getTargetNodeName(unsigned Opcode) const {
   switch ((LoongArchISD::NodeType)Opcode) {
   case LoongArchISD::FIRST_NUMBER:
     break;
 
 #define NODE_NAME_CASE(node)                                                   \
   case LoongArchISD::node:                                                     \
     return "LoongArchISD::" #node;
 
     // TODO: Add more target-dependent nodes later.
     NODE_NAME_CASE(RET)
     NODE_NAME_CASE(SLL_W)
     NODE_NAME_CASE(SRA_W)
     NODE_NAME_CASE(SRL_W)
+    NODE_NAME_CASE(BSTRPICK)
   }
 #undef NODE_NAME_CASE
   return nullptr;
 }
 
 //===----------------------------------------------------------------------===//
 //                     Calling Convention Implementation
 //===----------------------------------------------------------------------===//

llvm/lib/Target/LoongArch/LoongArchInstrInfo.td

 // LoongArch specific DAG Nodes.
 //===----------------------------------------------------------------------===//
 
 // Target-dependent type requirements.
 def SDT_LoongArchIntBinOpW : SDTypeProfile<1, 2, [
   SDTCisSameAs<0, 1>, SDTCisSameAs<0, 2>, SDTCisVT<0, i64>
 ]>;
 
+def SDT_LoongArchBStrPick: SDTypeProfile<1, 3, [
+  SDTCisInt<0>, SDTCisSameAs<0, 1>, SDTCisInt<2>, SDTCisSameAs<2, 3>
+]>;
+
 // TODO: Add LoongArch specific DAG Nodes
 // Target-dependent nodes.
 def loongarch_ret : SDNode<"LoongArchISD::RET", SDTNone,
                            [SDNPHasChain, SDNPOptInGlue, SDNPVariadic]>;
 def loongarch_sll_w : SDNode<"LoongArchISD::SLL_W", SDT_LoongArchIntBinOpW>;
 def loongarch_sra_w : SDNode<"LoongArchISD::SRA_W", SDT_LoongArchIntBinOpW>;
 def loongarch_srl_w : SDNode<"LoongArchISD::SRL_W", SDT_LoongArchIntBinOpW>;
+def loongarch_bstrpick
+    : SDNode<"LoongArchISD::BSTRPICK", SDT_LoongArchBStrPick>;
 
 //===----------------------------------------------------------------------===//
 // Operand and SDNode transformation definitions.
 //===----------------------------------------------------------------------===//
 
 class ImmAsmOperand<string prefix, int width, string suffix>
     : AsmOperandClass {
   let Name = prefix # "Imm" # width # suffix;
def : Pat<(select GPR:$cond, GPR:$t, GPR:$f),
           (OR (MASKEQZ GPR:$t, GPR:$cond), (MASKNEZ GPR:$f, GPR:$cond))>;
 
 /// Branches and jumps
 
 let isBarrier = 1, isReturn = 1, isTerminator = 1 in
 def PseudoRET : Pseudo<(outs), (ins), [(loongarch_ret)]>,
                 PseudoInstExpansion<(JIRL R0, R1, 0)>;
 
+/// BSTRPICK
+
+let Predicates = [IsLA32] in
+def : Pat<(loongarch_bstrpick GPR:$rj, uimm5:$msbd, uimm5:$lsbd),
+          (BSTRPICK_W GPR:$rj, uimm5:$msbd, uimm5:$lsbd)>;
+
+let Predicates = [IsLA64] in
+def : Pat<(loongarch_bstrpick GPR:$rj, uimm6:$msbd, uimm6:$lsbd),
+          (BSTRPICK_D GPR:$rj, uimm6:$msbd, uimm6:$lsbd)>;
+
 //===----------------------------------------------------------------------===//
 // Assembler Pseudo Instructions
 //===----------------------------------------------------------------------===//
 
 def : InstAlias<"nop", (ANDI R0, R0, 0)>;
 def : InstAlias<"move $dst, $src", (OR GPR:$dst, GPR:$src, R0)>;
 
 //===----------------------------------------------------------------------===//

llvm/test/CodeGen/LoongArch/bstrpick_d.ll

+; RUN: llc --mtriple=loongarch64 < %s | FileCheck %s
+
+define i64 @lshr40_and255(i64 %a) {
+; CHECK-LABEL: lshr40_and255:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    bstrpick.d $a0, $a0, 47, 40
+; CHECK-NEXT:    jirl $zero, $ra, 0
+  %shr = lshr i64 %a, 40
+  %and = and i64 %shr, 255
+  ret i64 %and
+}
+
+define i64 @ashr50_and511(i64 %a) {
+; CHECK-LABEL: ashr50_and511:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    bstrpick.d $a0, $a0, 58, 50
+; CHECK-NEXT:    jirl $zero, $ra, 0
+  %shr = ashr i64 %a, 50
+  %and = and i64 %shr, 511
+  ret i64 %and
+}
+
+define i64 @zext_i32_to_i64(i32 %a) {
+; CHECK-LABEL: zext_i32_to_i64:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    bstrpick.d $a0, $a0, 31, 0
+; CHECK-NEXT:    jirl $zero, $ra, 0
+  %res = zext i32 %a to i64
+  ret i64 %res
+}
+
+define i64 @and8191(i64 %a) {
+; CHECK-LABEL: and8191:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    bstrpick.d $a0, $a0, 12, 0
+; CHECK-NEXT:    jirl $zero, $ra, 0
+  %and = and i64 %a, 8191
+  ret i64 %and
+}
+
+;; Check that andi but not bstrpick.d is generated.
+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
+}

llvm/test/CodeGen/LoongArch/bstrpick_w.ll

+; RUN: llc --mtriple=loongarch32 < %s | FileCheck %s
+
+define i32 @lshr40_and255(i32 %a) {
+; CHECK-LABEL: lshr40_and255:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    bstrpick.w $a0, $a0, 17, 10
+; CHECK-NEXT:    jirl $zero, $ra, 0
+  %shr = lshr i32 %a, 10
+  %and = and i32 %shr, 255
+  ret i32 %and
+}
+
+define i32 @ashr50_and511(i32 %a) {
+; CHECK-LABEL: ashr50_and511:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    bstrpick.w $a0, $a0, 28, 20
+; CHECK-NEXT:    jirl $zero, $ra, 0
+  %shr = ashr i32 %a, 20
+  %and = and i32 %shr, 511
+  ret i32 %and
+}
+
+define i32 @zext_i16_to_i32(i16 %a) {
+; CHECK-LABEL: zext_i16_to_i32:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    bstrpick.w $a0, $a0, 15, 0
+; CHECK-NEXT:    jirl $zero, $ra, 0
+  %res = zext i16 %a to i32
+  ret i32 %res
+}
+
+define i32 @and8191(i32 %a) {
+; CHECK-LABEL: and8191:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    bstrpick.w $a0, $a0, 12, 0
+; CHECK-NEXT:    jirl $zero, $ra, 0
+  %and = and i32 %a, 8191
+  ret i32 %and
+}
+
+;; Check that andi but not bstrpick.d is generated.
+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
+}

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

 ; LA64-NEXT:    jirl $zero, $ra, 0
   %lshr = lshr i8 %x, %y
   ret i8 %lshr
 }
 
 define i16 @lshr_i16(i16 %x, i16 %y) {
 ; LA32-LABEL: lshr_i16:
 ; LA32:       # %bb.0:
-; LA32-NEXT:    lu12i.w $a2, 15
-; LA32-NEXT:    ori $a2, $a2, 4095
-; LA32-NEXT:    and $a0, $a0, $a2
+; LA32-NEXT:    bstrpick.w $a0, $a0, 15, 0
 ; LA32-NEXT:    srl.w $a0, $a0, $a1
 ; LA32-NEXT:    jirl $zero, $ra, 0
 ;
 ; LA64-LABEL: lshr_i16:
 ; LA64:       # %bb.0:
-; LA64-NEXT:    lu12i.w $a2, 15
-; LA64-NEXT:    ori $a2, $a2, 4095
-; LA64-NEXT:    and $a0, $a0, $a2
+; LA64-NEXT:    bstrpick.d $a0, $a0, 15, 0
 ; LA64-NEXT:    srl.d $a0, $a0, $a1
 ; LA64-NEXT:    jirl $zero, $ra, 0
   %lshr = lshr i16 %x, %y
   ret i16 %lshr
 }
 
 define i32 @lshr_i32(i32 %x, i32 %y) {
 ; LA32-LABEL: lshr_i32: