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

[Mehdi Amini]

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

>From c170ed8bb9370fc4be5d141deab5fbe7a72ddff3 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  |  9 ++++---
 .../Dialect/ArmSME/enable-arm-streaming.mlir  | 26 +++++++++----------
 mlir/test/Dialect/ArmSME/enable-arm-za.mlir   | 12 ++++-----
 3 files changed, 25 insertions(+), 22 deletions(-)

diff --git a/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp b/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
index eafdc1de5ef34..c5c90762eaa92 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
@@ -53,7 +53,7 @@ using namespace mlir::arm_sme;
 namespace {
 
 constexpr StringLiteral
-    kEnableArmStreamingIgnoreAttr("enable_arm_streaming_ignore");
+    kEnableArmStreamingIgnoreAttr("llvm.enable_arm_streaming_ignore");
 
 template <typename... Ops>
 constexpr auto opList() {
@@ -128,14 +128,17 @@ 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
     // streaming-mode (see section B1.1.1, IDGNQM of spec [1]). It may be worth
     // supporting this later.
     if (zaMode != ArmZaMode::Disabled)
-      function->setAttr(stringifyArmZaMode(zaMode), unitAttr);
+      function->setAttr((Twine("llvm.") + stringifyArmZaMode(zaMode)).str(),
+                        unitAttr);
   }
 };
 } // namespace
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>
diff --git a/mlir/test/Dialect/ArmSME/enable-arm-za.mlir b/mlir/test/Dialect/ArmSME/enable-arm-za.mlir
index d3325513a8482..71155594e4b7c 100644
--- a/mlir/test/Dialect/ArmSME/enable-arm-za.mlir
+++ b/mlir/test/Dialect/ArmSME/enable-arm-za.mlir
@@ -9,16 +9,16 @@
 func.func private @declaration()
 
 // ENABLE-ZA-LABEL: @arm_new_za
-// ENABLE-ZA-SAME: attributes {arm_new_za, arm_streaming}
+// ENABLE-ZA-SAME: attributes {llvm.arm_new_za, llvm.arm_streaming}
 // IN-ZA-LABEL: @arm_new_za
-// IN-ZA-SAME: attributes {arm_in_za, arm_streaming}
+// IN-ZA-SAME: attributes {llvm.arm_in_za, llvm.arm_streaming}
 // OUT-ZA-LABEL: @arm_new_za
-// OUT-ZA-SAME: attributes {arm_out_za, arm_streaming}
+// OUT-ZA-SAME: attributes {llvm.arm_out_za, llvm.arm_streaming}
 // INOUT-ZA-LABEL: @arm_new_za
-// INOUT-ZA-SAME: attributes {arm_inout_za, arm_streaming}
+// INOUT-ZA-SAME: attributes {llvm.arm_inout_za, llvm.arm_streaming}
 // PRESERVES-ZA-LABEL: @arm_new_za
-// PRESERVES-ZA-SAME: attributes {arm_preserves_za, arm_streaming}
+// PRESERVES-ZA-SAME: attributes {llvm.arm_preserves_za, llvm.arm_streaming}
 // DISABLE-ZA-LABEL: @arm_new_za
 // DISABLE-ZA-NOT: arm_new_za
-// DISABLE-ZA-SAME: attributes {arm_streaming}
+// DISABLE-ZA-SAME: attributes {llvm.arm_streaming}
 func.func @arm_new_za() { return }