[Mlir-commits] [mlir] Add more ZA modes (PR #77361)
Mats Petersson
llvmlistbot at llvm.org
Tue Jan 9 13:17:47 PST 2024
https://github.com/Leporacanthicus updated https://github.com/llvm/llvm-project/pull/77361
>From bf2ab497297fdbfa2f86d12715ae72b7328b90cf Mon Sep 17 00:00:00 2001
From: Mats Petersson <mats.petersson at arm.com>
Date: Thu, 28 Dec 2023 14:07:49 +0000
Subject: [PATCH 1/2] Add more ZA modes
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".
---
.../mlir/Dialect/ArmSME/Transforms/Passes.td | 16 +++++++++++---
mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td | 2 ++
mlir/lib/Target/LLVMIR/ModuleImport.cpp | 6 ++++++
mlir/test/Dialect/ArmSME/enable-arm-za.mlir | 8 +++++++
.../LLVMIR/Import/function-attributes.ll | 21 +++++++++++++++++++
5 files changed, 50 insertions(+), 3 deletions(-)
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" {
>From ea2f7f31d8d7eed59f9fd34126b893f15cd63559 Mon Sep 17 00:00:00 2001
From: Mats Petersson <mats.petersson at arm.com>
Date: Tue, 9 Jan 2024 21:14:58 +0000
Subject: [PATCH 2/2] Insert this into the final commit message:
Adds the arm_shared_za and arm_preserves_za attributes to the existing
arm_new_za attribute.
For more details see:
https://arm-software.github.io/acle/main/acle.html#sme-attributes-relating-to-za
---
mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td | 9 +++++++++
mlir/test/Dialect/ArmSME/enable-arm-za.mlir | 2 +-
2 files changed, 10 insertions(+), 1 deletion(-)
diff --git a/mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td b/mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td
index b33778967f6ac4..26e60142443a2c 100644
--- a/mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td
@@ -28,6 +28,15 @@ def ArmStreamingMode : I32EnumAttr<"ArmStreamingMode", "Armv9 Streaming SVE mode
let genSpecializedAttr = 0;
}
+// 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
+// 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">,
diff --git a/mlir/test/Dialect/ArmSME/enable-arm-za.mlir b/mlir/test/Dialect/ArmSME/enable-arm-za.mlir
index b21d22833ccf86..0aa00f75c3a567 100644
--- a/mlir/test/Dialect/ArmSME/enable-arm-za.mlir
+++ b/mlir/test/Dialect/ArmSME/enable-arm-za.mlir
@@ -20,5 +20,5 @@ func.func private @declaration()
// 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
+// NO-ARM-STREAMING-NOT: arm_preserves_za
func.func @arm_new_za() { return }
More information about the Mlir-commits
mailing list