[Generic][AIE2] Combiner for shufflevectors that use build vector

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

@@ -27,6 +27,7 @@
 #include "llvm/CodeGenTypes/LowLevelType.h"
 #include "llvm/IR/InstrTypes.h"
 #include <functional>
+#include <optional>
 
 namespace llvm {
 
@@ -256,8 +257,21 @@ class CombinerHelper {
   /// concat_vectors.
   ///
   /// \pre MI.getOpcode() == G_SHUFFLE_VECTOR.
-  bool matchCombineShuffleVector(MachineInstr &MI,
-                                 SmallVectorImpl<Register> &Ops);
+  using GeneratorType = std::function<std::optional<int32_t>()>;
+
+  bool matchCombineShuffleVector(MachineInstr &MI, GeneratorType Generator,
+                                 const size_t TargetDstSize);
+
+  /// Create G_UNMERGE_VECTOR instructions until the source has reached a
+  /// target vector size.
+  ///
+  /// Requires that the destination fits evenly in the source register. It
+  /// allows you to pass which of the different destination sized slices
+  /// you require.
+  Register createUnmergeValue(MachineInstr &MI, const Register SrcReg,
+                              const Register DstReg, uint8_t DestinationIndex,
+                              const uint32_t Start, const uint32_t End);
+
   /// Replace \p MI with a concat_vectors with \p Ops.
   void applyCombineShuffleVector(MachineInstr &MI,
                                  const ArrayRef<Register> Ops);
@@ -341,6 +355,15 @@ class CombinerHelper {
   applyCombineUnmergeMergeToPlainValues(MachineInstr &MI,
                                         SmallVectorImpl<Register> &Operands);
 
+  /// Transform <ty, ...> G_SHUFFLE_VECTOR(G_MERGE ty X Y Z) -> G_MERGE ty X,Y,Z
+  bool
+  matchCombineShuffleVectorBuildVector(MachineInstr &MI,
+                                       SmallVectorImpl<Register> &Operands);
+
+  void
+  applyCombineShuffleVectorBuildVector(MachineInstr &MI,
+                                       SmallVectorImpl<Register> &Operands);
+
   /// Transform G_UNMERGE Constant -> Constant1, Constant2, ...
   bool matchCombineUnmergeConstant(MachineInstr &MI,
                                    SmallVectorImpl<APInt> &Csts);

llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h

@@ -4,6 +4,9 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
+// Modifications (c) Copyright 2024 Advanced Micro Devices, Inc. or its
+// affiliates
+//
 //===----------------------------------------------------------------------===//
 /// \file
 /// Declares convenience wrapper classes for interpreting MachineInstr instances
@@ -240,6 +243,22 @@ class GUnmerge : public GenericMachineInstr {
   }
 };
 
+/// Represents a G_SHUFFLE_VECTOR.
+class GShuffleVector : public GenericMachineInstr {
+public:
+  /// Returns the number of source registers.
+  unsigned getNumSources() const { return getNumOperands() - 2; }
+  /// Returns the I'th source register.
+  Register getSourceReg(unsigned I) const {
+    assert(I + 1 <= getNumSources());
+    return getReg(I + 1);
+  }
+
+  static bool classof(const MachineInstr *MI) {
+    return MI->getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR;
+  }
+};
+
 /// Represents G_BUILD_VECTOR, G_CONCAT_VECTORS or G_MERGE_VALUES.
 /// All these have the common property of generating a single value from
 /// multiple sources.

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

@@ -756,7 +756,7 @@ def fneg_fneg_fold: GICombineRule <
   (apply (GIReplaceReg $dst, $src))
 >;
 
-// Fold (unmerge(merge x, y, z)) -> z, y, z.
+// Fold (unmerge(merge x, y, z)) -> x, y, z.
 def unmerge_merge_matchinfo : GIDefMatchData<"SmallVector<Register, 8>">;
 def unmerge_merge : GICombineRule<
   (defs root:$d, unmerge_merge_matchinfo:$info),
@@ -765,6 +765,16 @@ def unmerge_merge : GICombineRule<
   (apply [{ Helper.applyCombineUnmergeMergeToPlainValues(*${d}, ${info}); }])
 >;
 
+// Fold (unmerge(merge x, y, z)) -> z, y, z.
+def shufflevector_merge_matchinfo : GIDefMatchData<"SmallVector<Register, 8>">;
+def shufflevector_merge : GICombineRule<
+  (defs root:$d, shufflevector_merge_matchinfo:$info),
+  (match (wip_match_opcode G_SHUFFLE_VECTOR): $d,
+  [{ return Helper.matchCombineShuffleVectorBuildVector(*${d}, ${info}); }]),
+  (apply [{ Helper.applyCombineShuffleVectorBuildVector(*${d}, ${info}); }])
+>;
+
+
 // Fold merge(unmerge).
 def merge_unmerge : GICombineRule<
   (defs root:$d, register_matchinfo:$matchinfo),