[llvm] [AArch64][SME2] Preserve ZT0 state around function calls (PR #78321)

[Sander de Smalen via llvm-commits]

================
@@ -7664,6 +7666,32 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
     });
   }
 
+  SDValue ZTFrameIdx;
+  MachineFrameInfo &MFI = MF.getFrameInfo();
+  bool PreserveZT0 = CallerAttrs.requiresPreservingZT0(CalleeAttrs);
+
+  // If the caller has ZT0 state which will not be preserved by the callee,
+  // spill ZT0 before the call.
+  if (PreserveZT0) {
+    unsigned ZTObj = MFI.CreateSpillStackObject(64, Align(16));
+    ZTFrameIdx = DAG.getFrameIndex(
+        ZTObj,
+        DAG.getTargetLoweringInfo().getFrameIndexTy(DAG.getDataLayout()));
+
+    Chain = DAG.getNode(AArch64ISD::SAVE_ZT, DL, DAG.getVTList(MVT::Other),
+                        {Chain, DAG.getConstant(0, DL, MVT::i32), ZTFrameIdx});
+  }
+
+  // If caller shares ZT0 but the callee is not shared ZA, we need to stop
+  // PSTATE.ZA before the call if there is no lazy-save active.
+  bool ToggleZA = !RequiresLazySave && CallerAttrs.sharesZT0() &&
----------------
sdesmalen-arm wrote:

```suggestion
  bool ToggleZA = !RequiresLazySave && CallerAttrs.hasZT0State() &&
```
It doesn't really matter how the caller has obtained ZT0 state (either because it shares it, or because it sets up a new state), so `hasZT0State()` is more appropriate here. Can you also add a test that has `"aarch64_new_zt0"` as the state?

https://github.com/llvm/llvm-project/pull/78321