[GISel] Rework trunc/shl combine in a generic trunc/shift combine

[Pierre-vh]

Summary

This combine only handled left shifts, but now it can handle right shifts as well. It handles right shifts conservatively and only truncates them to the size returned by TLI.

AMDGPU benefits from always lowering shifts to 32 bits for instance, but AArch64 would rather keep them at 64 bits.

[Pierre-vh]

Remove dead var

[arsenm]

The generic combiner already has trunc(shl), this should go there along with it (preferably with fewer hardcoded sizes)

[Pierre-vh]

Move to generic combine, rework

[Pierre-vh]

Docs

[foad]

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

414

Interesting - this seems like an abuse of getPreferredShiftAmountTy, but I guess it works out OK in practice for the one case we care about (converting 64-bit shifts to 32-bit shifts).

[arsenm]

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

414

This is definitely not what getPreferredShiftAmountTy is for, which is only supposed to help produce already legal shift RHS operands

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

2278

This is what getPreferredShiftAmountTy is for

2292

This usage of getPreferredShiftAmountTy doesn't make sense, but you need something to pick an intermediate type here. Without adding anything, I guess you could bisect types until you found a legal shift? Creating the minimum required shift also "should" work out, but may create more legalizer and optimizer work. It would be nice if we could directly inspect the list of legal shift types, but the API doesn't have that.

Does just reusing the truncating type work as well?

[Pierre-vh]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

2292

Reusing the truncating type doesn't work as well, and trying to guess a legal one doesn't either for AArch64 because 16 bits shifts seem legal there, but are suboptimal compared to 64 (word size) shifts.

My first version actually did a simplified bisecting to find a good type, but AArch64 suffered and that's why I didn't do it.

I think a new TLI hook would be best so targets can just choose whatever they prefer. I'll post a version with that and we can discuss there.

[Pierre-vh]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

2292

Nevermind, I found that AArch64 "suffers" because the combine added in D109419 gets confused by the added trunc on the right shift operands. I'm looking into fixing it.

[Pierre-vh]

• Remove misuse of getPreferredShiftAmountTy
• Only truncate right shifts to 32 bits for now. 16 bits doesn't seem to benefit any targets, I only observed regressions in our backend and no clear gains.
• Don't truncate right shifts if the trunc has any store users to prevent blocking truncstore combine.
  • I looked into adding support for the patterns generated by this new combine to the truncstore combine, but it's a complex combine so it's not immediately clear what needs to change. I also didn't observe any regressions caused by adding this restriction so it's not high-priority IMO.

[Pierre-vh]

ping

[arsenm]

Not sure why you deleted the test

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

2269–2283

Start with lowercase letter

2320–2323

I don't understand special casing this, I'd rather just ignore it

[tschuett]

All combines are always on GISEL opcodes? There are no AMDGPU native combines?

[arsenm]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

2320

Do this first then?

[arsenm]

In D136319#3931108, @tschuett wrote:

All combines are always on GISEL opcodes? There are no AMDGPU native combines?

Yes. You can also define target specific "generic" opcodes

[tschuett]

You have define generic opcodes in the *amdgpu* namespace that represent AMDGPU native instructions?

[arsenm]

In D136319#3931139, @tschuett wrote:

You have define generic opcodes in the *amdgpu* namespace that represent AMDGPU native instructions?

Yes. They are regbankselectable and have the same restrictions as G_* opcodes

[tschuett]

Thanks!

[Pierre-vh]

Comments

[Pierre-vh]

In D136319#3931063, @arsenm wrote:

Not sure why you deleted the test

Test cases are in the new file, not sure why it shows as deleted

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

2320–2323

I spent a while trying to get the truncstore combine to work properly in such cases but it's really not easy. The matching logic is a bit complex.

There's also a comment in that combine's matcher indicating it may be moved to a separate pass in the future to improve its matching logic, so I think it's better to just leave this special case in for now since:

• Fixing the truncstore combine seems to be hard
• This seems to be a niche case - this special case doesn't look like it causes any harm, on the contrary, it allows the combine to not interfere with truncstore.

