[GlobalISel] Extending G_SBFX/G_UBFX combine to consider G_SEXT_INREG as well

[DavidFerencSzabo]

Summary

Currently simple C code like

typedef struct { int a:1; int b:4; int c:3; } s_t;
int f(s_t s) { return s.b; }

will produce the following assembly with O3

lsl     w8, w0, #3
sxtb    w8, w8
sbfx    w0, w8, #4, #8
ret

while GCC produce this

sbfx    x0, x0, 1, 4
ret

Extended the implementation of matchBitfieldExtractFromShr to consider cases where the shifts fed to a G_SEXT_INREG.

[foad]

This should be implemented by teaching computeNumSignBits about G_SBFX.

[DavidFerencSzabo]

This problem aroused in our down stream target, where G_SEXT_INREG(G_SBFX) were also produced, but for AArch64 it seems that not be the case, therefore removed the change in matchRedundantSExtInReg.

[DavidFerencSzabo]

In D141204#4035658, @foad wrote:

This should be implemented by teaching computeNumSignBits about G_SBFX.

Thanks for the feedback. Upon inspection it turned out that AArch64 does not require that part of the change, hence I removed it.

[foad]

Please upload diffs with full context, either with git diff -U9999999 or by using the arcanist tool (arc diff).

[DavidFerencSzabo]

[GlobalISel] Extending G_SBFX/G_UBFX combine to consider G_SEXT_INREG as well

[foad]

I wonder if a simpler approach would be just to combine the ashr (sextinreg x, 8), 4 into sextinreg (ashr x, 4), 4 (i.e. canonicalize ashr-of-sext into sext-of-ashr) and then let other combines simplify it further?

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

4399

Update the comment to say what else you are trying to match.

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

194

I don't think you need any of these COPYs of G_CONSTANTs.

211

This just seems wrong. You are combining (sextinreg ($w0 << 3), 8) >> 4 to ubfx $w0, 1, 4.

If $w0 was 0x000000FF then the original code gives 0x0FFFFFFF but the combined code gives 0x0000000F.

[arsenm]

Would probably be simpler to combine shifts and ext_inreg, plus should try to use new tablegen combiners

llvm/include/llvm/CodeGen/GlobalISel/MIPatternMatch.h

inline UnaryOp_match<SrcTy, TargetOpcode::COPY> m_Copy(SrcTy &&Src) {
   return UnaryOp_match<SrcTy, TargetOpcode::COPY>(std::forward<SrcTy>(Src));
 }
 
 template <typename SrcTy>
 inline UnaryOp_match<SrcTy, TargetOpcode::G_FSQRT> m_GFSqrt(const SrcTy &Src) {
   return UnaryOp_match<SrcTy, TargetOpcode::G_FSQRT>(Src);
 }
 
