[AIE2] Replace 8x8->8x8 tranpose shuffle vector with vshuffle

Author: ValentijnvdBeek

Diff for: llvm/lib/Target/AIE/AIE2PostLegalizerCustomCombiner.cpp

@@ -21,6 +21,8 @@
 #include "llvm/CodeGen/GlobalISel/CombinerInfo.h"
 #include "llvm/CodeGen/GlobalISel/GIMatchTableExecutorImpl.h"
 #include "llvm/CodeGen/GlobalISel/GISelKnownBits.h"
+#include "llvm/CodeGen/GlobalISel/LegalizerHelper.h"
+#include "llvm/CodeGen/GlobalISel/LegalizerInfo.h"
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/InitializePasses.h"
 

Diff for: llvm/lib/Target/AIE/AIE2PreLegalizerCombiner.cpp

@@ -178,11 +178,51 @@ CombinerHelper::GeneratorType sectionGenerator(const int32_t From,
 
 bool AIE2PreLegalizerCombinerImpl::tryCombineShuffleVector(
     MachineInstr &MI) const {
+  const Register DstReg = MI.getOperand(0).getReg();
+  const LLT DstTy = MRI.getType(DstReg);
+  const LLT SrcTy = MRI.getType(MI.getOperand(1).getReg());
+  const unsigned DstNumElts = DstTy.isVector() ? DstTy.getNumElements() : 1;
+  const unsigned SrcNumElts = SrcTy.isVector() ? SrcTy.getNumElements() : 1;
+  MachineIRBuilder MIB(MI);
+  MachineRegisterInfo &MRI = *MIB.getMRI();
+
   if (Helper.tryCombineShuffleVector(MI))
     return true;
 
+  const LLT V64S8 = LLT::fixed_vector(64, 8);
+  CombinerHelper::GeneratorType FourPartitions =
+      sectionGenerator(0, DstNumElts, 4, 1);
+  if (Helper.matchCombineShuffleVector(MI, FourPartitions, DstNumElts)) {
+    if (DstTy != V64S8)
+      return false;
+
+    const Register Src1 = MI.getOperand(1).getReg();
+    const Register Src2 = MI.getOperand(2).getReg();
+    const Register ShuffleModeReg =
+        MRI.createGenericVirtualRegister(LLT::scalar(32));
+
+    // This combiner only cares about the lower bits, so we can pad the
+    // vector to cover the case where two separate vectors are shuffled.
+    // together
+    MIB.buildConstant(ShuffleModeReg, 35);
+
+    if (SrcTy == V64S8) {
+      MIB.buildInstr(AIE2::G_AIE_VSHUFFLE, {DstReg},
+                     {Src1, Src2, ShuffleModeReg});
+    } else {
+      // We reuse the same register since we ignore the high part of the vector
+      const Register TmpRegister = MRI.createGenericVirtualRegister(V64S8);
+      MIB.buildConcatVectors(TmpRegister, {Src1, Src2});
+      MIB.buildInstr(AIE2::G_AIE_VSHUFFLE, {DstReg},
+                     {TmpRegister, TmpRegister, ShuffleModeReg});
+    }
+
+    MI.eraseFromParent();
+    return true;
+  }
   return false;
 }
+
 bool AIE2PreLegalizerCombinerImpl::tryCombineAll(MachineInstr &MI) const {
   if (tryCombineAllImpl(MI))
     return true;

Diff for: llvm/lib/Target/AIE/AIEInstrGISel.td

@@ -96,6 +96,12 @@ def G_AIE_BROADCAST_VECTOR : AIEGenericInstruction {
   let hasSideEffects = false;
 }
 
+def G_AIE_VSHUFFLE : AIEGenericInstruction {
+  let OutOperandList = (outs type0:$dst);
+  let InOperandList = (ins type0:$src1, type0:$src2, type1:$mode);
+  let hasSideEffects = false;
+}
+
 // Create a larger vector by padding undefined values in the high bits
 def G_AIE_PAD_VECTOR_UNDEF : AIEGenericInstruction {
   let OutOperandList = (outs type0:$dst);