The detailed reason is that the truncstore combine calculates the number of stores it should be merging together based on the type of the source it matches.

• If we always ignore the trunc, it breaks the combine on AArch64 for 16 bit sources, as the source is always going to be a trunc of a w register. (See trunc_i16_to_i8 testcase).
  • It would always take the 32 bit value as the source, but in those cases the 16 bit source is the correct one.
• Conditionally ignoring the trunc (e.g. only when the src was matched from a shift) can cause premature matching of the combine.
  • e.g. if this (right-shift trunc) combine kicks in on trunc_i64_to_i16, it causes the shift's source to be 32 bits. We then see that we need 2 or 4 stores (depending on which source we use). As the combiner works top-down within a block, we'll find 2 stores before finding 4 of them, and the matcher has to conservatively assume it won't find more. It'll end up merging 2 stores instead of 4.

TL;DR: I would rather leave the special case in for now as removing isn't worth the effort IMO (and it doesn't seem to improve codegen quality significantly either)

[Pierre-vh]

ping

[Pierre-vh]

@arsenm I accidentally forgot to leave in the special case to not fold in the presence of G_STORE users. That's why some aarch64 test was failing, see my previous comments.
Is it still good to land even with the special-case left in?

[arsenm]

In D136319#3980696, @Pierre-vh wrote:

@arsenm I accidentally forgot to leave in the special case to not fold in the presence of G_STORE users. That's why some aarch64 test was failing, see my previous comments.
Is it still good to land even with the special-case left in?

I forgot about that. Can you file an issue to pick it up later?

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

public:
   bool matchCombineFAbsOfFNeg(MachineInstr &MI, BuildFnTy &MatchInfo);
 
   /// Transform trunc ([asz]ext x) to x or ([asz]ext x) or (trunc x).
   bool matchCombineTruncOfExt(MachineInstr &MI,
                               std::pair<Register, unsigned> &MatchInfo);
   void applyCombineTruncOfExt(MachineInstr &MI,
                               std::pair<Register, unsigned> &MatchInfo);
 
-  /// Transform trunc (shl x, K) to shl (trunc x),
-  /// K => K < VT.getScalarSizeInBits().
-  bool matchCombineTruncOfShl(MachineInstr &MI,
-                              std::pair<Register, Register> &MatchInfo);
-  void applyCombineTruncOfShl(MachineInstr &MI,
-                              std::pair<Register, Register> &MatchInfo);
+  /// Transform trunc (shl x, K) to shl (trunc x), K
+  ///    if K < VT.getScalarSizeInBits().
+  ///
+  /// Transforms trunc ([al]shr x, K) to (trunc ([al]shr (MidVT (trunc x)), K))
+  ///    if K <= (MidVT.getScalarSizeInBits() - VT.getScalarSizeInBits())
+  /// MidVT is obtained by finding a legal type between the trunc's src and dst
+  /// types.
+  bool matchCombineTruncOfShift(MachineInstr &MI,
+                                std::pair<MachineInstr *, LLT> &MatchInfo);
+  void applyCombineTruncOfShift(MachineInstr &MI,
+                                std::pair<MachineInstr *, LLT> &MatchInfo);
 
   /// Transform G_MUL(x, -1) to G_SUB(0, x)
   void applyCombineMulByNegativeOne(MachineInstr &MI);
 
   /// Return true if any explicit use operand on \p MI is defined by a
   /// G_IMPLICIT_DEF.
   bool matchAnyExplicitUseIsUndef(MachineInstr &MI);
 

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

 def trunc_ext_fold_matchinfo : GIDefMatchData<"std::pair<Register, unsigned>">;
 def trunc_ext_fold: GICombineRule <
   (defs root:$root, trunc_ext_fold_matchinfo:$matchinfo),
   (match (wip_match_opcode G_TRUNC):$root,
          [{ return Helper.matchCombineTruncOfExt(*${root}, ${matchinfo}); }]),
   (apply [{ Helper.applyCombineTruncOfExt(*${root}, ${matchinfo}); }])
 >;
 
