[GlobalISel] Combine zext(trunc x) to x

[Petar.Avramovic]

Summary

Combine zext(trunc x) to x when truncated bits are known to be zero.

[arsenm]

Maybe should add a vector test, although known bits vector support may be too weak to test it

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

2312

getScalarSizeInBits

[foad]

Looks good but I think it needs more test coverage. Does it work on vector types? (You might need to use getScalarSizeInBits instead of getSizeInBits?) If it does work on vectors then that needs tests.

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

2310

Maybe add a test case to check that this combine does nothing if the inner and outer types don't match?

[Petar.Avramovic]

Adding vectors.

[foad]

LGTM.

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

60–61 ↗

(On Diff #321761)

In this case you could just use a trunc from s64 to s32. No need for the zext. That could be a follow-up though.

[foad]

Reverse ping.

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

public:
   // Transform G_PTR_ADD (G_PTRTOINT C1), C2 -> C1 + C2
   bool matchCombineConstPtrAddToI2P(MachineInstr &MI, int64_t &NewCst);
   bool applyCombineConstPtrAddToI2P(MachineInstr &MI, int64_t &NewCst);
 
   /// Transform anyext(trunc(x)) to x.
   bool matchCombineAnyExtTrunc(MachineInstr &MI, Register &Reg);
   bool applyCombineAnyExtTrunc(MachineInstr &MI, Register &Reg);
 
+  /// Transform zext(trunc(x)) to x.
+  bool matchCombineZextTrunc(MachineInstr &MI, Register &Reg);
+
   /// Transform [asz]ext([asz]ext(x)) to [asz]ext x.
   bool matchCombineExtOfExt(MachineInstr &MI,
                             std::tuple<Register, unsigned> &MatchInfo);
   bool applyCombineExtOfExt(MachineInstr &MI,
                             std::tuple<Register, unsigned> &MatchInfo);
 
   /// Transform fneg(fneg(x)) to x.
   bool matchCombineFNegOfFNeg(MachineInstr &MI, Register &Reg);

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

 def anyext_trunc_fold_matchinfo : GIDefMatchData<"Register">;
 def anyext_trunc_fold: GICombineRule <
   (defs root:$root, anyext_trunc_fold_matchinfo:$matchinfo),
   (match (wip_match_opcode G_ANYEXT):$root,
          [{ return Helper.matchCombineAnyExtTrunc(*${root}, ${matchinfo}); }]),
   (apply [{ return Helper.applyCombineAnyExtTrunc(*${root}, ${matchinfo}); }])
 >;
 
+// Fold (zext (trunc x)) -> x if the source type is same as the destination type
+// and truncated bits are known to be zero.
+def zext_trunc_fold_matchinfo : GIDefMatchData<"Register">;
+def zext_trunc_fold: GICombineRule <
+  (defs root:$root, zext_trunc_fold_matchinfo:$matchinfo),
+  (match (wip_match_opcode G_ZEXT):$root,
+         [{ return Helper.matchCombineZextTrunc(*${root}, ${matchinfo}); }]),
+  (apply [{ return Helper.replaceSingleDefInstWithReg(*${root}, ${matchinfo}); }])
+>;
+
 // Fold ([asz]ext ([asz]ext x)) -> ([asz]ext x).
 def ext_ext_fold_matchinfo : GIDefMatchData<"std::tuple<Register, unsigned>">;
 def ext_ext_fold: GICombineRule <
   (defs root:$root, ext_ext_fold_matchinfo:$matchinfo),
   (match (wip_match_opcode G_ANYEXT, G_SEXT, G_ZEXT):$root,
          [{ return Helper.matchCombineExtOfExt(*${root}, ${matchinfo}); }]),
   (apply [{ return Helper.applyCombineExtOfExt(*${root}, ${matchinfo}); }])
 >;
def identity_combines : GICombineGroup<[select_same_val, right_identity_zero,
                                         binop_same_val, binop_left_to_zero,
                                         binop_right_to_zero, p2i_to_i2p,
                                         i2p_to_p2i, anyext_trunc_fold,
                                         fneg_fneg_fold, right_identity_one]>;
 
 def const_combines : GICombineGroup<[constant_fp_op, const_ptradd_to_i2p]>;
 
 def known_bits_simplifications : GICombineGroup<[
-  redundant_and, redundant_sext_inreg, redundant_or, urem_pow2_to_mask]>;
+  redundant_and, redundant_sext_inreg, redundant_or, urem_pow2_to_mask,
+  zext_trunc_fold]>;
 
 def width_reduction_combines : GICombineGroup<[reduce_shl_of_extend]>;
 
 def select_combines : GICombineGroup<[select_undef_cmp, select_constant_cmp]>;
 
 def trivial_combines : GICombineGroup<[copy_prop, mul_to_shl, add_p2i_to_ptradd,
                                        mul_by_neg_one]>;
 

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

 bool CombinerHelper::applyCombineAnyExtTrunc(MachineInstr &MI, Register &Reg) {
   assert(MI.getOpcode() == TargetOpcode::G_ANYEXT && "Expected a G_ANYEXT");
   Register DstReg = MI.getOperand(0).getReg();
   MI.eraseFromParent();
   replaceRegWith(MRI, DstReg, Reg);
   return true;
 }
 
