[GlobalISel] Combine shr(shl x, c1), c2 to G_SBFX/G_UBFX

[foad]

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

544

This comment suggests that "n" is the same value on the LHS and the RHS, which is not true.

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

4155

You don't want a "one use" check for a constant.

4161

I think it would be better to combine this and the previous mi_match into one large match expression.

4163–4164

Since ShlAmt is signed don't you also need to check it's >= 0? (TBH I'm not sure how C++ compares int64_t vs unsigned.)

Anyway you might be able to simplify some of the bounds checking. I think all you need to check is: 0 <= ShlAmt <= ShrAmt < Size. Then the checks on Width and Pos should be unnecessary by construction.

4169

This is just Size - ShrAmt!

[gargaroff]

Address comments

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

544

Good point. Fixed!

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

4169

Good catch, thanks for spotting!

[foad]

How does this new combine interact with matchAshrShlToSextInreg? If n==k, aren't they both trying to combine the same thing with different results?

[gargaroff]

In D107330#2922162, @foad wrote:

How does this new combine interact with matchAshrShlToSextInreg? If n==k, aren't they both trying to combine the same thing with different results?

In that case they would indeed try to combine the same thing. Looks like SextInreg combine has a higher precedence, as no other tests were affected. Should we skip this new combine if n==k?

[foad]

Looks like SextInreg combine has a higher precedence, as no other tests were affected. Should we skip this new combine if n==k?

I don't know! Is there a well defined way of ensuring that SextInreg has higher precedence? Or is it just luck?

(Actually from an engineering perspective I would prefer that we just have a single combine, that generates sext_inreg if signed&&n==k and sbfx/ubfx otherwise. But I'm not sure that other globalisel folks would agree.)

[gargaroff]

Skip combine if ShlAmt == ShrAmt

[gargaroff]

AFAIK the order of the combines depends on the order in Combine.td. But I'm not sure if that's something we should rely on here and I would argue that G_SEXT_INREG should be preferred over G_SBFX. I have updated the patch to skip the combine if ShlAmt == ShrAmt.
But just in case, let's wait until tomorrow to give the others a chance to voice their opinion.

[foad]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

4161

This will also skip the unsigned n==k case, which you probably do want to turn into ubfx?

[gargaroff]

do not skip ubfx combine if shift sizes match

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

4161

Yep, was a bit too trigger happy :) Fixed now.

[gargaroff]

Ping. I'm leaving for vacation this weekend. Could we get this landed before that?

[foad]

LGTM apart from the question inline.

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

4145

Lots of combines use isLegalOrBeforeLegalizer, i.e. they do apply before legalization. Is there a reason you have made this one different, so it does not apply before legalization?

[gargaroff]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

4145

For that I followed the other bitfield extract combine. I'm guessing because there is only widenScalar implemented for G_SBFX and G_UBFX, so you wouldn't even be able to lower it if it would be illegal for your target.

Once at least the lower action is implemented, this check can be changed to make this combine also run before the legalizer.

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

public:
   bool matchICmpToTrueFalseKnownBits(MachineInstr &MI, int64_t &MatchInfo);
 
   bool matchBitfieldExtractFromSExtInReg(
       MachineInstr &MI, std::function<void(MachineIRBuilder &)> &MatchInfo);
   /// Match: and (lshr x, cst), mask -> ubfx x, cst, width
   bool matchBitfieldExtractFromAnd(
       MachineInstr &MI, std::function<void(MachineIRBuilder &)> &MatchInfo);
 
+  /// Match: shr (shl x, n), k -> sbfx/ubfx x, pos, width
+  bool matchBitfieldExtractFromShr(
+      MachineInstr &MI, std::function<void(MachineIRBuilder &)> &MatchInfo);
+
   /// Reassociate pointer calculations with G_ADD involved, to allow better
   /// addressing mode usage.
   bool matchReassocPtrAdd(MachineInstr &MI,
                           std::function<void(MachineIRBuilder &)> &MatchInfo);
 
 
   /// Do constant folding when opportunities are exposed after MIR building.
   bool matchConstantFold(MachineInstr &MI, APInt &MatchInfo);

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

 def funnel_shift_combines : GICombineGroup<[funnel_shift_to_rotate]>;
 
 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 [{ Helper.applyBuildFn(*${root}, ${info}); }])>;
 
