[mlir] Add pass to enable Armv9 Streaming SVE mode

[c-rhodes]

Summary

This patch adds a pass 'enable-arm-streaming' that enables the Armv9
Scalable Matrix Extension (SME) Streaming SVE (SSVE) mode [1] by adding
either of the following attributes to 'func.func' ops:

• arm_streaming (default)
• arm_locally_streaming

PATCH [2 / 2] in series for RFC: https://discourse.llvm.org/t/rfc-supporting-armv9-scalable-matrix-extension-sme-streaming-sve-ssve-mode-in-mlir/70678

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

[awarzynski]

Great stuff, thank you! LGTM

[c-rhodes]

Updated namespace to remove dependent pass.

[c-rhodes]

I'll land this tomorrow unless there's any further comments by then.

[dcaballe]

Thanks for moving forward with the implementation! I think we should move this outside of the Func dialect. Since an SME dialect is in the horizon, I would suggest that we move the implementation to Dialect/SME/Transforms and the lit tests to similar Dialect/SME folders.
We currently don't have a good place for target specific transformations without a dialect and introducing something like {lib/include}/Target/... would probably need much more discussion.

[c-rhodes]

Move to ArmSME dialect.

[c-rhodes]

In D150934#4366544, @dcaballe wrote:

Thanks for moving forward with the implementation! I think we should move this outside of the Func dialect. Since an SME dialect is in the horizon, I would suggest that we move the implementation to Dialect/SME/Transforms and the lit tests to similar Dialect/SME folders.

Thanks for reviewing. I've moved the implementation to a minimal SME dialect.

We currently don't have a good place for target specific transformations without a dialect and introducing something like {lib/include}/Target/... would probably need much more discussion.

I think a single dialect for each target (e.g. Arm, x86, ...) much like target in LLVM backend would be better than having dialects for each target feature. There's a fair amount of boilerplate in these dialects.

[dcaballe]

LGTM, thanks!

I think a single dialect for each target (e.g. Arm, x86, ...) much like target in LLVM backend would be better than having dialects for each target feature. There's a fair amount of boilerplate in these dialects.

That's fair. Others can correct me if I'm wrong but I think the current "trend" in MLIR (as this is something that has been changing over time) leans more towards more minimalist dialects with a small and well-defined scopes (e.g., 'func' dialect). I think this is, to some extent, motivated by the infra as, for example, registering ops that are not used will have a negative impact in compile time. We also have mechanisms to make a dialect legal/illegal after a conversion which I think also fosters this kind of small logical grouping (e.g., it would be tedious if we had a single ARM dialect and we had to mark as legal/illegal a large part of it, one instruction at a time). There are also clear exceptions to this, such as the Vector dialect, where we have multiple layers of abstraction altogether. I think the idea of "sub-dialects" or some kind of logical grouping within a dialect was discussed at some point. That would also make sense for target-like dialects and their features. At the end of the day, all of this is flexible and it mostly depends on how we use the operations in practice.

mlir/include/mlir/Dialect/CMakeLists.txt

7 ↗

