[Mlir-commits] [mlir] e947e76 - [mlir][ArmSME] Extend streaming-mode pass to support enabling ZA

[Cullen Rhodes]

Author: Cullen Rhodes
Date: 2023-06-16T09:26:42Z
New Revision: e947e760585c2a457775d8723e1503f7d5b56984

URL: https://github.com/llvm/llvm-project/commit/e947e760585c2a457775d8723e1503f7d5b56984
DIFF: https://github.com/llvm/llvm-project/commit/e947e760585c2a457775d8723e1503f7d5b56984.diff

LOG: [mlir][ArmSME] Extend streaming-mode pass to support enabling ZA

This patch extends the 'enable-arm-streaming' pass with a new option to
enable the ZA storage array by adding the 'arm_za' attribute to
'func.func' ops.

A later patch will insert `llvm.aarch64.sme.za.enable` at the beginning
of 'func.func' ops and `llvm.aarch64.sme.za.disable` before
`func.return` statements when lowering to LLVM dialect.

Currently the pass only supports enabling ZA with streaming-mode on but
the SME LDR, STR and ZERO instructions can access ZA when not in
streaming-mode (section B1.1.1, IDGNQM [1]), so it may be worth making
these options independent in the future.

N.B. This patch is generally useful in the context of SME enablement in
MLIR, but it will help enable writing an integration test for rewrite
pattern that lowers `vector.transfer_write` -> `zero {za}` (D152508).

[1] https://developer.arm.com/documentation/ddi0616/aa

Reviewed By: awarzynski, dcaballe

Differential Revision: https://reviews.llvm.org/D152695

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.h
    mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td
    mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
    mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.h b/mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.h
index b51b0a3950a9a..00ac5376ed7d9 100644
--- a/mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.h
@@ -27,7 +27,8 @@ enum class ArmStreaming { Default = 0, Locally = 1 };
 
 /// Pass to enable Armv9 Streaming SVE mode.
 std::unique_ptr<Pass>
-createEnableArmStreamingPass(const ArmStreaming mode = ArmStreaming::Default);
+createEnableArmStreamingPass(const ArmStreaming mode = ArmStreaming::Default,
+                             const bool enableZA = false);
 
 //===----------------------------------------------------------------------===//
 // Registration

diff  --git a/mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td b/mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td
index 8c9455d6a1788..7bc39e0534b8a 100644
--- a/mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td
@@ -33,6 +33,8 @@ def EnableArmStreaming
 						   "Streaming mode is internal to the function, callee "
 						   "manages PSTATE.SM on entry/exit.")
           )}]>,
+    Option<"enableZA", "enable-za", "bool", /*default=*/"false",
+           "Enable ZA storage array.">,
   ];
   let dependentDialects = ["func::FuncDialect"];
 }

diff  --git a/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp b/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
index 0f35b2248521d..97c38b5463495 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp
@@ -13,6 +13,8 @@
 //   * 'arm_streaming' (default)
 //   * 'arm_locally_streaming'
 //
+// It can also optionally enable the ZA storage array.
+//
 // Streaming-mode is part of the interface (ABI) for functions with the
 // first attribute and it's the responsibility of the caller to manage
 // PSTATE.SM on entry/exit to functions with this attribute [3]. The LLVM
@@ -49,11 +51,15 @@ using namespace mlir::arm_sme;
 
 static constexpr char kArmStreamingAttr[] = "arm_streaming";
 static constexpr char kArmLocallyStreamingAttr[] = "arm_locally_streaming";
+static constexpr char kArmZAAttr[] = "arm_za";
 
 namespace {
 struct EnableArmStreamingPass
     : public arm_sme::impl::EnableArmStreamingBase<EnableArmStreamingPass> {
-  EnableArmStreamingPass(ArmStreaming mode) { this->mode = mode; }
+  EnableArmStreamingPass(ArmStreaming mode, bool enableZA) {
+    this->mode = mode;
+    this->enableZA = enableZA;
+  }
   void runOnOperation() override {
     std::string attr;
     switch (mode) {
@@ -65,11 +71,19 @@ struct EnableArmStreamingPass
       break;
     }
     getOperation()->setAttr(attr, UnitAttr::get(&getContext()));
+
+    // 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 (enableZA)
+      getOperation()->setAttr(kArmZAAttr, UnitAttr::get(&getContext()));
   }
 };
 } // namespace
 
 std::unique_ptr<Pass>
-mlir::arm_sme::createEnableArmStreamingPass(const ArmStreaming mode) {
-  return std::make_unique<EnableArmStreamingPass>(mode);
+mlir::arm_sme::createEnableArmStreamingPass(const ArmStreaming mode,
+                                            const bool enableZA) {
+  return std::make_unique<EnableArmStreamingPass>(mode, enableZA);
 }

diff  --git a/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir b/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir
index 0c24f8c5389bd..f5cc83192f9f6 100644
--- a/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir
+++ b/mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir
@@ -1,8 +1,11 @@
 // RUN: mlir-opt %s -enable-arm-streaming -verify-diagnostics | FileCheck %s
 // RUN: mlir-opt %s -enable-arm-streaming=mode=locally -verify-diagnostics | FileCheck %s -check-prefix=CHECK-LOCALLY
+// RUN: mlir-opt %s -enable-arm-streaming=enable-za -verify-diagnostics | FileCheck %s -check-prefix=CHECK-ENABLE-ZA
 
 // CHECK-LABEL: @arm_streaming
 // CHECK-SAME: attributes {arm_streaming}
 // CHECK-LOCALLY-LABEL: @arm_streaming
 // CHECK-LOCALLY-SAME: attributes {arm_locally_streaming}
+// CHECK-ENABLE-ZA-LABEL: @arm_streaming
+// CHECK-ENABLE-ZA-SAME: attributes {arm_streaming, arm_za}
 func.func @arm_streaming() { return }