Add more ZA modes (#77361)

Leporacanthicus · web-flow · commit 21e1bf2d0001 · 2024-01-11T18:49:52.000Z
Add more ZA modes Adds the arm_shared_za and arm_preserves_za attributes to the existing arm_new_za attribute. The functionality already exists in LLVM, so just "linking the pieces together". For more details see: https://arm-software.github.io/acle/main/acle.html#sme-attributes-relating-to-za
diff --git a/mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td b/mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td
@@ -28,13 +28,25 @@ def ArmStreamingMode : I32EnumAttr<"ArmStreamingMode", "Armv9 Streaming SVE mode
   let genSpecializedAttr = 0;
 }
 
-// TODO: Add other ZA modes.
 // https://arm-software.github.io/acle/main/acle.html#sme-attributes-relating-to-za
+// See also the LLVM definitions: https://llvm.org/docs/AArch64SME.html
+//
+// Various frontends (e.g. Flang) that build on top of this may restrict or
+// enforce how these attributes are used, both individually and in terms of
+// combinations that are allowed.
+//
+// The MLIR interface here does not make any attempt to perform any checking,
+// it is up to the higher level to ensure that these attributes are used in a
+// way that both makes sense and is legal according to the Arm architecture.
 def ArmZaMode : I32EnumAttr<"ArmZaMode", "Armv9 ZA storage mode",
     [
       I32EnumAttrCase<"Disabled", 0, "disabled">,
       // A function's ZA state is created on entry and destroyed on exit.
       I32EnumAttrCase<"NewZA", 1, "arm_new_za">,
+      // A function that preserves ZA state.
+      I32EnumAttrCase<"PreservesZA", 2, "arm_preserves_za">,
+      // A function that uses ZA state as input and/or output
+      I32EnumAttrCase<"SharedZA", 3, "arm_shared_za">,
     ]>{
   let cppNamespace = "mlir::arm_sme";
   let genSpecializedAttr = 0;
@@ -79,7 +91,15 @@ def EnableArmStreaming
                  clEnumValN(mlir::arm_sme::ArmZaMode::NewZA,
                             "new-za",
                             "The function has ZA state. The ZA state is "
-                            "created on entry and destroyed on exit.")
+                            "created on entry and destroyed on exit."),
+                 clEnumValN(mlir::arm_sme::ArmZaMode::PreservesZA,
+                            "preserves-za",
+                            "The function preserves ZA state. The ZA state is "
+                            "saved on entry and restored on exit."),
+                 clEnumValN(mlir::arm_sme::ArmZaMode::SharedZA,
+                            "shared-za",
+                            "The function uses ZA state. The ZA state may "
+                            "be used for input and/or output.")
            )}]>,
     Option<"onlyIfRequiredByOps", "only-if-required-by-ops", "bool",
            /*default=*/"false",
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -1420,6 +1420,8 @@ def LLVM_LLVMFuncOp : LLVM_Op<"func", [
     OptionalAttr<UnitAttr>:$arm_locally_streaming,
     OptionalAttr<UnitAttr>:$arm_streaming_compatible,
     OptionalAttr<UnitAttr>:$arm_new_za,
+    OptionalAttr<UnitAttr>:$arm_preserves_za,
+    OptionalAttr<UnitAttr>:$arm_shared_za,
     OptionalAttr<StrAttr>:$section,
     OptionalAttr<UnnamedAddr>:$unnamed_addr,
     OptionalAttr<I64Attr>:$alignment,
diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
@@ -1641,6 +1641,8 @@ static constexpr std::array ExplicitAttributes{
     StringLiteral("aarch64_pstate_sm_body"),
     StringLiteral("aarch64_pstate_sm_compatible"),
     StringLiteral("aarch64_pstate_za_new"),
+    StringLiteral("aarch64_pstate_za_preserved"),
+    StringLiteral("aarch64_pstate_za_shared"),
     StringLiteral("vscale_range"),
     StringLiteral("frame-pointer"),
     StringLiteral("target-features"),
@@ -1717,6 +1719,11 @@ void ModuleImport::processFunctionAttributes(llvm::Function *func,
 
   if (func->hasFnAttribute("aarch64_pstate_za_new"))
     funcOp.setArmNewZa(true);
+  else if (func->hasFnAttribute("aarch64_pstate_za_shared"))
+    funcOp.setArmSharedZa(true);
+  // PreservedZA can be used with either NewZA or SharedZA.
+  if (func->hasFnAttribute("aarch64_pstate_za_preserved"))
+    funcOp.setArmPreservesZa(true);
 
   llvm::Attribute attr = func->getFnAttribute(llvm::Attribute::VScaleRange);
   if (attr.isValid()) {
diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -1102,6 +1102,10 @@ LogicalResult ModuleTranslation::convertOneFunction(LLVMFuncOp func) {
 
   if (func.getArmNewZa())
     llvmFunc->addFnAttr("aarch64_pstate_za_new");
+  else if (func.getArmSharedZa())
+    llvmFunc->addFnAttr("aarch64_pstate_za_shared");
+  if (func.getArmPreservesZa())
+    llvmFunc->addFnAttr("aarch64_pstate_za_preserved");
 
   if (auto targetFeatures = func.getTargetFeatures())
     llvmFunc->addFnAttr("target-features", targetFeatures->getFeaturesString());
diff --git a/mlir/test/Dialect/ArmSME/enable-arm-za.mlir b/mlir/test/Dialect/ArmSME/enable-arm-za.mlir
@@ -1,16 +1,24 @@
 // RUN: mlir-opt %s -enable-arm-streaming=za-mode=new-za -convert-arm-sme-to-llvm | FileCheck %s -check-prefix=ENABLE-ZA
 // RUN: mlir-opt %s -enable-arm-streaming -convert-arm-sme-to-llvm | FileCheck %s -check-prefix=DISABLE-ZA
+// RUN: mlir-opt %s -enable-arm-streaming=za-mode=shared-za -convert-arm-sme-to-llvm | FileCheck %s -check-prefix=SHARED-ZA
+// RUN: mlir-opt %s -enable-arm-streaming=za-mode=preserves-za -convert-arm-sme-to-llvm | FileCheck %s -check-prefix=PRESERVES-ZA
 // RUN: mlir-opt %s -convert-arm-sme-to-llvm | FileCheck %s -check-prefix=NO-ARM-STREAMING
 
 // CHECK-LABEL: @declaration
 func.func private @declaration()
 
 // ENABLE-ZA-LABEL: @arm_new_za
 // ENABLE-ZA-SAME: attributes {arm_new_za, arm_streaming}
+// SHARED-ZA-LABEL: @arm_new_za
+// SHARED-ZA-SAME: attributes {arm_shared_za, arm_streaming}
+// PRESERVES-ZA-LABEL: @arm_new_za
+// PRESERVES-ZA-SAME: attributes {arm_preserves_za, arm_streaming}
 // DISABLE-ZA-LABEL: @arm_new_za
 // DISABLE-ZA-NOT: arm_new_za
 // DISABLE-ZA-SAME: attributes {arm_streaming}
 // NO-ARM-STREAMING-LABEL: @arm_new_za
 // NO-ARM-STREAMING-NOT: arm_new_za
 // NO-ARM-STREAMING-NOT: arm_streaming
+// NO-ARM-STREAMING-NOT: arm_shared_za
+// NO-ARM-STREAMING-NOT: arm_preserves_za
 func.func @arm_new_za() { return }
diff --git a/mlir/test/Target/LLVMIR/Import/function-attributes.ll b/mlir/test/Target/LLVMIR/Import/function-attributes.ll
@@ -220,6 +220,27 @@ define void @streaming_compatible_func() "aarch64_pstate_sm_compatible" {
 
 // -----
 
+; CHECK-LABEL: @arm_new_za_func
+; CHECK-SAME: attributes {arm_new_za}
+define void @arm_new_za_func() "aarch64_pstate_za_new" {
+  ret void
+}
+
+
+; CHECK-LABEL: @arm_preserves_za_func
+; CHECK-SAME: attributes {arm_preserves_za}
+define void @arm_preserves_za_func() "aarch64_pstate_za_preserved" {
+  ret void
+}
+
+; CHECK-LABEL: @arm_shared_za_func
+; CHECK-SAME: attributes {arm_shared_za}
+define void @arm_shared_za_func() "aarch64_pstate_za_shared" {
+  ret void
+}
+
+// -----
+
 ; CHECK-LABEL: @section_func
 ; CHECK-SAME: attributes {section = ".section.name"}
 define void @section_func() section ".section.name" {
diff --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir
@@ -2330,6 +2330,29 @@ llvm.func @streaming_compatible_func() attributes {arm_streaming_compatible} {
 
 // -----
 
+// CHECK-LABEL: @new_za_func
+// CHECK: #[[ATTR:[0-9]*]]
+llvm.func @new_za_func() attributes {arm_new_za} {
+  llvm.return
+}
+// CHECK #[[ATTR]] = { "aarch64_pstate_za_new" }
+
+// CHECK-LABEL: @shared_za_func
+// CHECK: #[[ATTR:[0-9]*]]
+llvm.func @shared_za_func() attributes {arm_shared_za } {
+  llvm.return
+}
+// CHECK #[[ATTR]] = { "aarch64_pstate_za_shared" }
+
+// CHECK-LABEL: @preserves_za_func
+// CHECK: #[[ATTR:[0-9]*]]
+llvm.func @preserves_za_func() attributes {arm_preserves_za} {
+  llvm.return
+}
+// CHECK #[[ATTR]] = { "aarch64_pstate_za_preserved" }
+
+// -----
+
 //
 // Zero-initialize operation.
 //
