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

[Sander de Smalen via llvm-commits]

================
@@ -275,4 +278,88 @@ TEST(SMEAttributes, Transitions) {
                  .requiresSMChange(SA(SA::SM_Compatible | SA::SM_Body),
                                    /*BodyOverridesInterface=*/true),
             true);
+
+  SA ZT0_In = SA(SA::encodeZT0State(SA::StateValue::In));
+  SA ZT0_InOut = SA(SA::encodeZT0State(SA::StateValue::InOut));
+  SA ZT0_Out = SA(SA::encodeZT0State(SA::StateValue::Out));
+  SA ZT0_Preserved = SA(SA::encodeZT0State(SA::StateValue::Preserved));
+  SA ZT0_New = SA(SA::encodeZT0State(SA::StateValue::New));
+
+  // ZT0 New -> Normal
+  ASSERT_TRUE(ZT0_New.requiresPreservingZT0(SA(SA::Normal)));
+  ASSERT_TRUE(ZT0_New.requiresDisablingZABeforeCall(SA(SA::Normal)));
+  ASSERT_TRUE(ZT0_New.requiresEnablingZAAfterCall(SA(SA::Normal)));
+
+  // ZT0 New -> ZT0 New
+  ASSERT_TRUE(ZT0_New.requiresPreservingZT0(ZT0_New));
+  ASSERT_TRUE(ZT0_New.requiresDisablingZABeforeCall(ZT0_New));
+  ASSERT_TRUE(ZT0_New.requiresEnablingZAAfterCall(ZT0_New));
+
+  // ZT0 New -> ZT0 Shared
+  ASSERT_FALSE(ZT0_New.requiresPreservingZT0(ZT0_In));
+  ASSERT_FALSE(ZT0_New.requiresDisablingZABeforeCall(ZT0_In));
+  ASSERT_FALSE(ZT0_New.requiresEnablingZAAfterCall(ZT0_In));
+
+  ASSERT_FALSE(ZT0_New.requiresPreservingZT0(ZT0_InOut));
+  ASSERT_FALSE(ZT0_New.requiresDisablingZABeforeCall(ZT0_InOut));
+  ASSERT_FALSE(ZT0_New.requiresEnablingZAAfterCall(ZT0_InOut));
+
+  ASSERT_FALSE(ZT0_New.requiresPreservingZT0(ZT0_Out));
+  ASSERT_FALSE(ZT0_New.requiresDisablingZABeforeCall(ZT0_Out));
+  ASSERT_FALSE(ZT0_New.requiresEnablingZAAfterCall(ZT0_Out));
+
+  ASSERT_FALSE(ZT0_New.requiresPreservingZT0(ZT0_Preserved));
+  ASSERT_FALSE(ZT0_New.requiresDisablingZABeforeCall(ZT0_Preserved));
+  ASSERT_FALSE(ZT0_New.requiresEnablingZAAfterCall(ZT0_Preserved));
+
+  // ZT0 Shared -> Normal
+  ASSERT_TRUE(ZT0_In.requiresPreservingZT0(SA(SA::Normal)));
+  ASSERT_TRUE(ZT0_In.requiresDisablingZABeforeCall(SA(SA::Normal)));
+  ASSERT_TRUE(ZT0_In.requiresEnablingZAAfterCall(SA(SA::Normal)));
+
+  ASSERT_TRUE(ZT0_InOut.requiresPreservingZT0(SA(SA::Normal)));
+  ASSERT_TRUE(ZT0_InOut.requiresDisablingZABeforeCall(SA(SA::Normal)));
+  ASSERT_TRUE(ZT0_InOut.requiresEnablingZAAfterCall(SA(SA::Normal)));
+
+  ASSERT_TRUE(ZT0_Out.requiresPreservingZT0(SA(SA::Normal)));
+  ASSERT_TRUE(ZT0_Out.requiresDisablingZABeforeCall(SA(SA::Normal)));
+  ASSERT_TRUE(ZT0_Out.requiresEnablingZAAfterCall(SA(SA::Normal)));
+
+  ASSERT_TRUE(ZT0_Preserved.requiresPreservingZT0(SA(SA::Normal)));
+  ASSERT_TRUE(ZT0_Preserved.requiresDisablingZABeforeCall(SA(SA::Normal)));
+  ASSERT_TRUE(ZT0_Preserved.requiresEnablingZAAfterCall(SA(SA::Normal)));
+
+  // ZT0 Shared -> ZT0 Shared
+  ASSERT_FALSE(ZT0_In.requiresPreservingZT0(ZT0_In));
+  ASSERT_FALSE(ZT0_In.requiresDisablingZABeforeCall(ZT0_In));
+  ASSERT_FALSE(ZT0_In.requiresEnablingZAAfterCall(ZT0_In));
+
+  ASSERT_FALSE(ZT0_InOut.requiresPreservingZT0(ZT0_In));
+  ASSERT_FALSE(ZT0_InOut.requiresDisablingZABeforeCall(ZT0_In));
+  ASSERT_FALSE(ZT0_InOut.requiresEnablingZAAfterCall(ZT0_In));
+
+  ASSERT_FALSE(ZT0_Out.requiresPreservingZT0(ZT0_In));
+  ASSERT_FALSE(ZT0_Out.requiresDisablingZABeforeCall(ZT0_In));
+  ASSERT_FALSE(ZT0_Out.requiresEnablingZAAfterCall(ZT0_In));
+
+  ASSERT_FALSE(ZT0_Preserved.requiresPreservingZT0(ZT0_In));
+  ASSERT_FALSE(ZT0_Preserved.requiresDisablingZABeforeCall(ZT0_In));
+  ASSERT_FALSE(ZT0_Preserved.requiresEnablingZAAfterCall(ZT0_In));
+
+  // ZT0 Shared -> ZT0 New
+  ASSERT_TRUE(ZT0_In.requiresPreservingZT0(ZT0_New));
+  ASSERT_TRUE(ZT0_In.requiresPreservingZT0(ZT0_New));
+  ASSERT_TRUE(ZT0_In.requiresEnablingZAAfterCall(ZT0_New));
+
+  ASSERT_TRUE(ZT0_InOut.requiresPreservingZT0(ZT0_New));
+  ASSERT_TRUE(ZT0_InOut.requiresPreservingZT0(ZT0_New));
+  ASSERT_TRUE(ZT0_InOut.requiresEnablingZAAfterCall(ZT0_New));
+
+  ASSERT_TRUE(ZT0_Out.requiresPreservingZT0(ZT0_New));
+  ASSERT_TRUE(ZT0_Out.requiresPreservingZT0(ZT0_New));
+  ASSERT_TRUE(ZT0_Out.requiresEnablingZAAfterCall(ZT0_New));
+
+  ASSERT_TRUE(ZT0_Preserved.requiresPreservingZT0(ZT0_New));
+  ASSERT_TRUE(ZT0_Preserved.requiresPreservingZT0(ZT0_New));
+  ASSERT_TRUE(ZT0_Preserved.requiresEnablingZAAfterCall(ZT0_New));
----------------
sdesmalen-arm wrote:

Similar comment for these cases, they're not really testing anything that isn't already tested by the by other tests. `ZT0_In/Out/InOut/Preserves/New` all return `true` for `hasZT0State()`, so the codepath and data will be the same for these functions.

I think it's more valuable to add some tests for ZA (because of the lazy-save mechanism) and perhaps even some in combination with ZT0.

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