+def bitfield_extract_from_shr : GICombineRule<
+  (defs root:$root, build_fn_matchinfo:$info),
+  (match (wip_match_opcode G_ASHR, G_LSHR):$root,
+    [{ return Helper.matchBitfieldExtractFromShr(*${root}, ${info}); }]),
+  (apply [{ Helper.applyBuildFn(*${root}, ${info}); }])>;
+
 def form_bitfield_extract : GICombineGroup<[bitfield_extract_from_sext_inreg,
-                                            bitfield_extract_from_and]>;
+                                            bitfield_extract_from_and,
+                                            bitfield_extract_from_shr]>;
 def reassoc_ptradd : GICombineRule<
   (defs root:$root, build_fn_matchinfo:$matchinfo),
   (match (wip_match_opcode G_PTR_ADD):$root,
     [{ return Helper.matchReassocPtrAdd(*${root}, ${matchinfo}); }]),
   (apply [{ Helper.applyBuildFnNoErase(*${root}, ${matchinfo}); }])>;
 
 def reassocs : GICombineGroup<[reassoc_ptradd]>;
 

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

bool CombinerHelper::matchBitfieldExtractFromAnd(
   MatchInfo = [=](MachineIRBuilder &B) {
     auto WidthCst = B.buildConstant(ExtractTy, Width);
     auto LSBCst = B.buildConstant(ExtractTy, LSBImm);
     B.buildInstr(TargetOpcode::G_UBFX, {Dst}, {ShiftSrc, LSBCst, WidthCst});
   };
   return true;
 }
 
