[GISel][CombinerHelper] Add concat_vectors(build_vector, build_vector) => build_vector

[qcolombet]

Summary

Teach the combiner helper how to flatten concat_vectors of build_vectors
into a build_vector.

Add this combine as part of AArch64 pre-legalizer combiner.

[arsenm]

include/llvm/CodeGen/GlobalISel/CombinerHelper.h

108

ArrayRef?

lib/CodeGen/GlobalISel/CombinerHelper.cpp

92–94

Should there be a legality check here somewhere? A concat_vectors of build_vector (or possibly build_vector_trunc) is probably the preferred form of sub-32-bit element vectors, but not for 32 or larger.

104–105

I would expect this to be an assert

test/CodeGen/AArch64/GlobalISel/prelegalizercombiner-concat-vectors.mir

26–28

Can you add some tests with pointer elements?

[qcolombet]

include/llvm/CodeGen/GlobalISel/CombinerHelper.h

108

Sure!

lib/CodeGen/GlobalISel/CombinerHelper.cpp

92–94

I think the goal of the helper is just to provide the combine functionality. Whether or not the combine should be applied is left to the caller.

What do you think?

104–105

I would have expected this too, but decided against it to match the style of the other match functions.
Given the (arguably strange) use of the try methods in tryCombine, an assert does not seem desirable.

Happy to change it though :).

test/CodeGen/AArch64/GlobalISel/prelegalizercombiner-concat-vectors.mir

26–28

Sure!

[qcolombet]

• Use ArrayRef instead of const SmallVectorImpl&
• Add tests with pointer types.

[arsenm]

lib/CodeGen/GlobalISel/CombinerHelper.cpp

92–94

That makes sense

104–105

I replaced similar checks in the LegalizationArtifactCombiner in r367376

[qcolombet]

lib/CodeGen/GlobalISel/CombinerHelper.cpp

104–105

Love it!
I'll fix that one in that PR and the other in another commit.

[qcolombet]

• Replace check with an assertion
• Update comments with pre-conditions

[arsenm]

LGTM with nit

lib/CodeGen/GlobalISel/CombinerHelper.cpp

113

I don't think this is needed since the verifier will catch it

126

I don't think this is needed since the verifier will catch it

139–148

Pull out repeated getType calls into a variable?

[qcolombet]

• Remove asserts already covered by the verifier
• Put .getType calls into a variable.

[arsenm]

LGTM

include/llvm/CodeGen/GlobalISel/CombinerHelper.h

public:
 
   /// Combine \p MI into a pre-indexed or post-indexed load/store operation if
   /// legal and the surrounding code makes it useful.
   bool tryCombineIndexedLoadStore(MachineInstr &MI);
 
   bool matchCombineBr(MachineInstr &MI);
   bool tryCombineBr(MachineInstr &MI);
 
+  /// If \p MI is G_CONCAT_VECTORS, try to combine it.
+  /// Returns true if MI changed.
+  /// Right now, we support:
+  /// - concat_vector(undef, undef) => undef
+  /// - concat_vector(build_vector(A, B), build_vector(C, D)) =>
+  ///   build_vector(A, B, C, D)
+  bool tryCombineConcatVectors(MachineInstr &MI);
+  /// Check if the G_CONCAT_VECTORS \p MI is undef or if it
+  /// can be flattened into a build_vector.
+  /// In the first case \p IsUndef will be true.
+  /// In the second case \p Ops will contain the operands needed
+  /// to produce the flattened build_vector.
+  bool matchCombineConcatVectors(MachineInstr &MI, bool &IsUndef,
+                                 SmallVectorImpl<Register> &Ops);
+  /// Replace MI with a flattened build_vector with \p Ops or an
+  /// implicit_def if IsUndef is true.
+  void applyCombineConcatVectors(MachineInstr &MI, bool IsUndef,
+                                 const ArrayRef<Register> Ops);
+
   /// Optimize memcpy intrinsics et al, e.g. constant len calls.
   /// /p MaxLen if non-zero specifies the max length of a mem libcall to inline.
   ///
   /// For example (pre-indexed):
   ///
   ///     $addr = G_GEP $base, $offset
   ///     [...]
   ///     $val = G_LOAD $addr

lib/CodeGen/GlobalISel/CombinerHelper.cpp

 }
 void CombinerHelper::applyCombineCopy(MachineInstr &MI) {
   Register DstReg = MI.getOperand(0).getReg();
   Register SrcReg = MI.getOperand(1).getReg();
   MI.eraseFromParent();
   replaceRegWith(MRI, DstReg, SrcReg);
 }
 
