fixup! Directly run the shuffle vector analysis code

Author: ValentijnvdBeek

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

@@ -594,12 +594,9 @@ class CombinerHelper {
 
   /// Combine extracts of two different arrays into one build vector into a
   /// shuffle vector.
-  bool matchCombineExtractToShuffle(
-      MachineInstr &MI, SmallVectorImpl<std::pair<Register, int>> &MatchInfo,
-      std::pair<Register, Register> &VectorRegisters);
-  void applyCombineExtractToShuffle(
-      MachineInstr &MI, SmallVectorImpl<std::pair<Register, int>> &MatchInfo,
-      std::pair<Register, Register> &VectorRegisters);
+  bool
+  matchCombineExtractToShuffle(MachineInstr &MI, SmallVectorImpl<Register> &Ops,
+                               std::pair<Register, Register> &VectorRegisters);
 
   bool matchExtractAllEltsFromBuildVector(
       MachineInstr &MI,

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

@@ -905,15 +905,15 @@ def extract_vec_elt_build_vec : GICombineRule<
     [{ return Helper.matchExtractVecEltBuildVec(*${root}, ${matchinfo}); }]),
   (apply [{ Helper.applyExtractVecEltBuildVec(*${root}, ${matchinfo}); }])>;
 
-def extract_vector_register_to_id_mapping_matchinfo :
-  GIDefMatchData<"SmallVector<std::pair<Register, int>>">;
+def extract_vector_register_sources_matchinfo :
+  GIDefMatchData<"SmallVector<Register>">;
 def vector_reg_pair_matchinfo :
   GIDefMatchData<"std::pair<Register, Register>">;
 def extract_vector_element_build_vector_to_shuffle_vector : GICombineRule<
-  (defs root:$root, extract_vector_register_to_id_mapping_matchinfo:$matchinfo, vector_reg_pair_matchinfo:$regpair),
+  (defs root:$root, extract_vector_register_sources_matchinfo:$matchinfo, vector_reg_pair_matchinfo:$regpair),
   (match (wip_match_opcode G_BUILD_VECTOR):$root,
     [{ return Helper.matchCombineExtractToShuffle(*${root}, ${matchinfo}, ${regpair}); }]),
-  (apply [{ Helper.applyCombineExtractToShuffle(*${root}, ${matchinfo}, ${regpair}); }])>;
+  (apply [{ Helper.applyCombineShuffleVector(*${root}, ${matchinfo}); }])>;
 
 // Fold away full elt extracts from a build_vector.
 def extract_all_elts_from_build_vector_matchinfo :
@@ -926,8 +926,8 @@ def extract_all_elts_from_build_vector : GICombineRule<
 
 def extract_vec_elt_combines : GICombineGroup<[
   extract_vec_elt_build_vec,
-  extract_all_elts_from_build_vector,
-  extract_vector_element_build_vector_to_shuffle_vector]>;
+  extract_vector_element_build_vector_to_shuffle_vector,
+  extract_all_elts_from_build_vector]>;
 
 def funnel_shift_from_or_shift : GICombineRule<
   (defs root:$root, build_fn_matchinfo:$info),

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