(On Diff #525100)

alphabetically sorted?

mlir/lib/Dialect/CMakeLists.txt

6 ↗

(On Diff #525100)

same

[c-rhodes]

In D150934#4368724, @dcaballe wrote:

LGTM, thanks!

I think a single dialect for each target (e.g. Arm, x86, ...) much like target in LLVM backend would be better than having dialects for each target feature. There's a fair amount of boilerplate in these dialects.

That's fair. Others can correct me if I'm wrong but I think the current "trend" in MLIR (as this is something that has been changing over time) leans more towards more minimalist dialects with a small and well-defined scopes (e.g., 'func' dialect). I think this is, to some extent, motivated by the infra as, for example, registering ops that are not used will have a negative impact in compile time. We also have mechanisms to make a dialect legal/illegal after a conversion which I think also fosters this kind of small logical grouping (e.g., it would be tedious if we had a single ARM dialect and we had to mark as legal/illegal a large part of it, one instruction at a time). There are also clear exceptions to this, such as the Vector dialect, where we have multiple layers of abstraction altogether. I think the idea of "sub-dialects" or some kind of logical grouping within a dialect was discussed at some point. That would also make sense for target-like dialects and their features. At the end of the day, all of this is flexible and it mostly depends on how we use the operations in practice.

Ah ok, I wasn't aware of most of that, thanks for explaining.

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

 namespace mlir {
 namespace bufferization {
 class BufferizeTypeConverter;
 } // namespace bufferization
 
 class RewritePatternSet;
 
 namespace func {
+// Options for Armv9 Streaming SVE mode. By default, streaming-mode is part of
+// the function interface (ABI) and the caller manages PSTATE.SM on entry/exit.
+// In a locally streaming function PSTATE.SM is kept internal and the callee
+// manages it on entry/exit.
+enum class ArmStreaming { Default = 0, Locally = 1 };
 
 #define GEN_PASS_DECL
 #include "mlir/Dialect/Func/Transforms/Passes.h.inc"
 
 /// Creates an instance of func bufferization pass.
 std::unique_ptr<Pass> createFuncBufferizePass();
 
 /// Pass to deduplicate functions.
 std::unique_ptr<Pass> createDuplicateFunctionEliminationPass();
 
+/// Pass to enable Armv9 Streaming SVE mode.
+std::unique_ptr<Pass>
+createEnableArmStreamingPass(const ArmStreaming mode = ArmStreaming::Default);
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/Func/Transforms/Passes.h.inc"
 
 } // namespace func
 } // namespace mlir
 
 #endif // MLIR_DIALECT_FUNC_TRANSFORMS_PASSES_H

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

def DuplicateFunctionEliminationPass : Pass<"duplicate-function-elimination",
   let description = [{
     Deduplicate functions that are equivalent in all aspects but their symbol
     name. The pass chooses one representative per equivalence class, erases
     the remainder, and updates function calls accordingly.
   }];
   let constructor = "mlir::func::createDuplicateFunctionEliminationPass()";
 }
 
+def EnableArmStreaming
+    : Pass<"enable-arm-streaming", "mlir::func::FuncOp"> {
+  let summary = "Enable Armv9 Streaming SVE mode";
+  let description = [{
+    Enables the Armv9 Streaming SVE mode [1] for func.func ops by annotating
+    them with attributes. See options for more details.
+
+    [1] https://developer.arm.com/documentation/ddi0616/aa
+  }];
+  let constructor = "mlir::func::createEnableArmStreamingPass()";
+  let options = [
+    Option<"mode", "mode", "mlir::func::ArmStreaming",
+          /*default=*/"mlir::func::ArmStreaming::Default",
+          "Select how streaming-mode is managed at the function-level.",
+          [{::llvm::cl::values(
+            clEnumValN(mlir::func::ArmStreaming::Default, "default", "Streaming mode is part of the function interface (ABI), caller manages PSTATE.SM on entry/exit."),
+            clEnumValN(mlir::func::ArmStreaming::Locally, "locally", "Streaming mode is internal to the function, callee manages PSTATE.SM on entry/exit.")
+          )}]>,
+  ];
+}
+
 #endif // MLIR_DIALECT_FUNC_TRANSFORMS_PASSES_TD

mlir/lib/Dialect/Func/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRFuncTransforms
   DecomposeCallGraphTypes.cpp
   DuplicateFunctionElimination.cpp
   FuncBufferize.cpp
   FuncConversions.cpp
   OneToNFuncConversions.cpp
+  EnableArmStreaming.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Func/Transforms
 
   DEPENDS
   MLIRFuncTransformsIncGen
 
   LINK_LIBS PUBLIC
   MLIRBufferizationDialect
   MLIRBufferizationTransforms
   MLIRFuncDialect
   MLIRIR
   MLIRMemRefDialect
   MLIRPass
   MLIRTransforms
   )

mlir/lib/Dialect/Func/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'
+//
+// 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' /
+// 'smstop sm' in the prologue / epilogue for functions with this
+// attribute.
+//
+// [1] https://developer.arm.com/documentation/ddi0616/aa
+// [2] https://llvm.org/docs/AArch64SME.html
+// [3] https://github.com/ARM-software/abi-aa/blob/main/aapcs64/aapcs64.rst#671pstatesm-interfaces
+// [4] https://developer.arm.com/documentation/ddi0602/2023-03/Base-Instructions/SMSTART--Enables-access-to-Streaming-SVE-mode-and-SME-architectural-state--an-alias-of-MSR--immediate--
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Func/Transforms/Passes.h"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+
+#define DEBUG_TYPE "enable-arm-streaming"
+
+namespace mlir {
+namespace func {
+#define GEN_PASS_DEF_ENABLEARMSTREAMING
+#include "mlir/Dialect/Func/Transforms/Passes.h.inc"
+} // namespace func
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::func;
+
+static constexpr char kArmStreamingAttr[] = "arm_streaming";
+static constexpr char kArmLocallyStreamingAttr[] = "arm_locally_streaming";
+
+namespace {
+struct EnableArmStreamingPass
+    : public func::impl::EnableArmStreamingBase<EnableArmStreamingPass> {
+  EnableArmStreamingPass(ArmStreaming mode) { this->mode = mode; }
+  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()));
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass>
+mlir::func::createEnableArmStreamingPass(const ArmStreaming mode) {
+  return std::make_unique<EnableArmStreamingPass>(mode);
+}

mlir/test/Dialect/Func/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
+
+// CHECK-LABEL: @arm_streaming
+// CHECK-SAME: attributes {arm_streaming}
+func.func @arm_streaming() { return }
+
+// CHECK-LOCALLY-LABEL: @arm_locally_streaming
+// CHECK-LOCALLY-SAME: attributes {arm_locally_streaming}
+func.func @arm_locally_streaming() { return }