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

[c-rhodes]

Summary

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

[awarzynski]

Thanks for adding this, LGTM!

mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp

75–78

[nit] I would reword a bit to highlight that:

• right now, ZA can only be enabled when in streaming mode (that's a "limitation" of this pass), however,
• it is also possible use ZA when not in streaming mode (and we could add support for that later).

I am just thinking that this is clear to somebody very familiar with the spec, but perhaps it's worth clarifying for somebody with less background knowledge?

[c-rhodes]

Rebase and address comments.

[c-rhodes]

In D152695#4413602, @awarzynski wrote:

Thanks for adding this, LGTM!

Thanks for reviewing!

[dcaballe]

Thanks!

[c-rhodes]

As I mentioned on D152694 it isn't feasible to use the aarch64_pstate_za_new attribute since the backend emits calls to the following SME support routines [1] for the lazy-save mechanism
[2]:

__arm_tpidr2_restore
__arm_tpidr2_save

These will soon be added to compiler-rt but there's currently no public implementation, and using this attribute would introduce an MLIR dependency on compiler-rt. Furthermore, this mechanism is for routines with ZA enabled calling other routines with it also enabled. We can choose not to enable ZA in the compiler when this is case.

Given this ZA has to be enabled manually with these intrinsics [3]:

llvm.aarch64.sme.za.enable
llvm.aarch64.sme.za.disable

I've updated this patch to add an arm_za attribute and in D153050 added the intrinsics, as well as patterns for inserting them during legalization to LLVM at the start and end of functions if the function has the 'arm_za' attribute.

These 2 patches could be 1 really but given the approach has changed somewhat I didn't want to update this drastically.

[1] https://github.com/ARM-software/abi-aa/blob/main/aapcs64/aapcs64.rst#sme-support-routines
[2] https://github.com/ARM-software/abi-aa/blob/main/aapcs64/aapcs64.rst#the-za-lazy-saving-scheme
[3] https://github.com/llvm/llvm-project/blob/main/llvm/test/CodeGen/AArch64/sme-toggle-pstateza.ll

[awarzynski]

LGTM, thanks!

These 2 patches could be 1 really but given the approach has changed somewhat I didn't want to update this drastically.

IMHO, totally fine to have this as a separate patch.

mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.h

 // manages it on entry/exit.
 enum class ArmStreaming { Default = 0, Locally = 1 };
 
 #define GEN_PASS_DECL
 #include "mlir/Dialect/ArmSME/Transforms/Passes.h.inc"
 
 /// 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
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/ArmSME/Transforms/Passes.h.inc"
 
 } // namespace arm_sme
 } // namespace mlir
 
 #endif // MLIR_DIALECT_ARMSME_TRANSFORMS_PASSES_H

mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td

Option<"mode", "mode", "mlir::arm_sme::ArmStreaming",
           [{::llvm::cl::values(
                 clEnumValN(mlir::arm_sme::ArmStreaming::Default, "default",
                "Streaming mode is part of the function interface "
                "(ABI), caller manages PSTATE.SM on entry/exit."),
                 clEnumValN(mlir::arm_sme::ArmStreaming::Locally, "locally",
                "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"];
 }
 
 #endif // MLIR_DIALECT_ARMSME_TRANSFORMS_PASSES_TD

mlir/lib/Dialect/ArmSME/Transforms/EnableArmStreaming.cpp

 //===- EnableArmStreaming.cpp - Enable Armv9 Streaming SVE mode -----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This pass enables the Armv9 Scalable Matrix Extension (SME) Streaming SVE
 // (SSVE) mode [1][2] by adding either of the following attributes to
 // 'func.func' ops:
 //
 //   * '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
 // backend will emit 'smstart sm' / 'smstop sm' [4] around calls to
 // streaming functions.
 //
 // In locally streaming functions PSTATE.SM is kept internal and managed by
 // the callee on entry/exit. The LLVM backend will emit 'smstart sm' /

 } // namespace arm_sme
 } // namespace mlir
 
 using namespace mlir;
 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) {
     case ArmStreaming::Default:
       attr = kArmStreamingAttr;
       break;
     case ArmStreaming::Locally:
       attr = kArmLocallyStreamingAttr;
       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);
 }

mlir/test/Dialect/ArmSME/enable-arm-streaming.mlir

 // 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 }