+template <typename SrcTy, unsigned Opcode> struct SextInRegOp_match {
+  SrcTy L;
+  int64_t &R;
+
+  SextInRegOp_match(const SrcTy &LHS, int64_t &RHS) : L(LHS), R(RHS) {}
+  template <typename OpTy>
+  bool match(const MachineRegisterInfo &MRI, OpTy &&Op) {
+    MachineInstr *TmpMI;
+    if (mi_match(Op, MRI, m_MInstr(TmpMI))) {
+      if (TmpMI->getOpcode() == Opcode && TmpMI->getNumOperands() == 3) {
+        R = TmpMI->getOperand(2).getImm();
+        return L.match(MRI, TmpMI->getOperand(1).getReg());
+      }
+    }
+    return false;
+  }
+};
+
+template <typename SrcTy>
+inline SextInRegOp_match<SrcTy, TargetOpcode::G_SEXT_INREG>
+m_GSExtInReg(const SrcTy &Src, int64_t &BitPos) {
+  return SextInRegOp_match<SrcTy, TargetOpcode::G_SEXT_INREG>(Src, BitPos);
+}
+
 // General helper for generic MI compares, i.e. G_ICMP and G_FCMP
 // TODO: Allow checking a specific predicate.
 template <typename Pred_P, typename LHS_P, typename RHS_P, unsigned Opcode>
 struct CompareOp_match {
   Pred_P P;
   LHS_P L;
   RHS_P R;
 

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

bool CombinerHelper::matchBitfieldExtractFromShr(
   LLT Ty = MRI.getType(Dst);
   LLT ExtractTy = getTargetLowering().getPreferredShiftAmountTy(Ty);
   if (!LI || !LI->isLegalOrCustom({ExtrOpcode, {Ty, ExtractTy}}))
     return false;
 
   Register ShlSrc;
   int64_t ShrAmt;
   int64_t ShlAmt;
-  const unsigned Size = Ty.getScalarSizeInBits();
+  int64_t Size = Ty.getScalarSizeInBits();
 
   // Try to match shr (shl x, c1), c2
-  if (!mi_match(Dst, MRI,
-                m_BinOp(Opcode,
-                        m_OneNonDBGUse(m_GShl(m_Reg(ShlSrc), m_ICst(ShlAmt))),
-                        m_ICst(ShrAmt))))
+  if (!mi_match(
+          Dst, MRI,
+          m_BinOp(
+              Opcode,
+              m_any_of(m_GSExtInReg(m_OneNonDBGUse(
+                                        m_GShl(m_Reg(ShlSrc), m_ICst(ShlAmt))),
+                                    Size),
+                       m_OneNonDBGUse(m_GShl(m_Reg(ShlSrc), m_ICst(ShlAmt)))),
     return false;
 
   // Make sure that the shift sizes can fit a bitfield extract
   if (ShlAmt < 0 || ShlAmt > ShrAmt || ShrAmt >= Size)
     return false;
 
   // Skip this combine if the G_SEXT_INREG combine could handle it
   if (Opcode == TargetOpcode::G_ASHR && ShlAmt == ShrAmt)
bool CombinerHelper::matchBitfieldExtractFromShrAnd(
   assert(Opcode == TargetOpcode::G_LSHR || Opcode == TargetOpcode::G_ASHR);
 
   const Register Dst = MI.getOperand(0).getReg();
   LLT Ty = MRI.getType(Dst);
   LLT ExtractTy = getTargetLowering().getPreferredShiftAmountTy(Ty);
   if (!getTargetLowering().isConstantUnsignedBitfieldExtractLegal(
           TargetOpcode::G_UBFX, Ty, ExtractTy))
     return false;
 
   // Try to match shr (and x, c1), c2
   Register AndSrc;
   int64_t ShrAmt;
   int64_t SMask;
   if (!mi_match(Dst, MRI,
                 m_BinOp(Opcode,
                         m_OneNonDBGUse(m_GAnd(m_Reg(AndSrc), m_ICst(SMask))),
                         m_ICst(ShrAmt))))

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

bb.0.entry:
     %0:_(s32) = COPY $w0
     %1:_(s32) = G_CONSTANT i32 5
     %2:_(s32) = G_CONSTANT i32 16
     %3:_(s32) = G_SHL %0, %1
     %4:_(s32) = G_ASHR %3, %2
     %5:_(s32) = G_MUL %3, %4
     $w0 = COPY %5(s32)
 ...
+
+---
+name:            test_ashr_sextinreg_shl_to_sbfx
+legalized: true
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: test_ashr_sextinreg_shl_to_sbfx
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 4
+    ; CHECK-NEXT: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK-NEXT: [[SBFX:%[0-9]+]]:_(s32) = G_SBFX [[COPY]], [[C1]](s32), [[C]]
+    ; CHECK-NEXT: $w0 = COPY [[SBFX]](s32)
+    %0:_(s32) = COPY $w0
+    %1:_(s32) = G_CONSTANT i32 3
+    %2:_(s32) = G_CONSTANT i32 4
+    %3:_(s32) = COPY %1:_(s32)
+    %4:_(s32) = G_SHL %0:_, %3:_(s32)
+    %5:_(s32) = COPY %2:_(s32)
+    %6:_(s32) = G_SEXT_INREG %4:_, 8
+    %7:_(s32) = G_ASHR %6:_, %5:_(s32)
+    $w0 = COPY %7:_(s32)
+...
+
+---
+name:            test_lshr_sextinreg_shl_to_sbfx
+legalized: true
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: test_lshr_sextinreg_shl_to_sbfx
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 4
+    ; CHECK-NEXT: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK-NEXT: [[UBFX:%[0-9]+]]:_(s32) = G_UBFX [[COPY]], [[C1]](s32), [[C]]
+    ; CHECK-NEXT: $w0 = COPY [[UBFX]](s32)
+    %0:_(s32) = COPY $w0
+    %1:_(s32) = G_CONSTANT i32 3
+    %2:_(s32) = G_CONSTANT i32 4
+    %3:_(s32) = COPY %1:_(s32)
+    %4:_(s32) = G_SHL %0:_, %3:_(s32)
+    %5:_(s32) = COPY %2:_(s32)
+    %6:_(s32) = G_SEXT_INREG %4:_, 8
+    %7:_(s32) = G_LSHR %6:_, %5:_(s32)
+    $w0 = COPY %7:_(s32)
+...