[mlir][llvm] Add aarch64_pstate_za_new LLVM function attribute

[c-rhodes]

Summary

This patch adds an optional attribute 'aarch64_pstate_za_new' to the
'llvm.func' op for the Armv9 Scalable Matrix Extension (SME) [1]. It
maps 1-1 with the LLVM function attribute of the same name [2] and
enables the SME ZA storage array.

The attribute has the following behaviour [3]:

The function’s ZA state is created on entry to the function and
destroyed on return from the function. That is, the function does not
use ZA to receive data from callers or to pass data back to callers.

ZA state being private to a function vs shared with the caller will be
simpler to begin enabling SME in MLIR.

[1] https://developer.arm.com/documentation/ddi0616/aa
[2] https://llvm.org/docs/AArch64SME.html#introduction
[3] https://arm-software.github.io/acle/main/acle.html#aarch64_pstate_za_new

[c-rhodes]

Patch was missing context.

[awarzynski]

LGTM, thanks!

[c-rhodes]

Use aarch64_pstate_za_new name everywhere rather than arm_za_new. Better to more closely match LLVM than C/C++ ACLE and also avoid confusion with Clang (thanks @sdesmalen for pointing this out).

[dcaballe]

Thanks!

[c-rhodes]

The backend emits calls to the following SME support routines [1]:

• __arm_tpidr2_restore
• __arm_tpidr2_save

for this attribute. There's currently no public implementation of these routines and it's not yet clear where they will live. I think for now ZA will have to be enabled with the:

• llvm.aarch64.sme.za.enable
• llvm.aarch64.sme.za.disable

intrinsics [2] manually. Currently looking into this.

[1] https://github.com/ARM-software/abi-aa/blob/main/aapcs64/aapcs64.rst#sme-support-routines
[2] https://github.com/llvm/llvm-project/blob/main/llvm/test/CodeGen/AArch64/sme-toggle-pstateza.ll

[c-rhodes]

Abandoning this in favour of D153050.

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

let arguments = (ins
     OptionalAttr<StrAttr>:$garbageCollector,
     OptionalAttr<ArrayAttr>:$passthrough,
     OptionalAttr<DictArrayAttr>:$arg_attrs,
     OptionalAttr<DictArrayAttr>:$res_attrs,
     OptionalAttr<I64Attr>:$function_entry_count,
     OptionalAttr<LLVM_MemoryEffectsAttr>:$memory,
     DefaultValuedAttr<Visibility, "mlir::LLVM::Visibility::Default">:$visibility_,
     OptionalAttr<UnitAttr>:$arm_streaming,
-    OptionalAttr<UnitAttr>:$arm_locally_streaming
+    OptionalAttr<UnitAttr>:$arm_locally_streaming,
+    OptionalAttr<UnitAttr>:$aarch64_pstate_za_new
   );
 
   let regions = (region AnyRegion:$body);
 
   let skipDefaultBuilders = 1;
 
   let builders = [
     OpBuilder<(ins "StringRef":$name, "Type":$type,

mlir/lib/Target/LLVMIR/ModuleImport.cpp

for (llvm::Attribute attr : funcAttrs) {
 
     StringRef attrName;
     if (attr.isStringAttribute())
       attrName = attr.getKindAsString();
     else
       attrName = llvm::Attribute::getNameFromAttrKind(attr.getKindAsEnum());
     auto keyAttr = StringAttr::get(context, attrName);
 
-    // Skip the aarch64_pstate_sm_<body|enabled> since the LLVMFuncOp has an
-    // explicit attribute.
+    // Skip attributes the LLVMFuncOp has an explicit attribute for.
     if (attrName == "aarch64_pstate_sm_enabled" ||
-        attrName == "aarch64_pstate_sm_body")
+        attrName == "aarch64_pstate_sm_body" ||
+        attrName == "aarch64_pstate_za_new")
       continue;
 
     if (attr.isStringAttribute()) {
       StringRef val = attr.getValueAsString();
       if (val.empty()) {
         passthroughs.push_back(keyAttr);
         continue;
       }
void ModuleImport::processFunctionAttributes(llvm::Function *func,
                                              LLVMFuncOp funcOp) {
   processMemoryEffects(func, funcOp);
   processPassthroughAttrs(func, funcOp);
 
   if (func->hasFnAttribute("aarch64_pstate_sm_enabled"))
     funcOp.setArmStreaming(true);
   else if (func->hasFnAttribute("aarch64_pstate_sm_body"))
     funcOp.setArmLocallyStreaming(true);
+
+  if (func->hasFnAttribute("aarch64_pstate_za_new"))
+    funcOp.setAarch64PstateZaNew(true);
 }
 
 DictionaryAttr
 ModuleImport::convertParameterAttribute(llvm::AttributeSet llvmParamAttrs,
                                         OpBuilder &builder) {
   SmallVector<NamedAttribute> paramAttrs;
   for (auto [llvmKind, mlirName] : getAttrKindToNameMapping()) {
     auto llvmAttr = llvmParamAttrs.getAttribute(llvmKind);

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

LogicalResult ModuleTranslation::convertOneFunction(LLVMFuncOp func) {
   if (auto gc = func.getGarbageCollector())
     llvmFunc->setGC(gc->str());
 
   if (auto armStreaming = func.getArmStreaming())
     llvmFunc->addFnAttr("aarch64_pstate_sm_enabled");
   else if (auto armLocallyStreaming = func.getArmLocallyStreaming())
     llvmFunc->addFnAttr("aarch64_pstate_sm_body");
 
+  if (auto armNewZA = func.getAarch64PstateZaNew())
+    llvmFunc->addFnAttr("aarch64_pstate_za_new");
+
   // First, create all blocks so we can jump to them.
   llvm::LLVMContext &llvmContext = llvmFunc->getContext();
   for (auto &bb : func) {
     auto *llvmBB = llvm::BasicBlock::Create(llvmContext);
     llvmBB->insertInto(llvmFunc);
     mapBlock(&bb, llvmBB);
   }
 

mlir/test/Target/LLVMIR/Import/function-attributes.ll

 
 // -----
 
 ; CHECK-LABEL: @arm_locally_streaming_func
 ; CHECK-SAME: attributes {arm_locally_streaming}
 define void @arm_locally_streaming_func() "aarch64_pstate_sm_body" {
   ret void
 }
+
+// -----
+
+; CHECK-LABEL: @arm_private_za_func
+; CHECK-SAME: attributes {aarch64_pstate_za_new}
+define void @arm_private_za_func() "aarch64_pstate_za_new" {
+  ret void
+}

mlir/test/Target/LLVMIR/llvmir.mlir

 
 // CHECK-LABEL: @arm_locally_streaming_func
 // CHECK: #[[ATTR:[0-9]*]]
 llvm.func @arm_locally_streaming_func() attributes {arm_locally_streaming} {
   llvm.return
 }
 
 // CHECK: attributes #[[ATTR]] = { "aarch64_pstate_sm_body" }
+
+// -----
+
+//
+// aarch64_pstate_za_new attribute.
+//
+
+// CHECK-LABEL: @arm_private_za_func
+// CHECK: #[[ATTR:[0-9]*]]
+llvm.func @arm_private_za_func() attributes {aarch64_pstate_za_new} {
+  llvm.return
+}
+
+// CHECK: attributes #[[ATTR]] = { "aarch64_pstate_za_new" }