[mlir][ArmSME] Add sve streaming compatible attribute (#75222)

Following the same path already used for ArmStreaming and
ArmLocallyStreaming.

This should correspond to clang's __arm_streaming_compatible attribute.

Author: tblah

mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td

@@ -20,6 +20,9 @@ def ArmStreamingMode : I32EnumAttr<"ArmStreamingMode", "Armv9 Streaming SVE mode
       // StreamingLocally: PSTATE.SM is kept internal and the callee manages it
       // on entry/exit.
       I32EnumAttrCase<"StreamingLocally", 2, "arm_locally_streaming">,
+      // StreamingCompatible: the function may be entered in either
+      // non-streaming mode (PSTATE.SM=0) or in streaming mode (PSTATE.SM=1)
+      I32EnumAttrCase<"StreamingCompatible", 3, "arm_streaming_compatible">,
     ]>{
   let cppNamespace = "mlir::arm_sme";
   let genSpecializedAttr = 0;
@@ -61,7 +64,11 @@ def EnableArmStreaming
                 clEnumValN(mlir::arm_sme::ArmStreamingMode::StreamingLocally,
                            "streaming-locally",
                            "Streaming mode is internal to the function, callee "
-                           "manages PSTATE.SM on entry/exit.")
+                           "manages PSTATE.SM on entry/exit."),
+                clEnumValN(mlir::arm_sme::ArmStreamingMode::StreamingCompatible,
+                           "streaming-compatible",
+                           "Function supports both streaming and non-streaming "
+                           "modes.")
           )}]>,
     Option<"zaMode", "za-mode", "mlir::arm_sme::ArmZaMode",
            /*default=*/"mlir::arm_sme::ArmZaMode::Disabled",

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

@@ -1411,6 +1411,7 @@ def LLVM_LLVMFuncOp : LLVM_Op<"func", [
     DefaultValuedAttr<Visibility, "mlir::LLVM::Visibility::Default">:$visibility_,
     OptionalAttr<UnitAttr>:$arm_streaming,
     OptionalAttr<UnitAttr>:$arm_locally_streaming,
+    OptionalAttr<UnitAttr>:$arm_streaming_compatible,
     OptionalAttr<UnitAttr>:$arm_new_za,
     OptionalAttr<StrAttr>:$section,
     OptionalAttr<UnnamedAddr>:$unnamed_addr,

mlir/lib/Target/LLVMIR/ModuleImport.cpp

@@ -1637,6 +1637,7 @@ static void processMemoryEffects(llvm::Function *func, LLVMFuncOp funcOp) {
 static constexpr std::array ExplicitAttributes{
     StringLiteral("aarch64_pstate_sm_enabled"),
     StringLiteral("aarch64_pstate_sm_body"),
+    StringLiteral("aarch64_pstate_sm_compatible"),
     StringLiteral("aarch64_pstate_za_new"),
     StringLiteral("vscale_range"),
     StringLiteral("frame-pointer"),
@@ -1709,6 +1710,8 @@ void ModuleImport::processFunctionAttributes(llvm::Function *func,
     funcOp.setArmStreaming(true);
   else if (func->hasFnAttribute("aarch64_pstate_sm_body"))
     funcOp.setArmLocallyStreaming(true);
+  else if (func->hasFnAttribute("aarch64_pstate_sm_compatible"))
+    funcOp.setArmStreamingCompatible(true);
 
   if (func->hasFnAttribute("aarch64_pstate_za_new"))
     funcOp.setArmNewZa(true);

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

@@ -965,6 +965,8 @@ LogicalResult ModuleTranslation::convertOneFunction(LLVMFuncOp func) {
     llvmFunc->addFnAttr("aarch64_pstate_sm_enabled");
   else if (func.getArmLocallyStreaming())
     llvmFunc->addFnAttr("aarch64_pstate_sm_body");
+  else if (func.getArmStreamingCompatible())
+    llvmFunc->addFnAttr("aarch64_pstate_sm_compatible");
 
   if (func.getArmNewZa())
     llvmFunc->addFnAttr("aarch64_pstate_za_new");

mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir

@@ -1,12 +1,15 @@
 // RUN: mlir-opt %s -enable-arm-streaming -verify-diagnostics | FileCheck %s
 // RUN: mlir-opt %s -enable-arm-streaming=streaming-mode=streaming-locally -verify-diagnostics | FileCheck %s -check-prefix=CHECK-LOCALLY
+// RUN: mlir-opt %s -enable-arm-streaming=streaming-mode=streaming-compatible -verify-diagnostics | FileCheck %s -check-prefix=CHECK-COMPATIBLE
 // RUN: mlir-opt %s -enable-arm-streaming=za-mode=new-za -verify-diagnostics | FileCheck %s -check-prefix=CHECK-ENABLE-ZA
 // RUN: mlir-opt %s -enable-arm-streaming=only-if-required-by-ops -verify-diagnostics | FileCheck %s -check-prefix=IF-REQUIRED
 
 // CHECK-LABEL: @arm_streaming
 // CHECK-SAME: attributes {arm_streaming}
 // CHECK-LOCALLY-LABEL: @arm_streaming
 // CHECK-LOCALLY-SAME: attributes {arm_locally_streaming}
+// CHECK-COMPATIBLE-LABEL: @arm_streaming
+// CHECK-COMPATIBLE-SAME: attributes {arm_streaming_compatible}
 // CHECK-ENABLE-ZA-LABEL: @arm_streaming
 // CHECK-ENABLE-ZA-SAME: attributes {arm_new_za, arm_streaming}
 func.func @arm_streaming() { return }
@@ -15,6 +18,8 @@ func.func @arm_streaming() { return }
 // CHECK-SAME: attributes {enable_arm_streaming_ignore}
 // CHECK-LOCALLY-LABEL: @not_arm_streaming
 // CHECK-LOCALLY-SAME: attributes {enable_arm_streaming_ignore}
+// CHECK-COMPATIBLE-LABEL: @not_arm_streaming
+// CHECK-COMPATIBLE-SAME: attributes {enable_arm_streaming_ignore}
 // CHECK-ENABLE-ZA-LABEL: @not_arm_streaming
 // CHECK-ENABLE-ZA-SAME: attributes {enable_arm_streaming_ignore}
 func.func @not_arm_streaming() attributes {enable_arm_streaming_ignore} { return }

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

@@ -212,6 +212,14 @@ define void @locally_streaming_func() "aarch64_pstate_sm_body" {
 
 // -----
 
+; CHECK-LABEL: @streaming_compatible_func
+; CHECK-SAME: attributes {arm_streaming_compatible}
+define void @streaming_compatible_func() "aarch64_pstate_sm_compatible" {
+  ret void
+}
+
+// -----
+
 ; CHECK-LABEL: @section_func
 ; CHECK-SAME: attributes {section = ".section.name"}
 define void @section_func() section ".section.name" {

mlir/test/Target/LLVMIR/llvmir.mlir

@@ -2305,6 +2305,20 @@ llvm.func @locally_streaming_func() attributes {arm_locally_streaming} {
 
 // -----
 
+//
+// arm_streaming_compatible attribute.
+//
+
+// CHECK-LABEL: @streaming_compatible_func
+// CHECK: #[[ATTR:[0-9]*]]
+llvm.func @streaming_compatible_func() attributes {arm_streaming_compatible} {
+  llvm.return
+}
+
+// CHECK: attributes #[[ATTR]] = { "aarch64_pstate_sm_compatible" }
+
+// -----
+
 //
 // Zero-initialize operation.
 //