[Mlir-commits] [mlir] [mlir][ArmSME] Add `enable_arm_streaming_ignore` attribute (PR #66911)

[Benjamin Maxwell]

https://github.com/MacDue updated https://github.com/llvm/llvm-project/pull/66911

>From af4c8031b9f2c091b157855c70d1fee452672da3 Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Wed, 20 Sep 2023 13:51:07 +0000
Subject: [PATCH 1/2] [mlir][ArmSME] Add `enable_arm_streaming_ignore`
 attribute

This attribute makes the `enable_arm_streaming` pass ignore a function
(i.e. not add the enable streaming/za attributes). The main use case for
this is to prevent helper functions within tests being made streaming
functions.
---
 mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp | 6 +++++-
 mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir        | 8 ++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp b/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
index 97c38b546349510..1d3a090e861013b 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
@@ -52,6 +52,8 @@ using namespace mlir::arm_sme;
 static constexpr char kArmStreamingAttr[] = "arm_streaming";
 static constexpr char kArmLocallyStreamingAttr[] = "arm_locally_streaming";
 static constexpr char kArmZAAttr[] = "arm_za";
+static constexpr char kEnableArmStreamingIgnoreAttr[] =
+    "enable_arm_streaming_ignore";
 
 namespace {
 struct EnableArmStreamingPass
@@ -61,7 +63,9 @@ struct EnableArmStreamingPass
     this->enableZA = enableZA;
   }
   void runOnOperation() override {
-    std::string attr;
+    if (getOperation()->getAttr(kEnableArmStreamingIgnoreAttr))
+      return;
+    StringRef attr;
     switch (mode) {
     case ArmStreaming::Default:
       attr = kArmStreamingAttr;
diff --git a/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir b/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir
index f5cc83192f9f6e1..e7bbe8c0047687d 100644
--- a/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir
+++ b/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir
@@ -9,3 +9,11 @@
 // CHECK-ENABLE-ZA-LABEL: @arm_streaming
 // CHECK-ENABLE-ZA-SAME: attributes {arm_streaming, arm_za}
 func.func @arm_streaming() { return }
+
+// CHECK-LABEL: @not_arm_streaming
+// CHECK-SAME: attributes {enable_arm_streaming_ignore}
+// CHECK-LOCALLY-LABEL: @not_arm_streaming
+// CHECK-LOCALLY-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 }

>From 18896874e3fba68172623ca26ff301278f8e3104 Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Thu, 21 Sep 2023 16:46:59 +0000
Subject: [PATCH 2/2] Fixup: Add CHECK-NOTs

---
 mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir b/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir
index e7bbe8c0047687d..f93fca8df1ee6f9 100644
--- a/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir
+++ b/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir
@@ -10,10 +10,18 @@
 // CHECK-ENABLE-ZA-SAME: attributes {arm_streaming, arm_za}
 func.func @arm_streaming() { return }
 
+// -----
+
 // CHECK-LABEL: @not_arm_streaming
-// CHECK-SAME: attributes {enable_arm_streaming_ignore}
+// CHECK-NOT: arm_streaming
+// CHECK-SAME: enable_arm_streaming_ignore
+
 // CHECK-LOCALLY-LABEL: @not_arm_streaming
-// CHECK-LOCALLY-SAME: attributes {enable_arm_streaming_ignore}
+// CHECK-NOT: arm_locally_streaming
+// CHECK-LOCALLY-SAME: enable_arm_streaming_ignore
+
 // CHECK-ENABLE-ZA-LABEL: @not_arm_streaming
-// CHECK-ENABLE-ZA-SAME: attributes {enable_arm_streaming_ignore}
+// CHECK-NOT: arm_streaming
+// CHECK-NOT: arm_za
+// CHECK-ENABLE-ZA-SAME: enable_arm_streaming_ignore
 func.func @not_arm_streaming() attributes {enable_arm_streaming_ignore} { return }