[AMDGPU][GlobalISel] Fold a chain of two shift instructions with constant operands

[mbrkusanin]

Summary

Sequence of same shift instructions with constant operands can be combined into
a single shift instruction.

[foad]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

1551–1553

Change to assert().

1562

I don't think Shl2Def can ever be null here, can it?

1588

"Imm" instead of "MatchInfo.Imm"

1588

Logically you want the size of the first operand here, not the size of the immediate. I'm not sure if they are guaranteed to be the same.

1595

Typo "than".

[mbrkusanin]

• Addressed review comments.
• Updated comment at the start of match function.

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

1562

I guess if the MIR is valid it should never be null. This match is the very similar to the one above (matchPtrAddImmedChain).

[foad]

LGTM, but I'm adding some generic GlobalISel reviewers in case they have comments.

[arsenm]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

1590

This is undefined, not 0. I don't think this needs to specially consider this case

[arsenm]

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

57

Can't this just reuse RegisterImmPair?

llvm/test/CodeGen/AMDGPU/GlobalISel/combine-shift-imm-chain.ll

4

These cases are all SGPRs, and don't have any illegal type tests. A few MIR cases wouldn't hurt either

[foad]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

1590

No, this is handling cases like X >> 20 >> 25 where the sum (20 + 25) exceeds the width of the type.

[foad]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

1596–1597

You could also handle G_USHLSAT and G_SSHLSAT the same way as G_ASHR.

[mbrkusanin]

• Using RegisterImmPair, removed ShiftChain struct.
• Supported G_SSHLSAT and G_USHLSAT.
• Added .mir tests

[foad]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

1592–1594

This is wrong. With 8-bit ints:
1 sshlsat 6 = 0x40
1 sshlsat 7 = 0x7f
I think this should clamp to ScalarSizeInBits - 1 just like ashr does.

[foad]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

1596–1597

Sorry, I was wrong about G_USHLSAT. You can't handle it quite the same way as G_ASHR. With 8-bit ints:
1 ushlsat 5 ushlsat 5 = 0xff
1 ushlsat 7 = 0x80
So they are not the same.

If the combined shift amount for ushlsat is >= 8 then you could replace it with "x == 0 ? 0 : 0xff" but it's probably not worth doing this.

[arsenm]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

1587

Shouldn't construct a new MachineIRBuilder here

[mbrkusanin]

• Fix handling of G_SSHLSAT and G_USHLSAT.
• Removed new MachineIRBuilder.

[foad]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

1599

I think you mean "saturating signed or unsigned left shift" here.

But this does not work for saturating unsigned left shift. There is no good replacement for "x ushlsat 20 ushlsat 25" where 20 + 25 exceeds the width of the type. I suggest catching this case in the match func and returning false.

[mbrkusanin]

• Avoid combine for G_USHLSAT when sum exceeds scalar size.

[foad]

LGTM modulo minor comment inline.

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

1577–1578

The comment could be improved. (What you wrote is also true of all the other shifts.) How about "there is no simple replacement for a saturating unsigned left shift that exceeds the scalar size"?

[mbrkusanin]

• Rebase
• Updated comments

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

struct IndexedLoadStoreMatchInfo {
   bool IsPre;
 };
 
 struct PtrAddChain {
   int64_t Imm;
   Register Base;
 };
 
