[AArch64][GlobalISe] Fold G_AND and G_LSHR into ubfm

[paquette]

Summary

This partially implements isBitfieldExtractOpFromAnd from AArch64ISelDAGtoDAG.

It only handles the simplest case. That is,

%mask = G_CONSTANT
%constant = G_CONSTANT
%lshr = G_LSHR %something, %constant
%dst = G_AND %lshr, %mask

becomes

%dst = ubfm %something, immr, imms

To separate the matching from selection, and to avoid adding this to earlySelect, this is done in the post-legalizer combiner. This adds a target-specific G_UBFM instruction which is selected to UBFMWri or UBFMXri depending on the destination size.

The only thing that bothers me about this approach is that in one of the currently-unimplemented cases AArch64ISelDAGtoDAG sometimes inserts a SUBREG_TO_REG. I'm not sure how we would accomplish the same thing in the combiner right now.

[aemerson]

It seems simpler to me to handle this in the selector directly. I’m not sure why a SUBREG_TO_REG is an issue for us in the combiner, but if this could also help O0 it’s better there for now.

If there’s a custom SD node that we want to combine to however that makes more sense to put in the post-legal combiner.

[paquette]

Sure, that sounds fine to me. The code is about the same either way.

[paquette]

• Do the combine in the selector
• Regenerate a testcase which had a codegen change, but also was generally wonked up

[arsenm]

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

1709–1717

Why not use mi_match?

[arsenm]

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

1680

I think we should introduce a generic code for these and combine to it. Currently at least 4 backends are reproducing all of the patterns to match it

[aemerson]

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

1680

A generic opcode for UBFM? Seems a bit specific to me.

[arsenm]

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

1680

It's not very specific, the only target specific detail here is that the fields need to be immediates. BFE and BFM are available on multiple targets, and every target with these shouldn't have to reproduce the same matching logic

[aemerson]

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

1680

Fair enough. A G_BITFIELD_EXTRACT operation could be useful enough to justify inclusion then.

llvm/lib/Target/AArch64/AArch64Combine.td

def AArch64PreLegalizerCombinerHelper: GICombinerHelper<
   "AArch64GenPreLegalizerCombinerHelper", [all_combines,
                                            elide_br_by_inverting_cond,
                                            fconstant_to_constant]> {
   let DisableRuleOption = "aarch64prelegalizercombiner-disable-rule";
   let StateClass = "AArch64PreLegalizerCombinerHelperState";
   let AdditionalArguments = [];
 }
 
-// Matchdata for combines which replace a G_SHUFFLE_VECTOR with a
-// target-specific opcode.
-def shuffle_matchdata : GIDefMatchData<"ShuffleVectorPseudo">;
+// Matchdata for combines which replace a generic instruction with a
+// target-specific generic instruction.
+def pseudo_matchdata : GIDefMatchData<"PseudoInstr">;
 
 def rev : GICombineRule<
-  (defs root:$root, shuffle_matchdata:$matchinfo),
+  (defs root:$root, pseudo_matchdata:$matchinfo),
   (match (wip_match_opcode G_SHUFFLE_VECTOR):$root,
          [{ return matchREV(*${root}, MRI, ${matchinfo}); }]),
-  (apply [{ applyShuffleVectorPseudo(*${root}, ${matchinfo}); }])
+  (apply [{ replaceWithPseudo(*${root}, ${matchinfo}); }])
 >;
 
 def zip : GICombineRule<
-  (defs root:$root, shuffle_matchdata:$matchinfo),
+  (defs root:$root, pseudo_matchdata:$matchinfo),
   (match (wip_match_opcode G_SHUFFLE_VECTOR):$root,
          [{ return matchZip(*${root}, MRI, ${matchinfo}); }]),
-  (apply [{ applyShuffleVectorPseudo(*${root}, ${matchinfo}); }])
+  (apply [{ replaceWithPseudo(*${root}, ${matchinfo}); }])
 >;
 
 def uzp : GICombineRule<
-  (defs root:$root, shuffle_matchdata:$matchinfo),
+  (defs root:$root, pseudo_matchdata:$matchinfo),
   (match (wip_match_opcode G_SHUFFLE_VECTOR):$root,
          [{ return matchUZP(*${root}, MRI, ${matchinfo}); }]),
-  (apply [{ applyShuffleVectorPseudo(*${root}, ${matchinfo}); }])
+  (apply [{ replaceWithPseudo(*${root}, ${matchinfo}); }])
 >;
 
 def dup: GICombineRule <
-  (defs root:$root, shuffle_matchdata:$matchinfo),
+  (defs root:$root, pseudo_matchdata:$matchinfo),
   (match (wip_match_opcode G_SHUFFLE_VECTOR):$root,
          [{ return matchDup(*${root}, MRI, ${matchinfo}); }]),
-  (apply [{ applyShuffleVectorPseudo(*${root}, ${matchinfo}); }])
+  (apply [{ replaceWithPseudo(*${root}, ${matchinfo}); }])
 >;
 
 def trn : GICombineRule<