+bool CombinerHelper::matchBitfieldExtractFromShr(
+    MachineInstr &MI, std::function<void(MachineIRBuilder &)> &MatchInfo) {
+  const unsigned Opcode = MI.getOpcode();
+  assert(Opcode == TargetOpcode::G_ASHR || Opcode == TargetOpcode::G_LSHR);
+
+  const Register Dst = MI.getOperand(0).getReg();
+
+  const unsigned ExtrOpcode = Opcode == TargetOpcode::G_ASHR
+                                  ? TargetOpcode::G_SBFX
+                                  : TargetOpcode::G_UBFX;
+
+  // Check if the type we would use for the extract is legal
+  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();
+
+  // 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))))
+    return false;
+
+  // Make sure that the shift sizes can fit a bitfield extract.
+  if (ShlAmt < 0 || ShlAmt > ShrAmt || ShrAmt >= Size)
+    return false;
+
+  // Calculate start position and width of the extract
+  const int64_t Pos = ShrAmt - ShlAmt;
+  const int64_t Width = Size - ShrAmt;
+
+  MatchInfo = [=](MachineIRBuilder &B) {
+    auto WidthCst = B.buildConstant(ExtractTy, Width);
+    auto PosCst = B.buildConstant(ExtractTy, Pos);
+    B.buildInstr(ExtrOpcode, {Dst}, {ShlSrc, PosCst, WidthCst});
+  };
+  return true;
+}
+
 bool CombinerHelper::reassociationCanBreakAddressingModePattern(
     MachineInstr &PtrAdd) {
   assert(PtrAdd.getOpcode() == TargetOpcode::G_PTR_ADD);
 
   Register Src1Reg = PtrAdd.getOperand(1).getReg();
   MachineInstr *Src1Def = getOpcodeDef(TargetOpcode::G_PTR_ADD, Src1Reg, MRI);
   if (!Src1Def)
     return false;

llvm/test/CodeGen/AArch64/GlobalISel/form-bitfield-extract-from-shr.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_ASHR/G_LSHR fed by a G_SHL into a G_SBFX/G_UBFX.
+
+---
+name:            apply_ashr_shl_to_sbfx
+legalized: true
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: apply_ashr_shl_to_sbfx
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 16
+    ; CHECK: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 11
+    ; CHECK: [[SBFX:%[0-9]+]]:_(s32) = G_SBFX [[COPY]], [[C1]](s32), [[C]]
+    ; CHECK: $w0 = COPY [[SBFX]](s32)
+    %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
+    $w0 = COPY %4(s32)
+...
+
+---
+name:            apply_ashr_shl_to_sbfx_lower_bound
+legalized: true
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: apply_ashr_shl_to_sbfx_lower_bound
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK: [[SBFX:%[0-9]+]]:_(s32) = G_SBFX [[COPY]], [[C]](s32), [[C]]
+    ; CHECK: $w0 = COPY [[SBFX]](s32)
+    %0:_(s32) = COPY $w0
+    %1:_(s32) = G_CONSTANT i32 30
+    %2:_(s32) = G_CONSTANT i32 31
+    %3:_(s32) = G_SHL %0, %1
+    %4:_(s32) = G_ASHR %3, %2
+    $w0 = COPY %4(s32)
+...
+
+---
+name:            apply_ashr_shl_to_sbfx_upper_bound
+legalized: true
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: apply_ashr_shl_to_sbfx_upper_bound
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 30
+    ; CHECK: [[SBFX:%[0-9]+]]:_(s32) = G_SBFX [[COPY]], [[C1]](s32), [[C]]
+    ; CHECK: $w0 = COPY [[SBFX]](s32)
+    %0:_(s32) = COPY $w0
+    %1:_(s32) = G_CONSTANT i32 1
+    %2:_(s32) = G_CONSTANT i32 31
+    %3:_(s32) = G_SHL %0, %1
+    %4:_(s32) = G_ASHR %3, %2
+    $w0 = COPY %4(s32)
+...
+
+---
+name:            apply_lshr_shl_to_ubfx
+legalized: true
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: apply_lshr_shl_to_ubfx
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 16
+    ; CHECK: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 11
+    ; CHECK: [[UBFX:%[0-9]+]]:_(s32) = G_UBFX [[COPY]], [[C1]](s32), [[C]]
+    ; CHECK: $w0 = COPY [[UBFX]](s32)
+    %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_LSHR %3, %2
+    $w0 = COPY %4(s32)
+...
+
+---
+name:            apply_lshr_shl_to_ubfx_lower_bound
+legalized: true
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: apply_lshr_shl_to_ubfx_lower_bound
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK: [[UBFX:%[0-9]+]]:_(s32) = G_UBFX [[COPY]], [[C]](s32), [[C]]
+    ; CHECK: $w0 = COPY [[UBFX]](s32)
+    %0:_(s32) = COPY $w0
+    %1:_(s32) = G_CONSTANT i32 30
+    %2:_(s32) = G_CONSTANT i32 31
+    %3:_(s32) = G_SHL %0, %1
+    %4:_(s32) = G_LSHR %3, %2
+    $w0 = COPY %4(s32)
+...
+
+---
+name:            apply_lshr_shl_to_ubfx_upper_bound
+legalized: true
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: apply_lshr_shl_to_ubfx_upper_bound
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 30
+    ; CHECK: [[UBFX:%[0-9]+]]:_(s32) = G_UBFX [[COPY]], [[C1]](s32), [[C]]
+    ; CHECK: $w0 = COPY [[UBFX]](s32)
+    %0:_(s32) = COPY $w0
+    %1:_(s32) = G_CONSTANT i32 1
+    %2:_(s32) = G_CONSTANT i32 31
+    %3:_(s32) = G_SHL %0, %1
+    %4:_(s32) = G_LSHR %3, %2
+    $w0 = COPY %4(s32)
+...
+
+---
+name:            dont_apply_pos_out_of_bounds
+legalized: true
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: dont_apply_pos_out_of_bounds
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 5
+    ; CHECK: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 2
+    ; CHECK: [[SHL:%[0-9]+]]:_(s32) = G_SHL [[COPY]], [[C]](s32)
+    ; CHECK: [[ASHR:%[0-9]+]]:_(s32) = G_ASHR [[SHL]], [[C1]](s32)
+    ; CHECK: $w0 = COPY [[ASHR]](s32)
+    %0:_(s32) = COPY $w0
+    %1:_(s32) = G_CONSTANT i32 5
+    %2:_(s32) = G_CONSTANT i32 2
+    %3:_(s32) = G_SHL %0, %1
+    %4:_(s32) = G_ASHR %3, %2
+    $w0 = COPY %4(s32)
+...
+
+---
+name:            dont_apply_no_constant
+legalized: true
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: dont_apply_no_constant
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[COPY1:%[0-9]+]]:_(s32) = COPY $w1
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 2
+    ; CHECK: [[SHL:%[0-9]+]]:_(s32) = G_SHL [[COPY]], [[COPY1]](s32)
+    ; CHECK: [[ASHR:%[0-9]+]]:_(s32) = G_ASHR [[SHL]], [[C]](s32)
+    ; CHECK: $w0 = COPY [[ASHR]](s32)
+    %0:_(s32) = COPY $w0
+    %1:_(s32) = COPY $w1
+    %2:_(s32) = G_CONSTANT i32 2
+    %3:_(s32) = G_SHL %0, %1
+    %4:_(s32) = G_ASHR %3, %2
+    $w0 = COPY %4(s32)
+...
+
+---
+name:            dont_apply_more_than_one_use
+legalized: true
+body:             |
+  bb.0.entry:
+    ; CHECK-LABEL: name: dont_apply_more_than_one_use
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 5
+    ; CHECK: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 16
+    ; CHECK: [[SHL:%[0-9]+]]:_(s32) = G_SHL [[COPY]], [[C]](s32)
+    ; CHECK: [[ASHR:%[0-9]+]]:_(s32) = G_ASHR [[SHL]], [[C1]](s32)
+    ; CHECK: [[MUL:%[0-9]+]]:_(s32) = G_MUL [[SHL]], [[ASHR]]
+    ; CHECK: $w0 = COPY [[MUL]](s32)
+    %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)
+...

