[AIEX] Concat/Unmerge PHI Combiner

llvm/lib/Target/AIE/AIECombine.td

@@ -4,7 +4,7 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
-// (c) Copyright 2023-2024 Advanced Micro Devices, Inc. or its affiliates
+// (c) Copyright 2023-2025 Advanced Micro Devices, Inc. or its affiliates
 //
 //===----------------------------------------------------------------------===//
 include "llvm/Target/GlobalISel/Combine.td"
@@ -282,6 +282,14 @@ def combine_global_load_store_increment : GICombineRule <
           [{ return matchGlobalPtrModOptimizer(*${root}, MRI, Helper, B.getTII(), GlobalCombiners); }]),
   (apply [{ applyLdStInc(*${root}, MRI, Helper, B, Observer, GlobalCombiners); }] )>;
 
+def concat_unmerge_matchdata : GIDefMatchData<"AIEConcatUnmergeCombineMatchData">;
+def combine_concat_unmerge_phis : GICombineRule <
+  (defs root:$root, concat_unmerge_matchdata:$matchinfo),
+  (match (wip_match_opcode G_CONCAT_VECTORS):$root,
+  [{return matchConcatUnmergePhis(*${root}, MRI, Helper, ${matchinfo});}]),
+  (apply [{ applyConcatUnmergePhis(*${root}, MRI, B,${matchinfo}, Observer);}])
+>;
+
 
 
 def AIE2PostLegalizerCustomCombiner
@@ -294,7 +302,8 @@ def AIE2PostLegalizerCustomCombiner
                                                           combine_add_vector_elt_undef,
                                                           combine_extract_concat,
                                                           combine_unmerge_concat,
-                                                          combine_upd_to_concat 
+                                                          combine_upd_to_concat, 
+                                                          combine_concat_unmerge_phis
                                                           ]> {
 }
 
@@ -304,5 +313,7 @@ def AIE2PPostLegalizerCustomCombiner
                                                             ptr_add_immed_chain, 
                                                             combine_offset_load_store_ptradd,
                                                             combine_offset_load_store_share_ptradd,
-                                                            combine_add_vector_elt_undef ]> {
+                                                            combine_add_vector_elt_undef,
+                                                            combine_concat_unmerge_phis
+                                                             ]> {
 }

llvm/lib/Target/AIE/AIECombinerHelper.cpp

@@ -4,7 +4,7 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
-// (c) Copyright 2023-2024 Advanced Micro Devices, Inc. or its affiliates
+// (c) Copyright 2023-2025 Advanced Micro Devices, Inc. or its affiliates
 //
 //===----------------------------------------------------------------------===//
 
@@ -563,6 +563,183 @@ static MachineInstr *findPostIncMatch(MachineInstr &MemI,
   return nullptr;
 }
 
