[AIE] Integrate LatencyAware and SWPAware

The enable flags are now integers, with 0 meaning false, 1 meaning true,
anything else meaning auto-select based on LoopClass.

This selection avoids running swpaware if latencyaware has run

Add some tuning based on LoopClass

Author: Martien de Jong

llvm/lib/Target/AIE/AIELoopClass.cpp

@@ -76,6 +76,8 @@ static const KernelFeatures Kernels[] = {
     {43, {{900, 300, 0, 2460, 0, 0, 3360}, {30, 30, 0, 30, 30}}},
     {44, {{900, 300, 0, 2520, 0, 0, 3360}, {30, 30, 0, 30, 30}}},
     {45, {{900, 300, 0, 2520, 0, 0, 3360}, {45, 45, 0, 45, 45}}},
+    {46, {{0, 0, 0, 0, 2160, 0, 120, 1080}, {0, 420, 420}}},
+    {47, {{0, 0, 0, 0, 360, 0, 240, 360}, {0, 60, 60}}},
 };
 
 std::vector<int> getLoopClassScores(const SlotStatistics &Stats) {

llvm/lib/Target/AIE/AIEWawRegRewriter.cpp

@@ -49,6 +49,14 @@ using namespace llvm;
 // This might be compatible with a future extension of the DEBUG rigging
 #define DEBUG_DETAIL(x) DEBUG_WITH_TYPE("aie-waw-reg-rewrite:2", x)
 
+enum class RewriteMode {
+  Basic,
+  Automatic,
+  LatencyAware,
+  SWPAware,
+  SWPAwareAutoBias,
+};
+
 static cl::opt<bool> AggressiveReAlloc(
     "aie-aggressive-realloc", cl::Hidden, cl::init(false),
     cl::desc("Aggressively de-allocate live-through registers to favor "
@@ -66,16 +74,21 @@ static cl::opt<unsigned>
                                 "considered for WAW rewriting"),
                        cl::init(3));
 
-static cl::opt<bool>
-    LatencyAware("aie-realloc-latencyaware", cl::Hidden, cl::init(true),
-                 cl::desc("Enable latency-aware allocation strategy"));
-
-static cl::opt<bool>
-    SWPAware("aie-realloc-swp-aware", cl::Hidden, cl::init(false),
-             cl::desc("Use assignment order based on interleaved swp stages"));
+static cl::opt<RewriteMode> RegRewriteMode(
+    "aie-reg-rewrite-mode", cl::Hidden, cl::init(RewriteMode::Automatic),
+    cl::desc("Set the rewriting mode"),
+    cl::values(clEnumValN(RewriteMode::Basic, "basic", "Basic"),
+               clEnumValN(RewriteMode::Automatic, "auto",
+                          "Automatic selection based on loop class"),
+               clEnumValN(RewriteMode::LatencyAware, "latencyaware",
+                          "Latency aware"),
+               clEnumValN(RewriteMode::SWPAware, "swpaware",
+                          "SWP aware with default bias"),
+               clEnumValN(RewriteMode::SWPAwareAutoBias, "swpaware-auto",
+                          "SWP aware with automatic bias")));
 
 static cl::opt<int> MinIIBias("aie-realloc-ii-bias", cl::Hidden, cl::init(0),
-                              cl::desc("MinII bias for swp-aware"));
+                              cl::desc("Set default MinII bias for swpaware"));
 
 namespace {
 
@@ -89,6 +102,44 @@ using OriginalAllocation =
 // Record success for a whole register class.
 using RegClassSuccess = std::map<const TargetRegisterClass *, bool>;
 
+RewriteMode selectMode(RewriteMode Mode, int LoopClass) {
+  if (Mode != RewriteMode::Automatic) {
+    return Mode;
+  }
+  switch (LoopClass) {
+  case 14:
+  case 29:
+  case 47:
+    return RewriteMode::SWPAwareAutoBias;
+  default:
+    return RewriteMode::LatencyAware;
+  }
+}
+
+bool runSWPAware(RewriteMode Mode, int LoopClass, int &Bias) {
+  Bias = MinIIBias;
+  switch (Mode) {
+  case RewriteMode::SWPAwareAutoBias:
+    break;
+  case RewriteMode::SWPAware:
+    return true;
+  default:
+    return false;
+  }
+
+  switch (LoopClass) {
+  case 14:
+    Bias = -1;
+    return true;
+  case 18:
+  case 29:
+    Bias = 1;
+    return true;
+  default:
+    return true;
+  }
+}
+
 ///
 /// This pass rewrites physical register assignments in critical parts of the
 /// code (like loops) to break WAW and WAR dependencies.
@@ -141,7 +192,8 @@ class AIEWawRegRewriter : public MachineFunctionPass {
       const std::map<const TargetRegisterClass *, bool> &RegClasses);
 
   /// Sort the candidates to mimic interleaving the pipeline stages
-  void sortSWPAware(OriginalAllocation &Candidates, MachineBasicBlock &MBB);
+  void sortSWPAware(OriginalAllocation &Candidates, MachineBasicBlock &MBB,
+                    const llvm::AIE::SlotStatistics &Statistics, int LoopClass);
 
   /// Pre-allocate all virtual registers in Candidates. The sole purpose of
   /// this is to prime the LRURegisters, so that the end of the loop is
@@ -383,21 +435,17 @@ RoundRobin AIEWawRegRewriter::computeLRURegisters(
   return LRURegisters;
 }
 
-void AIEWawRegRewriter::sortSWPAware(OriginalAllocation &Candidates,
-                                     MachineBasicBlock &MBB) {
-
+void AIEWawRegRewriter::sortSWPAware(
+    OriginalAllocation &Candidates, MachineBasicBlock &MBB,
+    const llvm::AIE::SlotStatistics &Statistics, int Bias) {
   // We estimate the length of the schedule based on latencies and the
   // minimum II based on slots. We then estimate the modulo cycle of each
   // instruction based on its depth and apply LRU in the order of the modulo
   // cycle.
   // Note that both the depth and the II are underestimations since we don't
   // account for them interfering. Hence the modulo cycle estimate won't be
   // too far off.
-  AIE::SlotStatistics Statistics = AIE::computeSlotStatistics(MBB, TII);
-  DEBUG_DETAIL(dbgs() << "Stats="; Statistics.dumpShort(); dbgs() << "\n");
-  DEBUG_DETAIL(dbgs() << "LoopClass=" << llvm::AIE::classifyLoop(Statistics)
-                      << "\n");
-  const int MinII = std::max(Statistics.getMinII() + MinIIBias, 1);
+  const int MinII = std::max(Statistics.getMinII() + Bias, 1);
 
   MachineSchedContext Context;
   Context.MF = MF;
@@ -456,7 +504,18 @@ bool AIEWawRegRewriter::renameMBBPhysRegs(const MachineBasicBlock *MBB) {
   IndexedMap<const MachineInstr *, VirtReg2IndexFunctor> LastVRegDef =
       getLastVRegDef(*MBB);
 
-  std::set<MCRegister> HighLatencyRegs = getHighOutputLatencyRegs(MBB);
+  auto &NonConstMBB = *(const_cast<MachineBasicBlock *>(MBB));
+  AIE::SlotStatistics Statistics = AIE::computeSlotStatistics(NonConstMBB, TII);
+  const int LoopClass = llvm::AIE::classifyLoop(Statistics);
+  LLVM_DEBUG(dbgs() << "Stats="; Statistics.dumpShort(); dbgs() << "\n");
+  LLVM_DEBUG(dbgs() << "LoopClass=" << LoopClass << "\n");
+
+  RewriteMode Mode = selectMode(RegRewriteMode, LoopClass);
+
+  std::set<MCRegister> HighLatencyRegs;
+  if (Mode == RewriteMode::LatencyAware) {
+    HighLatencyRegs = getHighOutputLatencyRegs(MBB);
+  }
 
   OriginalAllocation Candidates;
 
@@ -533,9 +592,9 @@ bool AIEWawRegRewriter::renameMBBPhysRegs(const MachineBasicBlock *MBB) {
     }
   }
 
-  if (SWPAware) {
-    auto &NCMBB = *(const_cast<MachineBasicBlock *>(MBB));
-    sortSWPAware(Candidates, NCMBB);
+  int Bias = MinIIBias;
+  if (runSWPAware(Mode, LoopClass, Bias)) {
+    sortSWPAware(Candidates, NonConstMBB, Statistics, Bias);
   }
 
   // Least-Recently-Used list of physical registers for assignments to VRegs.
@@ -730,10 +789,6 @@ AIEWawRegRewriter::getLastVRegDef(const MachineBasicBlock &MBB) const {
 
 std::set<MCRegister> AIEWawRegRewriter::getHighOutputLatencyRegs(
     const MachineBasicBlock *MBB) const {
-
-  if (!LatencyAware)
-    return {};
-
   auto *ItinData = MF->getSubtarget().getInstrItineraryData();
   std::set<MCRegister> HighLatRegisters;
   for (const MachineInstr &MI : *MBB) {

llvm/test/CodeGen/AIE/aie2/end-to-end/TanhTemplated-swp.ll

@@ -4,9 +4,11 @@
 ; 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
 
-; RUN: llc -O2 -mtriple=aie2 %s -o - | FileCheck %s
+; RUN: llc -O2 -mtriple=aie2 \
+; RUN:    --aie-reg-rewrite-mode=latencyaware \
+; RUN:    %s -o - | FileCheck %s
 
 ; The test is meant as a quick way to spot QoR regressions. In this test, the
 ; code can only be pipelined (Pre-SWP) because of the removal of some WAW

llvm/test/CodeGen/AIE/aie2p/end-to-end/add-att-broadcasting.ll

@@ -5,7 +5,7 @@
 ; SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 ;
 ; (c) Copyright 2025 Advanced Micro Devices, Inc. or its affiliates
-; RUN: llc -mtriple=aie2p --aie-realloc-latencyaware=true %s -o - | FileCheck %s
+; RUN: llc -mtriple=aie2p --aie-reg-rewrite-mode=latencyaware %s -o - | FileCheck %s
 
 ; Test postSWP capabilities related to AddAttributeBroadcasting
 

llvm/test/CodeGen/AIE/aie2p/end-to-end/conv2d-dw-bf16.ll

@@ -5,7 +5,7 @@
 ; SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 ;
 ; (c) Copyright 2025 Advanced Micro Devices, Inc. or its affiliates
-; RUN: llc -mtriple=aie2p --aie-realloc-latencyaware=true %s -o - | FileCheck %s
+; RUN: llc -mtriple=aie2p --aie-reg-rewrite-mode=latencyaware %s -o - | FileCheck %s
 
 ; Test postSWP capabilities related to conv2d_dw_bf16.
 

llvm/test/CodeGen/AIE/aie2p/end-to-end/gemm-bfp16.ll

@@ -5,7 +5,7 @@
 ; SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 ;
 ; (c) Copyright 2025 Advanced Micro Devices, Inc. or its affiliates
-; RUN: llc -mtriple=aie2p --aie-realloc-latencyaware=true %s -o - | FileCheck %s
+; RUN: llc -mtriple=aie2p --aie-reg-rewrite-mode=latencyaware %s -o - | FileCheck %s
 
 ; Test postSWP capabilities related to gemm_bfp16.
 

llvm/test/CodeGen/AIE/aie2p/ra/waw_reg_renaming.mir

@@ -7,11 +7,13 @@
 #
 # (c) Copyright 2025 Advanced Micro Devices, Inc. or its affiliates
 
-# Basic test for the WAW register renaming pass. Check AIE2 tests for more coverage.
-# Additional pass options are disabled because we are testing the basic mechanism.
+# Basic test for the WAW register renaming pass. Check AIE2 tests for more
+# coverage. Additional pass options are disabled because we are testing the
+# basic mechanism.
 
-# RUN: llc -mtriple=aie2p -verify-machineinstrs --aie-realloc-latencyaware=false --start-before=greedy \
-# RUN:   --stop-after=virtregrewriter %s -o - | FileCheck %s
+# RUN: llc -mtriple=aie2p -verify-machineinstrs --aie-reg-rewrite-mode=basic \
+# RUN:   --start-before=greedy --stop-after=virtregrewriter \
+# RUN:   %s -o - | FileCheck %s
 
 # Make sure VLD and VMAX define different X registers.
 ---

llvm/test/CodeGen/AIE/aie2p/ra/waw_reg_renaming_latencyaware.mir

@@ -13,8 +13,10 @@
 # try to not allocate it as def of any other instruction. In the example below, no
 # other instructions should output to the same registers used by the loads.
 
-# RUN: llc -mtriple=aie2p -verify-machineinstrs --aie-realloc-latencyaware=true\
-# RUN: --start-before=greedy --stop-after=virtregrewriter %s -o - | FileCheck %s
+# RUN: llc -mtriple=aie2p -verify-machineinstrs \
+# RUN:     --aie-reg-rewrite-mode=latencyaware \
+# RUN:     --start-before=greedy --stop-after=virtregrewriter \
+# RUN:     %s -o - | FileCheck %s
 
 ---
 name:            conv2d_dw_bf16