Index: mlir/include/mlir/Dialect/Func/Transforms/Passes.h =================================================================== --- mlir/include/mlir/Dialect/Func/Transforms/Passes.h +++ mlir/include/mlir/Dialect/Func/Transforms/Passes.h @@ -24,6 +24,11 @@ class RewritePatternSet; +// 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" @@ -34,6 +39,10 @@ /// Pass to deduplicate functions. std::unique_ptr createDuplicateFunctionEliminationPass(); +/// Pass to enable Armv9 Streaming SVE mode. +std::unique_ptr createEnableArmStreamingPass(); +std::unique_ptr createEnableArmStreamingPass(const ArmStreaming mode); + //===----------------------------------------------------------------------===// // Registration //===----------------------------------------------------------------------===// Index: mlir/include/mlir/Dialect/Func/Transforms/Passes.td =================================================================== --- mlir/include/mlir/Dialect/Func/Transforms/Passes.td +++ mlir/include/mlir/Dialect/Func/Transforms/Passes.td @@ -51,4 +51,25 @@ let constructor = "mlir::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::createEnableArmStreamingPass()"; + let options = [ + Option<"mode", "mode", "mlir::ArmStreaming", + /*default=*/"mlir::ArmStreaming::Default", + "Select how streaming-mode is managed at the function-level.", + [{::llvm::cl::values( + clEnumValN(::mlir::ArmStreaming::Default, "default", "Streaming mode is part of the function interface (ABI), caller manages PSTATE.SM on entry/exit."), + clEnumValN(::mlir::ArmStreaming::Locally, "locally", "Streaming mode is internal to the function, callee manages PSTATE.SM on entry/exit.") + )}]>, + ]; +} + #endif // MLIR_DIALECT_FUNC_TRANSFORMS_PASSES_TD Index: mlir/lib/Dialect/Func/Transforms/CMakeLists.txt =================================================================== --- mlir/lib/Dialect/Func/Transforms/CMakeLists.txt +++ mlir/lib/Dialect/Func/Transforms/CMakeLists.txt @@ -4,6 +4,7 @@ FuncBufferize.cpp FuncConversions.cpp OneToNFuncConversions.cpp + EnableArmStreaming.cpp ADDITIONAL_HEADER_DIRS ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Func/Transforms Index: mlir/lib/Dialect/Func/Transforms/EnableArmStreaming.cpp =================================================================== --- /dev/null +++ mlir/lib/Dialect/Func/Transforms/EnableArmStreaming.cpp @@ -0,0 +1,79 @@ +//===- 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 { +#define GEN_PASS_DEF_ENABLEARMSTREAMING +#include "mlir/Dialect/Func/Transforms/Passes.h.inc" +} // 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 impl::EnableArmStreamingBase { + EnableArmStreamingPass() = default; + EnableArmStreamingPass(const EnableArmStreamingPass &pass) = default; + EnableArmStreamingPass(mlir::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 mlir::createEnableArmStreamingPass() { + return std::make_unique(); +} + +std::unique_ptr +mlir::createEnableArmStreamingPass(const mlir::ArmStreaming mode) { + return std::make_unique(mode); +} Index: mlir/test/Dialect/Func/enable-arm-streaming.mlir =================================================================== --- /dev/null +++ mlir/test/Dialect/Func/enable-arm-streaming.mlir @@ -0,0 +1,10 @@ +// 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 }