-  (defs root:$root, shuffle_matchdata:$matchinfo),
+  (defs root:$root, pseudo_matchdata:$matchinfo),
   (match (wip_match_opcode G_SHUFFLE_VECTOR):$root,
          [{ return matchTRN(*${root}, MRI, ${matchinfo}); }]),
-  (apply [{ applyShuffleVectorPseudo(*${root}, ${matchinfo}); }])
+  (apply [{ replaceWithPseudo(*${root}, ${matchinfo}); }])
 >;
 
 def ext: GICombineRule <
-  (defs root:$root, shuffle_matchdata:$matchinfo),
+  (defs root:$root, pseudo_matchdata:$matchinfo),
   (match (wip_match_opcode G_SHUFFLE_VECTOR):$root,
          [{ return matchEXT(*${root}, MRI, ${matchinfo}); }]),
   (apply [{ applyEXT(*${root}, ${matchinfo}); }])
 >;
 
+def and_to_ubfm : GICombineRule <
+  (defs root:$root, pseudo_matchdata:$matchinfo),
+  (match (wip_match_opcode G_AND):$root,
+         [{ return matchUBFMFromAnd(*${root}, MRI, ${matchinfo}); }]),
+  (apply [{ return replaceWithPseudo(*${root}, ${matchinfo}); }])
+>;
+
 // Combines which replace a G_SHUFFLE_VECTOR with a target-specific pseudo
 // instruction.
 def shuffle_vector_pseudos : GICombineGroup<[dup, rev, ext, zip, uzp, trn]>;
 
+// Combines which replace an instruction with a bitfield extract.
+def bitfield_combines : GICombineGroup<[and_to_ubfm]>;
+
 def AArch64PostLegalizerCombinerHelper
     : GICombinerHelper<"AArch64GenPostLegalizerCombinerHelper",
                        [erase_undef_store, combines_for_extload,
-                        sext_trunc_sextload, shuffle_vector_pseudos]> {
+                        shuffle_vector_pseudos, bitfield_combines]> {
   let DisableRuleOption = "aarch64postlegalizercombiner-disable-rule";
 }

llvm/lib/Target/AArch64/AArch64InstrGISel.td

 
 // Represents an ext instruction. Produced post-legalization from
 // G_SHUFFLE_VECTORs with appropriate masks.
 def G_EXT: AArch64GenericInstruction {
   let OutOperandList = (outs type0:$dst);
   let InOperandList = (ins type0:$v1, type0:$v2, untyped_imm_0:$imm);
 }
 
+// Represents a ubfm instruction.
+def G_UBFM: AArch64GenericInstruction {
+  let OutOperandList = (outs type0:$dst);
+  let InOperandList = (ins type0:$src,
+                       untyped_imm_0:$immr,
+                       untyped_imm_0:$imms);
+  let hasSideEffects = 0;
+}
+
 def : GINodeEquiv<G_REV16, AArch64rev16>;
 def : GINodeEquiv<G_REV32, AArch64rev32>;
 def : GINodeEquiv<G_REV64, AArch64rev64>;
 def : GINodeEquiv<G_UZP1, AArch64uzp1>;
 def : GINodeEquiv<G_UZP2, AArch64uzp2>;
 def : GINodeEquiv<G_ZIP1, AArch64zip1>;
 def : GINodeEquiv<G_ZIP2, AArch64zip2>;
 def : GINodeEquiv<G_DUP, AArch64dup>;
 def : GINodeEquiv<G_TRN1, AArch64trn1>;
 def : GINodeEquiv<G_TRN2, AArch64trn2>;
 def : GINodeEquiv<G_EXT, AArch64ext>;

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