llvm/test/CodeGen/AArch64/arm64-rev.ll

 ; CHECK-LABEL: test_rev_x_srl32_shift:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    ubfx x8, x0, #2, #29
 ; CHECK-NEXT:    rev32 x0, x8
 ; CHECK-NEXT:    ret
 ;
 ; GISEL-LABEL: test_rev_x_srl32_shift:
 ; GISEL:       // %bb.0: // %entry
-; GISEL-NEXT:    lsl x8, x0, #33
-; GISEL-NEXT:    lsr x8, x8, #35
+; GISEL-NEXT:    ubfx x8, x0, #2, #29
 ; GISEL-NEXT:    rev x8, x8
 ; GISEL-NEXT:    lsr x0, x8, #32
 ; GISEL-NEXT:    ret
 entry:
   %0 = shl i64 %a, 33
   %1 = lshr i64 %0, 35
   %2 = tail call i64 @llvm.bswap.i64(i64 %1)
   %3 = lshr i64 %2, 32

llvm/test/CodeGen/AMDGPU/GlobalISel/llvm.amdgcn.ubfe.ll

 ; GFX6:       ; %bb.0:
 ; GFX6-NEXT:    s_load_dwordx2 s[4:5], s[0:1], 0x9
 ; GFX6-NEXT:    s_load_dwordx2 s[0:1], s[0:1], 0xb
 ; GFX6-NEXT:    s_mov_b32 s6, -1
 ; GFX6-NEXT:    s_mov_b32 s7, 0xf000
 ; GFX6-NEXT:    s_waitcnt lgkmcnt(0)
 ; GFX6-NEXT:    s_load_dword s0, s[0:1], 0x0
 ; GFX6-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX6-NEXT:    s_lshl_b32 s0, s0, 31
-; GFX6-NEXT:    s_lshr_b32 s0, s0, 31
+; GFX6-NEXT:    s_bfe_u32 s0, s0, 0x10000
 ; GFX6-NEXT:    s_bfe_u32 s0, s0, 0x1001f
 ; GFX6-NEXT:    v_mov_b32_e32 v0, s0
 ; GFX6-NEXT:    buffer_store_dword v0, off, s[4:7], 0
 ; GFX6-NEXT:    s_endpgm
   %x = load i32, i32 addrspace(1)* %in, align 4
   %shl = shl i32 %x, 31
   %shr = lshr i32 %shl, 31
   %bfe = call i32 @llvm.amdgcn.ubfe.i32(i32 %shr, i32 31, i32 1)

 define amdgpu_kernel void @shl_lshr(i32 addrspace(1)* %out, i32 %a) #0 {
 ; GFX6-LABEL: shl_lshr:
 ; GFX6:       ; %bb.0:
 ; GFX6-NEXT:    s_load_dwordx2 s[4:5], s[0:1], 0x9
 ; GFX6-NEXT:    s_load_dword s0, s[0:1], 0xb
 ; GFX6-NEXT:    s_mov_b32 s6, -1
 ; GFX6-NEXT:    s_mov_b32 s7, 0xf000
 ; GFX6-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX6-NEXT:    s_lshl_b32 s0, s0, 9
-; GFX6-NEXT:    s_lshr_b32 s0, s0, 11
+; GFX6-NEXT:    s_bfe_u32 s0, s0, 0x150002
 ; GFX6-NEXT:    v_mov_b32_e32 v0, s0
 ; GFX6-NEXT:    buffer_store_dword v0, off, s[4:7], 0
 ; GFX6-NEXT:    s_endpgm
   %b = shl i32 %a, 9
   %c = lshr i32 %b, 11
   store i32 %c, i32 addrspace(1)* %out, align 8
   ret void
 }
 
 declare i32 @llvm.amdgcn.ubfe.i32(i32, i32, i32) #1
 declare i64 @llvm.amdgcn.ubfe.i64(i64, i32, i32) #1
 
 attributes #0 = { nounwind }
 attributes #1 = { nounwind readnone }