[AArch64][SME] Set fn attributes correctly on __arm_tpidr2_save call.

[sdesmalen]

Summary

Changing the way the attribute is created changes the way the attribute
is added. Now it seems to be correctly added as a function attribute, whereas
before it was not. This led ISel ignoring attributes and setting up a lazy-save
buffer for the call to __arm_tpidr2_save.

This patch also marks the intrinsic call as 'preserves_za' which together with the
above prevents the code-generator from setting up a lazy-save buffer.

[sdesmalen]

Gentle ping

[aemerson]

llvm/lib/Target/AArch64/SMEABIPass.cpp

68

No objections to the patch but I'm a bit confused why this didn't work? Seems like a footgun to me.

[sdesmalen]

llvm/lib/Target/AArch64/SMEABIPass.cpp

68

I was similarly surprised. Perhaps there is some subtle difference that by adding it with the above interface, the attribute was somehow not a function attribute?
But then I don't know what kind of attribute it would be, or what the attribute applies to, given that it's still attached to a function.

llvm/lib/Target/AArch64/SMEABIPass.cpp

 //===----------------------------------------------------------------------===//
 // Utility functions
 //===----------------------------------------------------------------------===//
 
 // Utility function to emit a call to __arm_tpidr2_save and clear TPIDR2_EL0.
 void emitTPIDR2Save(Module *M, IRBuilder<> &Builder) {
   auto *TPIDR2SaveTy =
       FunctionType::get(Builder.getVoidTy(), {}, /*IsVarArgs=*/false);
-
   auto Attrs =
-      AttributeList::get(M->getContext(), 0, {"aarch64_pstate_sm_compatible"});
+      AttributeList()
+          .addFnAttribute(M->getContext(), "aarch64_pstate_sm_compatible")
+          .addFnAttribute(M->getContext(), "aarch64_pstate_za_preserved");
   FunctionCallee Callee =
       M->getOrInsertFunction("__arm_tpidr2_save", TPIDR2SaveTy, Attrs);
-  Builder.CreateCall(Callee);
+  CallInst *Call = Builder.CreateCall(Callee);
+  Call->setCallingConv(
+      CallingConv::AArch64_SME_ABI_Support_Routines_PreserveMost_From_X0);
 
   // A save to TPIDR2 should be followed by clearing TPIDR2_EL0.
   Function *WriteIntr =
       Intrinsic::getDeclaration(M, Intrinsic::aarch64_sme_set_tpidr2);
   Builder.CreateCall(WriteIntr->getFunctionType(), WriteIntr,
                      Builder.getInt64(0));
 }
 

llvm/test/CodeGen/AArch64/sme-new-za-function.ll

 ; RUN: opt -S -mtriple=aarch64-linux-gnu -aarch64-sme-abi %s | FileCheck %s
 ; RUN: opt -S -mtriple=aarch64-linux-gnu -aarch64-sme-abi -aarch64-sme-abi %s | FileCheck %s
 
 declare void @shared_za_callee() "aarch64_pstate_za_shared"
 
 define void @private_za() "aarch64_pstate_za_new" {
 ; CHECK-LABEL: @private_za(
 ; CHECK-NEXT:  prelude:
 ; CHECK-NEXT:    [[TPIDR2:%.*]] = call i64 @llvm.aarch64.sme.get.tpidr2()
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i64 [[TPIDR2]], 0
 ; CHECK-NEXT:    br i1 [[CMP]], label [[SAVE_ZA:%.*]], label [[TMP0:%.*]]
 ; CHECK:       save.za:
-; CHECK-NEXT:    call void @__arm_tpidr2_save()
+; CHECK-NEXT:    call aarch64_sme_preservemost_from_x0 void @__arm_tpidr2_save()
 ; CHECK-NEXT:    call void @llvm.aarch64.sme.set.tpidr2(i64 0)
 ; CHECK-NEXT:    br label [[TMP0]]
 ; CHECK:       0:
 ; CHECK-NEXT:    call void @llvm.aarch64.sme.za.enable()
 ; CHECK-NEXT:    call void @shared_za_callee()
 ; CHECK-NEXT:    call void @llvm.aarch64.sme.za.disable()
 ; CHECK-NEXT:    ret void
 ;
   call void @shared_za_callee()
   ret void
 }
 
 define i32 @private_za_multiple_exit(i32 %a, i32 %b, i64 %cond) "aarch64_pstate_za_new" {
 ; CHECK-LABEL: @private_za_multiple_exit(
 ; CHECK-NEXT:  prelude:
 ; CHECK-NEXT:    [[TPIDR2:%.*]] = call i64 @llvm.aarch64.sme.get.tpidr2()
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i64 [[TPIDR2]], 0
 ; CHECK-NEXT:    br i1 [[CMP]], label [[SAVE_ZA:%.*]], label [[ENTRY:%.*]]
 ; CHECK:       save.za:
-; CHECK-NEXT:    call void @__arm_tpidr2_save()
+; CHECK-NEXT:    call aarch64_sme_preservemost_from_x0 void @__arm_tpidr2_save()
 ; CHECK-NEXT:    call void @llvm.aarch64.sme.set.tpidr2(i64 0)
 ; CHECK-NEXT:    br label [[ENTRY]]
 ; CHECK:       entry:
 ; CHECK-NEXT:    call void @llvm.aarch64.sme.za.enable()
 ; CHECK-NEXT:    [[TOBOOL:%.*]] = icmp eq i64 [[COND:%.*]], 1
 ; CHECK-NEXT:    br i1 [[TOBOOL]], label [[IF_ELSE:%.*]], label [[IF_END:%.*]]
 ; CHECK:       if.else:
 ; CHECK-NEXT:    [[ADD:%.*]] = add i32 [[A:%.*]], [[B:%.*]]
if.else:
   %add = add i32 %a, %b
   ret i32 %add
 
 if.end:
   %sub = sub i32 %a, %b
   ret i32 %sub
 }
 
-; CHECK: declare "aarch64_pstate_sm_compatible" void @__arm_tpidr2_save()
+; CHECK: declare void @__arm_tpidr2_save() #[[ATTR:[0-9]+]]
+; CHECK: attributes #[[ATTR]] = { "aarch64_pstate_sm_compatible" "aarch64_pstate_za_preserved" }