bool AArch64InstructionSelector::earlySelectSHL(
   // for a special case of UBFM. Otherwise, we fall back to the imported
   // selector which will match the register variant.
   assert(I.getOpcode() == TargetOpcode::G_SHL && "unexpected op");
   const auto &MO = I.getOperand(2);
   auto VRegAndVal = getConstantVRegVal(MO.getReg(), MRI);
   if (!VRegAndVal)
     return false;
 
   const LLT DstTy = MRI.getType(I.getOperand(0).getReg());
   if (DstTy.isVector())
     return false;
   bool Is64Bit = DstTy.getSizeInBits() == 64;
   auto Imm1Fn = Is64Bit ? selectShiftA_64(MO) : selectShiftA_32(MO);
   auto Imm2Fn = Is64Bit ? selectShiftB_64(MO) : selectShiftB_32(MO);
   MachineIRBuilder MIB(I);
 
   if (!Imm1Fn || !Imm2Fn)
bool AArch64InstructionSelector::earlySelectSHL(
   return constrainSelectedInstRegOperands(*NewI, TII, TRI, RBI);
 }
 
 bool AArch64InstructionSelector::contractCrossBankCopyIntoStore(
     MachineInstr &I, MachineRegisterInfo &MRI) {
   assert(I.getOpcode() == TargetOpcode::G_STORE && "Expected G_STORE");
   // If we're storing a scalar, it doesn't matter what register bank that
   // scalar is on. All that matters is the size.
   //
   // So, if we see something like this (with a 32-bit scalar as an example):
   //
   // %x:gpr(s32) = ... something ...
   // %y:fpr(s32) = COPY %x:gpr(s32)
   // G_STORE %y:fpr(s32)
   //
   // We can fix this up into something like this:
   //
   // G_STORE %x:gpr(s32)
   //
   // And then continue the selection process normally.
   Register DefDstReg = getSrcRegIgnoringCopies(I.getOperand(0).getReg(), MRI);
   if (!DefDstReg.isValid())
     return false;
   LLT DefDstTy = MRI.getType(DefDstReg);
   Register StoreSrcReg = I.getOperand(0).getReg();
if (selectImpl(I, *CoverageInfo))
     return true;
 
   LLT Ty =
       I.getOperand(0).isReg() ? MRI.getType(I.getOperand(0).getReg()) : LLT{};
 
   MachineIRBuilder MIB(I);
 
   switch (Opcode) {
+  case AArch64::G_UBFM: {
+#ifndef NDEBUG
+    // Verify the range of G_UBFM.
+    unsigned Size = Ty.getSizeInBits();
+    int64_t ImmR = I.getOperand(2).getImm();
+    int64_t ImmS = I.getOperand(3).getImm();
+    assert(ImmR < Size && "ImmR is too large!");
+    assert(ImmS < Size && "ImmS is too large!");
+    assert(ImmR >= 0 && "ImmR is too small!");
+    assert(ImmS >= 0 && "ImmS is too small!");
+#endif
+    unsigned Opc =
+        Ty.getSizeInBits() == 32 ? AArch64::UBFMWri : AArch64::UBFMXri;
+    I.setDesc(TII.get(Opc));
+    return constrainSelectedInstRegOperands(I, TII, TRI, RBI);
+  }
   case TargetOpcode::G_BRCOND: {
     if (Ty.getSizeInBits() > 32) {
       // We shouldn't need this on AArch64, but it would be implemented as an
       // EXTRACT_SUBREG followed by a TBNZW because TBNZX has no encoding if the
       // bit being tested is < 32.
       LLVM_DEBUG(dbgs() << "G_BRCOND has type: " << Ty
                         << ", expected at most 32-bits");
       return false;

llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerCombiner.cpp

 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "aarch64-postlegalizer-combiner"
 
 using namespace llvm;
 using namespace MIPatternMatch;
 
-/// Represents a pseudo instruction which replaces a G_SHUFFLE_VECTOR.
-///
-/// Used for matching target-supported shuffles before codegen.
-struct ShuffleVectorPseudo {
-  unsigned Opc; ///< Opcode for the instruction. (E.g. G_ZIP1)
-  Register Dst; ///< Destination register.
-  SmallVector<SrcOp, 2> SrcOps; ///< Source registers.
-  ShuffleVectorPseudo(unsigned Opc, Register Dst,
+/// Represents an instruction which should be used to replace some other
+/// instruction in a combine.
+struct PseudoInstr {
+  unsigned Opc; /// Instruction opcode
+  Register Dst; /// Destination register
+  SmallVector<SrcOp, 2> SrcOps; /// Source operands
+  PseudoInstr(unsigned Opc, Register Dst,
                       std::initializer_list<SrcOp> SrcOps)
       : Opc(Opc), Dst(Dst), SrcOps(SrcOps){};
-  ShuffleVectorPseudo() {}
+  PseudoInstr() {}
 };
 
 /// \returns The splat index of a G_SHUFFLE_VECTOR \p MI when \p MI is a splat.
 /// If \p MI is not a splat, returns None.
 static Optional<int> getSplatIndex(MachineInstr &MI) {
   assert(MI.getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR &&
          "Only G_SHUFFLE_VECTOR can have a splat index!");
   ArrayRef<int> Mask = MI.getOperand(3).getShuffleMask();
for (unsigned i = 0; i != NumElts; i += 2) {
     Idx += 1;
   }
   return true;
 }
 
 /// \return true if a G_SHUFFLE_VECTOR instruction \p MI can be replaced with a
 /// G_REV instruction. Returns the appropriate G_REV opcode in \p Opc.
 static bool matchREV(MachineInstr &MI, MachineRegisterInfo &MRI,
-                     ShuffleVectorPseudo &MatchInfo) {
+                     PseudoInstr &MatchInfo) {
   assert(MI.getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR);
   ArrayRef<int> ShuffleMask = MI.getOperand(3).getShuffleMask();
   Register Dst = MI.getOperand(0).getReg();
   Register Src = MI.getOperand(1).getReg();
   LLT Ty = MRI.getType(Dst);
   unsigned EltSize = Ty.getScalarSizeInBits();
 
   // Element size for a rev cannot be 64.
   if (EltSize == 64)
     return false;
 
   unsigned NumElts = Ty.getNumElements();
 
   // Try to produce G_REV64
   if (isREVMask(ShuffleMask, EltSize, NumElts, 64)) {
-    MatchInfo = ShuffleVectorPseudo(AArch64::G_REV64, Dst, {Src});
+    MatchInfo = PseudoInstr(AArch64::G_REV64, Dst, {Src});
     return true;
   }
 
   // TODO: Produce G_REV32 and G_REV16 once we have proper legalization support.
   // This should be identical to above, but with a constant 32 and constant
   // 16.
   return false;
 }
 
 /// \return true if a G_SHUFFLE_VECTOR instruction \p MI can be replaced with
 /// a G_TRN1 or G_TRN2 instruction.
 static bool matchTRN(MachineInstr &MI, MachineRegisterInfo &MRI,
-                     ShuffleVectorPseudo &MatchInfo) {
+                     PseudoInstr &MatchInfo) {
   assert(MI.getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR);
   unsigned WhichResult;
   ArrayRef<int> ShuffleMask = MI.getOperand(3).getShuffleMask();
   Register Dst = MI.getOperand(0).getReg();
   unsigned NumElts = MRI.getType(Dst).getNumElements();
   if (!isTRNMask(ShuffleMask, NumElts, WhichResult))
     return false;
   unsigned Opc = (WhichResult == 0) ? AArch64::G_TRN1 : AArch64::G_TRN2;
   Register V1 = MI.getOperand(1).getReg();
   Register V2 = MI.getOperand(2).getReg();
-  MatchInfo = ShuffleVectorPseudo(Opc, Dst, {V1, V2});
+  MatchInfo = PseudoInstr(Opc, Dst, {V1, V2});
   return true;
 }
 
 /// \return true if a G_SHUFFLE_VECTOR instruction \p MI can be replaced with
 /// a G_UZP1 or G_UZP2 instruction.
 ///
 /// \param [in] MI - The shuffle vector instruction.
 /// \param [out] MatchInfo - Either G_UZP1 or G_UZP2 on success.
 static bool matchUZP(MachineInstr &MI, MachineRegisterInfo &MRI,
-                     ShuffleVectorPseudo &MatchInfo) {
+                     PseudoInstr &MatchInfo) {
   assert(MI.getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR);
   unsigned WhichResult;
   ArrayRef<int> ShuffleMask = MI.getOperand(3).getShuffleMask();
   Register Dst = MI.getOperand(0).getReg();
   unsigned NumElts = MRI.getType(Dst).getNumElements();
   if (!isUZPMask(ShuffleMask, NumElts, WhichResult))
     return false;
   unsigned Opc = (WhichResult == 0) ? AArch64::G_UZP1 : AArch64::G_UZP2;
   Register V1 = MI.getOperand(1).getReg();
   Register V2 = MI.getOperand(2).getReg();
-  MatchInfo = ShuffleVectorPseudo(Opc, Dst, {V1, V2});
+  MatchInfo = PseudoInstr(Opc, Dst, {V1, V2});
   return true;
 }
 
 static bool matchZip(MachineInstr &MI, MachineRegisterInfo &MRI,
-                     ShuffleVectorPseudo &MatchInfo) {
+                     PseudoInstr &MatchInfo) {
   assert(MI.getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR);
   unsigned WhichResult;
   ArrayRef<int> ShuffleMask = MI.getOperand(3).getShuffleMask();
   Register Dst = MI.getOperand(0).getReg();
   unsigned NumElts = MRI.getType(Dst).getNumElements();
   if (!isZipMask(ShuffleMask, NumElts, WhichResult))
     return false;
   unsigned Opc = (WhichResult == 0) ? AArch64::G_ZIP1 : AArch64::G_ZIP2;
   Register V1 = MI.getOperand(1).getReg();
   Register V2 = MI.getOperand(2).getReg();
-  MatchInfo = ShuffleVectorPseudo(Opc, Dst, {V1, V2});
+  MatchInfo = PseudoInstr(Opc, Dst, {V1, V2});
   return true;
 }
 
 /// Helper function for matchDup.
 static bool matchDupFromInsertVectorElt(int Lane, MachineInstr &MI,
                                         MachineRegisterInfo &MRI,
-                                        ShuffleVectorPseudo &MatchInfo) {
+                                        PseudoInstr &MatchInfo) {
   if (Lane != 0)
     return false;
 
   // Try to match a vector splat operation into a dup instruction.
   // We're looking for this pattern:
   //
   // %scalar:gpr(s64) = COPY $x0
   // %undef:fpr(<2 x s64>) = G_IMPLICIT_DEF
if (!getOpcodeDef(TargetOpcode::G_IMPLICIT_DEF, InsMI->getOperand(1).getReg(),
                     MRI))
     return false;
 
   // Match the index constant 0.
   int64_t Index = 0;
   if (!mi_match(InsMI->getOperand(3).getReg(), MRI, m_ICst(Index)) || Index)
     return false;
 
-  MatchInfo = ShuffleVectorPseudo(AArch64::G_DUP, MI.getOperand(0).getReg(),
+  MatchInfo = PseudoInstr(AArch64::G_DUP, MI.getOperand(0).getReg(),
                                   {InsMI->getOperand(2).getReg()});
   return true;
 }
 
 /// Helper function for matchDup.
 static bool matchDupFromBuildVector(int Lane, MachineInstr &MI,
                                     MachineRegisterInfo &MRI,
-                                    ShuffleVectorPseudo &MatchInfo) {
+                                    PseudoInstr &MatchInfo) {
   assert(Lane >= 0 && "Expected positive lane?");
   // Test if the LHS is a BUILD_VECTOR. If it is, then we can just reference the
   // lane's definition directly.
   auto *BuildVecMI = getOpcodeDef(TargetOpcode::G_BUILD_VECTOR,
                                   MI.getOperand(1).getReg(), MRI);
   if (!BuildVecMI)
     return false;
   Register Reg = BuildVecMI->getOperand(Lane + 1).getReg();
   MatchInfo =
-      ShuffleVectorPseudo(AArch64::G_DUP, MI.getOperand(0).getReg(), {Reg});
+      PseudoInstr(AArch64::G_DUP, MI.getOperand(0).getReg(), {Reg});
   return true;
 }
 
 static bool matchDup(MachineInstr &MI, MachineRegisterInfo &MRI,
-                     ShuffleVectorPseudo &MatchInfo) {
+                     PseudoInstr &MatchInfo) {
   assert(MI.getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR);
   auto MaybeLane = getSplatIndex(MI);
   if (!MaybeLane)
     return false;
   int Lane = *MaybeLane;
   // If this is undef splat, generate it via "just" vdup, if possible.
   if (Lane < 0)
     Lane = 0;
   if (matchDupFromInsertVectorElt(Lane, MI, MRI, MatchInfo))
     return true;
   if (matchDupFromBuildVector(Lane, MI, MRI, MatchInfo))
     return true;
   return false;
 }
 
 static bool matchEXT(MachineInstr &MI, MachineRegisterInfo &MRI,
-                     ShuffleVectorPseudo &MatchInfo) {
+                     PseudoInstr &MatchInfo) {
   assert(MI.getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR);
   Register Dst = MI.getOperand(0).getReg();
   auto ExtInfo = getExtMask(MI.getOperand(3).getShuffleMask(),
                             MRI.getType(Dst).getNumElements());
   if (!ExtInfo)
     return false;
   bool ReverseExt;
   uint64_t Imm;
   std::tie(ReverseExt, Imm) = *ExtInfo;
   Register V1 = MI.getOperand(1).getReg();
   Register V2 = MI.getOperand(2).getReg();
   if (ReverseExt)
     std::swap(V1, V2);
   uint64_t ExtFactor = MRI.getType(V1).getScalarSizeInBits() / 8;
   Imm *= ExtFactor;
-  MatchInfo = ShuffleVectorPseudo(AArch64::G_EXT, Dst, {V1, V2, Imm});
+  MatchInfo = PseudoInstr(AArch64::G_EXT, Dst, {V1, V2, Imm});
   return true;
 }
 
-/// Replace a G_SHUFFLE_VECTOR instruction with a pseudo.
-/// \p Opc is the opcode to use. \p MI is the G_SHUFFLE_VECTOR.
-static bool applyShuffleVectorPseudo(MachineInstr &MI,
-                                     ShuffleVectorPseudo &MatchInfo) {
+/// Replace a \p MI with a pseudo instruction defined in \p MatchInfo.
+static bool replaceWithPseudo(MachineInstr &MI, PseudoInstr &MatchInfo) {
   MachineIRBuilder MIRBuilder(MI);
   MIRBuilder.buildInstr(MatchInfo.Opc, {MatchInfo.Dst}, MatchInfo.SrcOps);
   MI.eraseFromParent();
   return true;
 }
 
 /// Replace a G_SHUFFLE_VECTOR instruction with G_EXT.
 /// Special-cased because the constant operand must be emitted as a G_CONSTANT
 /// for the imported tablegen patterns to work.
-static bool applyEXT(MachineInstr &MI, ShuffleVectorPseudo &MatchInfo) {
+static bool applyEXT(MachineInstr &MI, PseudoInstr &MatchInfo) {
   MachineIRBuilder MIRBuilder(MI);
   // Tablegen patterns expect an i32 G_CONSTANT as the final op.
   auto Cst =
       MIRBuilder.buildConstant(LLT::scalar(32), MatchInfo.SrcOps[2].getImm());
   MIRBuilder.buildInstr(MatchInfo.Opc, {MatchInfo.Dst},
                         {MatchInfo.SrcOps[0], MatchInfo.SrcOps[1], Cst});
   MI.eraseFromParent();
   return true;
 }
 
+/// Return true if a G_AND and a G_LSHR can be replaced with a G_UBFM.
+static bool matchUBFMFromAnd(MachineInstr &MI, MachineRegisterInfo &MRI,
+                             PseudoInstr &MatchInfo) {
+  assert(MI.getOpcode() == TargetOpcode::G_AND);
+  // Look for the following:
+  // %mask = G_CONSTANT
+  // %constant = G_CONSTANT
+  // %lshr = G_LSHR %something, %constant
+  // %dst = G_AND %lshr, %mask
+  //
+  // And produce
+  //
+  // %dst = G_UBFM %something, immr, imms
+  //
+  // TODO: Handle other cases from isBitfieldExtractOpFromAnd in
+  // AArch64ISelDAGToDAG.
+  Register Dst = MI.getOperand(0).getReg();
+  auto DstTy = MRI.getType(Dst);
+  if (DstTy.isVector())
+    return false;
+
+  // UBFM only supports 32-bit and 64-bit registers.
+  unsigned DstSize = DstTy.getSizeInBits();
+  if (DstSize != 32 && DstSize != 64)
+    return false;
+  Register LHS = MI.getOperand(1).getReg();
+  Register RHS = MI.getOperand(2).getReg();
+
+  // Look for a mask on the G_AND's RHS.
+  // The immediate is a mask of the low bits iff imm & (imm+1) == 0
+  auto MaybeLowBitMask = getConstantVRegValWithLookThrough(RHS, MRI);
+  if (!MaybeLowBitMask)
+    return false;
+  uint64_t LowBitMask = MaybeLowBitMask->Value;
+  if (LowBitMask & (LowBitMask + 1))
+    return false;
+
+  // Look for %lshr = G_LSHR %something, %constant.
+  MachineInstr *Lshr =
+      getOpcodeDef(TargetOpcode::G_LSHR, LHS, MRI);
+  if (!Lshr)
+    return false;
+
+  auto MaybeLshrImm =
+      getConstantVRegValWithLookThrough(Lshr->getOperand(2).getReg(), MRI);
+  if (!MaybeLshrImm)
+    return false;
+  // Check that the immediates we want to pass to the UBFM are legal. Both must
+  // be in the range [0, DstSize).
+  uint64_t ImmR = MaybeLshrImm->Value;
+  if (ImmR >= DstSize)
+    return false;
+  uint64_t ImmS = ImmR +
+                  (DstSize == 32 ? countTrailingOnes<uint32_t>(LowBitMask)
+                                 : countTrailingOnes<uint64_t>(LowBitMask)) -
+                  1;
+  if (ImmS >= DstSize)
+    return false;
+  MatchInfo = PseudoInstr(AArch64::G_UBFM, Dst,
+                          {Lshr->getOperand(1).getReg(), ImmR, ImmS});
+  return true;
+}
+
 #define AARCH64POSTLEGALIZERCOMBINERHELPER_GENCOMBINERHELPER_DEPS
 #include "AArch64GenPostLegalizeGICombiner.inc"
 #undef AARCH64POSTLEGALIZERCOMBINERHELPER_GENCOMBINERHELPER_DEPS
 
 namespace {
 #define AARCH64POSTLEGALIZERCOMBINERHELPER_GENCOMBINERHELPER_H
 #include "AArch64GenPostLegalizeGICombiner.inc"
 #undef AARCH64POSTLEGALIZERCOMBINERHELPER_GENCOMBINERHELPER_H

llvm/test/CodeGen/AArch64/GlobalISel/postlegalizer-combiner-ubfm.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# Check combines which produce a G_UBFM.
+#
+# RUN: llc -mtriple aarch64 -run-pass=aarch64-postlegalizer-combiner -verify-machineinstrs %s -o - | FileCheck %s
+
+...
+---
+name:            and_to_ubfm_s32
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0
+
+    ; CHECK-LABEL: name: and_to_ubfm_s32
+    ; CHECK: liveins: $w0
+    ; CHECK: %copy:_(s32) = COPY $w0
+    ; CHECK: %and:_(s32) = G_UBFM %copy, 22, 22
+    ; CHECK: $w0 = COPY %and(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %copy:_(s32) = COPY $w0
+    %immr:_(s64) = G_CONSTANT i64 22
+    %mask:_(s32) = G_CONSTANT i32 1
+    %lshr:_(s32) = G_LSHR %copy, %immr(s64)
+    %and:_(s32) = G_AND %lshr, %mask
+    $w0 = COPY %and(s32)
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            and_to_ubfm_s64
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0
+    ; CHECK-LABEL: name: and_to_ubfm_s64
+    ; CHECK: liveins: $x0
+    ; CHECK: %copy:_(s64) = COPY $x0
+    ; CHECK: %and:_(s64) = G_UBFM %copy, 22, 22
+    ; CHECK: $x0 = COPY %and(s64)
+    ; CHECK: RET_ReallyLR implicit $x0
+    %copy:_(s64) = COPY $x0
+    %immr:_(s64) = G_CONSTANT i64 22
+    %mask:_(s64) = G_CONSTANT i64 1
+    %lshr:_(s64) = G_LSHR %copy, %immr(s64)
+    %and:_(s64) = G_AND %lshr, %mask
+    $x0 = COPY %and(s64)
+    RET_ReallyLR implicit $x0
+
+...
+---
+name:            too_large_immr_s32
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0
+    ; We can't combine here because both of the immediates passed to UBFM must
+    ; be smaller than the size of the register.
+    ;
+    ; In this case, immr is too large.
+    ;
+    ; CHECK-LABEL: name: too_large_immr_s32
+    ; CHECK: liveins: $w0
+    ; CHECK: %copy:_(s32) = COPY $w0
+    ; CHECK: %immr:_(s32) = G_CONSTANT i32 40
+    ; CHECK: %mask:_(s32) = G_CONSTANT i32 1
+    ; CHECK: %lshr:_(s32) = G_LSHR %copy, %immr(s32)
+    ; CHECK: %and:_(s32) = G_AND %lshr, %mask
+    ; CHECK: $w0 = COPY %and(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %copy:_(s32) = COPY $w0
+    %immr:_(s32) = G_CONSTANT i32 40
+    %mask:_(s32) = G_CONSTANT i32 1
+    %lshr:_(s32) = G_LSHR %copy, %immr(s32)
+    %and:_(s32) = G_AND %lshr, %mask
+    $w0 = COPY %and(s32)
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            too_large_imms_s32
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0
+    ; We can't combine here because both of the immediates passed to UBFM must
+    ; be smaller than the size of the register.
+    ;
+    ; Trailing ones of the mask: 2
+    ; immr = 31
+    ;
+    ; -> imms = 31 + 2 - 1 = 32, which is too large.
+    ;
+    ; CHECK-LABEL: name: too_large_imms_s32
+    ; CHECK: liveins: $w0
+    ; CHECK: %copy:_(s32) = COPY $w0
+    ; CHECK: %immr:_(s64) = G_CONSTANT i64 31
+    ; CHECK: %mask:_(s32) = G_CONSTANT i32 3
+    ; CHECK: %lshr:_(s32) = G_LSHR %copy, %immr(s64)
+    ; CHECK: %and:_(s32) = G_AND %lshr, %mask
+    ; CHECK: $w0 = COPY %and(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %copy:_(s32) = COPY $w0
+    %immr:_(s64) = G_CONSTANT i64 31
+    %mask:_(s32) = G_CONSTANT i32 3
+    %lshr:_(s32) = G_LSHR %copy, %immr(s64)
+    %and:_(s32) = G_AND %lshr, %mask
+    $w0 = COPY %and(s32)
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            too_large_immr_s64
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0
+    ; CHECK-LABEL: name: too_large_immr_s64
+    ; CHECK: liveins: $x0
+    ; CHECK: %copy:_(s64) = COPY $x0
+    ; CHECK: %immr:_(s64) = G_CONSTANT i64 64
+    ; CHECK: %mask:_(s64) = G_CONSTANT i64 1
+    ; CHECK: %lshr:_(s64) = G_LSHR %copy, %immr(s64)
+    ; CHECK: %and:_(s64) = G_AND %lshr, %mask
+    ; CHECK: $x0 = COPY %and(s64)
+    ; CHECK: RET_ReallyLR implicit $x0
+    %copy:_(s64) = COPY $x0
+    %immr:_(s64) = G_CONSTANT i64 64
+    %mask:_(s64) = G_CONSTANT i64 1
+    %lshr:_(s64) = G_LSHR %copy, %immr(s64)
+    %and:_(s64) = G_AND %lshr, %mask
+    $x0 = COPY %and(s64)
+    RET_ReallyLR implicit $x0
+
+...
+---
+name:            too_large_imms_s64
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0
+    ; We can't combine here because both of the immediates passed to UBFM must
+    ; be smaller than the size of the register.
+    ;
+    ; Trailing ones of the mask: 3
+    ; immr = 62
+    ;
+    ; -> imms = 62 + 3 = 65, which is too large.
+    ;
+    ; CHECK-LABEL: name: too_large_imms_s64
+    ; CHECK: liveins: $x0
+    ; CHECK: %copy:_(s64) = COPY $x0
+    ; CHECK: %immr:_(s64) = G_CONSTANT i64 62
+    ; CHECK: %mask:_(s64) = G_CONSTANT i64 7
+    ; CHECK: %lshr:_(s64) = G_LSHR %copy, %immr(s64)
+    ; CHECK: %and:_(s64) = G_AND %lshr, %mask
+    ; CHECK: $x0 = COPY %and(s64)
+    ; CHECK: RET_ReallyLR implicit $x0
+    %copy:_(s64) = COPY $x0
+    %immr:_(s64) = G_CONSTANT i64 62
+    %mask:_(s64) = G_CONSTANT i64 7
+    %lshr:_(s64) = G_LSHR %copy, %immr(s64)
+    %and:_(s64) = G_AND %lshr, %mask
+    $x0 = COPY %and(s64)
+    RET_ReallyLR implicit $x0
+
+...
+---
+name:            bad_low_bit_mask
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0
+    ; mask = 2
+    ; 2 & (2 + 1) != 0, so this is not a low-bit mask.
+    ;
+    ; CHECK-LABEL: name: bad_low_bit_mask
+    ; CHECK: liveins: $w0
+    ; CHECK: %copy:_(s32) = COPY $w0
+    ; CHECK: %mask:_(s32) = G_CONSTANT i32 2
+    ; CHECK: %immr:_(s64) = G_CONSTANT i64 20
+    ; CHECK: %lshr:_(s32) = G_LSHR %copy, %immr(s64)
+    ; CHECK: %and:_(s32) = G_AND %lshr, %mask
+    ; CHECK: $w0 = COPY %and(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %copy:_(s32) = COPY $w0
+    %mask:_(s32) = G_CONSTANT i32 2
+    %immr:_(s64) = G_CONSTANT i64 20
+    %lshr:_(s32) = G_LSHR %copy, %immr(s64)
+    %and:_(s32) = G_AND %lshr, %mask
+    $w0 = COPY %and(s32)
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            dont_fold_negative_immr
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0
+    ; immr = -10, which is not allowed for a ubfm.
+    ;
+    ; CHECK-LABEL: name: dont_fold_negative_immr
+    ; CHECK: liveins: $w0
+    ; CHECK: %copy:_(s32) = COPY $w0
+    ; CHECK: %immr:_(s32) = G_CONSTANT i32 -10
+    ; CHECK: %mask:_(s32) = G_CONSTANT i32 1
+    ; CHECK: %lshr:_(s32) = G_LSHR %copy, %immr(s32)
+    ; CHECK: %and:_(s32) = G_AND %lshr, %mask
+    ; CHECK: $w0 = COPY %and(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %copy:_(s32) = COPY $w0
+    %immr:_(s32) = G_CONSTANT i32 -10
+    %mask:_(s32) = G_CONSTANT i32 1
+    %lshr:_(s32) = G_LSHR %copy, %immr(s32)
+    %and:_(s32) = G_AND %lshr, %mask
+    $w0 = COPY %and(s32)
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            dont_fold_negative_imms
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0
+    ; Imms = 0 + trailing ones(0) - 1 = -1, which is not allowed for a ubfm.
+    ;
+    ; CHECK-LABEL: name: dont_fold_negative_imms
+    ; CHECK: liveins: $w0
+    ; CHECK: %copy:_(s32) = COPY $w0
+    ; CHECK: %immr:_(s32) = G_CONSTANT i32 0
+    ; CHECK: %mask:_(s32) = G_CONSTANT i32 0
+    ; CHECK: %lshr:_(s32) = G_LSHR %copy, %immr(s32)
+    ; CHECK: %and:_(s32) = G_AND %lshr, %mask
+    ; CHECK: $w0 = COPY %and(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %copy:_(s32) = COPY $w0
+    %immr:_(s32) = G_CONSTANT i32 0
+    %mask:_(s32) = G_CONSTANT i32 0
+    %lshr:_(s32) = G_LSHR %copy, %immr(s32)
+    %and:_(s32) = G_AND %lshr, %mask
+    $w0 = COPY %and(s32)
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            zero
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0
+    ; 0 is a valid value for immr and imms.
+    ; immr = 0
+    ; imms = 0 + trailing ones(1) - 1 = 0
+    ;
+    ; CHECK-LABEL: name: zero
+    ; CHECK: liveins: $w0
+    ; CHECK: %copy:_(s32) = COPY $w0
+    ; CHECK: %and:_(s32) = G_UBFM %copy, 0, 0
+    ; CHECK: $w0 = COPY %and(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %copy:_(s32) = COPY $w0
+    %immr:_(s32) = G_CONSTANT i32 0
+    %mask:_(s32) = G_CONSTANT i32 1
+    %lshr:_(s32) = G_LSHR %copy, %immr(s32)
+    %and:_(s32) = G_AND %lshr, %mask
+    $w0 = COPY %and(s32)
+    RET_ReallyLR implicit $w0

llvm/test/CodeGen/AArch64/GlobalISel/select-ubfm.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+#
+# Test G_UBFM selection.
+#
+# RUN: llc -mtriple aarch64 -run-pass=instruction-select -verify-machineinstrs %s -o - | FileCheck %s
+
+...
+---
+name:            ubfm_s64
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0
+    ; CHECK-LABEL: name: ubfm_s64
+    ; CHECK: liveins: $x0
+    ; CHECK: %copy:gpr64 = COPY $x0
+    ; CHECK: %ubfm:gpr64 = UBFMXri %copy, 0, 0
+    ; CHECK: $x0 = COPY %ubfm
+    ; CHECK: RET_ReallyLR implicit $x0
+    %copy:gpr(s64) = COPY $x0
+    %ubfm:gpr(s64) = G_UBFM %copy, 0, 0
+    $x0 = COPY %ubfm(s64)
+    RET_ReallyLR implicit $x0
+
+...
+---
+name:            ubfm_s32
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0
+    ; CHECK-LABEL: name: ubfm_s32
+    ; CHECK: liveins: $w0
+    ; CHECK: %copy:gpr32 = COPY $w0
+    ; CHECK: %ubfm:gpr32 = UBFMWri %copy, 0, 0
+    ; CHECK: $w0 = COPY %ubfm
+    ; CHECK: RET_ReallyLR implicit $w0
+    %copy:gpr(s32) = COPY $w0
+    %ubfm:gpr(s32) = G_UBFM %copy, 0, 0
+    $w0 = COPY %ubfm(s32)
+    RET_ReallyLR implicit $w0
+
+...