From 59107e5f71647a23048e8c27f832716731030842 Mon Sep 17 00:00:00 2001
From: Marian Lukac <Marian.Lukac@arm.com>
Date: Tue, 15 Apr 2025 17:48:53 +0000
Subject: [PATCH] [AArch64] Fix FPMR handling when switching streaming mode

---
 .../Target/AArch64/AArch64ISelLowering.cpp    |  3 +++
 llvm/test/CodeGen/AArch64/sme-write-fpmr.ll   | 23 +++++++++++++++++++
 2 files changed, 26 insertions(+)
 create mode 100644 llvm/test/CodeGen/AArch64/sme-write-fpmr.ll

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 830ec6886e6bc..20dfda34fb14d 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -8880,12 +8880,15 @@ void AArch64TargetLowering::AdjustInstrPostInstrSelection(MachineInstr &MI,
         MI.removeOperand(I);
 
     // The SVE vector length can change when entering/leaving streaming mode.
+    // FPMR is set to 0 when entering/leaving streaming mode.
     if (MI.getOperand(0).getImm() == AArch64SVCR::SVCRSM ||
         MI.getOperand(0).getImm() == AArch64SVCR::SVCRSMZA) {
       MI.addOperand(MachineOperand::CreateReg(AArch64::VG, /*IsDef=*/false,
                                               /*IsImplicit=*/true));
       MI.addOperand(MachineOperand::CreateReg(AArch64::VG, /*IsDef=*/true,
                                               /*IsImplicit=*/true));
+      MI.addOperand(MachineOperand::CreateReg(AArch64::FPMR, /*IsDef=*/true,
+                                              /*IsImplicit=*/true));
     }
   }
 
diff --git a/llvm/test/CodeGen/AArch64/sme-write-fpmr.ll b/llvm/test/CodeGen/AArch64/sme-write-fpmr.ll
new file mode 100644
index 0000000000000..074d65831a584
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/sme-write-fpmr.ll
@@ -0,0 +1,23 @@
+; RUN: llc -mattr=+sme -stop-after=finalize-isel < %s | FileCheck %s
+
+target triple = "aarch64"
+
+; Check that we don't define FPMR for 'smstart za' and 'smstop za'
+define void @smstart_za() "aarch64_new_za" nounwind {
+  ; CHECK-LABEL:    name: smstart_za
+  ; CHECK-NOT:        implicit-def {{[^,]*}}$fpmr
+  ret void
+}
+
+; Check that we do define FPMR for 'smstart sm' and 'smstop sm'
+define void @smstart_sm() nounwind {
+  ; CHECK-LABEL: name: smstart_sm
+  ; CHECK:          MSRpstatesvcrImm1 1, 1,
+  ; CHECK-SAME:       implicit-def {{[^,]*}}$fpmr
+  ; CHECK:          MSRpstatesvcrImm1 1, 0,
+  ; CHECK-SAME:       implicit-def {{[^,]*}}$fpmr
+  call void @require_sm()
+  ret void
+}
+
+declare void @require_sm() "aarch64_pstate_sm_enabled"