[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

Author: c-rhodes

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

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

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

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

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

mlir/test/Target/LLVMIR/arm-ssve.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" }