Skip to content

[mlir][ArmSME] Prefix SME streaming attributes with "llvm." for propagation#190864

Open
joker-eph wants to merge 1 commit intollvm:mainfrom
joker-eph:fix_discardable_attr_sme
Open

[mlir][ArmSME] Prefix SME streaming attributes with "llvm." for propagation#190864
joker-eph wants to merge 1 commit intollvm:mainfrom
joker-eph:fix_discardable_attr_sme

Conversation

@joker-eph
Copy link
Copy Markdown
Collaborator

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.

@llvmbot
Copy link
Copy Markdown
Member

llvmbot commented Apr 7, 2026

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

Changes

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.


Full diff: https://github.com/llvm/llvm-project/pull/190864.diff

2 Files Affected:

  • (modified) mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp (+3-1)
  • (modified) mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir (+13-13)
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>

@llvmbot
Copy link
Copy Markdown
Member

llvmbot commented Apr 7, 2026

@llvm/pr-subscribers-mlir-sme

Author: Mehdi Amini (joker-eph)

Changes

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.


Full diff: https://github.com/llvm/llvm-project/pull/190864.diff

2 Files Affected:

  • (modified) mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp (+3-1)
  • (modified) mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir (+13-13)
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>

…gation

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.
@joker-eph joker-eph force-pushed the fix_discardable_attr_sme branch from 4a5ffd4 to c170ed8 Compare April 7, 2026 22:36
@hockyy
Copy link
Copy Markdown
Contributor

hockyy commented Apr 8, 2026

lgtm

Copy link
Copy Markdown
Contributor

@banach-space banach-space left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks!

Btw, I am working on introducing a dedicated buildbot for this sort of stuff. As usual, these things take time.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants