[llvm] [AArch64][SME]Check streaming mode when using SME2 instruction in fra… (PR #109680)

[via llvm-commits]

https://github.com/CarolineConcatto created https://github.com/llvm/llvm-project/pull/109680

…me lowering

SME instructions can only be used in streaming mode. PTRUE for predicated counter and the ld/st pair can be used when:
  sve2.1  is available or
  sme2 available in function in streaming mode.
Previously the frame lowering only checking if sme2 available when building the machine instruction.
This fix checks if sme2 is available and is subtarget in streaming mode

>From 2ce1be225eb09b89b8efceeea0338b20b072e9d8 Mon Sep 17 00:00:00 2001
From: Caroline Concatto <caroline.concatto at arm.com>
Date: Mon, 23 Sep 2024 13:37:41 +0000
Subject: [PATCH] [AArch64][SME]Check streaming mode when using SME2
 instruction in frame lowering

SME instructions can only be used in streaming mode.
PTRUE for predicated counter and the ld/st pair can be used when:
  sve2.1  is available or
  sme2 available in function in streaming mode.
Previously the frame lowering only checking if sme2 available when building
the machine instruction.
This fix checks if sme2 is available and is subtarget in streaming mode
---
 llvm/lib/Target/AArch64/AArch64FrameLowering.cpp   | 14 +++++++++++---
 .../AArch64/sve-callee-save-restore-pairs.ll       |  3 ++-
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
index ad20e76d0fe2e0..9849e35bcb266b 100644
--- a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
@@ -2955,6 +2955,14 @@ unsigned findFreePredicateReg(BitVector &SavedRegs) {
   return AArch64::NoRegister;
 }
 
+bool isSve2p1OrSme2InStreaming(const AArch64Subtarget &Subtarget) {
+  if (Subtarget.hasSVE2p1())
+    return true;
+  if (Subtarget.hasSME2() && Subtarget.isStreaming())
+    return true;
+  return false;
+}
+
 static void computeCalleeSaveRegisterPairs(
     MachineFunction &MF, ArrayRef<CalleeSavedInfo> CSI,
     const TargetRegisterInfo *TRI, SmallVectorImpl<RegPairInfo> &RegPairs,
@@ -3326,7 +3334,7 @@ bool AArch64FrameLowering::spillCalleeSavedRegisters(
                               MF.getSubtarget<AArch64Subtarget>();
       AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
       unsigned PnReg = AFI->getPredicateRegForFillSpill();
-      assert(((Subtarget.hasSVE2p1() || Subtarget.hasSME2()) && PnReg != 0) &&
+      assert((PnReg != 0 && isSve2p1OrSme2InStreaming(Subtarget)) &&
              "Expects SVE2.1 or SME2 target and a predicate register");
 #ifdef EXPENSIVE_CHECKS
       auto IsPPR = [](const RegPairInfo &c) {
@@ -3504,7 +3512,7 @@ bool AArch64FrameLowering::restoreCalleeSavedRegisters(
       [[maybe_unused]] const AArch64Subtarget &Subtarget =
                               MF.getSubtarget<AArch64Subtarget>();
       unsigned PnReg = AFI->getPredicateRegForFillSpill();
-      assert(((Subtarget.hasSVE2p1() || Subtarget.hasSME2()) && PnReg != 0) &&
+      assert((PnReg != 0 && isSve2p1OrSme2InStreaming(Subtarget)) &&
              "Expects SVE2.1 or SME2 target and a predicate register");
 #ifdef EXPENSIVE_CHECKS
       assert(!(PPRBegin < ZPRBegin) &&
@@ -3718,7 +3726,7 @@ void AArch64FrameLowering::determineCalleeSaves(MachineFunction &MF,
                     SavedRegs.test(CSRegs[i ^ 1]));
   }
 
-  if (HasPairZReg && (Subtarget.hasSVE2p1() || Subtarget.hasSME2())) {
+  if (HasPairZReg && isSve2p1OrSme2InStreaming(Subtarget)) {
     AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
     // Find a suitable predicate register for the multi-vector spill/fill
     // instructions.
diff --git a/llvm/test/CodeGen/AArch64/sve-callee-save-restore-pairs.ll b/llvm/test/CodeGen/AArch64/sve-callee-save-restore-pairs.ll
index 470c0dd45782c9..9b5cd3251c22a2 100644
--- a/llvm/test/CodeGen/AArch64/sve-callee-save-restore-pairs.ll
+++ b/llvm/test/CodeGen/AArch64/sve-callee-save-restore-pairs.ll
@@ -3,7 +3,8 @@
 ; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve2 -verify-machineinstrs < %s | FileCheck %s --check-prefixes=NOPAIR
 ; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sme2 -verify-machineinstrs -force-streaming < %s | FileCheck %s --check-prefixes=PAIR
 ; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve2p1 -verify-machineinstrs < %s | FileCheck %s --check-prefixes=PAIR
-
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sme2 -mattr=+sve2 -verify-machineinstrs < %s | FileCheck %s --check-prefixes=NOPAIR
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sme2  -mattr=+sve2 -verify-machineinstrs -force-streaming < %s | FileCheck %s --check-prefixes=PAIR
 
 declare void @my_func()