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

[Tomer Shafir via llvm-commits]

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

This patch adds an assert that AES fused pairs are tied to more explicitly avoid regressions.

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 regalloc. Notice that in practice all subtargets effectively fuse only on the tied case. Thus, adding an explicit assert to avoid the case of wrong fusions. If a subtarget would be able to fuse on the untied case, we should re-address and maybe distinguish 2 subtarget features for the 2 cases.

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

>From a525a0762217fafbe1fc42c6876094862cf0759e 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] Assert AES fused pairs are tied (NFC)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This patch adds an assert that AES fused pairs are tied to more explicitly avoid regressions.

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 regalloc. Notice that in practice all subtargets effectively fuse only on the tied case. Thus, adding an explicit assert to avoid the case of wrong fusions. If a subtarget would be able to fuse on the untied case, we should re-address and maybe distinguish 2 subtarget features for the 2 cases.

Plus adding a test runline for latest `apple-m5`.
---
 .../lib/Target/AArch64/AArch64MacroFusion.cpp | 25 +++++++++++++++++--
 .../CodeGen/AArch64/misched-fusion-aes.ll     |  1 +
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64MacroFusion.cpp b/llvm/lib/Target/AArch64/AArch64MacroFusion.cpp
index 55214f3f749e4..27e15a9071d69 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,27 @@ 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;
+    assert(
+        (SecondMI.getOpcode() == AArch64::AESMCrrTied ||
+         FirstMI->getOperand(0).getReg() == SecondMI.getOperand(0).getReg()) &&
+        "Expected tied AES pair");
+    return true;
   // 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;
+    assert(
+        (SecondMI.getOpcode() == AArch64::AESIMCrrTied ||
+         FirstMI->getOperand(0).getReg() == SecondMI.getOperand(0).getReg()) &&
+        "Expected tied AES pair");
+    return true;
   }
 
   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)