[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

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);
 
+  /// 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 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}); }])>;
 

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

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::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

def build_vector_to_dup : GICombineRule<
   (defs root:$root),
   (match (wip_match_opcode G_BUILD_VECTOR):$root,
           [{ return matchBuildVectorToDup(*${root}, MRI); }]),
   (apply [{ return applyBuildVectorToDup(*${root}, MRI, B); }])
 >;
 
 def build_vector_lowering : GICombineGroup<[build_vector_to_dup]>;
 
+def bitfield_extract_from_sext_inreg : GICombineRule<
+  (defs root:$root, build_fn_matchinfo:$info),
+  (match (wip_match_opcode G_SEXT_INREG):$root,
+    [{ return matchBitfieldExtractFromSExtInReg(*${root}, MRI, ${info}); }]),
+  (apply [{ return Helper.applyBuildFn(*${root}, ${info}); }])>;
+
+def form_bitfield_extract : GICombineGroup<[bitfield_extract_from_sext_inreg]>;
+
 // Post-legalization combines which should happen at all optimization levels.
 // (E.g. ones that facilitate matching for the selector) For example, matching
 // pseudos.
 def AArch64PostLegalizerLoweringHelper
     : GICombinerHelper<"AArch64GenPostLegalizerLoweringHelper",
                        [shuffle_vector_lowering, vashr_vlshr_imm,
                         icmp_lowering, build_vector_lowering]> {
   let DisableRuleOption = "aarch64postlegalizerlowering-disable-rule";
 }
 
 // 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/lib/Target/AArch64/GISel/AArch64PostLegalizerCombiner.cpp

 ///
 //===----------------------------------------------------------------------===//
 
 #include "AArch64TargetMachine.h"
 #include "llvm/CodeGen/GlobalISel/Combiner.h"
 #include "llvm/CodeGen/GlobalISel/CombinerHelper.h"
 #include "llvm/CodeGen/GlobalISel/CombinerInfo.h"
 #include "llvm/CodeGen/GlobalISel/GISelKnownBits.h"
+#include "llvm/CodeGen/GlobalISel/MIPatternMatch.h"
 #include "llvm/CodeGen/GlobalISel/MachineIRBuilder.h"
 #include "llvm/CodeGen/GlobalISel/Utils.h"
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetOpcodes.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "aarch64-postlegalizer-combiner"
 
 using namespace llvm;
+using namespace MIPatternMatch;
 
 /// This combine tries do what performExtractVectorEltCombine does in SDAG.
 /// Rewrite for pairwise fadd pattern
 ///   (s32 (g_extract_vector_elt
 ///           (g_fadd (vXs32 Other)
 ///                  (g_vector_shuffle (vXs32 Other) undef <1,X,...> )) 0))
 /// ->
 ///   (s32 (g_fadd (g_extract_vector_elt (vXs32 Other) 0)
bool applyAArch64MulConstCombine(
     MachineInstr &MI, MachineRegisterInfo &MRI, MachineIRBuilder &B,
     std::function<void(MachineIRBuilder &B, Register DstReg)> &ApplyFn) {
   B.setInstrAndDebugLoc(MI);
   ApplyFn(B, MI.getOperand(0).getReg());
   MI.eraseFromParent();
   return true;
 }
 
+/// Form a G_SBFX from a G_SEXT_INREG fed by a right shift.
+static bool matchBitfieldExtractFromSExtInReg(
+    MachineInstr &MI, MachineRegisterInfo &MRI,
+    std::function<void(MachineIRBuilder &)> &MatchInfo) {
+  assert(MI.getOpcode() == TargetOpcode::G_SEXT_INREG);
+  Register Dst = MI.getOperand(0).getReg();
+  Register Src = MI.getOperand(1).getReg();
+  int64_t Width = MI.getOperand(2).getImm();
+  LLT Ty = MRI.getType(Src);
+  assert((Ty == LLT::scalar(32) || Ty == LLT::scalar(64)) &&
+         "Unexpected type for G_SEXT_INREG?");
+  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;
+}
+
 #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/form-bitfield-extract-from-sextinreg.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