[AIE2P] Extend allocation filter to include only Target MIs

Now we filter by register class and usage. Basically, we exclude here
instructions like copies and non-2D/3D ones.

Co-Authored-By: Krishnam Tibrewala <[email protected]>

Author: andcarminati

llvm/lib/Target/AIE/AIEBaseTargetMachine.cpp

@@ -93,6 +93,11 @@ cl::opt<bool>
     EnableStagedRA("aie-staged-ra", cl::Hidden, cl::init(true),
                    cl::desc("Enable multi-stage register allocation"));
 
+cl::opt<bool> EnableFineGrainedStagedRA(
+    "aie-staged-ra-fine-grained-alloc", cl::Hidden, cl::init(true),
+    cl::desc("Enable multi-stage register allocation with fine-grained "
+             "selection of live intervals"));
+
 cl::opt<bool>
     EnableWAWRegRewrite("aie-wawreg-rewrite",
                         cl::desc("Enable the WAW Register Renaming in loops"),

llvm/lib/Target/AIE/aie2p/AIE2PTargetMachine.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 2024 Advanced Micro Devices, Inc. or its affiliates
+// (c) Copyright 2024-2025 Advanced Micro Devices, Inc. or its affiliates
 //
 //===----------------------------------------------------------------------===//
 //
@@ -15,6 +15,7 @@
 #include "AIE2PTargetMachine.h"
 #include "AIE2PTargetTransformInfo.h"
 #include "llvm/CodeGen/TargetLoweringObjectFileImpl.h"
+#include "llvm/CodeGen/TargetRegisterInfo.h"
 
 using namespace llvm;
 extern cl::opt<bool> EnableStagedRA;
@@ -25,6 +26,7 @@ extern cl::opt<bool> EnableAddressChaining;
 extern cl::opt<bool> EnableGlobalPtrModOptimizer;
 extern cl::opt<bool> EnableWAWRegRewrite;
 extern cl::opt<bool> EnableAIEIfConversion;
+extern cl::opt<bool> EnableFineGrainedStagedRA;
 
 void AIE2PTargetMachine::anchor() {}
 
@@ -70,17 +72,44 @@ void AIE2PPassConfig::addPreRegBankSelect() {
   }
 }
 
+static bool isRegUsedBy2DOr3DInstruction(const MachineRegisterInfo &MRI,
+                                         const Register &R) {
+
+  return std::any_of(
+      MRI.use_nodbg_instructions(R).begin(),
+      MRI.use_nodbg_instructions(R).end(), [&](const MachineInstr &MI) {
+        auto &TII = *static_cast<const AIEBaseInstrInfo *>(
+            MI.getMF()->getSubtarget().getInstrInfo());
+
+        // We should recognize both cases, with and without splitting. A 2D/3D
+        // instruction will always be split os splittable.
+        return TII.getOpcodeWithTupleOperands(MI.getOpcode()).has_value() ||
+               TII.getOpcodeWithAtomicOperands(MI.getOpcode()).has_value();
+      });
+}
+
 static bool onlyAllocate3DRegisters(const TargetRegisterInfo &TRI,
                                     const MachineRegisterInfo &MRI,
                                     const Register &R) {
-  return AIE2P::eDSRegClass.hasSubClassEq(MRI.getRegClass(R));
+
+  const TargetRegisterClass *RegClass = MRI.getRegClass(R);
+  if (!AIE2P::eDSRegClass.hasSubClassEq(RegClass))
+    return false;
+  return EnableFineGrainedStagedRA ? isRegUsedBy2DOr3DInstruction(MRI, R)
+                                   : true;
 }
+
 static bool onlyAllocate3D2DRegisters(const TargetRegisterInfo &TRI,
                                       const MachineRegisterInfo &MRI,
                                       const Register &R) {
-  return AIE2P::eDSRegClass.hasSubClassEq(MRI.getRegClass(R)) ||
-         AIE2P::eDRegClass.hasSubClassEq(MRI.getRegClass(R));
+  const TargetRegisterClass *RegClass = MRI.getRegClass(R);
+  if (!AIE2P::eDSRegClass.hasSubClassEq(RegClass) &&
+      !AIE2P::eDRegClass.hasSubClassEq(RegClass))
+    return false;
+  return EnableFineGrainedStagedRA ? isRegUsedBy2DOr3DInstruction(MRI, R)
+                                   : true;
 }
+
 static bool onlyAllocateMRegisters(const TargetRegisterInfo &TRI,
                                    const MachineRegisterInfo &MRI,
                                   const Register &R) {