[GlobalISel] Combine G_SEXT_INREG + right shift -> G_SBFX

[paquette]

Summary

Only handling the constant case right now, and no handling for truncs.

This is only done post-legalization for now. Once the legalizer knows how to decompose these back into shifts, this requirement can probably be removed.

[aemerson]

LGTM.

[paquette]

I made this AArch64-specific for now.

Since G_UBFX/G_SBFX requires custom legalization, the legalization hooks in the CombinerHelper won't work for AArch64. Not sure what the best way to go about fixing that is.

llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h

CombinerHelper(GISelChangeObserver &Observer, MachineIRBuilder &B,
                  const LegalizerInfo *LI = nullptr);
 
   GISelKnownBits *getKnownBits() const {
     return KB;
   }
 
   const TargetLowering &getTargetLowering() const;
 
+  /// \returns true if the combine is running prior to legalization.
+  bool isBeforeLegalizer() const { return !LI; }
+
+  /// \returns true if \p Query is legal.
+  /// \note Assumes that the combine is running post-legalization.
+  bool isLegal(const LegalityQuery &Query) const;
+
   /// \return true if the combine is running prior to legalization, or if \p
   /// Query is legal on the target.
   bool isLegalOrBeforeLegalizer(const LegalityQuery &Query) const;
 
   /// MachineRegisterInfo::replaceRegWith() and inform the observer of the changes
   void replaceRegWith(MachineRegisterInfo &MRI, Register FromReg, Register ToReg) const;
 
   /// Replace a single register operand with a new register and inform the
public:
   ///  sN *a = ...
   ///  sM val = a[0] | (a[1] << N) | (a[2] << 2N) | (a[3] << 3N) ...
   /// \endcode
   ///
   /// And check if the tree can be replaced with a M-bit load + possibly a
   /// bswap.
   bool matchLoadOrCombine(MachineInstr &MI,
                           std::function<void(MachineIRBuilder &)> &MatchInfo);
-  bool applyLoadOrCombine(MachineInstr &MI,
-                          std::function<void(MachineIRBuilder &)> &MatchInfo);
 
   bool matchExtendThroughPhis(MachineInstr &MI, MachineInstr *&ExtMI);
   bool applyExtendThroughPhis(MachineInstr &MI, MachineInstr *&ExtMI);
 
   bool matchExtractVecEltBuildVec(MachineInstr &MI, Register &Reg);
   void applyExtractVecEltBuildVec(MachineInstr &MI, Register &Reg);
 
   bool matchExtractAllEltsFromBuildVector(
       MachineInstr &MI,
       SmallVectorImpl<std::pair<Register, MachineInstr *>> &MatchInfo);
   void applyExtractAllEltsFromBuildVector(
       MachineInstr &MI,
       SmallVectorImpl<std::pair<Register, MachineInstr *>> &MatchInfo);
 
