[llvm] [AArch64][MacroFusion] Assert AES fused pairs are tied (NFC) (PR #196484)

[Tomer Shafir via llvm-commits]

https://github.com/tomershafir updated https://github.com/llvm/llvm-project/pull/196484

>From 449f5d59e3fcf9ee2979fb4dc348ed914c472295 Mon Sep 17 00:00:00 2001
From: tomershafir <tomer.shafir8 at gmail.com>
Date: Fri, 8 May 2026 10:52:03 +0300
Subject: [PATCH] [AArch64][MacroFusion] Fuse only tied AES pairs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This patch adds an ad-hoc check to macro fusion to only fuse AES pairs that are tied, to more explicitly avoid regressions.

In the case of full compilation - ISel captures every RAW dependent AESE/D+AES[I]MC pair (by data-dependence DAG), and applies a constraint that the pair must write to the same dest, i.e the second instruction is tied (a thing that cannot be expressed in SSA IR). However, I think it’s not clear or enforced on macro fusion. The tests in `llvm/test/CodeGen/AArch64/misched-fusion-aes.ll` may not catch an ISel change that would happen to pass, satisfying the register allocation being filechecked. Notice that in practice all subtargets effectively fuse only tied pairs. Thus, adding an explicit check to avoid inaccurate fusions. If it appears in the future that a subtarget can fuse untied pairs, we should re-address and maybe distinguish 2 subtarget features for the 2 cases.

Plus adding a test runline for latest `apple-m5`.

Note: it cannot be an assert, e.g. a partial pipeline like in `llvm/test/CodeGen/AArch64/misched-fusion-crypto-eor.mir` would fail.
---
 .../lib/Target/AArch64/AArch64MacroFusion.cpp | 19 +++++++++++++++++--
 .../CodeGen/AArch64/misched-fusion-aes.ll     |  1 +
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64MacroFusion.cpp b/llvm/lib/Target/AArch64/AArch64MacroFusion.cpp
index 55214f3f749e4..b9153203c311d 100644
--- a/llvm/lib/Target/AArch64/AArch64MacroFusion.cpp
+++ b/llvm/lib/Target/AArch64/AArch64MacroFusion.cpp
@@ -118,6 +118,11 @@ static bool isArithmeticCbzPair(const MachineInstr *FirstMI,
 }
 
 /// AES crypto encoding or decoding.
+///
+/// ISel should always match RAW dependent AES pairs via tied pseudos,
+/// regardless of the program order, so matched pairs here should always have
+/// the same dest register (AESMCrrTied pre-RA, AESMCrr post-RA after pseudo
+/// expansion).
 static bool isAESPair(const MachineInstr *FirstMI,
                       const MachineInstr &SecondMI) {
   // Assume the 1st instr to be a wildcard if it is unspecified.
@@ -125,11 +130,21 @@ static bool isAESPair(const MachineInstr *FirstMI,
   // AES encode.
   case AArch64::AESMCrr:
   case AArch64::AESMCrrTied:
-    return FirstMI == nullptr || FirstMI->getOpcode() == AArch64::AESErr;
+    if (FirstMI == nullptr)
+      return true;
+    if (FirstMI->getOpcode() != AArch64::AESErr)
+      return false;
+    return SecondMI.getOpcode() == AArch64::AESMCrrTied ||
+           FirstMI->getOperand(0).getReg() == SecondMI.getOperand(0).getReg();
   // AES decode.
   case AArch64::AESIMCrr:
   case AArch64::AESIMCrrTied:
-    return FirstMI == nullptr || FirstMI->getOpcode() == AArch64::AESDrr;
+    if (FirstMI == nullptr)
+      return true;
+    if (FirstMI->getOpcode() != AArch64::AESDrr)
+      return false;
+    return SecondMI.getOpcode() == AArch64::AESIMCrrTied ||
+           FirstMI->getOperand(0).getReg() == SecondMI.getOperand(0).getReg();
   }
 
   return false;
diff --git a/llvm/test/CodeGen/AArch64/misched-fusion-aes.ll b/llvm/test/CodeGen/AArch64/misched-fusion-aes.ll
index dc6fa9128e933..cd1e46f187b09 100644
--- a/llvm/test/CodeGen/AArch64/misched-fusion-aes.ll
+++ b/llvm/test/CodeGen/AArch64/misched-fusion-aes.ll
@@ -20,6 +20,7 @@
 ; RUN: llc %s -o - -mtriple=aarch64-unknown -mcpu=ampere1    | FileCheck %s
 ; RUN: llc %s -o - -mtriple=aarch64-unknown -mcpu=ampere1a   | FileCheck %s
 ; RUN: llc %s -o - -mtriple=aarch64-unknown -mcpu=ampere1b   | FileCheck %s
+; RUN: llc %s -o - -mtriple=aarch64-unknown -mcpu=apple-m5   | FileCheck %s
 
 declare <16 x i8> @llvm.aarch64.crypto.aese(<16 x i8> %d, <16 x i8> %k)
 declare <16 x i8> @llvm.aarch64.crypto.aesmc(<16 x i8> %d)