[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:

nit: While the way you've written it is perhaps a bit more future proof in case we ever need to add more state, writing it as below does make it a bit more readable (at least for me):
```
bool ToggleZA = CallerAttrs.hasZT0State() && !CallerAttrs.hasZAState() &&
                CalleeAttrs.hasPrivateZAInterface();
assert((!ToggleZA || !RequiresLazySave) &&
       "Lazy-save should have PSTATE.SM=1 on entry to the function");
```

The fact that this does not require a lazy-save is then more of a consequence.

Is it worth creating a new method to the SMEAttributes class, e.g. something like `bool requiresZAToggle(const SMEAttributes &Callee) const` ?

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