-// Fold trunc (shl x, K) -> shl (trunc x), K => K < VT.getScalarSizeInBits().
-def trunc_shl_matchinfo : GIDefMatchData<"std::pair<Register, Register>">;
-def trunc_shl: GICombineRule <
-  (defs root:$root, trunc_shl_matchinfo:$matchinfo),
+// Under certain conditions, transform:
+//  trunc (shl x, K)     -> shl (trunc x), K//
+//  trunc ([al]shr x, K) -> (trunc ([al]shr (trunc x), K))
+def trunc_shift_matchinfo : GIDefMatchData<"std::pair<MachineInstr*, LLT>">;
+def trunc_shift: GICombineRule <
+  (defs root:$root, trunc_shift_matchinfo:$matchinfo),
   (match (wip_match_opcode G_TRUNC):$root,
-         [{ return Helper.matchCombineTruncOfShl(*${root}, ${matchinfo}); }]),
-  (apply [{ Helper.applyCombineTruncOfShl(*${root}, ${matchinfo}); }])
+         [{ return Helper.matchCombineTruncOfShift(*${root}, ${matchinfo}); }]),
+  (apply [{ Helper.applyCombineTruncOfShift(*${root}, ${matchinfo}); }])
 >;
 
 // Transform (mul x, -1) -> (sub 0, x)
 def mul_by_neg_one: GICombineRule <
   (defs root:$root),
   (match (wip_match_opcode G_MUL):$root,
          [{ return Helper.matchConstantOp(${root}->getOperand(2), -1); }]),
   (apply [{ Helper.applyCombineMulByNegativeOne(*${root}); }])
def all_combines : GICombineGroup<[trivial_combines, insert_vec_elt_combines,
     combine_indexed_load_store, undef_combines, identity_combines, phi_combines,
     simplify_add_to_sub, hoist_logic_op_with_same_opcode_hands,
     reassocs, ptr_add_immed_chain,
     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, unmerge_cst, unmerge_dead_to_trunc,
-    unmerge_zext_to_zext, merge_unmerge, trunc_ext_fold, trunc_shl,
+    unmerge_zext_to_zext, merge_unmerge, trunc_ext_fold, trunc_shift,
     const_combines, xor_of_and_with_same_reg, ptr_add_with_zero,
     shift_immed_chain, shift_of_shifted_logic_chain, load_or_combine,
     truncstore_merge, div_rem_to_divrem, funnel_shift_combines,
     form_bitfield_extract, constant_fold, fabs_fneg_fold,
     intdiv_combines, mulh_combines, redundant_neg_operands,
     and_or_disjoint_mask, fma_combines, fold_binop_into_select,
     sub_add_reg, select_to_minmax, redundant_binop_in_equality,
     fsub_to_fneg]>;
 
 // A combine group used to for prelegalizer combiners at -O0. The combines in
 // this group have been selected based on experiments to balance code size and
 // compile time performance.
 def optnone_combines : GICombineGroup<[trivial_combines,
     ptr_add_immed_chain, combines_for_extload,
     not_cmp_fold, opt_brcond_by_inverting_cond]>;

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

void CombinerHelper::applyCombineTruncOfExt(
   Builder.setInstrAndDebugLoc(MI);
   if (SrcTy.getSizeInBits() < DstTy.getSizeInBits())
     Builder.buildInstr(SrcExtOp, {DstReg}, {SrcReg});
   else
     Builder.buildTrunc(DstReg, SrcReg);
   MI.eraseFromParent();
 }
 
