diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td --- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td @@ -1567,7 +1567,9 @@ OptionalAttr:$res_attrs, OptionalAttr:$function_entry_count, OptionalAttr:$memory, - DefaultValuedAttr:$visibility_ + DefaultValuedAttr:$visibility_, + OptionalAttr:$arm_streaming, + OptionalAttr:$arm_locally_streaming ); let regions = (region AnyRegion:$body); diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp --- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp +++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp @@ -1542,6 +1542,12 @@ attrName = llvm::Attribute::getNameFromAttrKind(attr.getKindAsEnum()); auto keyAttr = StringAttr::get(context, attrName); + // Skip the aarch64_pstate_sm_ 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 @@ 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 --- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp @@ -881,6 +881,11 @@ 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 --- 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, %dst : memref, %size : index) attributes {passthrough = ["aarch64_pstate_sm_enabled"]} { +func.func @streaming_kernel_copy(%src : memref, %dst : memref, %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 --- a/mlir/test/Target/LLVMIR/Import/function-attributes.ll +++ b/mlir/test/Target/LLVMIR/Import/function-attributes.ll @@ -193,3 +193,19 @@ 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 --- 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 --- a/mlir/test/Target/LLVMIR/llvmir.mlir +++ b/mlir/test/Target/LLVMIR/llvmir.mlir @@ -2209,3 +2209,31 @@ memory = #llvm.memory_effects} // 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" }