+  /// Form a G_SBFX from a G_SEXT_INREG fed by a right shift.
+  bool matchBitfieldExtractFromSExtInReg(
+      MachineInstr &MI, std::function<void(MachineIRBuilder &)> &MatchInfo);
+
+  /// Use a function which takes in a MachineIRBuilder to perform a combine.
+  bool applyBuildFn(MachineInstr &MI,
+                    std::function<void(MachineIRBuilder &)> &MatchInfo);
+
   /// Try to transform \p MI by using all of the above
   /// combine functions. Returns true if changed.
   bool tryCombine(MachineInstr &MI);
 
 private:
   // Memcpy family optimization helpers.
   bool optimizeMemcpy(MachineInstr &MI, Register Dst, Register Src,
                       unsigned KnownLen, Align DstAlign, Align SrcAlign,

llvm/include/llvm/Target/GlobalISel/Combine.td

 def apply;
 /// All arguments of the apply operator must be subclasses of GIApplyKind, or
 /// sub-dags whose operator is GIApplyKindWithArgs, or an MIR block
 /// (deprecated).
 class GIApplyKind;
 class GIApplyKindWithArgs;
 
 def register_matchinfo: GIDefMatchData<"Register">;
+def build_fn_matchinfo :
+GIDefMatchData<"std::function<void(MachineIRBuilder &)>">;
 
 def copy_prop : GICombineRule<
   (defs root:$d),
   (match (COPY $d, $s):$mi,
          [{ return Helper.matchCombineCopy(*${mi}); }]),
   (apply [{ Helper.applyCombineCopy(*${mi}); }])>;
 
 def extending_loads : GICombineRule<

 
 def regs_small_vec : GIDefMatchData<"SmallVector<Register, 4>">;
 def combine_insert_vec_elts_build_vector : GICombineRule<
   (defs root:$root, regs_small_vec:$info),
   (match (wip_match_opcode G_INSERT_VECTOR_ELT):$root,
     [{ return Helper.matchCombineInsertVecElts(*${root}, ${info}); }]),
   (apply [{ return Helper.applyCombineInsertVecElts(*${root}, ${info}); }])>;
 
-def load_or_combine_matchdata :
-GIDefMatchData<"std::function<void(MachineIRBuilder &)>">;
 def load_or_combine : GICombineRule<
-  (defs root:$root, load_or_combine_matchdata:$info),
+  (defs root:$root, build_fn_matchinfo:$info),
   (match (wip_match_opcode G_OR):$root,
     [{ return Helper.matchLoadOrCombine(*${root}, ${info}); }]),
-  (apply [{ return Helper.applyLoadOrCombine(*${root}, ${info}); }])>;
+  (apply [{ return Helper.applyBuildFn(*${root}, ${info}); }])>;
+
+def bitfield_extract_from_sext_inreg : GICombineRule<
+  (defs root:$root, build_fn_matchinfo:$info),
+  (match (wip_match_opcode G_SEXT_INREG):$root,
+    [{ return Helper.matchBitfieldExtractFromSExtInReg(*${root}, ${info}); }]),
+  (apply [{ return Helper.applyBuildFn(*${root}, ${info}); }])>;
 
 def extend_through_phis_matchdata: GIDefMatchData<"MachineInstr*">;
 def extend_through_phis : GICombineRule<
   (defs root:$root, extend_through_phis_matchdata:$matchinfo),
   (match (wip_match_opcode G_PHI):$root,
     [{ return Helper.matchExtendThroughPhis(*${root}, ${matchinfo}); }]),
   (apply [{ return Helper.applyExtendThroughPhis(*${root}, ${matchinfo}); }])>;
 
def extract_all_elts_from_build_vector : GICombineRule<
   (match (wip_match_opcode G_BUILD_VECTOR):$root,
     [{ return Helper.matchExtractAllEltsFromBuildVector(*${root}, ${matchinfo}); }]),
   (apply [{ Helper.applyExtractAllEltsFromBuildVector(*${root}, ${matchinfo}); }])>;
 
 def extract_vec_elt_combines : GICombineGroup<[
   extract_vec_elt_build_vec,
   extract_all_elts_from_build_vector]>;
 
+def form_bitfield_extract : GICombineGroup<[
+  bitfield_extract_from_sext_inreg]>;
+
 // FIXME: These should use the custom predicate feature once it lands.
 def undef_combines : GICombineGroup<[undef_to_fp_zero, undef_to_int_zero,
                                      undef_to_negative_one,
                                      binop_left_undef_to_zero,
                                      propagate_undef_any_op,
                                      propagate_undef_all_ops,
                                      propagate_undef_shuffle_mask,
                                      erase_undef_store]>;

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

if (!BigEndian && !LittleEndian)
       return None;
   }
 
   assert((BigEndian != LittleEndian) &&
          "Pattern cannot be both big and little endian!");
   return BigEndian;
 }
 
