[AArch64][SME] Enable TPIDR2 lazy-save for za_preserved

[MattDevereau]

Summary

This change makes callees with the arm_preserves_za
type attribute comply with the dormant state requirements
when it's caller has the arm_shared_za type attribute

https://github.com/ARM-software/abi-aa/blob/5e67092434b50c04f8ad178a9c272ce3c6ada7fd/aapcs64/aapcs64.rst?plain=1#L1381

[sdesmalen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

7354

requiresLazySave should return true, even when the function has the aarch64_pstate_za_preserved attribute. That also makes this code below a bit simpler.

7366

Streaming-properties are orthogonal to ZA, so don't belong in this condition.

llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h

80

Could you keep the original name please?

[MattDevereau]

llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h

80

Can I do it as a separate change at least? It doesn't match the rest of the class's interface at all.

[MattDevereau]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

7366

It's possible for

if (RequiresSMChange)
  PStateSM = getPStateSM(DAG, Chain, CallerAttrs, DL, MVT::i64);

to call LowerCall an additional time and generate an extra redundant store into the TPIDR2 save slices parameter, which is why this is here. I will look for a better way of handling this though.

[sdesmalen]

llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h

80

The preserving ZA is not really an interface in the sense that the ABI defines e.g. a SharedZA/PrivateZA/Streaming interface. The 'interface' part of the name refers to the caller having to be aware of the callee's interface when generating code for the call.

For example, a private-ZA interface requires the caller to set up the lazy-save mechanism when the caller has ZA state.
Or a streaming interface requires the caller to ensure that streaming-SVE mode is enabled before the call.

The __arm_preserves_za attribute (or aarch64_pstate_za_preserved as it's named in LLVM IR) is more of a 'hint' to the compiler to generate more efficient code. It is a property that can be ignored by the compiler and the generated code would still be correct.

Note that I pushed a change yesterday to rename hasNewZAInterface() in rG0a32a999aea07e373e1b3c9e6f0d20d3f7f1a7b8, because 'new ZA' is not part of the interface (externally visible to callers), but rather an implementation property of the function's body/definition.

As such, I'd rather stick with preservesZA as the name of this interface.

[sdesmalen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

7342

Am I reading this wrong, or should this have been named NumZaSaveSlicesAddr ?

7767–7809

Can you write this as:

if (RequiresLazySave) {
  if (!CalleeAttrs.preservesZA()) {
    // code to restore ZA
  }
  Result = DAG.getNode(...aarch64_sme_set_tpidr2...);
}

llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h

38

nit: maybe add a comment that we use this for SME ABI routines.

84

nit: this interface seems redundant.

[sdesmalen]

llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp

34

Perhaps it's worth adding a new SMEAttrs constructor, such that you can use that both here and in AArch64ISelLowering.cpp

SMEAttrs(StringRef funcname) {  /* handle SME ABI functions here */ }

[MattDevereau]

llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp

34

What benefit is this supposed to give? If you just use a StringRef param in a constructor you will have to make many assumptions to satisfy the mutex asserts in set and probably just end up doing Bitmask |= Za_NoLazySave with whatever comes out of that constructor anyway. Isn't getting this functionality from AArch64ISelLowering.cpp is what getCalleeAttrsFromExternalFunction is providing anyway?

[sdesmalen]

llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp

34

My suggestion is to avoid having to write this if (FuncName == "__arm_tpidr2_save" || .. code in multiple places.

you will have to make many assumptions to satisfy the mutex asserts in set and probably just end up doing Bitmask |= Za_NoLazySave with whatever comes out of that constructor anyway

What assumptions will you have to make?

[MattDevereau]

llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp

34

In the end I realised what you meant and I uploaded another patch set with the StringRef constructor

[sdesmalen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

7324

Do we really need the if-statement? For example, can we just write:

else if (auto *ES = ...)
  CalleeAttrs = SMEAttrs(ES->getSymbol());

If the default for SMEAttrs(<symbol>) is that it doesn't have any streaming- or ZA properties then this is equivalent to what we had before this patch

llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp

26

Does this need a new assert to say that NewZA and SharedZA cannot be combined with NoLazySave?

llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h

90

I'd rather we don't add an operator bool(), because it's not clear to me what it is supposed to represent that the existing interfaces don't represent already.

[sdesmalen]

llvm/test/CodeGen/AArch64/sme-lazy-save-call.ll

172

Could you also add the "aarch64_pstate_sm_compatible" attribute so that the compiler must call __arm_sme_state() as part of the call-sequence and test that we don't set up the lazy-save mechanism for that __arm_sme_state() call?

nit: Could you maybe also give this test a better name than foo? :)

[MattDevereau]

llvm/test/CodeGen/AArch64/sme-lazy-save-call.ll

172

Done & nit done.

[sdesmalen]

LGTM

llvm/test/CodeGen/AArch64/sme-lazy-save-call.ll

216

nit: can you add "aarch64_pstate_za_preserved" next to the declaration of @private_za_preserved_callee similar to how you've done it for @za_shared_caller_za_preserved_callee, and then remove this attribute?

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

SDValue AArch64TargetLowering::getPStateSM(SelectionDAG &DAG, SDValue Chain,
   return DAG.getNode(ISD::AND, DL, MVT::i64, CallResult.first.getOperand(0),
                      Mask);
 }
 
 static std::optional<SMEAttrs> getCalleeAttrsFromExternalFunction(SDValue V) {
   if (auto *ES = dyn_cast<ExternalSymbolSDNode>(V)) {
     StringRef S(ES->getSymbol());
     if (S == "__arm_sme_state" || S == "__arm_tpidr2_save")
-      return SMEAttrs(SMEAttrs::SM_Compatible | SMEAttrs::ZA_Preserved);
+      return SMEAttrs(SMEAttrs::SM_Compatible | SMEAttrs::ZA_Preserved |
+                      SMEAttrs::ZA_NoLazySave);
     if (S == "__arm_tpidr2_restore")
-      return SMEAttrs(SMEAttrs::SM_Compatible | SMEAttrs::ZA_Shared);
+      return SMEAttrs(SMEAttrs::SM_Compatible | SMEAttrs::ZA_Shared |
+                      SMEAttrs::ZA_NoLazySave);
   }
   return std::nullopt;
 }
 
 SDValue AArch64TargetLowering::LowerINTRINSIC_VOID(SDValue Op,
                                                    SelectionDAG &DAG) const {
   unsigned IntNo = Op.getConstantOperandVal(1);
   SDLoc DL(Op);
AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
 
   // Determine whether we need any streaming mode changes.
   SMEAttrs CalleeAttrs, CallerAttrs(MF.getFunction());
   if (CLI.CB)
     CalleeAttrs = SMEAttrs(*CLI.CB);
   else if (std::optional<SMEAttrs> Attrs =
                getCalleeAttrsFromExternalFunction(CLI.Callee))
     CalleeAttrs = *Attrs;
 
   bool RequiresLazySave = CallerAttrs.requiresLazySave(CalleeAttrs);
-
-  MachineFrameInfo &MFI = MF.getFrameInfo();
   if (RequiresLazySave) {
+    SDValue NumZaSaveSlices;
+    if (!CalleeAttrs.preservesZA()) {
       // Set up a lazy save mechanism by storing the runtime live slices
-    // (worst-case N*N) to the TPIDR2 stack object.
-    SDValue N = DAG.getNode(AArch64ISD::RDSVL, DL, MVT::i64,
+      // (worst-case SVL*SVL) to the TPIDR2 stack object.
+      SDValue SVL = DAG.getNode(AArch64ISD::RDSVL, DL, MVT::i64,
                                 DAG.getConstant(1, DL, MVT::i32));
-    SDValue NN = DAG.getNode(ISD::MUL, DL, MVT::i64, N, N);
-    unsigned TPIDR2Obj = FuncInfo->getLazySaveTPIDR2Obj();
+      NumZaSaveSlices = DAG.getNode(ISD::MUL, DL, MVT::i64, SVL, SVL);
+    } else if (CalleeAttrs.preservesZA()) {
+      NumZaSaveSlices = DAG.getConstant(0, DL, MVT::i64);
+    }
 
+    unsigned TPIDR2Obj = FuncInfo->getLazySaveTPIDR2Obj();
     MachinePointerInfo MPI = MachinePointerInfo::getStack(MF, TPIDR2Obj);
     SDValue TPIDR2ObjAddr = DAG.getFrameIndex(TPIDR2Obj,
         DAG.getTargetLoweringInfo().getFrameIndexTy(DAG.getDataLayout()));
     SDValue BufferPtrAddr =
         DAG.getNode(ISD::ADD, DL, TPIDR2ObjAddr.getValueType(), TPIDR2ObjAddr,
                     DAG.getConstant(8, DL, TPIDR2ObjAddr.getValueType()));
-    Chain = DAG.getTruncStore(Chain, DL, NN, BufferPtrAddr, MPI, MVT::i16);
+    Chain = DAG.getTruncStore(Chain, DL, NumZaSaveSlices, BufferPtrAddr, MPI,
+                              MVT::i16);
     Chain = DAG.getNode(
         ISD::INTRINSIC_VOID, DL, MVT::Other, Chain,
         DAG.getConstant(Intrinsic::aarch64_sme_set_tpidr2, DL, MVT::i32),
         TPIDR2ObjAddr);
   }
 
   SDValue PStateSM;
   std::optional<bool> RequiresSMChange =
       CallerAttrs.requiresSMChange(CalleeAttrs);
   if (RequiresSMChange)
     PStateSM = getPStateSM(DAG, Chain, CallerAttrs, DL, MVT::i64);
 
   // Adjust the stack pointer for the new arguments...
   // These operations are automatically eliminated by the prolog/epilog pass
   if (!IsSibCall)
     Chain = DAG.getCALLSEQ_START(Chain, IsTailCall ? 0 : NumBytes, 0, DL);
 
   SDValue StackPtr = DAG.getCopyFromReg(Chain, DL, AArch64::SP,
                                         getPointerTy(DAG.getDataLayout()));
 
   SmallVector<std::pair<unsigned, SDValue>, 8> RegsToPass;
   SmallSet<unsigned, 8> RegsUsed;
   SmallVector<SDValue, 8> MemOpChains;
   auto PtrVT = getPointerTy(DAG.getDataLayout());
 
   if (IsVarArg && CLI.CB && CLI.CB->isMustTailCall()) {
     const auto &Forwards = FuncInfo->getForwardedMustTailRegParms();
     for (const auto &F : Forwards) {
case CCValAssign::Indirect:
         assert(!Outs[i].Flags.isInConsecutiveRegsLast());
         while (!Outs[i + NumParts - 1].Flags.isInConsecutiveRegsLast())
           ++NumParts;
         StoreSize *= NumParts;
       }
 
       Type *Ty = EVT(VA.getValVT()).getTypeForEVT(*DAG.getContext());
       Align Alignment = DAG.getDataLayout().getPrefTypeAlign(Ty);
+      MachineFrameInfo &MFI = MF.getFrameInfo();
       int FI = MFI.CreateStackObject(StoreSize, Alignment, false);
       if (isScalable)
         MFI.setStackID(FI, TargetStackID::ScalableVector);
 
       MachinePointerInfo MPI = MachinePointerInfo::getFixedStack(MF, FI);
       SDValue Ptr = DAG.getFrameIndex(
           FI, DAG.getTargetLoweringInfo().getFrameIndexTy(DAG.getDataLayout()));
       SDValue SpillSlot = Ptr;
if (!Ins.empty())
     InGlue = Result.getValue(Result->getNumValues() - 1);
 
   if (RequiresSMChange) {
     assert(PStateSM && "Expected a PStateSM to be set");
     Result = changeStreamingMode(DAG, DL, !*RequiresSMChange, Result, InGlue,
                                  PStateSM, false);
   }
 
-  if (RequiresLazySave) {
+  if (RequiresLazySave && !CalleeAttrs.preservesZA()) {
     // Unconditionally resume ZA.
     Result = DAG.getNode(
         AArch64ISD::SMSTART, DL, MVT::Other, Result,
         DAG.getTargetConstant((int32_t)(AArch64SVCR::SVCRZA), DL, MVT::i32),
         DAG.getConstant(0, DL, MVT::i64), DAG.getConstant(1, DL, MVT::i64));
 
     // Conditionally restore the lazy save using a pseudo node.
     unsigned FI = FuncInfo->getLazySaveTPIDR2Obj();
     SDValue RegMask = DAG.getRegisterMask(
         TRI->SMEABISupportRoutinesCallPreservedMaskFromX0());
     SDValue RestoreRoutine = DAG.getTargetExternalSymbol(
         "__arm_tpidr2_restore", getPointerTy(DAG.getDataLayout()));
     SDValue TPIDR2_EL0 = DAG.getNode(
         ISD::INTRINSIC_W_CHAIN, DL, MVT::i64, Result,
         DAG.getConstant(Intrinsic::aarch64_sme_get_tpidr2, DL, MVT::i32));
 
     // Copy the address of the TPIDR2 block into X0 before 'calling' the
     // RESTORE_ZA pseudo.
     SDValue Glue;
     SDValue TPIDR2Block = DAG.getFrameIndex(
         FI, DAG.getTargetLoweringInfo().getFrameIndexTy(DAG.getDataLayout()));
     Result = DAG.getCopyToReg(Result, DL, AArch64::X0, TPIDR2Block, Glue);
     Result = DAG.getNode(AArch64ISD::RESTORE_ZA, DL, MVT::Other,
                          {Result, TPIDR2_EL0,
                           DAG.getRegister(AArch64::X0, MVT::i64),
                           RestoreRoutine,
                           RegMask,
                           Result.getValue(1)});
 
     // Finally reset the TPIDR2_EL0 register to 0.
     Result = DAG.getNode(
         ISD::INTRINSIC_VOID, DL, MVT::Other, Result,
         DAG.getConstant(Intrinsic::aarch64_sme_set_tpidr2, DL, MVT::i32),
         DAG.getConstant(0, DL, MVT::i64));
   }
 
+  if (RequiresLazySave && CalleeAttrs.preservesZA()) {
+    Result = DAG.getNode(
+        ISD::INTRINSIC_VOID, DL, MVT::Other, Result,
+        DAG.getConstant(Intrinsic::aarch64_sme_set_tpidr2, DL, MVT::i32),
+        DAG.getConstant(0, DL, MVT::i64));
+  }
+
   if (RequiresSMChange || RequiresLazySave) {
     for (unsigned I = 0; I < InVals.size(); ++I) {
       // The smstart/smstop is chained as part of the call, but when the
       // resulting chain is discarded (which happens when the call is not part
       // of a chain, e.g. a call to @llvm.cos()), we need to ensure the
       // smstart/smstop is chained to the result value. We can do that by doing
       // a vreg -> vreg copy.
       Register Reg = MF.getRegInfo().createVirtualRegister(

llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.h

public:
   enum Mask {
     Normal = 0,
     SM_Enabled = 1 << 0,    // aarch64_pstate_sm_enabled
     SM_Compatible = 1 << 1, // aarch64_pstate_sm_compatible
     SM_Body = 1 << 2,       // aarch64_pstate_sm_body
     ZA_Shared = 1 << 3,     // aarch64_pstate_sm_shared
     ZA_New = 1 << 4,        // aarch64_pstate_sm_new
     ZA_Preserved = 1 << 5,  // aarch64_pstate_sm_preserved
+    ZA_NoLazySave = 1 << 6, // Don't emit TPI_DR2 lazy saves
     All = ZA_Preserved - 1
   };
 
   SMEAttrs(unsigned Mask = Normal) : Bitmask(0) { set(Mask); }
   SMEAttrs(const Function &F) : SMEAttrs(F.getAttributes()) {}
   SMEAttrs(const CallBase &CB);
   SMEAttrs(const AttributeList &L);
 
public:
   std::optional<bool>
   requiresSMChange(const SMEAttrs &Callee,
                    bool BodyOverridesInterface = false) const;
 
   // Interfaces to query PSTATE.ZA
   bool hasNewZABody() const { return Bitmask & ZA_New; }
   bool hasSharedZAInterface() const { return Bitmask & ZA_Shared; }
   bool hasPrivateZAInterface() const { return !hasSharedZAInterface(); }
   bool preservesZA() const { return Bitmask & ZA_Preserved; }
   bool hasZAState() const {
     return hasNewZABody() || hasSharedZAInterface();
   }
+  bool hasNoLazySave() const { return Bitmask & ZA_NoLazySave; }
   bool requiresLazySave(const SMEAttrs &Callee) const {
     return hasZAState() && Callee.hasPrivateZAInterface() &&
-           !Callee.preservesZA();
+           !Callee.hasNoLazySave();
   }
 };
 
 } // namespace llvm
 
 #endif // LLVM_LIB_TARGET_AARCH64_UTILS_AARCH64SMEATTRIBUTES_H

llvm/lib/Target/AArch64/Utils/AArch64SMEAttributes.cpp

void SMEAttrs::set(unsigned M, bool Enable) {
   else
     Bitmask &= ~M;
 
   assert(!(hasStreamingInterface() && hasStreamingCompatibleInterface()) &&
          "SM_Enabled and SM_Compatible are mutually exclusive");
   assert(!(hasNewZABody() && hasSharedZAInterface()) &&
          "ZA_New and ZA_Shared are mutually exclusive");
   assert(!(hasNewZABody() && preservesZA()) &&
          "ZA_New and ZA_Preserved are mutually exclusive");
 }
 
 SMEAttrs::SMEAttrs(const CallBase &CB) {
   *this = SMEAttrs(CB.getAttributes());
-  if (auto *F = CB.getCalledFunction())
+  if (auto *F = CB.getCalledFunction()) {
     set(SMEAttrs(*F).Bitmask);
+    StringRef FuncName = F->getName();
+    if (FuncName == "__arm_tpidr2_save" || FuncName == "__arm_sme_state" ||
+        FuncName == "__arm_tpidr2_restore")
+      Bitmask |= ZA_NoLazySave;
+  }
 }
 
 SMEAttrs::SMEAttrs(const AttributeList &Attrs) {
   Bitmask = 0;
   if (Attrs.hasFnAttr("aarch64_pstate_sm_enabled"))
     Bitmask |= SM_Enabled;
   if (Attrs.hasFnAttr("aarch64_pstate_sm_compatible"))
     Bitmask |= SM_Compatible;

llvm/test/CodeGen/AArch64/sme-lazy-save-call.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=aarch64 -mattr=+sme < %s | FileCheck %s
 
 declare void @private_za_callee()
+declare void @private_za_preserved_callee() #0
 declare float @llvm.cos.f32(float)
 
 ; Test lazy-save mechanism for a single callee.
 define void @test_lazy_save_1_callee() nounwind "aarch64_pstate_za_shared" {
 ; CHECK-LABEL: test_lazy_save_1_callee:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    stp x29, x30, [sp, #-16]! // 16-byte Folded Spill
 ; CHECK-NEXT:    mov x29, sp

 ; CHECK-NEXT:    ldp d9, d8, [sp, #48] // 16-byte Folded Reload
 ; CHECK-NEXT:    ldp d11, d10, [sp, #32] // 16-byte Folded Reload
 ; CHECK-NEXT:    ldp d13, d12, [sp, #16] // 16-byte Folded Reload
 ; CHECK-NEXT:    ldp d15, d14, [sp], #96 // 16-byte Folded Reload
 ; CHECK-NEXT:    ret
   call void @private_za_callee()
   ret void
 }
+
+
+; Test lazy-save mechanism for a single callee with aarch64_pstate_za_preserved.
+define void @foo() nounwind "aarch64_pstate_za_shared" {
+; CHECK-LABEL: foo:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    stp x29, x30, [sp, #-16]! // 16-byte Folded Spill
+; CHECK-NEXT:    mov x29, sp
+; CHECK-NEXT:    sub sp, sp, #16
+; CHECK-NEXT:    rdsvl x8, #1
+; CHECK-NEXT:    mov x9, sp
+; CHECK-NEXT:    msub x8, x8, x8, x9
+; CHECK-NEXT:    mov sp, x8
+; CHECK-NEXT:    stur x8, [x29, #-16]
+; CHECK-NEXT:    sub x8, x29, #16
+; CHECK-NEXT:    sturh wzr, [x29, #-8]
+; CHECK-NEXT:    msr TPIDR2_EL0, x8
+; CHECK-NEXT:    bl private_za_preserved_callee
+; CHECK-NEXT:    msr TPIDR2_EL0, xzr
+; CHECK-NEXT:    mov sp, x29
+; CHECK-NEXT:    ldp x29, x30, [sp], #16 // 16-byte Folded Reload
+; CHECK-NEXT:    ret
+  call void @private_za_preserved_callee()
+  ret void
+}
+
+attributes #0 = {
+  "aarch64_pstate_za_preserved"
+}