Skip to content

Add more ZA modes #77361

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jan 11, 2024
Merged

Add more ZA modes #77361

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
89686b4
Add more ZA modes
Leporacanthicus Dec 28, 2023
3869805
Support LLVMIR target and add tests
Leporacanthicus Jan 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 22 additions & 2 deletions mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Keep the link to the ACLE

// 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">,
Comment on lines +46 to +47
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

// 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;
Expand Down Expand Up @@ -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",
Expand Down
2 changes: 2 additions & 0 deletions mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
7 changes: 7 additions & 0 deletions mlir/lib/Target/LLVMIR/ModuleImport.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"),
Expand Down Expand Up @@ -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 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()) {
Expand Down
4 changes: 4 additions & 0 deletions mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1081,6 +1081,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());
Expand Down
8 changes: 8 additions & 0 deletions mlir/test/Dialect/ArmSME/enable-arm-za.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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 }
21 changes: 21 additions & 0 deletions mlir/test/Target/LLVMIR/Import/function-attributes.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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" {
Expand Down
23 changes: 23 additions & 0 deletions mlir/test/Target/LLVMIR/llvmir.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2319,6 +2319,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" }
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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


// 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" }
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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


// 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" }
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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


// -----

//
// Zero-initialize operation.
//
Expand Down