[Mlir-commits] [mlir] 7d46590 - [mlir][llvm] Add arm_streaming LLVM function attributes
Cullen Rhodes
llvmlistbot at llvm.org
Thu May 25 02:21:44 PDT 2023
Author: Cullen Rhodes
Date: 2023-05-25T09:20:35Z
New Revision: 7d4659095aa535f1f2bbc01a094b86907863b1ce
URL: https://github.com/llvm/llvm-project/commit/7d4659095aa535f1f2bbc01a094b86907863b1ce
DIFF: https://github.com/llvm/llvm-project/commit/7d4659095aa535f1f2bbc01a094b86907863b1ce.diff
LOG: [mlir][llvm] Add arm_streaming LLVM function attributes
This patch adds two optional attributes to 'llvm.func' op for the Armv9
Streaming SVE (SSVE) mode [1] that map 1-1 with LLVM function attributes [2]:
* arm_streaming -> aarch64_pstate_sm_enabled
* arm_locally_streaming -> aarch64_pstate_sm_body
Streaming-mode is part of the interface (ABI) for functions with the
first attribute and it's the responsibility of the caller to manage
PSTATE.SM on entry/exit to functions with this attribute [3]. The LLVM
backend will emit 'smstart sm' / 'smstop sm' [4] around calls to
streaming functions.
In locally streaming functions PSTATE.SM is kept internal and managed by
the callee on entry/exit. The LLVM backend will emit 'smstart sm' /
'smstop sm' in the prologue / epilogue for functions with this
attribute.
The integration test for SSVE has been updated to no longer use the
passthrough mechanism that's intended for prototyping.
PATCH [1 / 2] in series for RFC: https://discourse.llvm.org/t/rfc-supporting-armv9-scalable-matrix-extension-sme-streaming-sve-ssve-mode-in-mlir/70678
[1] https://developer.arm.com/documentation/ddi0616/aa
[2] https://llvm.org/docs/AArch64SME.html#introduction
[3] https://github.com/ARM-software/abi-aa/blob/main/aapcs64/aapcs64.rst#671pstatesm-interfaces
[4] https://developer.arm.com/documentation/ddi0602/2023-03/Base-Instructions/SMSTART--Enables-access-to-Streaming-SVE-mode-and-SME-architectural-state--an-alias-of-MSR--immediate--
Reviewed By: awarzynski, dcaballe, WanderAway
Differential Revision: https://reviews.llvm.org/D150932
Added:
Modified:
mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
mlir/lib/Target/LLVMIR/ModuleImport.cpp
mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-ssve.mlir
mlir/test/Target/LLVMIR/Import/function-attributes.ll
mlir/test/Target/LLVMIR/llvmir.mlir
Removed:
mlir/test/Target/LLVMIR/arm-ssve.mlir
################################################################################
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index 232e180ab8975..3218701cd5c3e 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -1567,7 +1567,9 @@ def LLVM_LLVMFuncOp : LLVM_Op<"func", [
OptionalAttr<DictArrayAttr>:$res_attrs,
OptionalAttr<I64Attr>:$function_entry_count,
OptionalAttr<LLVM_MemoryEffectsAttr>:$memory,
- DefaultValuedAttr<Visibility, "mlir::LLVM::Visibility::Default">:$visibility_
+ DefaultValuedAttr<Visibility, "mlir::LLVM::Visibility::Default">:$visibility_,
+ OptionalAttr<UnitAttr>:$arm_streaming,
+ OptionalAttr<UnitAttr>:$arm_locally_streaming
);
let regions = (region AnyRegion:$body);
diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
index d9f115c6d77a7..2e67db55abc6b 100644
--- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
@@ -1542,6 +1542,12 @@ static void processPassthroughAttrs(llvm::Function *func, LLVMFuncOp funcOp) {
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.
+ if (attrName == "aarch64_pstate_sm_enabled" ||
+ attrName == "aarch64_pstate_sm_body")
+ continue;
+
if (attr.isStringAttribute()) {
StringRef val = attr.getValueAsString();
if (val.empty()) {
@@ -1574,6 +1580,11 @@ 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);
}
DictionaryAttr
diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index ca52b59641cec..d1d23b6c6f647 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -881,6 +881,11 @@ 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");
+
// First, create all blocks so we can jump to them.
llvm::LLVMContext &llvmContext = llvmFunc->getContext();
for (auto &bb : func) {
diff --git a/mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-ssve.mlir b/mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-ssve.mlir
index fce890fd93a25..116bf9ffe5e72 100644
--- a/mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-ssve.mlir
+++ b/mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-ssve.mlir
@@ -9,7 +9,7 @@
// NOTE: To run this test, your CPU must support SME.
// VLA memcopy in streaming mode.
-func.func @streaming_kernel_copy(%src : memref<?xi64>, %dst : memref<?xi64>, %size : index) attributes {passthrough = ["aarch64_pstate_sm_enabled"]} {
+func.func @streaming_kernel_copy(%src : memref<?xi64>, %dst : memref<?xi64>, %size : index) attributes {arm_streaming} {
%c0 = arith.constant 0 : index
%c2 = arith.constant 2 : index
%vscale = vector.vscale
diff --git a/mlir/test/Target/LLVMIR/Import/function-attributes.ll b/mlir/test/Target/LLVMIR/Import/function-attributes.ll
index e2696ad951b74..60cdd5864518f 100644
--- a/mlir/test/Target/LLVMIR/Import/function-attributes.ll
+++ b/mlir/test/Target/LLVMIR/Import/function-attributes.ll
@@ -193,3 +193,19 @@ define hidden void @hidden() {
define protected void @protected() {
ret void
}
+
+// -----
+
+; CHECK-LABEL: @streaming_func
+; CHECK-SAME: attributes {arm_streaming}
+define void @streaming_func() "aarch64_pstate_sm_enabled" {
+ ret void
+}
+
+// -----
+
+; CHECK-LABEL: @locally_streaming_func
+; CHECK-SAME: attributes {arm_locally_streaming}
+define void @locally_streaming_func() "aarch64_pstate_sm_body" {
+ ret void
+}
diff --git a/mlir/test/Target/LLVMIR/arm-ssve.mlir b/mlir/test/Target/LLVMIR/arm-ssve.mlir
deleted file mode 100644
index 91bf3e6daf517..0000000000000
--- a/mlir/test/Target/LLVMIR/arm-ssve.mlir
+++ /dev/null
@@ -1,11 +0,0 @@
-// RUN: mlir-translate -mlir-to-llvmir -split-input-file %s | FileCheck %s
-
-// Attribute to enable streaming-mode.
-
-// CHECK-LABEL: @streaming_callee
-// CHECK: #[[ATTR:[0-9]*]]
-llvm.func @streaming_callee() attributes {passthrough = ["aarch64_pstate_sm_enabled"]} {
- llvm.return
-}
-
-// CHECK: attributes #[[ATTR]] = { "aarch64_pstate_sm_enabled" }
diff --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir
index 6d340bc57fcd1..1c1581c6c6705 100644
--- a/mlir/test/Target/LLVMIR/llvmir.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir.mlir
@@ -2209,3 +2209,31 @@ llvm.func @readwrite_func() attributes {
memory = #llvm.memory_effects<other = readwrite, argMem = readwrite, inaccessibleMem = readwrite>}
// CHECK: attributes #[[ATTR]] = { memory(readwrite) }
+
+// -----
+
+//
+// arm_streaming attribute.
+//
+
+// CHECK-LABEL: @streaming_func
+// CHECK: #[[ATTR:[0-9]*]]
+llvm.func @streaming_func() attributes {arm_streaming} {
+ llvm.return
+}
+
+// CHECK: attributes #[[ATTR]] = { "aarch64_pstate_sm_enabled" }
+
+// -----
+
+//
+// arm_locally_streaming attribute.
+//
+
+// CHECK-LABEL: @locally_streaming_func
+// CHECK: #[[ATTR:[0-9]*]]
+llvm.func @locally_streaming_func() attributes {arm_locally_streaming} {
+ llvm.return
+}
+
+// CHECK: attributes #[[ATTR]] = { "aarch64_pstate_sm_body" }
More information about the Mlir-commits
mailing list