[llvm] [llvm-exegesis] Implement the loop repetition mode for AArch64 (PR #154751)

[Sjoerd Meijer via llvm-commits]

https://github.com/sjoerdmeijer updated https://github.com/llvm/llvm-project/pull/154751

>From 727c2713b0d10e4c3d6d547793cdda2189331f32 Mon Sep 17 00:00:00 2001
From: Sjoerd Meijer <smeijer at nvidia.com>
Date: Thu, 21 Aug 2025 02:54:31 -0700
Subject: [PATCH] [llvm-exegesis] Implement the loop repetition mode for
 AArch64

Subject says it all: implement the loop iterator decrement and jump
function functions, and reserve X19 for the loop counter.
---
 .../llvm-exegesis/AArch64/loop-register.s     | 17 +++++++++++
 .../llvm-exegesis/lib/AArch64/Target.cpp      | 30 +++++++++++++++++++
 2 files changed, 47 insertions(+)
 create mode 100644 llvm/test/tools/llvm-exegesis/AArch64/loop-register.s

diff --git a/llvm/test/tools/llvm-exegesis/AArch64/loop-register.s b/llvm/test/tools/llvm-exegesis/AArch64/loop-register.s
new file mode 100644
index 0000000000000..e1ffe462b50c1
--- /dev/null
+++ b/llvm/test/tools/llvm-exegesis/AArch64/loop-register.s
@@ -0,0 +1,17 @@
+REQUIRES: aarch64-registered-target, asserts
+
+RUN: llvm-exegesis -mcpu=neoverse-v2 --use-dummy-perf-counters --mode=latency --debug-only=print-gen-assembly --opcode-name=ADDVv4i16v -repetition-mode=loop 2>&1 | FileCheck %s
+
+CHECK:        0:  {{.*}}  str     x19, [sp, #-16]!
+CHECK-NEXT:   4:  {{.*}}  movi    d[[REG:[0-9]+]], #0000000000000000
+CHECK-NEXT:   8:  {{.*}}  mov     x19, #10000
+CHECK-NEXT:   c:  {{.*}}  nop
+CHECK-NEXT:   10: {{.*}}  nop
+CHECK-NEXT:   14: {{.*}}  nop
+CHECK-NEXT:   18: {{.*}}  nop
+CHECK-NEXT:   1c: {{.*}}  nop
+CHECK-NEXT:   20: {{.*}}  addv    h[[REG]], v[[REG]].4h
+CHECK-NEXT:   24: {{.*}}  subs    x19, x19, #1
+CHECK-NEXT:   28: {{.*}}  b.ne    #-8
+CHECK-NEXT:   2c: {{.*}}  ldr     x19, [sp], #16
+CHECK-NEXT:   30: {{.*}}  ret
diff --git a/llvm/tools/llvm-exegesis/lib/AArch64/Target.cpp b/llvm/tools/llvm-exegesis/lib/AArch64/Target.cpp
index 3a0021e3c132d..b5fc8c3a2bd37 100644
--- a/llvm/tools/llvm-exegesis/lib/AArch64/Target.cpp
+++ b/llvm/tools/llvm-exegesis/lib/AArch64/Target.cpp
@@ -8,6 +8,7 @@
 #include "../Target.h"
 #include "AArch64.h"
 #include "AArch64RegisterInfo.h"
+#include "llvm/CodeGen/MachineInstrBuilder.h"
 
 #if defined(__aarch64__) && defined(__linux__)
 #include <sys/prctl.h> // For PR_PAC_* constants
@@ -109,6 +110,10 @@ static MCInst loadFPImmediate(MCRegister Reg, unsigned RegBitWidth,
 
 namespace {
 
+// Use X19 as the loop counter register since it's a callee-saved register
+// that's available for temporary use.
+constexpr const MCPhysReg kDefaultLoopCounterReg = AArch64::X19;
+
 class ExegesisAArch64Target : public ExegesisTarget {
 public:
   ExegesisAArch64Target()
@@ -141,6 +146,31 @@ class ExegesisAArch64Target : public ExegesisTarget {
     errs() << "setRegTo is not implemented, results will be unreliable\n";
     return {};
   }
+  MCRegister getDefaultLoopCounterRegister(const Triple &) const override {
+    return kDefaultLoopCounterReg;
+  }
+
+void decrementLoopCounterAndJump(MachineBasicBlock &MBB,
+                                 MachineBasicBlock &TargetMBB,
+                                 const MCInstrInfo &MII,
+                                 MCRegister LoopRegister) const override {
+     // subs LoopRegister, LoopRegister, #1
+    BuildMI(&MBB, DebugLoc(), MII.get(AArch64::SUBSXri))
+        .addDef(LoopRegister)
+        .addUse(LoopRegister)
+        .addImm(1)  // Subtract 1
+        .addImm(0); // No shift amount
+    // b.ne TargetMBB
+    BuildMI(&MBB, DebugLoc(), MII.get(AArch64::Bcc))
+        .addImm(AArch64CC::NE)
+        .addMBB(&TargetMBB);
+  }
+
+  // Registers that should not be selected for use in snippets.
+  const MCPhysReg UnavailableRegisters[1] = {kDefaultLoopCounterReg};
+  ArrayRef<MCPhysReg> getUnavailableRegisters() const override {
+    return UnavailableRegisters;
+  }
 
   bool matchesArch(Triple::ArchType Arch) const override {
     return Arch == Triple::aarch64 || Arch == Triple::aarch64_be;