[GlobalISel] Add a X, Y = G_UNMERGE(G_ZEXT Z) -> X = G_ZEXT Z; Y = 0 combine

Add a combiner helper to transform unmerge of zext into one zext and
a constant 0

Differential Revision: https://reviews.llvm.org/D87427

Author: Quentin Colombet

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

@@ -262,6 +262,10 @@ class CombinerHelper {
   bool matchCombineUnmergeWithDeadLanesToTrunc(MachineInstr &MI);
   bool applyCombineUnmergeWithDeadLanesToTrunc(MachineInstr &MI);
 
+  /// Transform X, Y = G_UNMERGE(G_ZEXT(Z)) -> X = G_ZEXT(Z); Y = G_CONSTANT 0
+  bool matchCombineUnmergeZExtToZExt(MachineInstr &MI);
+  bool applyCombineUnmergeZExtToZExt(MachineInstr &MI);
+
   /// Transform IntToPtr(PtrToInt(x)) to x if cast is in the same address space.
   bool matchCombineI2PToP2I(MachineInstr &MI, Register &Reg);
   bool applyCombineI2PToP2I(MachineInstr &MI, Register &Reg);

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

@@ -429,6 +429,14 @@ def unmerge_dead_to_trunc : GICombineRule<
   (apply [{ return Helper.applyCombineUnmergeWithDeadLanesToTrunc(*${d}); }])
 >;
 
+// Transform x,y = unmerge(zext(z)) -> x = zext z; y = 0.
+def unmerge_zext_to_zext : GICombineRule<
+  (defs root:$d),
+  (match (wip_match_opcode G_UNMERGE_VALUES): $d,
+  [{ return Helper.matchCombineUnmergeZExtToZExt(*${d}); }]),
+  (apply [{ return Helper.applyCombineUnmergeZExtToZExt(*${d}); }])
+>;
+
 // FIXME: These should use the custom predicate feature once it lands.
 def undef_combines : GICombineGroup<[undef_to_fp_zero, undef_to_int_zero,
                                      undef_to_negative_one,
@@ -460,4 +468,5 @@ def all_combines : GICombineGroup<[trivial_combines, ptr_add_immed_chain,
     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_merge, fabs_fabs_fold, unmerge_cst, unmerge_dead_to_trunc,
+    unmerge_zext_to_zext]>;

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

@@ -1687,6 +1687,67 @@ bool CombinerHelper::applyCombineUnmergeWithDeadLanesToTrunc(MachineInstr &MI) {
   return true;
 }
 
+bool CombinerHelper::matchCombineUnmergeZExtToZExt(MachineInstr &MI) {
+  assert(MI.getOpcode() == TargetOpcode::G_UNMERGE_VALUES &&
+         "Expected an unmerge");
+  Register Dst0Reg = MI.getOperand(0).getReg();
+  LLT Dst0Ty = MRI.getType(Dst0Reg);
+  // G_ZEXT on vector applies to each lane, so it will
+  // affect all destinations. Therefore we won't be able
+  // to simplify the unmerge to just the first definition.
+  if (Dst0Ty.isVector())
+    return false;
+  Register SrcReg = MI.getOperand(MI.getNumDefs()).getReg();
+  LLT SrcTy = MRI.getType(SrcReg);
+  if (SrcTy.isVector())
+    return false;
+
+  Register ZExtSrcReg;
+  if (!mi_match(SrcReg, MRI, m_GZExt(m_Reg(ZExtSrcReg))))
+    return false;
+
+  // Finally we can replace the first definition with
+  // a zext of the source if the definition is big enough to hold
+  // all of ZExtSrc bits.
+  LLT ZExtSrcTy = MRI.getType(ZExtSrcReg);
+  return ZExtSrcTy.getSizeInBits() <= Dst0Ty.getSizeInBits();
+}
+
+bool CombinerHelper::applyCombineUnmergeZExtToZExt(MachineInstr &MI) {
+  assert(MI.getOpcode() == TargetOpcode::G_UNMERGE_VALUES &&
+         "Expected an unmerge");
+
+  Register Dst0Reg = MI.getOperand(0).getReg();
+
+  MachineInstr *ZExtInstr =
+      MRI.getVRegDef(MI.getOperand(MI.getNumDefs()).getReg());
+  assert(ZExtInstr && ZExtInstr->getOpcode() == TargetOpcode::G_ZEXT &&
+         "Expecting a G_ZEXT");
+
+  Register ZExtSrcReg = ZExtInstr->getOperand(1).getReg();
+  LLT Dst0Ty = MRI.getType(Dst0Reg);
+  LLT ZExtSrcTy = MRI.getType(ZExtSrcReg);
+
+  Builder.setInstrAndDebugLoc(MI);
+
+  if (Dst0Ty.getSizeInBits() > ZExtSrcTy.getSizeInBits()) {
+    Builder.buildZExt(Dst0Reg, ZExtSrcReg);
+  } else {
+    assert(Dst0Ty.getSizeInBits() == ZExtSrcTy.getSizeInBits() &&
+           "ZExt src doesn't fit in destination");
+    replaceRegWith(MRI, Dst0Reg, ZExtSrcReg);
+  }
+
+  Register ZeroReg;
+  for (unsigned Idx = 1, EndIdx = MI.getNumDefs(); Idx != EndIdx; ++Idx) {
+    if (!ZeroReg)
+      ZeroReg = Builder.buildConstant(Dst0Ty, 0).getReg(0);
+    replaceRegWith(MRI, MI.getOperand(Idx).getReg(), ZeroReg);
+  }
+  MI.eraseFromParent();
+  return true;
+}
+
 bool CombinerHelper::matchCombineShiftToUnmerge(MachineInstr &MI,
                                                 unsigned TargetShiftSize,
                                                 unsigned &ShiftVal) {

llvm/test/CodeGen/AArch64/GlobalISel/combine-unmerge.mir

@@ -369,3 +369,110 @@ body:             |
     %1:_(<2 x s16>),%2:_(<2 x s16>) = G_UNMERGE_VALUES %0(s64)
     $w0 = COPY %1(<2 x s16>)
 ...
+
+# Transform unmerge(zext) into zext.
+# In that test, the source of the zext is same size as the first definition
+# of the unmerge. Therefore a we can just reuse the input of the zext for
+# this definition.
+---
+name:            test_combine_unmerge_zext_to_zext_same_size
+body:             |
+  bb.1:
+    ; CHECK-LABEL: name: test_combine_unmerge_zext_to_zext_same_size
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
+    ; CHECK: $w0 = COPY [[COPY]](s32)
+    ; CHECK: $w1 = COPY [[C]](s32)
+    %0:_(s32) = COPY $w0
+    %3:_(s64) = G_ZEXT %0(s32)
+    %1:_(s32),%2:_(s32) = G_UNMERGE_VALUES %3(s64)
+    $w0 = COPY %1(s32)
+    $w1 = COPY %2(s32)
+...
+
+# Transform unmerge(zext) into zext.
+# In that test, the source of the zext is smaller than the first definition
+# of the unmerge. Therefore a G_ZEXT is required.
+---
+name:            test_combine_unmerge_zext_to_zext
+body:             |
+  bb.1:
+    ; CHECK-LABEL: name: test_combine_unmerge_zext_to_zext
+    ; CHECK: [[COPY:%[0-9]+]]:_(s8) = COPY $b0
+    ; CHECK: [[ZEXT:%[0-9]+]]:_(s16) = G_ZEXT [[COPY]](s8)
+    ; CHECK: [[C:%[0-9]+]]:_(s16) = G_CONSTANT i16 0
+    ; CHECK: $h0 = COPY [[ZEXT]](s16)
+    ; CHECK: $h1 = COPY [[C]](s16)
+    ; CHECK: $h2 = COPY [[C]](s16)
+    ; CHECK: $h3 = COPY [[C]](s16)
+    %0:_(s8) = COPY $b0
+    %3:_(s64) = G_ZEXT %0(s8)
+    %1:_(s16),%2:_(s16),%4:_(s16),%5:_(s16) = G_UNMERGE_VALUES %3(s64)
+    $h0 = COPY %1(s16)
+    $h1 = COPY %2(s16)
+    $h2 = COPY %4(s16)
+    $h3 = COPY %5(s16)
+...
+
+# Check that we don't apply the unmerge(zext) to zext transformation
+# when the first destination of the unmerge is smaller than the source
+# of the zext.
+---
+name:            test_dont_combine_unmerge_zext_to_zext_src_bigger
+body:             |
+  bb.1:
+    ; CHECK-LABEL: name: test_dont_combine_unmerge_zext_to_zext_src_bigger
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[ZEXT:%[0-9]+]]:_(s64) = G_ZEXT [[COPY]](s32)
+    ; CHECK: [[UV:%[0-9]+]]:_(s16), [[UV1:%[0-9]+]]:_(s16), [[UV2:%[0-9]+]]:_(s16), [[UV3:%[0-9]+]]:_(s16) = G_UNMERGE_VALUES [[ZEXT]](s64)
+    ; CHECK: $h0 = COPY [[UV]](s16)
+    ; CHECK: $h1 = COPY [[UV1]](s16)
+    ; CHECK: $h2 = COPY [[UV2]](s16)
+    ; CHECK: $h3 = COPY [[UV3]](s16)
+    %0:_(s32) = COPY $w0
+    %3:_(s64) = G_ZEXT %0(s32)
+    %1:_(s16),%2:_(s16),%4:_(s16),%5:_(s16) = G_UNMERGE_VALUES %3(s64)
+    $h0 = COPY %1(s16)
+    $h1 = COPY %2(s16)
+    $h2 = COPY %4(s16)
+    $h3 = COPY %5(s16)
+...
+
+# Check that we don't apply the unmerge(zext) to zext transformation
+# when the input zext deals with a vector type.
+---
+name:            test_dont_combine_unmerge_zext_to_zext_src_vector
+body:             |
+  bb.1:
+    ; CHECK-LABEL: name: test_dont_combine_unmerge_zext_to_zext_src_vector
+    ; CHECK: [[COPY:%[0-9]+]]:_(<2 x s16>) = COPY $w0
+    ; CHECK: [[ZEXT:%[0-9]+]]:_(<2 x s32>) = G_ZEXT [[COPY]](<2 x s16>)
+    ; CHECK: [[UV:%[0-9]+]]:_(s32), [[UV1:%[0-9]+]]:_(s32) = G_UNMERGE_VALUES [[ZEXT]](<2 x s32>)
+    ; CHECK: $w0 = COPY [[UV]](s32)
+    ; CHECK: $w1 = COPY [[UV1]](s32)
+    %0:_(<2 x s16>) = COPY $w0
+    %3:_(<2 x s32>) = G_ZEXT %0(<2 x s16>)
+    %1:_(s32),%2:_(s32) = G_UNMERGE_VALUES %3(<2 x s32>)
+    $w0 = COPY %1(s32)
+    $w1 = COPY %2(s32)
+...
+
+# Check that we don't apply the unmerge(zext) to zext transformation
+# when the destination type is a vector type.
+# We could actually handle this case but we would need to insert a cast.
+---
+name:            test_dont_combine_unmerge_zext_to_zext_dst_vector
+body:             |
+  bb.1:
+    ; CHECK-LABEL: name: test_dont_combine_unmerge_zext_to_zext_dst_vector
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[ZEXT:%[0-9]+]]:_(s64) = G_ZEXT [[COPY]](s32)
+    ; CHECK: [[UV:%[0-9]+]]:_(<2 x s16>), [[UV1:%[0-9]+]]:_(<2 x s16>) = G_UNMERGE_VALUES [[ZEXT]](s64)
+    ; CHECK: $w0 = COPY [[UV]](<2 x s16>)
+    ; CHECK: $w1 = COPY [[UV1]](<2 x s16>)
+    %0:_(s32) = COPY $w0
+    %3:_(s64) = G_ZEXT %0(s32)
+    %1:_(<2 x s16>),%2:_(<2 x s16>) = G_UNMERGE_VALUES %3(s64)
+    $w0 = COPY %1(<2 x s16>)
+    $w1 = COPY %2(<2 x s16>)
+...