+bool CombinerHelper::matchCombineZextTrunc(MachineInstr &MI, Register &Reg) {
+  assert(MI.getOpcode() == TargetOpcode::G_ZEXT && "Expected a G_ZEXT");
+  Register DstReg = MI.getOperand(0).getReg();
+  Register SrcReg = MI.getOperand(1).getReg();
+  LLT DstTy = MRI.getType(DstReg);
+  if (mi_match(SrcReg, MRI,
+               m_GTrunc(m_all_of(m_Reg(Reg), m_SpecificType(DstTy))))) {
+    return KB->getKnownBits(Reg).countMinLeadingZeros() >=
+           DstTy.getSizeInBits() - MRI.getType(SrcReg).getSizeInBits();
+  }
+  return false;
+}
+
 bool CombinerHelper::matchCombineExtOfExt(
     MachineInstr &MI, std::tuple<Register, unsigned> &MatchInfo) {
   assert((MI.getOpcode() == TargetOpcode::G_ANYEXT ||
           MI.getOpcode() == TargetOpcode::G_SEXT ||
           MI.getOpcode() == TargetOpcode::G_ZEXT) &&
          "Expected a G_[ASZ]EXT");
   Register SrcReg = MI.getOperand(1).getReg();
   MachineInstr *SrcMI = MRI.getVRegDef(SrcReg);

llvm/test/CodeGen/AMDGPU/GlobalISel/shl-ext-reduce.ll

 ; GFX7-NEXT:    s_lshl_b32 s0, s0, 2
 ; GFX7-NEXT:    s_and_b32 s0, s0, 0xffff
 ; GFX7-NEXT:    ; return to shader part epilog
 ;
 ; GFX8-LABEL: s_shl_i32_zext_i16:
 ; GFX8:       ; %bb.0:
 ; GFX8-NEXT:    s_and_b32 s0, s0, 0xffff
 ; GFX8-NEXT:    s_and_b32 s0, s0, 0x3fff
-; GFX8-NEXT:    s_bfe_u32 s1, 2, 0x100000
-; GFX8-NEXT:    s_lshl_b32 s0, s0, s1
-; GFX8-NEXT:    s_bfe_u32 s0, s0, 0x100000
+; GFX8-NEXT:    s_lshl_b32 s0, s0, 2
 ; GFX8-NEXT:    ; return to shader part epilog
 ;
 ; GFX9-LABEL: s_shl_i32_zext_i16:
 ; GFX9:       ; %bb.0:
 ; GFX9-NEXT:    s_and_b32 s0, s0, 0xffff
 ; GFX9-NEXT:    s_and_b32 s0, s0, 0x3fff
-; GFX9-NEXT:    s_bfe_u32 s1, 2, 0x100000
-; GFX9-NEXT:    s_lshl_b32 s0, s0, s1
-; GFX9-NEXT:    s_bfe_u32 s0, s0, 0x100000
+; GFX9-NEXT:    s_lshl_b32 s0, s0, 2
 ; GFX9-NEXT:    ; return to shader part epilog
   %and = and i16 %x, 16383
   %ext = zext i16 %and to i32
   %shl = shl i32 %ext, 2
   ret i32 %shl
 }
 
 define i32 @v_shl_i32_zext_i16(i16 %x) {