[llvm] [AArch64] Add SME clustering (PR #211483)
Tomer Shafir via llvm-commits
llvm-commits at lists.llvm.org
Thu Jul 23 01:06:35 PDT 2026
https://github.com/tomershafir updated https://github.com/llvm/llvm-project/pull/211483
>From 2214ca47e5080699c63a63d6a1a35c9d2e44f278 Mon Sep 17 00:00:00 2001
From: tomershafir <tomer.shafir8 at gmail.com>
Date: Thu, 23 Jul 2026 10:22:05 +0300
Subject: [PATCH] [AArch64] Add SME clustering on Apple CPU
This patch adds a subtarget feature that controls scheduling SME isntruction back to back. Enabled on Apple CPU.
---
llvm/lib/Target/AArch64/AArch64Features.td | 5 +
.../lib/Target/AArch64/AArch64MacroFusion.cpp | 69 ++++++++
llvm/lib/Target/AArch64/AArch64Processors.td | 4 +-
llvm/lib/Target/AArch64/AArch64Subtarget.h | 2 +-
.../CodeGen/AArch64/misched-fusion-sme.mir | 164 ++++++++++++++++++
.../TableGen/aarch64-apple-tuning-features.td | 2 +
6 files changed, 244 insertions(+), 2 deletions(-)
create mode 100644 llvm/test/CodeGen/AArch64/misched-fusion-sme.mir
diff --git a/llvm/lib/Target/AArch64/AArch64Features.td b/llvm/lib/Target/AArch64/AArch64Features.td
index 7cb7cd32f673b..7e278cdbd29c0 100644
--- a/llvm/lib/Target/AArch64/AArch64Features.td
+++ b/llvm/lib/Target/AArch64/AArch64Features.td
@@ -856,6 +856,11 @@ def FeatureFuseAddSub2RegAndConstOne : SubtargetFeature<
"CPU fuses (a + b + 1) and (a - b - 1)",
[], InlineIgnore>;
+def FeatureFuseSME : SubtargetFeature<
+ "fuse-sme", "HasFuseSME", "true",
+ "CPU fuses SME matrix operations",
+ [], InlineIgnore>;
+
def FeatureDisableLatencySchedHeuristic : SubtargetFeature<
"disable-latency-sched-heuristic", "DisableLatencySchedHeuristic", "true",
"Disable latency scheduling heuristic",
diff --git a/llvm/lib/Target/AArch64/AArch64MacroFusion.cpp b/llvm/lib/Target/AArch64/AArch64MacroFusion.cpp
index 31b3e097de8ca..468194945c028 100644
--- a/llvm/lib/Target/AArch64/AArch64MacroFusion.cpp
+++ b/llvm/lib/Target/AArch64/AArch64MacroFusion.cpp
@@ -537,6 +537,72 @@ static bool isAddSub2RegAndConstOnePair(const MachineInstr *FirstMI,
return false;
}
+static bool definesRegInClass(const MachineInstr &MI,
+ const TargetRegisterInfo *TRI,
+ const TargetRegisterClass &Class) {
+ return llvm::any_of(Class, [&MI, TRI](MCPhysReg Reg) {
+ return MI.definesRegister(Reg, TRI);
+ });
+}
+
+static bool readsRegInClass(const MachineInstr &MI,
+ const TargetRegisterInfo *TRI,
+ const TargetRegisterClass &Class) {
+ return llvm::any_of(
+ Class, [&MI, TRI](MCPhysReg Reg) { return MI.readsRegister(Reg, TRI); });
+}
+
+static bool isFusableSMEOp(const MachineInstr &MI, const TargetInstrInfo &TII,
+ const TargetRegisterInfo *TRI) {
+ const bool ReadOrWriteZA = MI.readsRegister(AArch64::ZA, TRI) ||
+ MI.definesRegister(AArch64::ZA, TRI);
+
+ // (read\write ZA or read\write Z)
+ if (!ReadOrWriteZA && !definesRegInClass(MI, TRI, AArch64::ZPRRegClass))
+ return false;
+
+ // (NOT load\store)
+ if (MI.mayLoad() || MI.mayStore())
+ return false;
+
+ // (NOT write P)
+ if (definesRegInClass(MI, TRI, AArch64::PPRRegClass))
+ return false;
+
+ // (NOT write GPR)
+ const bool WriteGPR = definesRegInClass(MI, TRI, AArch64::GPR32RegClass) ||
+ definesRegInClass(MI, TRI, AArch64::GPR64RegClass);
+ if (WriteGPR)
+ return false;
+
+ // (NOT read\write NZCV)
+ if (MI.readsRegister(AArch64::NZCV, TRI) ||
+ MI.definesRegister(AArch64::NZCV, TRI))
+ return false;
+
+ const bool ReadGPR = readsRegInClass(MI, TRI, AArch64::GPR32RegClass) &&
+ readsRegInClass(MI, TRI, AArch64::GPR64RegClass);
+
+ // ( (NOT read GPR) or read\write ZA )
+ if (ReadGPR && !ReadOrWriteZA)
+ return false;
+
+ return true;
+}
+
+static bool isSMEPair(const MachineInstr *FirstMI, const MachineInstr &SecondMI,
+ const TargetInstrInfo &TII,
+ const TargetRegisterInfo *TRI) {
+ if (!isFusableSMEOp(SecondMI, TII, TRI))
+ return false;
+ // Assume the 1st instr to be a wildcard if it is unspecified.
+ if (FirstMI == nullptr)
+ return true;
+ if (isFusableSMEOp(*FirstMI, TII, TRI))
+ return true;
+ return false;
+}
+
/// \brief Check if the instr pair, FirstMI and SecondMI, should be fused
/// together. Given SecondMI, when FirstMI is unspecified, then check if
/// SecondMI may be part of a fused pair at all.
@@ -576,6 +642,9 @@ static bool shouldScheduleAdjacent(const TargetInstrInfo &TII,
if (ST.hasFuseAddSub2RegAndConstOne() &&
isAddSub2RegAndConstOnePair(FirstMI, SecondMI))
return true;
+ const TargetRegisterInfo *TRI = TSI.getRegisterInfo();
+ if (ST.hasFuseSME() && isSMEPair(FirstMI, SecondMI, TII, TRI))
+ return true;
return false;
}
diff --git a/llvm/lib/Target/AArch64/AArch64Processors.td b/llvm/lib/Target/AArch64/AArch64Processors.td
index f8b344c61ba1e..b5eaccb1b0a0d 100644
--- a/llvm/lib/Target/AArch64/AArch64Processors.td
+++ b/llvm/lib/Target/AArch64/AArch64Processors.td
@@ -447,8 +447,10 @@ def TuneAppleA17 : SubtargetFeature<"apple-a17", "ARMProcFamily", "AppleA17",
def TuneAppleM4 : SubtargetFeature<"apple-m4", "ARMProcFamily", "AppleM4",
"Apple M4",
+ !listconcat(
!listremove(TuneAppleA17.Implies,
- [FeatureStorePairSuppress])>;
+ [FeatureStorePairSuppress]),
+ [FeatureFuseSME])>;
def TuneAppleM5 : SubtargetFeature<"apple-m5", "ARMProcFamily", "AppleM5",
"Apple M5", TuneAppleM4.Implies>;
diff --git a/llvm/lib/Target/AArch64/AArch64Subtarget.h b/llvm/lib/Target/AArch64/AArch64Subtarget.h
index b8303224a8976..b999bf7f3b3e9 100644
--- a/llvm/lib/Target/AArch64/AArch64Subtarget.h
+++ b/llvm/lib/Target/AArch64/AArch64Subtarget.h
@@ -266,7 +266,7 @@ class AArch64Subtarget final : public AArch64GenSubtargetInfo {
return hasArithmeticBccFusion() || hasArithmeticCbzFusion() ||
hasFuseAES() || hasFuseArithmeticLogic() || hasFuseCmpCSel() ||
hasFuseFCmpFCSel() || hasFuseCmpCSet() || hasFuseAdrpAdd() ||
- hasFuseLiterals();
+ hasFuseLiterals() || hasFuseSME();
}
unsigned getEpilogueVectorizationMinVF() const {
diff --git a/llvm/test/CodeGen/AArch64/misched-fusion-sme.mir b/llvm/test/CodeGen/AArch64/misched-fusion-sme.mir
new file mode 100644
index 0000000000000..bafbaf5e2722d
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/misched-fusion-sme.mir
@@ -0,0 +1,164 @@
+# REQUIRES: asserts
+
+# RUN: llc -o /dev/null %s -mtriple=arm64-apple-macosx -mattr=+fuse-sme -passes=machine-scheduler -misched-print-dags -print-before=machine-scheduler 2>&1 | FileCheck %s --check-prefixes=CHECK,FUSE
+# RUN: llc -o /dev/null %s -mtriple=arm64-apple-macosx -mcpu=apple-m4 -passes=machine-scheduler -misched-print-dags -print-before=machine-scheduler 2>&1 | FileCheck %s --check-prefixes=CHECK,FUSE
+# RUN: llc -o /dev/null %s -mtriple=arm64-apple-macosx -mcpu=apple-m5 -passes=machine-scheduler -misched-print-dags -print-before=machine-scheduler 2>&1 | FileCheck %s --check-prefixes=CHECK,FUSE
+
+# RUN: llc -o /dev/null %s -mtriple=arm64-apple-macosx -passes=machine-scheduler -misched-print-dags -print-before=machine-scheduler 2>&1 | FileCheck %s --check-prefixes=CHECK,NOFUSE
+# RUN: llc -o /dev/null %s -mtriple=arm64-apple-macosx -mcpu=apple-m4 -mattr=-fuse-sme -passes=machine-scheduler -misched-print-dags -print-before=machine-scheduler 2>&1 | FileCheck %s --check-prefixes=CHECK,NOFUSE
+
+# CHECK-LABEL: fuse_fmla
+# CHECK: SU({{[0-9]+}}): $za = FMLA_VG2_M2ZZ_S $za(tied-def 0), $w8, 0, $z0_z1, $z2
+# CHECK: Successors:
+# FUSE: SU({{[0-9]+}}): Ord Latency={{[0-9]+}} Cluster
+# NOFUSE-NOT: SU({{[0-9]+}}): Ord Latency={{[0-9]+}} Cluster
+# CHECK: SU({{[0-9]+}}): $za = FMLA_VG2_M2ZZ_S $za(tied-def 0), $w8, 7, $z0_z1, $z2
+---
+name: fuse_fmla
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $za, $w8, $z0_z1, $z2, $w2
+ $za = FMLA_VG2_M2ZZ_S $za, $w8, 0, $z0_z1, $z2
+ $w3 = ORRWri $w2, 4096
+ $za = FMLA_VG2_M2ZZ_S $za, $w8, 7, $z0_z1, $z2
+ RET undef $lr
+...
+
+# CHECK-LABEL: fuse_za_write
+# CHECK: SU({{[0-9]+}}): ZERO_M 255, implicit-def $za
+# CHECK: Successors:
+# FUSE: SU({{[0-9]+}}): Ord Latency={{[0-9]+}} Cluster
+# NOFUSE-NOT: SU({{[0-9]+}}): Ord Latency={{[0-9]+}} Cluster
+# CHECK: SU({{[0-9]+}}): $za = FMLA_VG2_M2ZZ_S $za(tied-def 0), $w8, 0, $z0_z1, $z2
+---
+name: fuse_za_write
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $w8, $z0_z1, $z2, $w2
+ ZERO_M 255, implicit-def $za
+ $w3 = ORRWri $w2, 4096
+ $za = FMLA_VG2_M2ZZ_S $za, $w8, 0, $z0_z1, $z2
+ RET undef $lr
+...
+
+# CHECK-LABEL: fuse_z_write
+# CHECK: SU({{[0-9]+}}): $z0 = ADD_ZZZ_S $z1, $z2
+# CHECK: Successors:
+# FUSE: SU({{[0-9]+}}): Ord Latency={{[0-9]+}} Cluster
+# NOFUSE-NOT: SU({{[0-9]+}}): Ord Latency={{[0-9]+}} Cluster
+# CHECK: SU({{[0-9]+}}): $z4 = ADD_ZZZ_S $z0, $z3
+---
+name: fuse_z_write
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $z1, $z2, $z3, $w2
+ $z0 = ADD_ZZZ_S $z1, $z2
+ $w3 = ORRWri $w2, 4096
+ $z4 = ADD_ZZZ_S $z0, $z3
+ RET undef $lr
+...
+
+# CHECK-LABEL: no_fuse_load
+# CHECK: SU({{[0-9]+}}): $z2 = LDR_ZXI $x0, 0
+# CHECK: Successors:
+# CHECK-NOT: SU({{[0-9]+}}): Ord Latency={{[0-9]+}} Cluster
+# CHECK: SU({{[0-9]+}}): $za = FMLA_VG2_M2ZZ_S
+---
+name: no_fuse_load
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $za, $w8, $z0_z1, $x0, $w2
+ $z2 = LDR_ZXI $x0, 0
+ $w3 = ORRWri $w2, 4096
+ $za = FMLA_VG2_M2ZZ_S $za, $w8, 0, $z0_z1, $z2
+ RET undef $lr
+...
+
+# CHECK-LABEL: no_fuse_store
+# CHECK: SU({{[0-9]+}}): $za = FMLA_VG2_M2ZZ_S
+# CHECK: Successors:
+# CHECK-NOT: SU({{[0-9]+}}): Ord Latency={{[0-9]+}} Cluster
+# CHECK: SU({{[0-9]+}}): STR_ZA
+---
+name: no_fuse_store
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $za, $w8, $z0_z1, $z2, $w12, $x0, $w2
+ $za = FMLA_VG2_M2ZZ_S $za, $w8, 0, $z0_z1, $z2
+ $w3 = ORRWri $w2, 4096
+ STR_ZA $za, $w12, 0, $x0, 0
+ RET undef $lr
+...
+
+# CHECK-LABEL: no_fuse_write_p
+# CHECK: SU({{[0-9]+}}): $z0 = ADD_ZZZ_S $z1, $z2
+# CHECK: Successors:
+# CHECK-NOT: SU({{[0-9]+}}): Ord Latency={{[0-9]+}} Cluster
+# CHECK: SU({{[0-9]+}}): $p1 = FCMEQ_PPzZZ_S $p0, $z0, $z3
+---
+name: no_fuse_write_p
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $p0, $z1, $z2, $z3, $w2
+ $z0 = ADD_ZZZ_S $z1, $z2
+ $w3 = ORRWri $w2, 4096
+ $p1 = FCMEQ_PPzZZ_S $p0, $z0, $z3
+ RET undef $lr
+...
+
+# CHECK-LABEL: no_fuse_write_gpr32
+# CHECK: SU({{[0-9]+}}): $z0 = ADD_ZZZ_S $z1, $z2
+# CHECK: Successors:
+# CHECK-NOT: SU({{[0-9]+}}): Ord Latency={{[0-9]+}} Cluster
+# CHECK: SU({{[0-9]+}}): $w0 = LASTB_RPZ_S $p0, $z0
+---
+name: no_fuse_write_gpr32
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $p0, $z1, $z2, $w2
+ $z0 = ADD_ZZZ_S $z1, $z2
+ $w3 = ORRWri $w2, 4096
+ $w0 = LASTB_RPZ_S $p0, $z0
+ RET undef $lr
+...
+
+# CHECK-LABEL: no_fuse_write_gpr64
+# CHECK: SU({{[0-9]+}}): $z0 = ADD_ZZZ_S $z1, $z2
+# CHECK: Successors:
+# CHECK-NOT: SU({{[0-9]+}}): Ord Latency={{[0-9]+}} Cluster
+# CHECK: SU({{[0-9]+}}): $x0 = LASTB_RPZ_D $p0, $z0
+---
+name: no_fuse_write_gpr64
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $p0, $z1, $z2, $w2
+ $z0 = ADD_ZZZ_S $z1, $z2
+ $w3 = ORRWri $w2, 4096
+ $x0 = LASTB_RPZ_D $p0, $z0
+ RET undef $lr
+...
+
+# CHECK-LABEL: no_fuse_read_gpr_not_za_z
+# CHECK: SU({{[0-9]+}}): $z2 = DUP_ZR_S $w0
+# CHECK: Successors:
+# CHECK-NOT: SU({{[0-9]+}}): Ord Latency={{[0-9]+}} Cluster
+# CHECK: SU({{[0-9]+}}): $za = FMLA_VG2_M2ZZ_S
+---
+name: no_fuse_read_gpr_not_za_z
+tracksRegLiveness: true
+body: |
+ bb.0:
+ liveins: $za, $w8, $z0_z1, $w0, $w2
+ $z2 = DUP_ZR_S $w0
+ $w3 = ORRWri $w2, 4096
+ $za = FMLA_VG2_M2ZZ_S $za, $w8, 0, $z0_z1, $z2
+ RET undef $lr
+...
diff --git a/llvm/test/TableGen/aarch64-apple-tuning-features.td b/llvm/test/TableGen/aarch64-apple-tuning-features.td
index fb0bd1e1c204f..d3e4f007b2512 100644
--- a/llvm/test/TableGen/aarch64-apple-tuning-features.td
+++ b/llvm/test/TableGen/aarch64-apple-tuning-features.td
@@ -250,6 +250,7 @@
// CHECK-NEXT: FeatureFuseCmpCSel,
// CHECK-NEXT: FeatureFuseFCmpFCSel,
// CHECK-NEXT: FeatureFuseLiterals,
+// CHECK-NEXT: FeatureFuseSME,
// CHECK-NEXT: FeatureMaxInterleaveFactor4,
// CHECK-NEXT: FeatureNoZCZeroingFPR64,
// CHECK-NEXT: FeatureZCRegMoveFPR128,
@@ -279,6 +280,7 @@
// CHECK-NEXT: FeatureFuseCmpCSel,
// CHECK-NEXT: FeatureFuseFCmpFCSel,
// CHECK-NEXT: FeatureFuseLiterals,
+// CHECK-NEXT: FeatureFuseSME,
// CHECK-NEXT: FeatureMaxInterleaveFactor4,
// CHECK-NEXT: FeatureNoZCZeroingFPR64,
// CHECK-NEXT: FeatureZCRegMoveFPR128,
More information about the llvm-commits
mailing list