-bool CombinerHelper::matchCombineTruncOfShl(
-    MachineInstr &MI, std::pair<Register, Register> &MatchInfo) {
+static LLT getMidVTForTruncRightShiftCombine(LLT ShiftTy, LLT TruncTy) {
+  const unsigned ShiftSize = ShiftTy.getScalarSizeInBits();
+  const unsigned TruncSize = TruncTy.getScalarSizeInBits();
+
+  // ShiftTy > 32 > TruncTy -> 32
+  if (ShiftSize > 32 && TruncSize < 32)
+    return ShiftTy.changeElementSize(32);
+
+  // TODO: We could also reduce to 16 bits, but that's more target-dependent.
+  //  Some targets like it, some don't, some only like it under certain
+  //  conditions/processor versions, etc.
+  //  A TL hook might be needed for this.
+
+  // Don't combine
+  return ShiftTy;
+}
+
+bool CombinerHelper::matchCombineTruncOfShift(
+    MachineInstr &MI, std::pair<MachineInstr *, LLT> &MatchInfo) {
   assert(MI.getOpcode() == TargetOpcode::G_TRUNC && "Expected a G_TRUNC");
   Register DstReg = MI.getOperand(0).getReg();
   Register SrcReg = MI.getOperand(1).getReg();
+
+  if (!MRI.hasOneNonDBGUse(SrcReg))
+    return false;
+
+  LLT SrcTy = MRI.getType(SrcReg);
   LLT DstTy = MRI.getType(DstReg);
-  Register ShiftSrc;
-  Register ShiftAmt;
 
-  if (MRI.hasOneNonDBGUse(SrcReg) &&
-      mi_match(SrcReg, MRI, m_GShl(m_Reg(ShiftSrc), m_Reg(ShiftAmt))) &&
-      isLegalOrBeforeLegalizer(
-          {TargetOpcode::G_SHL,
-           {DstTy, getTargetLowering().getPreferredShiftAmountTy(DstTy)}})) {
-    KnownBits Known = KB->getKnownBits(ShiftAmt);
-    unsigned Size = DstTy.getSizeInBits();
-    if (Known.countMaxActiveBits() <= Log2_32(Size)) {
-      MatchInfo = std::make_pair(ShiftSrc, ShiftAmt);
-      return true;
-    }
+  MachineInstr *SrcMI = getDefIgnoringCopies(SrcReg, MRI);
+  const auto &TL = getTargetLowering();
+
+  LLT NewShiftTy;
+  switch (SrcMI->getOpcode()) {
+  default:
+    return false;
+  case TargetOpcode::G_SHL: {
+    NewShiftTy = DstTy;
+
+    // Make sure new shift amount is legal.
+    KnownBits Known = KB->getKnownBits(SrcMI->getOperand(2).getReg());
+    if (Known.getMaxValue().uge(NewShiftTy.getScalarSizeInBits()))
+      return false;
+    break;
   }
+  case TargetOpcode::G_LSHR:
+  case TargetOpcode::G_ASHR: {
+    // For right shifts, we conservatively do not do the transform if the TRUNC
+    // has any STORE users. The reason is that if we change the type of the
+    // shift, we may break the truncstore combine.
+    //
+    // TODO: Fix truncstore combine to handle (trunc(lshr (trunc x), k)).
+    for (auto &User : MRI.use_instructions(DstReg))
+      if (User.getOpcode() == TargetOpcode::G_STORE)
+        return false;
+
+    NewShiftTy = getMidVTForTruncRightShiftCombine(SrcTy, DstTy);
+    if (NewShiftTy == SrcTy)
+      return false;
+
+    // Make sure we won't lose information by truncating the high bits.
+    KnownBits Known = KB->getKnownBits(SrcMI->getOperand(2).getReg());
+    if (Known.getMaxValue().ugt(NewShiftTy.getScalarSizeInBits() -
+                                DstTy.getScalarSizeInBits()))
       return false;
+    break;
+  }
   }
 
-void CombinerHelper::applyCombineTruncOfShl(
-    MachineInstr &MI, std::pair<Register, Register> &MatchInfo) {
-  assert(MI.getOpcode() == TargetOpcode::G_TRUNC && "Expected a G_TRUNC");
-  Register DstReg = MI.getOperand(0).getReg();
-  Register SrcReg = MI.getOperand(1).getReg();
-  LLT DstTy = MRI.getType(DstReg);
-  MachineInstr *SrcMI = MRI.getVRegDef(SrcReg);
+  if (!isLegalOrBeforeLegalizer(
+          {SrcMI->getOpcode(),
+           {NewShiftTy, TL.getPreferredShiftAmountTy(NewShiftTy)}}))
+    return false;
 
-  Register ShiftSrc = MatchInfo.first;
-  Register ShiftAmt = MatchInfo.second;
+  MatchInfo = std::make_pair(SrcMI, NewShiftTy);
+  return true;
+}
+
+void CombinerHelper::applyCombineTruncOfShift(
+    MachineInstr &MI, std::pair<MachineInstr *, LLT> &MatchInfo) {
   Builder.setInstrAndDebugLoc(MI);
-  auto TruncShiftSrc = Builder.buildTrunc(DstTy, ShiftSrc);
-  Builder.buildShl(DstReg, TruncShiftSrc, ShiftAmt, SrcMI->getFlags());
-  MI.eraseFromParent();
+
+  MachineInstr *ShiftMI = MatchInfo.first;
+  LLT NewShiftTy = MatchInfo.second;
+
+  Register Dst = MI.getOperand(0).getReg();
+  LLT DstTy = MRI.getType(Dst);
+
+  Register ShiftAmt = ShiftMI->getOperand(2).getReg();
+  Register ShiftSrc = ShiftMI->getOperand(1).getReg();
+  ShiftSrc = Builder.buildTrunc(NewShiftTy, ShiftSrc).getReg(0);
+
+  Register NewShift =
+      Builder
+          .buildInstr(ShiftMI->getOpcode(), {NewShiftTy}, {ShiftSrc, ShiftAmt})
+          .getReg(0);
+
+  if (NewShiftTy == DstTy)
+    replaceRegWith(MRI, Dst, NewShift);
+  else
+    Builder.buildTrunc(Dst, NewShift);
+
+  eraseInst(MI);
 }
 
 bool CombinerHelper::matchAnyExplicitUseIsUndef(MachineInstr &MI) {
   return any_of(MI.explicit_uses(), [this](const MachineOperand &MO) {
     return MO.isReg() &&
            getOpcodeDef(TargetOpcode::G_IMPLICIT_DEF, MO.getReg(), MRI);
   });
 }

llvm/test/CodeGen/AMDGPU/GlobalISel/combine-trunc-shift.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -march=amdgcn -run-pass=amdgpu-postlegalizer-combiner -verify-machineinstrs %s -o - | FileCheck %s
+
+---
+name: trunc_s32_shl_s64_5
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $vgpr0_vgpr1
+
+    ; CHECK-LABEL: name: trunc_s32_shl_s64_5
+    ; CHECK: liveins: $vgpr0_vgpr1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(s64) = COPY $vgpr0_vgpr1
+    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK-NEXT: [[TRUNC:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
+    ; CHECK-NEXT: [[SHL:%[0-9]+]]:_(s32) = G_SHL [[TRUNC]], [[C]](s32)
+    ; CHECK-NEXT: $vgpr0 = COPY [[SHL]](s32)
+    %0:_(s64) = COPY $vgpr0_vgpr1
+    %1:_(s32) = G_CONSTANT i32 1
+    %2:_(s64) = G_SHL %0:_, %1
+    %3:_(s32) = G_TRUNC %2
+    $vgpr0 = COPY %3
+...
+
+---
+name: trunc_s16_shl_s32_5
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $vgpr0
+
+    ; CHECK-LABEL: name: trunc_s16_shl_s32_5
+    ; CHECK: liveins: $vgpr0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(s32) = COPY $vgpr0
+    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK-NEXT: [[SHL:%[0-9]+]]:_(s32) = G_SHL [[COPY]], [[C]](s32)
+    ; CHECK-NEXT: [[TRUNC:%[0-9]+]]:_(s16) = G_TRUNC [[SHL]](s32)
+    ; CHECK-NEXT: S_ENDPGM 0, implicit [[TRUNC]](s16)
+    %0:_(s32) = COPY $vgpr0
+    %1:_(s32) = G_CONSTANT i32 1
+    %2:_(s32) = G_SHL %0:_, %1
+    %3:_(s16) = G_TRUNC %2
+    S_ENDPGM 0, implicit %3
+
+...
+
+---
+name: trunc_s16_shl_s64_5
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $vgpr0_vgpr1
+
+    ; CHECK-LABEL: name: trunc_s16_shl_s64_5
+    ; CHECK: liveins: $vgpr0_vgpr1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(s64) = COPY $vgpr0_vgpr1
+    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK-NEXT: [[SHL:%[0-9]+]]:_(s64) = G_SHL [[COPY]], [[C]](s32)
+    ; CHECK-NEXT: [[TRUNC:%[0-9]+]]:_(s16) = G_TRUNC [[SHL]](s64)
+    ; CHECK-NEXT: S_ENDPGM 0, implicit [[TRUNC]](s16)
+    %0:_(s64) = COPY $vgpr0_vgpr1
+    %1:_(s32) = G_CONSTANT i32 1
+    %2:_(s64) = G_SHL %0:_, %1
+    %3:_(s16) = G_TRUNC %2
+    S_ENDPGM 0, implicit %3
+
+...
+
+---
+name:            s16_trunc_s64_lshr_16
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $vgpr0
+    ; CHECK-LABEL: name: s16_trunc_s64_lshr_16
+    ; CHECK: liveins: $vgpr0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(s32) = COPY $vgpr0
+    ; CHECK-NEXT: %amt:_(s32) = G_CONSTANT i32 16
+    ; CHECK-NEXT: [[LSHR:%[0-9]+]]:_(s32) = G_LSHR [[COPY]], %amt(s32)
+    ; CHECK-NEXT: %trunc:_(s16) = G_TRUNC [[LSHR]](s32)
+    ; CHECK-NEXT: %foo:_(s16) = G_CONSTANT i16 55
+    ; CHECK-NEXT: %keep:_(s32) = G_MERGE_VALUES %trunc(s16), %foo(s16)
+    ; CHECK-NEXT: $vgpr0 = COPY %keep(s32)
+    %0:_(s32) = COPY $vgpr0
+    %src:_(s64) = G_ZEXT %0
+    %amt:_(s32) = G_CONSTANT i32 16
+    %shift:_(s64) = G_LSHR %src, %amt
+    %trunc:_(s16) = G_TRUNC %shift
+    %foo:_(s16) = G_CONSTANT i16 55
+    %keep:_(s32) = G_MERGE_VALUES %trunc, %foo
+    $vgpr0 = COPY %keep
+...
+
+---
+name:            s16_trunc_s64_ashr_16
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $vgpr0
+    ; CHECK-LABEL: name: s16_trunc_s64_ashr_16
+    ; CHECK: liveins: $vgpr0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(s32) = COPY $vgpr0
+    ; CHECK-NEXT: %amt:_(s32) = G_CONSTANT i32 16
+    ; CHECK-NEXT: [[ASHR:%[0-9]+]]:_(s32) = G_ASHR [[COPY]], %amt(s32)
+    ; CHECK-NEXT: %trunc:_(s16) = G_TRUNC [[ASHR]](s32)
+    ; CHECK-NEXT: %foo:_(s16) = G_CONSTANT i16 55
+    ; CHECK-NEXT: %keep:_(s32) = G_MERGE_VALUES %trunc(s16), %foo(s16)
+    ; CHECK-NEXT: $vgpr0 = COPY %keep(s32)
+    %0:_(s32) = COPY $vgpr0
+    %src:_(s64) = G_ZEXT %0
+    %amt:_(s32) = G_CONSTANT i32 16
+    %shift:_(s64) = G_ASHR %src, %amt
+    %trunc:_(s16) = G_TRUNC %shift
+    %foo:_(s16) = G_CONSTANT i16 55
+    %keep:_(s32) = G_MERGE_VALUES %trunc, %foo
+    $vgpr0 = COPY %keep
+...
+
+---
+name:            s16_trunc_s64_lshr_17_nofold
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $vgpr0
+    ; CHECK-LABEL: name: s16_trunc_s64_lshr_17_nofold
+    ; CHECK: liveins: $vgpr0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(s32) = COPY $vgpr0
+    ; CHECK-NEXT: %src:_(s64) = G_ZEXT [[COPY]](s32)
+    ; CHECK-NEXT: %amt:_(s32) = G_CONSTANT i32 17
+    ; CHECK-NEXT: %shift:_(s64) = G_LSHR %src, %amt(s32)
+    ; CHECK-NEXT: %trunc:_(s16) = G_TRUNC %shift(s64)
+    ; CHECK-NEXT: %foo:_(s16) = G_CONSTANT i16 55
+    ; CHECK-NEXT: %keep:_(s32) = G_MERGE_VALUES %trunc(s16), %foo(s16)
+    ; CHECK-NEXT: $vgpr0 = COPY %keep(s32)
+    %0:_(s32) = COPY $vgpr0
+    %src:_(s64) = G_ZEXT %0
+    %amt:_(s32) = G_CONSTANT i32 17
+    %shift:_(s64) = G_LSHR %src, %amt
+    %trunc:_(s16) = G_TRUNC %shift
+    %foo:_(s16) = G_CONSTANT i16 55
+    %keep:_(s32) = G_MERGE_VALUES %trunc, %foo
+    $vgpr0 = COPY %keep
+...
+
+---
+name:            s26_trunc_s64_lshr_6
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $vgpr0
+    ; CHECK-LABEL: name: s26_trunc_s64_lshr_6
+    ; CHECK: liveins: $vgpr0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(s32) = COPY $vgpr0
+    ; CHECK-NEXT: %amt:_(s32) = G_CONSTANT i32 6
+    ; CHECK-NEXT: [[LSHR:%[0-9]+]]:_(s32) = G_LSHR [[COPY]], %amt(s32)
+    ; CHECK-NEXT: %trunc:_(s26) = G_TRUNC [[LSHR]](s32)
+    ; CHECK-NEXT: %foo:_(s26) = G_CONSTANT i26 55
+    ; CHECK-NEXT: %keep0:_(s26) = G_ADD %trunc, %foo
+    ; CHECK-NEXT: %keep1:_(s32) = G_ANYEXT %keep0(s26)
+    ; CHECK-NEXT: $vgpr0 = COPY %keep1(s32)
+    %0:_(s32) = COPY $vgpr0
+    %src:_(s64) = G_ZEXT %0
+    %amt:_(s32) = G_CONSTANT i32 6
+    %shift:_(s64) = G_LSHR %src, %amt
+    %trunc:_(s26) = G_TRUNC %shift
+    %foo:_(s26) = G_CONSTANT i26 55
+    %keep0:_(s26) = G_ADD %trunc, %foo
+    %keep1:_(s32) = G_ANYEXT %keep0
+    $vgpr0 = COPY %keep1
+...
+
+---
+name:            s26_trunc_s64_lshr_7_nofold
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $vgpr0
+    ; CHECK-LABEL: name: s26_trunc_s64_lshr_7_nofold
+    ; CHECK: liveins: $vgpr0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(s32) = COPY $vgpr0
+    ; CHECK-NEXT: %src:_(s64) = G_ZEXT [[COPY]](s32)
+    ; CHECK-NEXT: %amt:_(s32) = G_CONSTANT i32 7
+    ; CHECK-NEXT: %shift:_(s64) = G_LSHR %src, %amt(s32)
+    ; CHECK-NEXT: %trunc:_(s26) = G_TRUNC %shift(s64)
+    ; CHECK-NEXT: %foo:_(s26) = G_CONSTANT i26 55
+    ; CHECK-NEXT: %keep0:_(s26) = G_ADD %trunc, %foo
+    ; CHECK-NEXT: %keep1:_(s32) = G_ANYEXT %keep0(s26)
+    ; CHECK-NEXT: $vgpr0 = COPY %keep1(s32)
+    %0:_(s32) = COPY $vgpr0
+    %src:_(s64) = G_ZEXT %0
+    %amt:_(s32) = G_CONSTANT i32 7
+    %shift:_(s64) = G_LSHR %src, %amt
+    %trunc:_(s26) = G_TRUNC %shift
+    %foo:_(s26) = G_CONSTANT i26 55
+    %keep0:_(s26) = G_ADD %trunc, %foo
+    %keep1:_(s32) = G_ANYEXT %keep0
+    $vgpr0 = COPY %keep1
+...

llvm/test/CodeGen/AMDGPU/GlobalISel/combine-trunc-shl.mir

-# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
-# RUN: llc -mtriple=amdgcn-amd-amdhsa -run-pass=amdgpu-postlegalizer-combiner -verify-machineinstrs %s -o - | FileCheck %s
-
----
-name: trunc_s32_shl_s64_5
-legalized:       true
-tracksRegLiveness: true
-body:             |
-  bb.0:
-    liveins: $vgpr0_vgpr1
-
-    ; CHECK-LABEL: name: trunc_s32_shl_s64_5
-    ; CHECK: liveins: $vgpr0_vgpr1
-    ; CHECK-NEXT: {{  $}}
-    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(s64) = COPY $vgpr0_vgpr1
-    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
-    ; CHECK-NEXT: [[TRUNC:%[0-9]+]]:_(s32) = G_TRUNC [[COPY]](s64)
-    ; CHECK-NEXT: [[SHL:%[0-9]+]]:_(s32) = G_SHL [[TRUNC]], [[C]](s32)
-    ; CHECK-NEXT: $vgpr0 = COPY [[SHL]](s32)
-    %0:_(s64) = COPY $vgpr0_vgpr1
-    %1:_(s32) = G_CONSTANT i32 1
-    %2:_(s64) = G_SHL %0:_, %1
-    %3:_(s32) = G_TRUNC %2
-    $vgpr0 = COPY %3
-...
-
----
-name: trunc_s16_shl_s32_5
-legalized:       true
-tracksRegLiveness: true
-body:             |
-  bb.0:
-    liveins: $vgpr0
-
-    ; CHECK-LABEL: name: trunc_s16_shl_s32_5
-    ; CHECK: liveins: $vgpr0
-    ; CHECK-NEXT: {{  $}}
-    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(s32) = COPY $vgpr0
-    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
-    ; CHECK-NEXT: [[SHL:%[0-9]+]]:_(s32) = G_SHL [[COPY]], [[C]](s32)
-    ; CHECK-NEXT: [[TRUNC:%[0-9]+]]:_(s16) = G_TRUNC [[SHL]](s32)
-    ; CHECK-NEXT: S_ENDPGM 0, implicit [[TRUNC]](s16)
-    %0:_(s32) = COPY $vgpr0
-    %1:_(s32) = G_CONSTANT i32 1
-    %2:_(s32) = G_SHL %0:_, %1
-    %3:_(s16) = G_TRUNC %2
-    S_ENDPGM 0, implicit %3
-
-...
-
----
-name: trunc_s16_shl_s64_5
-legalized:       true
-tracksRegLiveness: true
-body:             |
-  bb.0:
-    liveins: $vgpr0_vgpr1
-
-    ; CHECK-LABEL: name: trunc_s16_shl_s64_5
-    ; CHECK: liveins: $vgpr0_vgpr1
-    ; CHECK-NEXT: {{  $}}
-    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(s64) = COPY $vgpr0_vgpr1
-    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
-    ; CHECK-NEXT: [[SHL:%[0-9]+]]:_(s64) = G_SHL [[COPY]], [[C]](s32)
-    ; CHECK-NEXT: [[TRUNC:%[0-9]+]]:_(s16) = G_TRUNC [[SHL]](s64)
-    ; CHECK-NEXT: S_ENDPGM 0, implicit [[TRUNC]](s16)
-    %0:_(s64) = COPY $vgpr0_vgpr1
-    %1:_(s32) = G_CONSTANT i32 1
-    %2:_(s64) = G_SHL %0:_, %1
-    %3:_(s16) = G_TRUNC %2
-    S_ENDPGM 0, implicit %3
-
-...