Add more ZA modes #77361

Leporacanthicus · 2024-01-08T19:14:21Z

This adds support for preserving and sharing ZA Modes to the MLIR infrastructure. The functionality already exists in LLVM, so just "linking the pieces together".

llvmbot · 2024-01-08T19:14:51Z

@llvm/pr-subscribers-mlir-llvm

@llvm/pr-subscribers-mlir-sme

Author: Mats Petersson (Leporacanthicus)

Changes

This adds support for preserving and sharing ZA Modes to the MLIR infrastructure. The functionality already exists in LLVM, so just "linking the pieces together".

Full diff: https://github.com/llvm/llvm-project/pull/77361.diff

5 Files Affected:

(modified) mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td (+13-3)
(modified) mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td (+2)
(modified) mlir/lib/Target/LLVMIR/ModuleImport.cpp (+6)
(modified) mlir/test/Dialect/ArmSME/enable-arm-za.mlir (+8)
(modified) mlir/test/Target/LLVMIR/Import/function-attributes.ll (+21)

diff --git a/mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td b/mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td
index 4266ac5b0c8cf6..b33778967f6ac4 100644
--- a/mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td
@@ -28,13 +28,15 @@ 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
 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 +81,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
index 9e65898154bd65..24541792d0c773 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -1413,6 +1413,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
index 905405e9398820..7c8d90316ff6ef 100644
--- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
@@ -1639,6 +1639,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"),
@@ -1715,6 +1717,10 @@ void ModuleImport::processFunctionAttributes(llvm::Function *func,
 
   if (func->hasFnAttribute("aarch64_pstate_za_new"))
     funcOp.setArmNewZa(true);
+  if (func->hasFnAttribute("aarch64_pstate_za_preserved"))
+    funcOp.setArmPreservesZa(true);
+  if (func->hasFnAttribute("aarch64_pstate_za_shared"))
+    funcOp.setArmSharedZa(true);
 
   llvm::Attribute attr = func->getFnAttribute(llvm::Attribute::VScaleRange);
   if (attr.isValid()) {
diff --git a/mlir/test/Dialect/ArmSME/enable-arm-za.mlir b/mlir/test/Dialect/ArmSME/enable-arm-za.mlir
index ba650b031e6110..b21d22833ccf86 100644
--- a/mlir/test/Dialect/ArmSME/enable-arm-za.mlir
+++ b/mlir/test/Dialect/ArmSME/enable-arm-za.mlir
@@ -1,5 +1,7 @@
 // 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
@@ -7,10 +9,16 @@ 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_conserves_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
index f76e7293809628..28ca4343062ca6 100644
--- 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, arm_streaming}
+define void @arm_new_za_func() "aarch64_pstate_sm_enabled" "aarch64_pstate_za_new" {
+  ret void
+}
+
+
+; CHECK-LABEL: @arm_preserves_za_func
+; CHECK-SAME: attributes {arm_preserves_za, arm_streaming}
+define void @arm_preserves_za_func() "aarch64_pstate_sm_enabled" "aarch64_pstate_za_preserved" {
+  ret void
+}
+
+; CHECK-LABEL: @arm_shared_za_func
+; CHECK-SAME: attributes {arm_shared_za, arm_streaming}
+define void @arm_shared_za_func() "aarch64_pstate_sm_enabled" "aarch64_pstate_za_shared" {
+  ret void
+}
+
+// -----
+
 ; CHECK-LABEL: @section_func
 ; CHECK-SAME: attributes {section = ".section.name"}
 define void @section_func() section ".section.name" {

MacDue · 2024-01-09T10:32:30Z

mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td

@@ -28,13 +28,15 @@ 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


Keep the link to the ACLE

MacDue · 2024-01-09T10:33:46Z

mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td

+      // A function that preserves ZA state.
+      I32EnumAttrCase<"PreservesZA", 2, "arm_preserves_za">,


It's a little odd to have 'PreservesZA' as a ZA mode, as it's an attribute that can be used alongside arm_new_za or arm_shared_za, not an independent mode.

You are referring to C semantics (as per https://arm-software.github.io/acle/main/acle.html#sme-attributes-relating-to-za). However, we are not required to follow that in MLIR and can focus on LLVM instead (i.e. https://llvm.org/docs/AArch64SME.html).

That's a very good point nonetheless that will be very relevant within Flang. Not sure we could capture it here.

banach-space

This adds support for preserving and sharing ZA Modes to the MLIR infrastructure. The functionality already exists in LLVM, so just "linking the pieces together".

Mats, I'd appreciate if we could be a bit more stricter with the documentation. Could you add links with definitions of these attributes? (I guess both ACLE and LLVM IR are relevant). And could you specify which attributes are already supported and which are added in this patch? Thank you :)

banach-space · 2024-01-09T12:48:05Z

mlir/test/Dialect/ArmSME/enable-arm-za.mlir

 // 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_conserves_za


That's unlikely what you meant ;-)

banach-space · 2024-01-09T12:50:26Z

mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td

+      // A function that preserves ZA state.
+      I32EnumAttrCase<"PreservesZA", 2, "arm_preserves_za">,


You are referring to C semantics (as per https://arm-software.github.io/acle/main/acle.html#sme-attributes-relating-to-za). However, we are not required to follow that in MLIR and can focus on LLVM instead (i.e. https://llvm.org/docs/AArch64SME.html).

That's a very good point nonetheless that will be very relevant within Flang. Not sure we could capture it here.

Leporacanthicus · 2024-01-09T21:20:06Z

I have pushed an update that I think resolves the comments provided. It includes a comment explaining that the legality of any attributes used is up to the higher level user of these attributes. I think that makes the sense.

Thanks for the comments, valuable as always.

banach-space

LGTM, but please wait for @MacDue to approve before landing :)

banach-space · 2024-01-10T14:17:56Z

mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td

 // 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 languages may restrict or enforce how these attributes


[nit]

Suggested change

// Various languages may restrict or enforce how these attributes

// Various frontends (e.g. Flang) that build on top of this may restrict or enforce how these attributes

MacDue

LGTM, just one thing:

mlir/lib/Target/LLVMIR/ModuleImport.cpp

MacDue · 2024-01-10T18:24:44Z

mlir/lib/Target/LLVMIR/ModuleImport.cpp

@@ -1715,6 +1717,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 eiether NewZA or SharedZA.


Typo: eiether -> either

c-rhodes

the attributes also need adding to mlir/lib/Target/LLVMIR/ModuleTranslation.cpp with corresponding tests in mlir/test/Target/LLVMIR/llvmir.mlir

c-rhodes · 2024-01-11T08:59:43Z

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

+; CHECK-SAME: attributes {arm_new_za, arm_streaming}
+define void @arm_new_za_func() "aarch64_pstate_sm_enabled" "aarch64_pstate_za_new" {
+  ret void
+}
+
+
+; CHECK-LABEL: @arm_preserves_za_func
+; CHECK-SAME: attributes {arm_preserves_za, arm_streaming}
+define void @arm_preserves_za_func() "aarch64_pstate_sm_enabled" "aarch64_pstate_za_preserved" {
+  ret void
+}
+
+; CHECK-LABEL: @arm_shared_za_func
+; CHECK-SAME: attributes {arm_shared_za, arm_streaming}
+define void @arm_shared_za_func() "aarch64_pstate_sm_enabled" "aarch64_pstate_za_shared" {
+  ret void
+}


for the purpose of these tests "aarch64_pstate_sm_enabled" isn't necessary

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

c-rhodes

LGTM cheers

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

definelicht

While modifying llvmir.mlir I noticed that there are some CHECKs here that are missing a : and thus aren't actually tested 😬 Could you fix this in a followup?

Please also add splits between the tests!

definelicht · 2024-06-11T13:35:46Z

mlir/test/Target/LLVMIR/llvmir.mlir

+llvm.func @new_za_func() attributes {arm_new_za} {
+  llvm.return
+}
+// CHECK #[[ATTR]] = { "aarch64_pstate_za_new" }


There is a : missing, so this is not actually checked 🙂

definelicht · 2024-06-11T13:35:51Z

mlir/test/Target/LLVMIR/llvmir.mlir

+llvm.func @shared_za_func() attributes {arm_shared_za } {
+  llvm.return
+}
+// CHECK #[[ATTR]] = { "aarch64_pstate_za_shared" }


There is a : missing, so this is not actually checked 🙂

definelicht · 2024-06-11T13:35:55Z

mlir/test/Target/LLVMIR/llvmir.mlir

+llvm.func @preserves_za_func() attributes {arm_preserves_za} {
+  llvm.return
+}
+// CHECK #[[ATTR]] = { "aarch64_pstate_za_preserved" }


There is a : missing, so this is not actually checked 🙂

Dinistro · 2024-06-11T13:42:32Z

@definelicht This should be addressed in #93476

Leporacanthicus requested review from banach-space, dcaballe and nicolasvasilache as code owners January 8, 2024 19:14

llvmbot added mlir:llvm mlir mlir:sme labels Jan 8, 2024

banach-space requested review from MacDue and c-rhodes January 8, 2024 20:24

MacDue reviewed Jan 9, 2024

View reviewed changes

banach-space reviewed Jan 9, 2024

View reviewed changes

banach-space approved these changes Jan 10, 2024

View reviewed changes

MacDue approved these changes Jan 10, 2024

View reviewed changes

mlir/lib/Target/LLVMIR/ModuleImport.cpp Outdated Show resolved Hide resolved

MacDue reviewed Jan 10, 2024

View reviewed changes

c-rhodes reviewed Jan 11, 2024

View reviewed changes

Leporacanthicus force-pushed the more-za-modes branch from 45cfde6 to 89686b4 Compare January 11, 2024 12:02

Support LLVMIR target and add tests

3869805

c-rhodes approved these changes Jan 11, 2024

View reviewed changes

Leporacanthicus merged commit 21e1bf2 into llvm:main Jan 11, 2024

definelicht reviewed Jun 11, 2024

View reviewed changes

		// A function that preserves ZA state.
		I32EnumAttrCase<"PreservesZA", 2, "arm_preserves_za">,

	// Various languages may restrict or enforce how these attributes
	// Various frontends (e.g. Flang) that build on top of this may restrict or enforce how these attributes

Add more ZA modes #77361

Add more ZA modes #77361

Uh oh!

Conversation

Leporacanthicus commented Jan 8, 2024

Uh oh!

llvmbot commented Jan 8, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Leporacanthicus commented Jan 9, 2024

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MacDue left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

MacDue Jan 10, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

c-rhodes left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

c-rhodes left a comment

Choose a reason for hiding this comment

Uh oh!

definelicht left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Dinistro commented Jun 11, 2024

Uh oh!

Uh oh!

llvmbot commented Jan 8, 2024 •

edited

Loading

MacDue Jan 10, 2024 •

edited

Loading

definelicht left a comment •

edited

Loading