[mlir][llvm] Add arm_streaming LLVM function attributes

[c-rhodes]

Summary

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

[awarzynski]

Thank you for working on this!

LGTM modulo one comment. I'll be away for a week, hence approving as is - I don't want to block it.

mlir/lib/Target/LLVMIR/ModuleImport.cpp

1545–1549

Do we need this? I don't see similar checks for other attributes. If you'd like to keep it, then perhaps worth adding a test that shows that {passthrough = ["aarch64_pstate_sm_enabled"]} does nothing?

[WanderAway]

LGTM. Maybe when the SME dialect eventually gets created (soon TM), we can migrate these attributes to the dialect itself.

[c-rhodes]

mlir/lib/Target/LLVMIR/ModuleImport.cpp

1545–1549

Do we need this? I don't see similar checks for other attributes. If you'd like to keep it, then perhaps worth adding a test that shows that {passthrough = ["aarch64_pstate_sm_enabled"]} does nothing?

We do, they'll be added as passthrough attributes otherwise. There's a similar check above for memory attribute.

[c-rhodes]

I'll land this tomorrow unless there's any further comments by then.

[dcaballe]

Sorry for the delay. LGTM as a first step.

LGTM. Maybe when the SME dialect eventually gets created (soon TM), we can migrate these attributes to the dialect itself.

+1

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

let arguments = (ins
     DefaultValuedAttr<CConv, "CConv::C">:$CConv,
     OptionalAttr<FlatSymbolRefAttr>:$personality,
     OptionalAttr<StrAttr>:$garbageCollector,
     OptionalAttr<ArrayAttr>:$passthrough,
     OptionalAttr<DictArrayAttr>:$arg_attrs,
     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);
 
   let skipDefaultBuilders = 1;
 
   let builders = [
     OpBuilder<(ins "StringRef":$name, "Type":$type,

mlir/lib/Target/LLVMIR/ModuleImport.cpp

for (llvm::Attribute attr : funcAttrs) {
 
     StringRef attrName;
     if (attr.isStringAttribute())
       attrName = attr.getKindAsString();
     else
       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()) {
         passthroughs.push_back(keyAttr);
         continue;
       }
       passthroughs.push_back(
           ArrayAttr::get(context, {keyAttr, StringAttr::get(context, val)}));
static void processPassthroughAttrs(llvm::Function *func, LLVMFuncOp funcOp) {
   if (!passthroughs.empty())
     funcOp.setPassthroughAttr(ArrayAttr::get(context, passthroughs));
 }
 
 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
 ModuleImport::convertParameterAttribute(llvm::AttributeSet llvmParamAttrs,
                                         OpBuilder &builder) {
   SmallVector<NamedAttribute> paramAttrs;
   for (auto [llvmKind, mlirName] : getAttrKindToNameMapping()) {
     auto llvmAttr = llvmParamAttrs.getAttribute(llvmKind);

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

if (func.getPersonality()) {
     if (llvm::Constant *pfunc = getLLVMConstant(ty, func.getPersonalityAttr(),
                                                 func.getLoc(), *this))
       llvmFunc->setPersonalityFn(pfunc);
   }
 
   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) {
     auto *llvmBB = llvm::BasicBlock::Create(llvmContext);
     llvmBB->insertInto(llvmFunc);
     mapBlock(&bb, llvmBB);
   }
 

mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-ssve.mlir

 // RUN: mlir-opt %s -test-lower-to-llvm | \
 // RUN: mlir-translate -mlir-to-llvmir | \
 // RUN: %lli_aarch64_cmd --march=aarch64 --mattr="+sve,+sme" \
 // RUN:      -force-streaming-compatible-sve \
 // RUN:      --entry-function=entry \
 // RUN:      --dlopen=%mlir_native_utils_lib_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 
 // 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
   %step = arith.muli %c2, %vscale : index
   scf.for %i = %c0 to %size step %step {
     %0 = vector.load %src[%i] : memref<?xi64>, vector<[2]xi64>
     vector.store %0, %dst[%i] : memref<?xi64>, vector<[2]xi64>
   }

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

 }
 
 // -----
 
 ; CHECK-LABEL: llvm.func protected @protected()
 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
+}

mlir/test/Target/LLVMIR/arm-ssve.mlir

-// 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" }

mlir/test/Target/LLVMIR/llvmir.mlir

 
 // -----
 
 // CHECK: declare void @readwrite_func() #[[ATTR:[0-9]+]]
 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" }