@@ -4225,9 +4225,12 @@ void CombinerHelper::applyExtractVecEltBuildVec(MachineInstr &MI,
 }
 
 bool CombinerHelper::matchCombineExtractToShuffle(
-    MachineInstr &MI, SmallVectorImpl<std::pair<Register, int>> &VecIndexPair,
+    MachineInstr &MI, SmallVectorImpl<Register> &Ops,
     std::pair<Register, Register> &VectorRegisters) {
   const GBuildVector *Build = cast<GBuildVector>(&MI);
+  const unsigned SrcNumElts =
+      MRI.getType(MI.getOperand(0).getReg()).getNumElements();
+
   // This combine tries to find all the build vectors whose source elements
   // all originate from a G_EXTRACT_VECTOR_ELT from one or two donor vectors.
   // One example where this may happen is for AI chips where there are a lot
@@ -4245,24 +4248,29 @@ bool CombinerHelper::matchCombineExtractToShuffle(
   // replace with:
   //   %vector = G_SHUFFLE_VECTOR %donor1, %donor2, shufflemask(0, 1, 2, 3)
   SmallSetVector<Register, 2> RegisterVector;
+  SmallVector<int, 32> VectorMask;
   const unsigned NumElements = Build->getNumSources();
   for (unsigned Index = 0; Index < NumElements; Index++) {
     Register SrcReg = peekThroughBitcast(Build->getSourceReg(Index), MRI);
     auto *ExtractInstr = getOpcodeDef<GExtractVectorElement>(SrcReg, MRI);
     if (!ExtractInstr)
       return false;
 
+    RegisterVector.insert(ExtractInstr->getVectorReg());
+
     // For shufflemasks we need to know exactly what index to place each element
     // so if it this build vector doesn't use exclusively constants than we
     // can't replace with a shufflevector
     auto Cst = getIConstantVRegVal(ExtractInstr->getIndexReg(), MRI);
     if (!Cst)
       return false;
+
     unsigned Idx = Cst->getZExtValue();
+    if (ExtractInstr->getVectorReg() != RegisterVector.front()) {
+      Idx += SrcNumElts;
+    }
 
-    Register VectorReg = ExtractInstr->getVectorReg();
-    RegisterVector.insert(VectorReg);
-    VecIndexPair.emplace_back(std::make_pair(VectorReg, Idx));
+    VectorMask.emplace_back(Idx);
   }
 
   // Create a pair so that we don't need to look for them later. This code is
@@ -4273,44 +4281,17 @@ bool CombinerHelper::matchCombineExtractToShuffle(
       std::make_pair(RegisterVector.front(), RegisterVector.back());
 
   // We check that they're the same type before running. We can also grow the
-  // smaller one to the target size, but there isn't an elegant way to do that
-  // until we have a good lowering for G_EXTRACT_SUBVECTOR.
+  // smaller one tro the target size, but there isn't an elegant way to do that
+  // until we have a good lowerng for G_EXTRACT_SUBVECTOR.
+  // Apparently if they are the same, they don't necessary have the same type?
   if (MRI.getType(VectorRegisters.first) != MRI.getType(VectorRegisters.second))
     return false;
 
-  return RegisterVector.size() <= 2;
-}
-
-void CombinerHelper::applyCombineExtractToShuffle(
-    MachineInstr &MI, SmallVectorImpl<std::pair<Register, int>> &MatchInfo,
-    std::pair<Register, Register> &VectorRegisters) {
-  assert(MI.getOpcode() == TargetOpcode::G_BUILD_VECTOR);
-
-  const Register FirstRegister = VectorRegisters.first;
-  const LLT FirstRegisterType = MRI.getType(FirstRegister);
-  const unsigned VectorSize = FirstRegisterType.getNumElements();
-  SmallVector<int, 32> ShuffleMask;
-  for (auto &Pair : MatchInfo) {
-    const Register VectorReg = Pair.first;
-    int Idx = Pair.second;
-
-    if (VectorReg != VectorRegisters.first) {
-      Idx += VectorSize;
-    }
-    ShuffleMask.emplace_back(Idx);
-  }
-
-  // We could reuse the same vector register and shuffle them both together
-  // but it is nicer for later optimizations to explicitly make it undef.
-  const GBuildVector *BuildVector = cast<GBuildVector>(&MI);
-  Register SecondRegister = VectorRegisters.second;
-  if (FirstRegister == SecondRegister) {
-    SecondRegister = Builder.buildUndef(FirstRegisterType).getReg(0);
-  }
+  if (RegisterVector.size() > 2)
+    return false;
 
-  Builder.buildShuffleVector(BuildVector->getOperand(0), FirstRegister,
-                             SecondRegister, ShuffleMask);
-  MI.eraseFromParent();
+  return analysePatternVectorMask(MI, Ops, MI.getOperand(0).getReg(),
+                                 VectorRegisters, VectorMask);
 }
 
 bool CombinerHelper::matchExtractAllEltsFromBuildVector(