[Mlir-commits] [mlir] [mlir][ArmSME] Prefix SME streaming attributes with "llvm." for propagation (PR #190864)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/190864

Function-level attributes enabling SME in the backend (e.g. arm_streaming, arm_locally_streaming) were dropped when converting func.func to llvm.func after ed37bdc, because attribute propagation during FuncToLLVM conversion only preserves attributes prefixed with "llvm.". Prefix these discardable attributes accordingly.

>From 4a5ffd4a6bfda81b627173ac0dbfc70d3c3e36ad Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Tue, 7 Apr 2026 15:31:24 -0700
Subject: [PATCH] [mlir][ArmSME] Prefix SME streaming attributes with "llvm."
 for propagation

Function-level attributes enabling SME in the backend (e.g.
arm_streaming, arm_locally_streaming) were dropped when converting
func.func to llvm.func after ed37bdc, because attribute propagation
during FuncToLLVM conversion only preserves attributes prefixed with
"llvm.". Prefix these discardable attributes accordingly.
---
 .../ArmSME/Transforms/EnableArmStreaming.cpp  |  4 ++-
 .../Dialect/ArmSME/enable-arm-streaming.mlir  | 26 +++++++++----------
 2 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp b/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
index eafdc1de5ef34..8d14bff3b5875 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
@@ -128,7 +128,9 @@ struct EnableArmStreamingPass
 
     auto unitAttr = UnitAttr::get(&getContext());
 
-    function->setAttr(stringifyArmStreamingMode(streamingMode), unitAttr);
+    function->setDiscardableAttr(
+        (Twine("llvm.") + stringifyArmStreamingMode(streamingMode)).str(),
+        unitAttr);
 
     // The pass currently only supports enabling ZA when in streaming-mode, but
     // ZA can be accessed by the SME LDR, STR and ZERO instructions when not in
diff --git a/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir b/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir
index 00b38b86f0d6e..087cc98edbdd6 100644
--- a/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir
+++ b/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir
@@ -6,42 +6,42 @@
 // RUN: mlir-opt %s -enable-arm-streaming=if-scalable-and-supported -verify-diagnostics | FileCheck %s -check-prefix=IF-SCALABLE
 
 // CHECK-LABEL: @arm_streaming
-// CHECK-SAME: attributes {arm_streaming}
+// CHECK-SAME: attributes {llvm.arm_streaming}
 // CHECK-LOCALLY-LABEL: @arm_streaming
-// CHECK-LOCALLY-SAME: attributes {arm_locally_streaming}
+// CHECK-LOCALLY-SAME: attributes {llvm.arm_locally_streaming}
 // CHECK-COMPATIBLE-LABEL: @arm_streaming
-// CHECK-COMPATIBLE-SAME: attributes {arm_streaming_compatible}
+// CHECK-COMPATIBLE-SAME: attributes {llvm.arm_streaming_compatible}
 // CHECK-ENABLE-ZA-LABEL: @arm_streaming
-// CHECK-ENABLE-ZA-SAME: attributes {arm_new_za, arm_streaming}
+// CHECK-ENABLE-ZA-SAME: attributes {llvm.arm_new_za, llvm.arm_streaming}
 func.func @arm_streaming() { return }
 
 // CHECK-LABEL: @not_arm_streaming
-// CHECK-SAME: attributes {enable_arm_streaming_ignore}
+// CHECK-SAME: attributes {llvm.enable_arm_streaming_ignore}
 // CHECK-LOCALLY-LABEL: @not_arm_streaming
-// CHECK-LOCALLY-SAME: attributes {enable_arm_streaming_ignore}
+// CHECK-LOCALLY-SAME: attributes {llvm.enable_arm_streaming_ignore}
 // CHECK-COMPATIBLE-LABEL: @not_arm_streaming
-// CHECK-COMPATIBLE-SAME: attributes {enable_arm_streaming_ignore}
+// CHECK-COMPATIBLE-SAME: attributes {llvm.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 }
+// CHECK-ENABLE-ZA-SAME: attributes {llvm.enable_arm_streaming_ignore}
+func.func @not_arm_streaming() attributes {llvm.enable_arm_streaming_ignore} { return }
 
 // CHECK-LABEL: @requires_arm_streaming
-// CHECK-SAME: attributes {arm_streaming}
+// CHECK-SAME: attributes {llvm.arm_streaming}
 // IF-REQUIRED: @requires_arm_streaming
-// IF-REQUIRED-SAME: attributes {arm_streaming}
+// IF-REQUIRED-SAME: attributes {llvm.arm_streaming}
 func.func @requires_arm_streaming() {
   %tile = arm_sme.get_tile : vector<[4]x[4]xi32>
   return
 }
 
 // CHECK-LABEL: @does_not_require_arm_streaming
-// CHECK-SAME: attributes {arm_streaming}
+// CHECK-SAME: attributes {llvm.arm_streaming}
 // IF-REQUIRED: @does_not_require_arm_streaming
 // IF-REQUIRED-NOT: arm_streaming
 func.func @does_not_require_arm_streaming() { return }
 
 // IF-SCALABLE-LABEL: @contains_scalable_vectors
-// IF-SCALABLE-SAME: attributes {arm_streaming}
+// IF-SCALABLE-SAME: attributes {llvm.arm_streaming}
 func.func @contains_scalable_vectors(%vec: vector<[4]xf32>) -> vector<[4]xf32> {
   %0 = arith.addf %vec, %vec : vector<[4]xf32>
   return %0 : vector<[4]xf32>