+bool CombinerHelper::tryCombineConcatVectors(MachineInstr &MI) {
+  bool IsUndef = false;
+  SmallVector<Register, 4> Ops;
+  if (matchCombineConcatVectors(MI, IsUndef, Ops)) {
+    applyCombineConcatVectors(MI, IsUndef, Ops);
+    return true;
+  }
+  return false;
+}
+
+bool CombinerHelper::matchCombineConcatVectors(MachineInstr &MI, bool &IsUndef,
+                                               SmallVectorImpl<Register> &Ops) {
+  if (MI.getOpcode() != TargetOpcode::G_CONCAT_VECTORS)
+    return false;
+  IsUndef = true;
+  MachineInstr *Undef = nullptr;
+
+  // Walk over all the operands of concat vectors and check if they are
+  // build_vector themselves or undef.
+  // Then collect their operands in Ops.
+  for (const MachineOperand &MO : MI.operands()) {
+    assert(MO.isReg() && "Concat non-register values");
+    // Skip the instruction definition.
+    if (MO.isDef())
+      continue;
+    Register Reg = MO.getReg();
+    MachineInstr *Def = MRI.getVRegDef(Reg);
+    assert(Def && "Operand not defined");
+    switch (Def->getOpcode()) {
+    case TargetOpcode::G_BUILD_VECTOR:
+      IsUndef = false;
+      // Remember the operands of the build_vector to fold
+      // them into the yet-to-build flattened concat vectors.
+      for (const MachineOperand &BuildVecMO : Def->operands()) {
+        assert(MO.isReg() && "Build a vector with non-register values");
+        // Skip the definition.
+        if (BuildVecMO.isDef())
+          continue;
+        Ops.push_back(BuildVecMO.getReg());
+      }
+      break;
+    case TargetOpcode::G_IMPLICIT_DEF:
+      // Keep one undef value for all the undef operands.
+      if (!Undef) {
+        assert(MRI.getType(Reg).isVector() &&
+               "We should be concatenating vectors");
+        Builder.setInsertPt(*MI.getParent(), MI);
+        Undef = Builder.buildUndef(MRI.getType(Reg).getScalarType());
+      }
+      assert(MRI.getType(Undef->getOperand(0).getReg()) ==
+                 MRI.getType(Reg).getScalarType() &&
+             "All undefs should have the same type");
+      // Break the undef vector in as many scalar elements as needed
+      // for the flattening.
+      for (unsigned EltIdx = 0, EltEnd = MRI.getType(Reg).getNumElements();
+           EltIdx != EltEnd; ++EltIdx)
+        Ops.push_back(Undef->getOperand(0).getReg());
+      break;
+    default:
+      return false;
+    }
+  }
+  return true;
+}
+void CombinerHelper::applyCombineConcatVectors(
+    MachineInstr &MI, bool IsUndef, const ArrayRef<Register> Ops) {
+  // We determined that the concat_vectors can be flatten.
+  // Generate the flattened build_vector.
+  Register DstReg = MI.getOperand(0).getReg();
+  Builder.setInsertPt(*MI.getParent(), MI);
+  Register NewDstReg = MRI.cloneVirtualRegister(DstReg);
+
+  // Note: IsUndef is sort of redundant. We could have determine it by
+  // checking that at all Ops are undef.  Alternatively, we could have
+  // generate a build_vector of undefs and rely on another combine to
+  // clean that up.  For now, given we already gather this information
+  // in tryCombineConcatVectors, just save compile time and issue the
+  // right thing.
+  if (IsUndef)
+    Builder.buildUndef(NewDstReg);
+  else
+    Builder.buildBuildVector(NewDstReg, Ops);
+  MI.eraseFromParent();
+  replaceRegWith(MRI, DstReg, NewDstReg);
+}
+
 namespace {
 
 /// Select a preference between two uses. CurrentUse is the current preference
 /// while *ForCandidate is attributes of the candidate under consideration.
 PreferredTuple ChoosePreferredUse(PreferredTuple &CurrentUse,
                                   const LLT &TyForCandidate,
                                   unsigned OpcodeForCandidate,
                                   MachineInstr *MIForCandidate) {

lib/Target/AArch64/AArch64PreLegalizerCombiner.cpp

bool AArch64PreLegalizerCombinerInfo::combine(GISelChangeObserver &Observer,
 
   switch (MI.getOpcode()) {
   default:
     return false;
   case TargetOpcode::COPY:
     return Helper.tryCombineCopy(MI);
   case TargetOpcode::G_BR:
     return Helper.tryCombineBr(MI);
+  case TargetOpcode::G_CONCAT_VECTORS:
+    return Helper.tryCombineConcatVectors(MI);
   case TargetOpcode::G_LOAD:
   case TargetOpcode::G_SEXTLOAD:
   case TargetOpcode::G_ZEXTLOAD: {
     bool Changed = false;
     Changed |= Helper.tryCombineExtendingLoads(MI);
     Changed |= Helper.tryCombineIndexedLoadStore(MI);
     return Changed;
   }

test/CodeGen/AArch64/GlobalISel/prelegalizercombiner-concat-vectors.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple aarch64-apple-ios  -run-pass=aarch64-prelegalizer-combiner %s -o - | FileCheck %s
+
+# Check that we canonicalize concat_vectors(build_vector, build_vector)
+# into build_vector.
+---
+name: concat_to_build_vector
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $x0, $x1, $x2, $x3
+
+    ; CHECK-LABEL: name: concat_to_build_vector
+    ; CHECK: liveins: $x0, $x1, $x2, $x3
+    ; CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY $x0
+    ; CHECK: [[COPY1:%[0-9]+]]:_(s64) = COPY $x1
+    ; CHECK: [[COPY2:%[0-9]+]]:_(s64) = COPY $x2
+    ; CHECK: [[COPY3:%[0-9]+]]:_(s64) = COPY $x3
+    ; CHECK: [[BUILD_VECTOR:%[0-9]+]]:_(<4 x s64>) = G_BUILD_VECTOR [[COPY]](s64), [[COPY1]](s64), [[COPY2]](s64), [[COPY3]](s64)
+    ; CHECK: RET_ReallyLR implicit [[BUILD_VECTOR]](<4 x s64>)
+    %0:_(s64) = COPY $x0
+    %1:_(s64) = COPY $x1
+    %2:_(s64) = COPY $x2
+    %3:_(s64) = COPY $x3
+    %4:_(<2 x s64>) = G_BUILD_VECTOR %0(s64), %1
+    %5:_(<2 x s64>) = G_BUILD_VECTOR %2(s64), %3
+    %6:_(<4 x s64>) = G_CONCAT_VECTORS %4(<2 x s64>), %5
+    RET_ReallyLR implicit %6
+...
+# Same test as concat_to_build_vector but with pointer types.
+---
+name: concat_to_build_vector_ptr
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $x0, $x1, $x2, $x3
+
+    ; CHECK-LABEL: name: concat_to_build_vector_ptr
+    ; CHECK: liveins: $x0, $x1, $x2, $x3
+    ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
+    ; CHECK: [[COPY1:%[0-9]+]]:_(p0) = COPY $x1
+    ; CHECK: [[COPY2:%[0-9]+]]:_(p0) = COPY $x2
+    ; CHECK: [[COPY3:%[0-9]+]]:_(p0) = COPY $x3
+    ; CHECK: [[BUILD_VECTOR:%[0-9]+]]:_(<4 x p0>) = G_BUILD_VECTOR [[COPY]](p0), [[COPY1]](p0), [[COPY2]](p0), [[COPY3]](p0)
+    ; CHECK: RET_ReallyLR implicit [[BUILD_VECTOR]](<4 x p0>)
+    %0:_(p0) = COPY $x0
+    %1:_(p0) = COPY $x1
+    %2:_(p0) = COPY $x2
+    %3:_(p0) = COPY $x3
+    %4:_(<2 x p0>) = G_BUILD_VECTOR %0(p0), %1
+    %5:_(<2 x p0>) = G_BUILD_VECTOR %2(p0), %3
+    %6:_(<4 x p0>) = G_CONCAT_VECTORS %4(<2 x p0>), %5
+    RET_ReallyLR implicit %6
+...
+# Check that we canonicalize concat_vectors(undef, undef) into undef.
+---
+name: concat_to_undef
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    ; CHECK-LABEL: name: concat_to_undef
+    ; CHECK: [[DEF:%[0-9]+]]:_(<4 x s64>) = G_IMPLICIT_DEF
+    ; CHECK: RET_ReallyLR implicit [[DEF]](<4 x s64>)
+    %4:_(<2 x s64>) = G_IMPLICIT_DEF
+    %5:_(<2 x s64>) = G_IMPLICIT_DEF
+    %6:_(<4 x s64>) = G_CONCAT_VECTORS %4(<2 x s64>), %5
+    RET_ReallyLR implicit %6
+...
+# Check that when combining concat_vectors(build_vector, undef) into
+# build_vector, we correctly break the undef vector into a sequence
+# of undef scalar.
+---
+name: concat_to_build_vector_with_undef
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $x0, $x1
+
+    ; CHECK-LABEL: name: concat_to_build_vector_with_undef
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY $x0
+    ; CHECK: [[COPY1:%[0-9]+]]:_(s64) = COPY $x1
+    ; CHECK: [[DEF:%[0-9]+]]:_(s64) = G_IMPLICIT_DEF
+    ; CHECK: [[BUILD_VECTOR:%[0-9]+]]:_(<4 x s64>) = G_BUILD_VECTOR [[COPY]](s64), [[COPY1]](s64), [[DEF]](s64), [[DEF]](s64)
+    ; CHECK: RET_ReallyLR implicit [[BUILD_VECTOR]](<4 x s64>)
+    %0:_(s64) = COPY $x0
+    %1:_(s64) = COPY $x1
+    %4:_(<2 x s64>) = G_BUILD_VECTOR %0(s64), %1
+    %5:_(<2 x s64>) = G_IMPLICIT_DEF
+    %6:_(<4 x s64>) = G_CONCAT_VECTORS %4(<2 x s64>), %5
+    RET_ReallyLR implicit %6
+...
+# Same as concat_to_build_vector_with_undef but with pointer types.
+---
+name: concat_to_build_vector_with_undef_ptr
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $x0, $x1
+
+    ; CHECK-LABEL: name: concat_to_build_vector_with_undef_ptr
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
+    ; CHECK: [[COPY1:%[0-9]+]]:_(p0) = COPY $x1
+    ; CHECK: [[DEF:%[0-9]+]]:_(p0) = G_IMPLICIT_DEF
+    ; CHECK: [[BUILD_VECTOR:%[0-9]+]]:_(<4 x p0>) = G_BUILD_VECTOR [[COPY]](p0), [[COPY1]](p0), [[DEF]](p0), [[DEF]](p0)
+    ; CHECK: RET_ReallyLR implicit [[BUILD_VECTOR]](<4 x p0>)
+    %0:_(p0) = COPY $x0
+    %1:_(p0) = COPY $x1
+    %4:_(<2 x p0>) = G_BUILD_VECTOR %0(p0), %1
+    %5:_(<2 x p0>) = G_IMPLICIT_DEF
+    %6:_(<4 x p0>) = G_CONCAT_VECTORS %4(<2 x p0>), %5
+    RET_ReallyLR implicit %6
+...
+# Check that we keep a concat_vectors as soon as one of the operand is
+# not undef or build_vector. I.e., we cannot flatten the concat_vectors.
+---
+name: concat_to_build_vector_negative_test
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $x0, $x1
+
+    ; CHECK-LABEL: name: concat_to_build_vector_negative_test
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY $x0
+    ; CHECK: [[DEF:%[0-9]+]]:_(<2 x s64>) = G_IMPLICIT_DEF
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK: [[IVEC:%[0-9]+]]:_(<2 x s64>) = G_INSERT_VECTOR_ELT [[DEF]], [[COPY]](s64), [[C]](s32)
+    ; CHECK: [[DEF1:%[0-9]+]]:_(<2 x s64>) = G_IMPLICIT_DEF
+    ; CHECK: [[CONCAT_VECTORS:%[0-9]+]]:_(<4 x s64>) = G_CONCAT_VECTORS [[IVEC]](<2 x s64>), [[DEF1]](<2 x s64>)
+    ; CHECK: RET_ReallyLR implicit [[CONCAT_VECTORS]](<4 x s64>)
+    %0:_(s64) = COPY $x0
+    %1:_(s64) = COPY $x1
+    %2:_(<2 x s64>) = G_IMPLICIT_DEF
+    %3:_(s32) = G_CONSTANT i32 1
+    %4:_(<2 x s64>) = G_INSERT_VECTOR_ELT %2, %0(s64), %3(s32)
+    %5:_(<2 x s64>) = G_IMPLICIT_DEF
+    %6:_(<4 x s64>) = G_CONCAT_VECTORS %4(<2 x s64>), %5
+    RET_ReallyLR implicit %6
+...