[llvm] [AArch64] Allow high-II schedules only when vectorised (PR #206079)

[William Huynh via llvm-commits]

https://github.com/saturn691 created https://github.com/llvm/llvm-project/pull/206079

MachinePipeliner has a cap of II=27 which is reasonable, but it rejects a subset of profitable loops. These are vectorised loops, that have not been interleaved by the Loop Vectorizer (i.e. loops written with intrinsics).

However, this is still a draft, as we need to uncap the II limit in code. Also this is codex code, so there may be a better way to do this.

Assisted-by: codex, reviewed and tested by me

>From 85c753625a4064a2e644f36b6350c17f129cc483 Mon Sep 17 00:00:00 2001
From: William Huynh <William.Huynh at arm.com>
Date: Fri, 26 Jun 2026 15:06:53 +0100
Subject: [PATCH] [AArch64] Allow high-II schedules only when vectorised

MachinePipeliner has a cap of II=27 which is reasonable, but it rejects
a subset of profitable loops. These are vectorised loops, that have not
been interleaved by the Loop Vectorizer (i.e. loops written with
intrinsics).

However, this is still a draft, as we need to uncap the II limit in code.
Also this is codex code, so there may be a better way to do this.

Assisted-by: codex, reviewed and tested by me
---
 llvm/lib/Target/AArch64/AArch64InstrInfo.cpp | 55 ++++++++++++++++++++
 1 file changed, 55 insertions(+)

diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index ca1a729dd1dd1..cea03d08a902a 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -36,6 +36,7 @@
 #include "llvm/CodeGen/MachineMemOperand.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/MachineOperand.h"
+#include "llvm/CodeGen/MachinePipeliner.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/RegisterScavenging.h"
 #include "llvm/CodeGen/StackMaps.h"
@@ -103,6 +104,11 @@ static cl::opt<unsigned> GatherOptSearchLimit(
     cl::desc("Restrict range of instructions to search for the "
              "machine-combiner gather pattern optimization"));
 
+static cl::opt<int> AArch64PipelinerMaxScalarII(
+    "aarch64-pipeliner-max-scalar-ii", cl::Hidden, cl::init(-1),
+    cl::desc("Reject scalar-only AArch64 MachinePipeliner schedules above this "
+             "initiation interval (-1 disables the check)"));
+
 AArch64InstrInfo::AArch64InstrInfo(const AArch64Subtarget &STI)
     : AArch64GenInstrInfo(STI, RI, AArch64::ADJCALLSTACKDOWN,
                           AArch64::ADJCALLSTACKUP, AArch64::CATCHRET),
@@ -11695,6 +11701,8 @@ class AArch64PipelinerLoopInfo : public TargetInstrInfo::PipelinerLoopInfo {
   /// The normalized condition used by createTripCountGreaterCondition()
   SmallVector<MachineOperand, 4> Cond;
 
+  bool containsFpOrVectorRegisters(SwingSchedulerDAG &SSD) const;
+
 public:
   AArch64PipelinerLoopInfo(MachineBasicBlock *LoopBB, MachineInstr *CondBranch,
                            MachineInstr *Comp, unsigned CompCounterOprNum,
@@ -11715,6 +11723,15 @@ class AArch64PipelinerLoopInfo : public TargetInstrInfo::PipelinerLoopInfo {
     return MI == Comp;
   }
 
+  bool shouldUseSchedule(SwingSchedulerDAG &SSD, SMSchedule &SMS) override {
+    if (AArch64PipelinerMaxScalarII >= 0 &&
+        SMS.getInitiationInterval() > AArch64PipelinerMaxScalarII &&
+        !containsFpOrVectorRegisters(SSD))
+      return false;
+
+    return true;
+  }
+
   std::optional<bool> createTripCountGreaterCondition(
       int TC, MachineBasicBlock &MBB,
       SmallVectorImpl<MachineOperand> &CondParam) override {
@@ -11735,6 +11752,44 @@ class AArch64PipelinerLoopInfo : public TargetInstrInfo::PipelinerLoopInfo {
 
   bool isMVEExpanderSupported() override { return true; }
 };
+
+bool AArch64PipelinerLoopInfo::containsFpOrVectorRegisters(
+    SwingSchedulerDAG &SSD) const {
+  auto IsFpOrVectorRC = [](const TargetRegisterClass *RC) {
+    return RC && (AArch64::FPR8RegClass.hasSubClassEq(RC) ||
+                  AArch64::FPR16RegClass.hasSubClassEq(RC) ||
+                  AArch64::FPR32RegClass.hasSubClassEq(RC) ||
+                  AArch64::FPR64RegClass.hasSubClassEq(RC) ||
+                  AArch64::FPR128RegClass.hasSubClassEq(RC) ||
+                  AArch64::PPRRegClass.hasSubClassEq(RC) ||
+                  AArch64::PNRRegClass.hasSubClassEq(RC) ||
+                  AArch64::ZPRRegClass.hasSubClassEq(RC) ||
+                  AArch64::ZPR2RegClass.hasSubClassEq(RC) ||
+                  AArch64::ZPR3RegClass.hasSubClassEq(RC) ||
+                  AArch64::ZPR4RegClass.hasSubClassEq(RC));
+  };
+
+  for (SUnit &SU : SSD.SUnits) {
+    const MachineInstr *MI = SU.getInstr();
+    for (const MachineOperand &MO : MI->operands()) {
+      if (!MO.isReg() || !MO.getReg())
+        continue;
+
+      Register Reg = MO.getReg();
+      const TargetRegisterClass *RC = nullptr;
+      if (Reg.isVirtual())
+        RC = MRI.getRegClassOrNull(Reg);
+      else if (Reg.isPhysical())
+        RC = TRI->getMinimalPhysRegClass(Reg.asMCReg());
+
+      if (IsFpOrVectorRC(RC))
+        return true;
+    }
+  }
+
+  return false;
+}
+
 } // namespace
 
 /// Clone an instruction from MI. The register of ReplaceOprNum-th operand