+bool CombinerHelper::isLegal(const LegalityQuery &Query) const {
+  assert(LI && "Expected to have LegalizerInfo?");
+  return LI->getAction(Query).Action == LegalizeActions::Legal;
+}
+
 bool CombinerHelper::isLegalOrBeforeLegalizer(
     const LegalityQuery &Query) const {
-  return !LI || LI->getAction(Query).Action == LegalizeActions::Legal;
+  return isBeforeLegalizer() || isLegal(Query);
 }
 
 void CombinerHelper::replaceRegWith(MachineRegisterInfo &MRI, Register FromReg,
                                     Register ToReg) const {
   Observer.changingAllUsesOfReg(MRI, FromReg);
 
   if (MRI.constrainRegAttrs(ToReg, FromReg))
     MRI.replaceRegWith(FromReg, ToReg);
void CombinerHelper::applyExtractAllEltsFromBuildVector(
   for (auto &Pair : SrcDstPairs) {
     auto *ExtMI = Pair.second;
     replaceRegWith(MRI, ExtMI->getOperand(0).getReg(), Pair.first);
     ExtMI->eraseFromParent();
   }
   MI.eraseFromParent();
 }
 
-bool CombinerHelper::applyLoadOrCombine(
+bool CombinerHelper::matchBitfieldExtractFromSExtInReg(
+    MachineInstr &MI, std::function<void(MachineIRBuilder &)> &MatchInfo) {
+  assert(MI.getOpcode() == TargetOpcode::G_SEXT_INREG);
+
+  // TODO: Need legalizer support for splitting these back into shifts +
+  // extends. For now, only allow the combine post-legalization.
+  if (isBeforeLegalizer())
+    return false;
+  Register Dst = MI.getOperand(0).getReg();
+  Register Src = MI.getOperand(1).getReg();
+  int64_t Width = MI.getOperand(2).getImm();
+
+  LLT Ty = MRI.getType(Src);
+  if (!isLegal({TargetOpcode::G_SBFX, {Ty}}))
+    return false;
+
+  // TODO: Some targets don't need a constant. Support that.
+  Register ShiftSrc;
+  int64_t ShiftImm;
+  if (!mi_match(Src, MRI,
+                m_OneNonDBGUse(
+                    m_any_of(m_GAShr(m_Reg(ShiftSrc), m_ICst(ShiftImm)),
+                             m_GLShr(m_Reg(ShiftSrc), m_ICst(ShiftImm))))))
+    return false;
+  if (ShiftImm < 0 || ShiftImm + Width > Ty.getSizeInBits())
+    return false;
+  MatchInfo = [=](MachineIRBuilder &B) {
+    auto Cst1 = B.buildConstant(Ty, ShiftImm);
+    auto Cst2 = B.buildConstant(Ty, ShiftImm + Width - 1);
+    B.buildInstr(TargetOpcode::G_SBFX, {Dst}, {ShiftSrc, Cst1, Cst2});
+  };
+  return true;
+}
+
+bool CombinerHelper::applyBuildFn(
     MachineInstr &MI, std::function<void(MachineIRBuilder &)> &MatchInfo) {
   Builder.setInstrAndDebugLoc(MI);
   MatchInfo(Builder);
   MI.eraseFromParent();
   return true;
 }
 
 bool CombinerHelper::tryCombine(MachineInstr &MI) {
   if (tryCombineCopy(MI))
     return true;
   if (tryCombineExtendingLoads(MI))
     return true;
   if (tryCombineIndexedLoadStore(MI))
     return true;
   return false;
 }

llvm/lib/Target/AArch64/AArch64Combine.td

 // Post-legalization combines which are primarily optimizations.
 def AArch64PostLegalizerCombinerHelper
     : GICombinerHelper<"AArch64GenPostLegalizerCombinerHelper",
                        [copy_prop, erase_undef_store, combines_for_extload,
                         sext_trunc_sextload,
                         hoist_logic_op_with_same_opcode_hands,
                         redundant_and, xor_of_and_with_same_reg,
                         extractvecelt_pairwise_add, redundant_or,
-                        mul_const, redundant_sext_inreg]> {
+                        mul_const, redundant_sext_inreg,
+                        form_bitfield_extract]> {
   let DisableRuleOption = "aarch64postlegalizercombiner-disable-rule";
 }

llvm/test/CodeGen/AArch64/GlobalISel/form-bitfield-extract.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple aarch64 -run-pass=aarch64-postlegalizer-combiner -verify-machineinstrs %s -o - | FileCheck %s
+
+# Check that we can fold a G_SEXT_INREG fed by a G_ASHR/G_LSHR into a G_SBFX.
+
+...
+---
+name:            sextinreg_ashr_to_sbfx
+tracksRegLiveness: true
+legalized: true
+body:             |
+  bb.0:
+    liveins: $w0
+    ; CHECK-LABEL: name: sextinreg_ashr_to_sbfx
+    ; CHECK: liveins: $w0
+    ; CHECK: %x:_(s32) = COPY $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 5
+    ; CHECK: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 14
+    ; CHECK: %sext_inreg:_(s32) = G_SBFX %x, [[C]], [[C1]]
+    ; CHECK: $w0 = COPY %sext_inreg(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %x:_(s32) = COPY $w0
+    %lsb:_(s32) = G_CONSTANT i32 5
+    %shift:_(s32) = G_ASHR %x, %lsb
+    %sext_inreg:_(s32) = G_SEXT_INREG %shift, 10
+    $w0 = COPY %sext_inreg
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            sextinreg_lshr_to_sbfx
+tracksRegLiveness: true
+legalized: true
+body:             |
+  bb.0:
+    liveins: $w0
+    ; CHECK-LABEL: name: sextinreg_lshr_to_sbfx
+    ; CHECK: liveins: $w0
+    ; CHECK: %x:_(s32) = COPY $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 5
+    ; CHECK: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 14
+    ; CHECK: %sext_inreg:_(s32) = G_SBFX %x, [[C]], [[C1]]
+    ; CHECK: $w0 = COPY %sext_inreg(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %x:_(s32) = COPY $w0
+    %lsb:_(s32) = G_CONSTANT i32 5
+    %shift:_(s32) = G_LSHR %x, %lsb
+    %sext_inreg:_(s32) = G_SEXT_INREG %shift, 10
+    $w0 = COPY %sext_inreg
+    RET_ReallyLR implicit $w0
+
+
+...
+---
+name:            dont_apply_no_constant
+tracksRegLiveness: true
+legalized: true
+body:             |
+  bb.0:
+    liveins: $w0
+    ; AArch64 needs a constant on the shift for this combine.
+
+    ; CHECK-LABEL: name: dont_apply_no_constant
+    ; CHECK: liveins: $w0
+    ; CHECK: %x:_(s32) = COPY $w0
+    ; CHECK: %y:_(s32) = COPY $w0
+    ; CHECK: %shift:_(s32) = G_LSHR %x, %y(s32)
+    ; CHECK: %sext_inreg:_(s32) = G_SEXT_INREG %shift, 10
+    ; CHECK: $w0 = COPY %sext_inreg(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %x:_(s32) = COPY $w0
+    %y:_(s32) = COPY $w0
+    %shift:_(s32) = G_LSHR %x, %y
+    %sext_inreg:_(s32) = G_SEXT_INREG %shift, 10
+    $w0 = COPY %sext_inreg
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            dont_apply_shift_imm_too_large
+tracksRegLiveness: true
+legalized: true
+body:             |
+  bb.0:
+    liveins: $w0
+
+    ; LSB must be in 0-31.
+
+    ; CHECK-LABEL: name: dont_apply_shift_imm_too_large
+    ; CHECK: liveins: $w0
+    ; CHECK: %x:_(s32) = COPY $w0
+    ; CHECK: %lsb:_(s32) = G_CONSTANT i32 32
+    ; CHECK: %shift:_(s32) = G_ASHR %x, %lsb(s32)
+    ; CHECK: %sext_inreg:_(s32) = G_SEXT_INREG %shift, 1
+    ; CHECK: $w0 = COPY %sext_inreg(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %x:_(s32) = COPY $w0
+    %lsb:_(s32) = G_CONSTANT i32 32
+    %shift:_(s32) = G_ASHR %x, %lsb
+    %sext_inreg:_(s32) = G_SEXT_INREG %shift, 1
+    $w0 = COPY %sext_inreg
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            dont_apply_negative_shift_imm
+tracksRegLiveness: true
+legalized: true
+body:             |
+  bb.0:
+    liveins: $w0
+
+    ; LSB must be in 0-31.
+
+    ; CHECK-LABEL: name: dont_apply_negative_shift_imm
+    ; CHECK: liveins: $w0
+    ; CHECK: %x:_(s32) = COPY $w0
+    ; CHECK: %lsb:_(s32) = G_CONSTANT i32 -1
+    ; CHECK: %shift:_(s32) = G_ASHR %x, %lsb(s32)
+    ; CHECK: %sext_inreg:_(s32) = G_SEXT_INREG %shift, 1
+    ; CHECK: $w0 = COPY %sext_inreg(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %x:_(s32) = COPY $w0
+    %lsb:_(s32) = G_CONSTANT i32 -1
+    %shift:_(s32) = G_ASHR %x, %lsb
+    %sext_inreg:_(s32) = G_SEXT_INREG %shift, 1
+    $w0 = COPY %sext_inreg
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            dont_apply_more_than_one_use
+tracksRegLiveness: true
+legalized: true
+body:             |
+  bb.0:
+    liveins: $w0
+    ; CHECK-LABEL: name: dont_apply_more_than_one_use
+    ; CHECK: liveins: $w0
+    ; CHECK: %x:_(s32) = COPY $w0
+    ; CHECK: %lsb:_(s32) = G_CONSTANT i32 1
+    ; CHECK: %shift:_(s32) = G_ASHR %x, %lsb(s32)
+    ; CHECK: %sext_inreg:_(s32) = G_SEXT_INREG %shift, 1
+    ; CHECK: %mul:_(s32) = G_MUL %shift, %sext_inreg
+    ; CHECK: $w0 = COPY %mul(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %x:_(s32) = COPY $w0
+    %lsb:_(s32) = G_CONSTANT i32 1
+    %shift:_(s32) = G_ASHR %x, %lsb
+    %sext_inreg:_(s32) = G_SEXT_INREG %shift, 1
+    %mul:_(s32) = G_MUL %shift, %sext_inreg
+    $w0 = COPY %mul
+    RET_ReallyLR implicit $w0