+/// Checking the following Concat-Unmerge-PHI pattern:
+/// bb.0
+/// concatIn0 = ...
+/// concatIn1 = ...
+/// bb.1
+/// 1 = phi 7, bb.1, concatIn0, bb.0
+/// 2 = phi 10, bb.1 concatIn1, bb.0
+/// 3 = G_CONCAT 1, 2
+/// ....
+/// 6 = ...
+/// 7, 8 = G_UNMERGE 6
+/// 9,10 = G_UNMERGE 6
+/// \return unmerge source (6) if the pattern is valid and the sub-vector index
+/// used [ 0 (7, 9) or 1 (8, 10)]. Start the pattern checking from \p ConcatI .
+/// Follow \p UseIdx (0 or 1) of G_CONCAT to identify the pattern. Collect all
+/// relevant Results in \p MatchData .
+static std::pair<MachineInstr *, unsigned>
+findUnmergeOrigin(MachineInstr &ConcatI, unsigned UseIdx,
+                  MachineRegisterInfo &MRI, CombinerHelper &Helper,
+                  AIEConcatUnmergeCombineMatchData &MatchData) {
+  auto *FirstMI = &*ConcatI.getParent()->begin();
+
+  auto FindPhiUnMergeComponents = [UseIdx, &FirstMI, &Helper, &MRI,
+                                   &MatchData](MachineInstr *PhiInst) {
+    // Find Unmerge and Concat Components of the PHI node
+    MachineInstr *UnMergeI = nullptr;
+    MachineOperand *UnMergeDefMO = nullptr;
+    for (unsigned I = 1; I < PhiInst->getNumOperands(); I += 2) {
+      MachineOperand *IncomingMO = &PhiInst->getOperand(I);
+      MachineOperand &MBB = PhiInst->getOperand(I + 1);
+
+      // Check if the basic block operand matches the given MBB
+      if (MBB.getMBB()->empty() ||
+          !Helper.dominates(*FirstMI, *MBB.getMBB()->begin())) {
+        // If MBB is empty, it can only be a fallthrough MBB, therefore, it
+        // dominates ConcatI's MBB
+
+        if (MatchData.ConcatMBB && MatchData.ConcatMBB != MBB.getMBB()) {
+          LLVM_DEBUG(dbgs()
+                     << "Both Parts of the Incoming Concat Values do not "
+                        "originate in the same MBB.");
+          return std::pair<MachineInstr *, MachineOperand *>{nullptr, nullptr};
+        }
+
+        // Concat Components
+        MatchData.ConcatMBB = MBB.getMBB();
+        if (UseIdx == 0)
+          MatchData.ConcatSub0 = IncomingMO;
+        else
+          MatchData.ConcatSub1 = IncomingMO;
+        continue;
+      }
+
+      // Unmerge Components
+      UnMergeDefMO = IncomingMO;
+      UnMergeI = MRI.getVRegDef(IncomingMO->getReg());
+      MatchData.UnmergeMBB = MBB.getMBB();
+    }
+    return std::pair{UnMergeI, UnMergeDefMO};
+  };
+
+  // Find PHI node
+  const auto OperandIdx = UseIdx + 1;
+  assert(ConcatI.getOperand(OperandIdx).isReg());
+  const Register Reg = ConcatI.getOperand(OperandIdx).getReg();
+  auto *PhiInst = MRI.getVRegDef(Reg);
+  if (!PhiInst || !PhiInst->isPHI())
+    return {nullptr, /*UnmergeSubVecDefIdx=*/0};
+
+  auto [UnMergeInst, UnMergeDefMO] = FindPhiUnMergeComponents(PhiInst);
+
+  if (!UnMergeInst ||
+      UnMergeInst->getOpcode() != TargetOpcode::G_UNMERGE_VALUES ||
+      !MatchData.ConcatMBB) {
+    // Could not find Unmerge or no non-dominating-MBB could be determined.
+    return {nullptr, /*UnmergeSubVecDefIdx=*/0};
+  }
+
+  MachineInstr *UnMergeSrc =
+      MRI.getVRegDef(UnMergeInst->getOperand(2).getReg());
+  if (!UnMergeSrc) {
+    // could not find source of UnMerge
+    return {nullptr, /*UnmergeSubVecDefIdx=*/0};
+  }
+
+  MatchData.UnmergeSourceOp = &UnMergeInst->getOperand(2);
+
+  auto GetUnmergeDefSubVectorIndex = [](const MachineInstr *UnMergeInst,
+                                        const auto UnmergeMOReg) {
+    unsigned UnMergeDefIdx = 0;
+    for (auto Def : UnMergeInst->defs()) {
+      if (Def.getReg() == UnmergeMOReg) {
+        return UnMergeDefIdx;
+      }
+      UnMergeDefIdx++;
+    }
+
+    llvm_unreachable("Lost PHIs MO in UnMerge");
+    return (unsigned)-1;
+  };
+
+  const auto SubVecDefIdx =
+      GetUnmergeDefSubVectorIndex(UnMergeInst, UnMergeDefMO->getReg());
+  return {UnMergeSrc, /*UnmergeSubVecDefIdx=*/SubVecDefIdx};
+}
+
+bool llvm::matchConcatUnmergePhis(MachineInstr &ConcatI,
+                                  MachineRegisterInfo &MRI,
+                                  CombinerHelper &Helper,
+                                  AIEConcatUnmergeCombineMatchData &MatchInfo) {
+  LLVM_DEBUG(dbgs() << "MF: " << ConcatI.getParent()->getParent()->getName()
+                    << "\n");
+
+  auto UnMergeOrigin0 =
+      findUnmergeOrigin(ConcatI, /*UseIdx=*/0, MRI, Helper, MatchInfo);
+  if (!UnMergeOrigin0.first)
+    return false;
+  auto UnMergeOrigin1 =
+      findUnmergeOrigin(ConcatI, /*UseIdx=*/1, MRI, Helper, MatchInfo);
+  if (!UnMergeOrigin1.first)
+    return false;
+
+  // If Unmerge Origins diverge, skip
+  if (UnMergeOrigin0.first != UnMergeOrigin1.first)
+    return false;
+
+  // If Concat behaves like a vector broadcast, i.e. Concat uses the same
+  // Sub-vector in both source operands, skip
+  if (UnMergeOrigin0.second == UnMergeOrigin1.second)
+    return false;
+
+  assert(MatchInfo.ConcatMBB);
+  assert(MatchInfo.ConcatSub0);
+  assert(MatchInfo.ConcatSub1);
+  assert(MatchInfo.UnmergeMBB);
+  assert(MatchInfo.UnmergeSourceOp);
+  return true;
+}
+
+void llvm::applyConcatUnmergePhis(MachineInstr &ConcatI,
+                                  MachineRegisterInfo &MRI, MachineIRBuilder &B,
+                                  AIEConcatUnmergeCombineMatchData &MatchInfo,
+                                  GISelChangeObserver &Observer) {
+  /// Set Insertion Point
+  if (MatchInfo.ConcatMBB->empty())
+    B.setMBB(*MatchInfo.ConcatMBB);
+  else
+    B.setInstr(MatchInfo.ConcatMBB->instr_back());
+
+  // Create new Concat Instruction
+  auto NewConcat = B.buildInstr(TargetOpcode::G_CONCAT_VECTORS);
+  const auto TargetVecType = MRI.getType(MatchInfo.UnmergeSourceOp->getReg());
+  NewConcat.addDef(MRI.createGenericVirtualRegister(TargetVecType));
+
+  NewConcat->addOperand(*MatchInfo.ConcatSub0);
+  NewConcat->addOperand(*MatchInfo.ConcatSub1);
+  LLVM_DEBUG(dbgs() << "In bb." << MatchInfo.ConcatMBB->getNumber()
+                    << " created " << *NewConcat.getInstr());
+
+  // set Insertion Point to top of phi-MBB
+  B.setInsertPt(*ConcatI.getParent(), ConcatI.getParent()->begin());
+
+  // Create new PHI node
+  auto NewPHI = B.buildInstr(TargetOpcode::G_PHI);
+  NewPHI.addDef(ConcatI.getOperand(0).getReg());
+
+  // add first PHI operand (newly create G_CONCAT)
+  NewPHI.addUse(NewConcat.getInstr()->getOperand(0).getReg());
+  NewPHI->addOperand(MachineOperand::CreateMBB(MatchInfo.ConcatMBB));
+
+  // Add second PHI operand (unmerge Components)
+  NewPHI.addUse(MatchInfo.UnmergeSourceOp->getReg());
+  NewPHI->addOperand(MachineOperand::CreateMBB(MatchInfo.UnmergeMBB));
+  LLVM_DEBUG(dbgs() << "Created New Instruction " << *NewPHI.getInstr());
+  ConcatI.removeFromParent();
+}
+
 bool llvm::matchGlobalPtrModOptimizer(MachineInstr &MemI,
                                       MachineRegisterInfo &MRI,
                                       CombinerHelper &Helper,

llvm/lib/Target/AIE/AIECombinerHelper.h

@@ -4,7 +4,7 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
-// (c) Copyright 2023-2024 Advanced Micro Devices, Inc. or its affiliates
+// (c) Copyright 2023-2025 Advanced Micro Devices, Inc. or its affiliates
 //
 //===----------------------------------------------------------------------===//
 
@@ -30,6 +30,23 @@ struct FrequentIndexResult {
   unsigned NonMatchingCount;
 };
 
+struct AIEConcatUnmergeCombineMatchData {
+  // New Concat Instruction
+  // MBB where the G_CONCAT of ConcatSub0 and ConcatSub1
+  // should be created.
+  MachineBasicBlock *ConcatMBB = nullptr;
+  // Phi Src Operand 0: Concat of ConcatSub0 and ConcatSub1
+  MachineOperand *ConcatSub0 = nullptr;
+  // Phi Src Operand 0: Concat of ConcatSub0 and ConcatSub1
+  MachineOperand *ConcatSub1 = nullptr;
+
+  // PHI components of the backedge
+  // Backedge MachineBasicBlock where Unmerge Instructions will be dropped
+  MachineBasicBlock *UnmergeMBB = nullptr;
+  // Instruction for the PHI Source operand 2
+  MachineOperand *UnmergeSourceOp = nullptr;
+};
+
 /// The mask is represented by a sawtooth function F with Period, Height and
 /// Amplitude, i.e., F(idx + Period) = F(idx) = Height + idx * Amplitude, where
 /// idx >= 0.
@@ -77,6 +94,41 @@ struct AIESingleDiffLaneBuildVectorMatchData {
 
 void foundPattern(MachineInstr &MemI);
 
+/// This is the matching function for Concat-Unmerge-PHI pattern.
+/// Convert the following:
+/// bb.0
+/// concatIn0 = ...
+/// concatIn1 = ...
+/// bb.1
+/// 1 = phi 7, bb.1, concatIn0, bb.0
+/// 2 = phi 10, bb.1 concatIn1, bb.0
+/// 3 = G_CONCAT 1, 2
+/// ....
+/// 6 = ...
+/// 7, 8 = G_UNMERGE 6
+/// 9,10 = G_UNMERGE 6
+///
+/// Into:
+/// bb.0
+/// concatIn0 = ...
+/// concatIn1 = ...
+/// 11 = G_CONCAT concatIn0, concatIn1
+/// bb.1
+/// 3 = phi 6, bb.1, 11, bb.0
+/// ...
+/// 6 =
+///
+/// \p ConcatI is the starting Point for this pattern.
+/// \return true if the pattern is found.
+bool matchConcatUnmergePhis(MachineInstr &ConcatI, MachineRegisterInfo &MRI,
+                            CombinerHelper &Helper,
+                            AIEConcatUnmergeCombineMatchData &MatchInfo);
+/// apply Concat-Unmerge-PHI Combiner
+void applyConcatUnmergePhis(MachineInstr &ConcatI, MachineRegisterInfo &MRI,
+                            MachineIRBuilder &B,
+                            AIEConcatUnmergeCombineMatchData &MatchInfo,
+                            GISelChangeObserver &Observer);
+
 bool matchGlobalPtrModOptimizer(MachineInstr &MemI, MachineRegisterInfo &MRI,
                                 CombinerHelper &Helper,
                                 const TargetInstrInfo &TII,