+struct ShiftChain {
+  int64_t Imm;
+  Register Base;
+};
+
 struct RegisterImmPair {
   Register Reg;
   int64_t Imm;
 };
 
 using OperandBuildSteps =
     SmallVector<std::function<void(MachineInstrBuilder &)>, 4>;
 struct InstructionBuildSteps {
public:
   ///     $addr = G_INDEXED_STORE $val, $base, $offset
   ///     [...]
   ///     $whatever = COPY $addr
   bool tryCombineMemCpyFamily(MachineInstr &MI, unsigned MaxLen = 0);
 
   bool matchPtrAddImmedChain(MachineInstr &MI, PtrAddChain &MatchInfo);
   bool applyPtrAddImmedChain(MachineInstr &MI, PtrAddChain &MatchInfo);
 
+  bool matchShiftImmedChain(MachineInstr &MI, ShiftChain &MatchInfo);
+  bool applyShiftImmedChain(MachineInstr &MI, ShiftChain &MatchInfo);
+
   /// Transform a multiply by a power-of-2 value to a left shift.
   bool matchCombineMulToShl(MachineInstr &MI, unsigned &ShiftVal);
   bool applyCombineMulToShl(MachineInstr &MI, unsigned &ShiftVal);
 
   // Transform a G_SHL with an extended source into a narrower shift if
   // possible.
   bool matchCombineShlOfExtend(MachineInstr &MI, RegisterImmPair &MatchData);
   bool applyCombineShlOfExtend(MachineInstr &MI,

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

 
 def ptr_add_immed_matchdata : GIDefMatchData<"PtrAddChain">;
 def ptr_add_immed_chain : GICombineRule<
   (defs root:$d, ptr_add_immed_matchdata:$matchinfo),
   (match (wip_match_opcode G_PTR_ADD):$d,
          [{ return Helper.matchPtrAddImmedChain(*${d}, ${matchinfo}); }]),
   (apply [{ Helper.applyPtrAddImmedChain(*${d}, ${matchinfo}); }])>;
 
+def shift_immed_matchdata : GIDefMatchData<"ShiftChain">;
+def shift_immed_chain : GICombineRule<
+  (defs root:$d, shift_immed_matchdata:$matchinfo),
+  (match (wip_match_opcode G_SHL, G_ASHR, G_LSHR):$d,
+         [{ return Helper.matchShiftImmedChain(*${d}, ${matchinfo}); }]),
+  (apply [{ Helper.applyShiftImmedChain(*${d}, ${matchinfo}); }])>;
+
 def mul_to_shl_matchdata : GIDefMatchData<"unsigned">;
 def mul_to_shl : GICombineRule<
   (defs root:$d, mul_to_shl_matchdata:$matchinfo),
   (match (G_MUL $d, $op1, $op2):$mi,
          [{ return Helper.matchCombineMulToShl(*${mi}, ${matchinfo}); }]),
   (apply [{ Helper.applyCombineMulToShl(*${mi}, ${matchinfo}); }])>;
 
 // shl ([asz]ext x), y => zext (shl x, y), if shift does not overflow int
def all_combines : GICombineGroup<[trivial_combines, ptr_add_immed_chain,
     identity_combines, simplify_add_to_sub,
     hoist_logic_op_with_same_opcode_hands,
     shl_ashr_to_sext_inreg, sext_inreg_of_load,
     width_reduction_combines, select_combines,
     known_bits_simplifications, ext_ext_fold,
     not_cmp_fold, opt_brcond_by_inverting_cond,
     unmerge_merge, fabs_fabs_fold, unmerge_cst, unmerge_dead_to_trunc,
     unmerge_zext_to_zext, trunc_ext_fold, trunc_shl,
-    const_combines, xor_of_and_with_same_reg, ptr_add_with_zero]>;
+    const_combines, xor_of_and_with_same_reg, ptr_add_with_zero,
+    shift_immed_chain]>;

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

bool CombinerHelper::applyPtrAddImmedChain(MachineInstr &MI,
   auto NewOffset = MIB.buildConstant(OffsetTy, MatchInfo.Imm);
   Observer.changingInstr(MI);
   MI.getOperand(1).setReg(MatchInfo.Base);
   MI.getOperand(2).setReg(NewOffset.getReg(0));
   Observer.changedInstr(MI);
   return true;
 }
 
+bool CombinerHelper::matchShiftImmedChain(MachineInstr &MI,
+                                          ShiftChain &MatchInfo) {
+  // We're trying to match the following pattern with any of G_SHL/G_ASHR/G_LSHR
+  // shift instructions:
+  //   %t1 = SHIFT %base, G_CONSTANT imm1
+  //   %root = SHIFT %t1, G_CONSTANT imm2
+  // -->
+  //   %root = SHIFT %base, G_CONSTANT (imm1 + imm2)
+
+  unsigned Opcode = MI.getOpcode();
+  assert((Opcode == TargetOpcode::G_SHL || Opcode == TargetOpcode::G_ASHR ||
+          Opcode == TargetOpcode::G_LSHR) &&
+         "Expected G_SHL, G_ASHR or G_LSHR");
+
+  Register Shl2 = MI.getOperand(1).getReg();
+  Register Imm1 = MI.getOperand(2).getReg();
+  auto MaybeImmVal = getConstantVRegValWithLookThrough(Imm1, MRI);
+  if (!MaybeImmVal)
+    return false;
+
+  MachineInstr *Shl2Def = MRI.getUniqueVRegDef(Shl2);
+  if (Shl2Def->getOpcode() != Opcode)
+    return false;
+
+  Register Base = Shl2Def->getOperand(1).getReg();
+  Register Imm2 = Shl2Def->getOperand(2).getReg();
+  auto MaybeImm2Val = getConstantVRegValWithLookThrough(Imm2, MRI);
+  if (!MaybeImm2Val)
+    return false;
+
+  // Pass the combined immediate to the apply function.
+  MatchInfo.Imm = MaybeImmVal->Value + MaybeImm2Val->Value;
+  MatchInfo.Base = Base;
+  return true;
+}
+
+bool CombinerHelper::applyShiftImmedChain(MachineInstr &MI,
+                                          ShiftChain &MatchInfo) {
+  unsigned Opcode = MI.getOpcode();
+  assert((Opcode == TargetOpcode::G_SHL || Opcode == TargetOpcode::G_ASHR ||
+          Opcode == TargetOpcode::G_LSHR) &&
+         "Expected G_SHL, G_ASHR or G_LSHR");
+
+  MachineIRBuilder MIB(MI);
+  LLT Ty = MRI.getType(MI.getOperand(1).getReg());
+  auto Imm = MatchInfo.Imm;
+
+  if (Imm >= Ty.getScalarSizeInBits()) {
+    // Any logical shift that exceeds scalar size will produce zero.
+    if (Opcode != TargetOpcode::G_ASHR) {
+      MIB.buildConstant(MI.getOperand(0), 0);
+      MI.eraseFromParent();
+      return true;
+    }
+    // Arithmetic shift larger than scalar size has no effect.
+    Imm = Ty.getScalarSizeInBits() - 1;
+  }
+
+  LLT ImmTy = MRI.getType(MI.getOperand(2).getReg());
+  Register NewImm = MIB.buildConstant(ImmTy, Imm).getReg(0);
+  Observer.changingInstr(MI);
+  MI.getOperand(1).setReg(MatchInfo.Base);
+  MI.getOperand(2).setReg(NewImm);
+  Observer.changedInstr(MI);
+  return true;
+}
+
 bool CombinerHelper::matchCombineMulToShl(MachineInstr &MI,
                                           unsigned &ShiftVal) {
   assert(MI.getOpcode() == TargetOpcode::G_MUL && "Expected a G_MUL");
   auto MaybeImmVal =
       getConstantVRegValWithLookThrough(MI.getOperand(2).getReg(), MRI);
   if (!MaybeImmVal || !isPowerOf2_64(MaybeImmVal->Value))
     return false;
   ShiftVal = Log2_64(MaybeImmVal->Value);

llvm/test/CodeGen/AMDGPU/GlobalISel/ashr.ll

 ; GFX6-NEXT:    v_bfe_i32 v0, v0, 0, 8
 ; GFX6-NEXT:    v_ashrrev_i32_e32 v0, 7, v0
 ; GFX6-NEXT:    s_setpc_b64 s[30:31]
 ;
 ; GFX8-LABEL: v_ashr_i8_7:
 ; GFX8:       ; %bb.0:
 ; GFX8-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX8-NEXT:    v_lshlrev_b16_e32 v0, 8, v0
-; GFX8-NEXT:    v_mov_b32_e32 v1, 7
-; GFX8-NEXT:    v_ashrrev_i16_sdwa v0, v1, sext(v0) dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:DWORD src1_sel:BYTE_1
+; GFX8-NEXT:    v_ashrrev_i16_e32 v0, 15, v0
 ; GFX8-NEXT:    s_setpc_b64 s[30:31]
 ;
 ; GFX9-LABEL: v_ashr_i8_7:
 ; GFX9:       ; %bb.0:
 ; GFX9-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX9-NEXT:    v_mov_b32_e32 v1, 7
 ; GFX9-NEXT:    v_ashrrev_i16_sdwa v0, v1, sext(v0) dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:DWORD src1_sel:BYTE_0
 ; GFX9-NEXT:    s_setpc_b64 s[30:31]

llvm/test/CodeGen/AMDGPU/GlobalISel/combine-shift-imm-chain.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -global-isel -march=amdgcn -verify-machineinstrs < %s | FileCheck %s
+
+define amdgpu_cs i32 @test_shl_1(i32 inreg %arg1) {
+; CHECK-LABEL: test_shl_1:
+; CHECK:       ; %bb.0: ; %.entry
+; CHECK-NEXT:    s_lshl_b32 s0, s0, 5
+; CHECK-NEXT:    ; return to shader part epilog
+.entry:
+  %z1 = shl i32 %arg1, 2
+  %z2 = shl i32 %z1, 3
+  ret i32 %z2
+}
+
+define amdgpu_cs i32 @test_shl_2(i32 inreg %arg1) {
+; CHECK-LABEL: test_shl_2:
+; CHECK:       ; %bb.0: ; %.entry
+; CHECK-NEXT:    s_lshl_b32 s0, s0, 10
+; CHECK-NEXT:    ; return to shader part epilog
+.entry:
+  %z1 = shl i32 %arg1, 1
+  %z2 = shl i32 %z1, 2
+  %z3 = shl i32 %z2, 3
+  %z4 = shl i32 %z3, 4
+  ret i32 %z4
+}
+
+define amdgpu_cs i32 @test_shl_i32(i32 inreg %arg1) {
+; CHECK-LABEL: test_shl_i32:
+; CHECK:       ; %bb.0: ; %.entry
+; CHECK-NEXT:    s_mov_b32 s0, 0
+; CHECK-NEXT:    ; return to shader part epilog
+.entry:
+  %z1 = shl i32 %arg1, 10
+  %z2 = shl i32 %z1, 10
+  %z3 = shl i32 %z2, 10
+  %z4 = shl i32 %z3, 10
+  ret i32 %z4
+}
+
+define amdgpu_cs i64 @test_shl_i64(i64 inreg %arg1) {
+; CHECK-LABEL: test_shl_i64:
+; CHECK:       ; %bb.0: ; %.entry
+; CHECK-NEXT:    s_mov_b32 s0, 0
+; CHECK-NEXT:    s_mov_b32 s1, 0
+; CHECK-NEXT:    ; return to shader part epilog
+.entry:
+  %z1 = shl i64 %arg1, 10
+  %z2 = shl i64 %z1, 10
+  %z3 = shl i64 %z2, 10
+  %z4 = shl i64 %z3, 10
+  %z5 = shl i64 %z4, 10
+  %z6 = shl i64 %z5, 10
+  %z7 = shl i64 %z6, 10
+  ret i64 %z7
+}
+
+define amdgpu_cs i32 @test_ashr_1(i32 inreg %arg1) {
+; CHECK-LABEL: test_ashr_1:
+; CHECK:       ; %bb.0: ; %.entry
+; CHECK-NEXT:    s_ashr_i32 s0, s0, 5
+; CHECK-NEXT:    ; return to shader part epilog
+.entry:
+  %z1 = ashr i32 %arg1, 2
+  %z2 = ashr i32 %z1, 3
+  ret i32 %z2
+}
+
+define amdgpu_cs i32 @test_ashr_2(i32 inreg %arg1) {
+; CHECK-LABEL: test_ashr_2:
+; CHECK:       ; %bb.0: ; %.entry
+; CHECK-NEXT:    s_ashr_i32 s0, s0, 10
+; CHECK-NEXT:    ; return to shader part epilog
+.entry:
+  %z1 = ashr i32 %arg1, 1
+  %z2 = ashr i32 %z1, 2
+  %z3 = ashr i32 %z2, 3
+  %z4 = ashr i32 %z3, 4
+  ret i32 %z4
+}
+
+define amdgpu_cs i32 @test_ashr_i32(i32 inreg %arg1) {
+; CHECK-LABEL: test_ashr_i32:
+; CHECK:       ; %bb.0: ; %.entry
+; CHECK-NEXT:    s_ashr_i32 s0, s0, 31
+; CHECK-NEXT:    ; return to shader part epilog
+.entry:
+  %z1 = ashr i32 %arg1, 10
+  %z2 = ashr i32 %z1, 10
+  %z3 = ashr i32 %z2, 10
+  %z4 = ashr i32 %z3, 10
+  ret i32 %z4
+}
+
+define amdgpu_cs i64 @test_ashr_i64(i64 inreg %arg1) {
+; CHECK-LABEL: test_ashr_i64:
+; CHECK:       ; %bb.0: ; %.entry
+; CHECK-NEXT:    s_ashr_i32 s0, s1, 31
+; CHECK-NEXT:    s_mov_b32 s1, s0
+; CHECK-NEXT:    ; return to shader part epilog
+.entry:
+  %z1 = ashr i64 %arg1, 10
+  %z2 = ashr i64 %z1, 10
+  %z3 = ashr i64 %z2, 10
+  %z4 = ashr i64 %z3, 10
+  %z5 = ashr i64 %z4, 10
+  %z6 = ashr i64 %z5, 10
+  %z7 = ashr i64 %z6, 10
+  ret i64 %z7
+}
+
+define amdgpu_cs i32 @test_lshr_1(i32 inreg %arg1) {
+; CHECK-LABEL: test_lshr_1:
+; CHECK:       ; %bb.0: ; %.entry
+; CHECK-NEXT:    s_lshr_b32 s0, s0, 5
+; CHECK-NEXT:    ; return to shader part epilog
+.entry:
+  %z1 = lshr i32 %arg1, 2
+  %z2 = lshr i32 %z1, 3
+  ret i32 %z2
+}
+
+define amdgpu_cs i32 @test_lshr_2(i32 inreg %arg1) {
+; CHECK-LABEL: test_lshr_2:
+; CHECK:       ; %bb.0: ; %.entry
+; CHECK-NEXT:    s_lshr_b32 s0, s0, 10
+; CHECK-NEXT:    ; return to shader part epilog
+.entry:
+  %z1 = lshr i32 %arg1, 1
+  %z2 = lshr i32 %z1, 2
+  %z3 = lshr i32 %z2, 3
+  %z4 = lshr i32 %z3, 4
+  ret i32 %z4
+}
+
+define amdgpu_cs i32 @test_lshr_i32(i32 inreg %arg1) {
+; CHECK-LABEL: test_lshr_i32:
+; CHECK:       ; %bb.0: ; %.entry
+; CHECK-NEXT:    s_mov_b32 s0, 0
+; CHECK-NEXT:    ; return to shader part epilog
+.entry:
+  %z1 = lshr i32 %arg1, 10
+  %z2 = lshr i32 %z1, 10
+  %z3 = lshr i32 %z2, 10
+  %z4 = lshr i32 %z3, 10
+  ret i32 %z4
+}
+
+define amdgpu_cs i64 @test_lshr_i64(i64 inreg %arg1) {
+; CHECK-LABEL: test_lshr_i64:
+; CHECK:       ; %bb.0: ; %.entry
+; CHECK-NEXT:    s_mov_b32 s0, 0
+; CHECK-NEXT:    s_mov_b32 s1, 0
+; CHECK-NEXT:    ; return to shader part epilog
+.entry:
+  %z1 = lshr i64 %arg1, 10
+  %z2 = lshr i64 %z1, 10
+  %z3 = lshr i64 %z2, 10
+  %z4 = lshr i64 %z3, 10
+  %z5 = lshr i64 %z4, 10
+  %z6 = lshr i64 %z5, 10
+  %z7 = lshr i64 %z6, 10
+  ret i64 %z7
+}