From 5b7e112954099374d6d07f676a9a57779c7a32f1 Mon Sep 17 00:00:00 2001
From: Sai Abhinay Anubola <saiabhinay.anubola@amd.com>
Date: Tue, 17 Sep 2024 15:21:27 +0530
Subject: [PATCH] Combine VInsert sequence in PreLegalizerCombiner

---
 .../Target/AIE/AIE2PreLegalizerCombiner.cpp   | 62 +++++++++++++++++++
 1 file changed, 62 insertions(+)

diff --git a/llvm/lib/Target/AIE/AIE2PreLegalizerCombiner.cpp b/llvm/lib/Target/AIE/AIE2PreLegalizerCombiner.cpp
index 0b3279e9de96..08cb36211136 100644
--- a/llvm/lib/Target/AIE/AIE2PreLegalizerCombiner.cpp
+++ b/llvm/lib/Target/AIE/AIE2PreLegalizerCombiner.cpp
@@ -71,6 +71,8 @@ class AIE2PreLegalizerCombinerImpl : public Combiner {
 
   bool tryToCombineSetExtract(MachineInstr &MI) const;
 
+  bool tryToCombineVectorInserts(MachineInstr &MI, int SclSrcBits) const;
+
   bool tryToCombineIntrinsic(MachineInstr &MI) const;
 
 private:
@@ -157,6 +159,63 @@ bool AIE2PreLegalizerCombinerImpl::tryToCombineSetExtract(
   return true;
 }
 
+/// Look for VINSERT sequence that can be rewritten as G_BUILD_VECTOR
+bool AIE2PreLegalizerCombinerImpl::tryToCombineVectorInserts(
+    MachineInstr &MI, int SclSrcBits) const {
+  const Register DstReg = MI.getOperand(0).getReg();
+  int DstRegLen = MRI.getType(DstReg).getElementCount().getKnownMinValue();
+  SmallVector<Register, 16> Regs;
+  std::map<int, Register> RegMap;
+  MachineInstr *CurMI = &MI;
+  int NumInserts = DstRegLen;
+
+  auto IsSet = [](const MachineInstr *MI) {
+    return MI->getOpcode() == AIE2::G_INTRINSIC &&
+           cast<GIntrinsic>(*MI).getIntrinsicID() ==
+               Intrinsic::aie2_set_I512_I128;
+  };
+
+  auto IsVInsert = [](const MachineInstr *MI) {
+    return MI->getOpcode() == AIE2::G_INTRINSIC &&
+           cast<GIntrinsic>(*MI).getIntrinsicID() ==
+               Intrinsic::aie2_vinsert8_I512;
+  };
+
+  while (NumInserts-- && IsVInsert(CurMI)) {
+    // In this case of G_INTRINSIC operand 1 is target intrinsic
+    const Register SrcReg = CurMI->getOperand(2).getReg();
+    const Register IdxReg = CurMI->getOperand(3).getReg();
+    const Register SclSrcReg = CurMI->getOperand(4).getReg();
+
+    // Collecting registers and their indices
+    auto Cst = getIConstantVRegValWithLookThrough(IdxReg, MRI);
+    if (!Cst ||
+        !RegMap.try_emplace(Cst->Value.getSExtValue(), SclSrcReg).second)
+      return false;
+
+    MachineInstr *SrcMI = MRI.getUniqueVRegDef(SrcReg);
+    assert(SrcMI && "Expected SSA.");
+    if (IsSet(SrcMI) && !tryToCombineSetExtract(*SrcMI))
+      break;
+    CurMI = getDefIgnoringCopies(SrcReg, MRI);
+  }
+
+  MachineIRBuilder MIRBuilder(MI);
+  // Collect registers in order for G_BUILD_VECTOR
+  for (int i = 0; i < DstRegLen; i++) {
+    auto It = RegMap.find(i);
+    if (It == RegMap.end())
+      return false;
+    Regs.push_back(It->second);
+  }
+  Register Dst128BitReg = MRI.createGenericVirtualRegister(
+      LLT::fixed_vector(DstRegLen, SclSrcBits));
+  MIRBuilder.buildBuildVectorTrunc(Dst128BitReg, Regs);
+  MIRBuilder.buildCopy(DstReg, Dst128BitReg);
+  MI.eraseFromParent();
+  return true;
+}
+
 bool AIE2PreLegalizerCombinerImpl::tryToCombineIntrinsic(
     MachineInstr &MI) const {
 
@@ -167,6 +226,9 @@ bool AIE2PreLegalizerCombinerImpl::tryToCombineIntrinsic(
   case Intrinsic::aie2_set_I512_I128: {
     return tryToCombineSetExtract(MI);
   }
+  case Intrinsic::aie2_vinsert8_I512: {
+    return tryToCombineVectorInserts(MI, 8);
+  }
   default:
     break;
   }