Skip to content

[mlir][ArmSME] Add test-lower-to-arm-sme pipeline #81732

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Feb 23, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -1,13 +1,8 @@
// RUN: mlir-opt %s \
// RUN: -transform-interpreter \
// RUN: -test-transform-dialect-erase-schedule \
// RUN: -transform-interpreter -test-transform-dialect-erase-schedule \
// RUN: -lower-vector-mask \
// RUN: -one-shot-bufferize="bufferize-function-boundaries" \
// RUN: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za" \
// RUN: -convert-vector-to-arm-sme -convert-arith-to-arm-sme \
// RUN: -allocate-arm-sme-tiles -convert-arm-sme-to-scf \
// RUN: -convert-arm-sme-to-llvm -cse -canonicalize \
// RUN: -test-lower-to-llvm | \
// RUN: -test-lower-to-arm-sme -test-lower-to-llvm | \
// RUN: %mcr_aarch64_cmd \
// RUN: -e=entry -entry-point-result=void \
// RUN: -march=aarch64 -mattr="+sve,+sme" \
Expand Down
Original file line number Diff line number Diff line change
@@ -1,12 +1,7 @@
// RUN: mlir-opt %s \
// RUN: -transform-interpreter -test-transform-dialect-erase-schedule \
// RUN: -one-shot-bufferize="bufferize-function-boundaries" -canonicalize \
// RUN: -convert-vector-to-arm-sme -allocate-arm-sme-tiles -convert-arm-sme-to-scf \
// RUN: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za only-if-required-by-ops" \
// RUN: -convert-vector-to-scf -cse -arm-sve-legalize-vector-storage \
// RUN: -convert-arm-sme-to-llvm \
// RUN: -convert-vector-to-llvm=enable-arm-sve \
// RUN: -cse -canonicalize -test-lower-to-llvm | \
// RUN: -one-shot-bufferize="bufferize-function-boundaries" \
// RUN: -test-lower-to-arm-sme -test-lower-to-llvm | \
// RUN: %mcr_aarch64_cmd \
// RUN: -e=main -entry-point-result=void \
// RUN: -march=aarch64 -mattr="+sve,+sme" \
Expand Down
Original file line number Diff line number Diff line change
@@ -1,12 +1,6 @@
// RUN: mlir-opt %s \
// RUN: -transform-interpreter -test-transform-dialect-erase-schedule \
// RUN: -canonicalize \
// RUN: -convert-vector-to-arm-sme -allocate-arm-sme-tiles -convert-arm-sme-to-scf \
// RUN: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za only-if-required-by-ops" \
// RUN: -convert-vector-to-scf -cse -arm-sve-legalize-vector-storage \
// RUN: -convert-arm-sme-to-llvm \
// RUN: -convert-vector-to-llvm=enable-arm-sve \
// RUN: -cse -canonicalize -test-lower-to-llvm | \
// RUN: -test-lower-to-arm-sme -test-lower-to-llvm | \
// RUN: %mcr_aarch64_cmd \
// RUN: -e=main -entry-point-result=void \
// RUN: -march=aarch64 -mattr="+sve,+sme" \
Expand Down
Original file line number Diff line number Diff line change
@@ -1,11 +1,7 @@
// RUN: mlir-opt %s \
// RUN: -transform-interpreter -test-transform-dialect-erase-schedule \
// RUN: -one-shot-bufferize="bufferize-function-boundaries" -canonicalize \
// RUN: -arm-sme-vector-legalization -canonicalize -cse \
// RUN: -convert-vector-to-arm-sme -allocate-arm-sme-tiles -convert-arm-sme-to-scf \
// RUN: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za only-if-required-by-ops" \
// RUN: -convert-vector-to-scf=full-unroll -convert-arm-sme-to-llvm \
// RUN: -test-lower-to-llvm | \
// RUN: -test-lower-to-arm-sme -test-lower-to-llvm | \
// RUN: %mcr_aarch64_cmd \
// RUN: -e=main -entry-point-result=void \
// RUN: -march=aarch64 -mattr="+sve,+sme" \
Expand Down
Original file line number Diff line number Diff line change
@@ -1,10 +1,5 @@
// RUN: mlir-opt %s \
// RUN: -convert-vector-to-arm-sme -convert-arith-to-arm-sme \
// RUN: -allocate-arm-sme-tiles -convert-arm-sme-to-scf \
// RUN: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za only-if-required-by-ops" \
// RUN: -convert-vector-to-scf -cse -arm-sve-legalize-vector-storage \
// RUN: -convert-arm-sme-to-llvm -convert-vector-to-llvm=enable-arm-sve -cse \
// RUN: -canonicalize -test-lower-to-llvm -verify-diagnostics | \
// RUN: -test-lower-to-arm-sme -test-lower-to-llvm -verify-diagnostics | \
// RUN: %mcr_aarch64_cmd \
// RUN: -e=main -entry-point-result=void \
// RUN: -march=aarch64 -mattr="+sve,+sme" \
Expand Down
Original file line number Diff line number Diff line change
@@ -1,9 +1,5 @@
// DEFINE: %{entry_point} = test_load_store_zaq0
// DEFINE: %{compile} = mlir-opt %s \
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za" \
// DEFINE: -convert-vector-to-arm-sme -convert-arm-sme-to-scf \
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
// DEFINE: -allocate-arm-sme-tiles -test-lower-to-llvm
// DEFINE: %{compile} = mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm
// DEFINE: %{run} = %mcr_aarch64_cmd \
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
// DEFINE: -e %{entry_point} -entry-point-result=void \
Expand Down
Original file line number Diff line number Diff line change
@@ -1,9 +1,5 @@
// DEFINE: %{entry_point} = entry
// DEFINE: %{compile} = mlir-opt %s \
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za" \
// DEFINE: -convert-vector-to-arm-sme -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
// DEFINE: -test-lower-to-llvm
// DEFINE: %{compile} = mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm
// DEFINE: %{run} = %mcr_aarch64_cmd \
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
// DEFINE: -e %{entry_point} -entry-point-result=void \
Expand Down
Original file line number Diff line number Diff line change
@@ -1,10 +1,4 @@
// RUN: mlir-opt %s -arm-sme-vector-legalization -cse -canonicalize \
// RUN: -convert-vector-to-arm-sme -allocate-arm-sme-tiles -convert-arm-sme-to-scf \
// RUN: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za only-if-required-by-ops" \
// RUN: -convert-vector-to-scf -cse -arm-sve-legalize-vector-storage \
// RUN: -convert-arm-sme-to-llvm \
// RUN: -convert-vector-to-llvm=enable-arm-sve \
// RUN: -cse -canonicalize -test-lower-to-llvm | \
// RUN: mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm | \
// RUN: %mcr_aarch64_cmd \
// RUN: -e=main -entry-point-result=void \
// RUN: -march=aarch64 -mattr="+sve,+sme" \
Expand Down
Original file line number Diff line number Diff line change
@@ -1,11 +1,7 @@
// DEFINE: %{opts} =
// DEFINE: %{entry} = main
// DEFINE: %{fusion_opts} = -arm-sme-outer-product-fusion
// DEFINE: %{compile} = mlir-opt %s \
// DEFINE: -convert-vector-to-arm-sme -convert-arith-to-arm-sme %{fusion_opts} \
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za only-if-required-by-ops" \
// DEFINE: -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
// DEFINE: -test-lower-to-llvm -o %t
// DEFINE: -test-lower-to-arm-sme=%{opts} -test-lower-to-llvm -o %t
// DEFINE: %{run} = %mcr_aarch64_cmd %t \
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
// DEFINE: -e %{entry} -entry-point-result=void \
Expand All @@ -18,7 +14,7 @@
// Check result is the same when outerproducts are not combined into widening
// variant.

// REDEFINE: %{fusion_opts} =
// REDEFINE: %{opts} = fuse-outer-products=false
// RUN: %{run} | FileCheck %s

func.func @main() {
Expand Down
Original file line number Diff line number Diff line change
@@ -1,10 +1,6 @@
// DEFINE: %{entry_point} = test_outerproduct_no_accumulator_4x4xf32
// DEFINE: %{compile} = mlir-opt %s \
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za" \
// DEFINE: -convert-vector-to-arm-sme -convert-arith-to-arm-sme \
// DEFINE: -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
// DEFINE: -test-lower-to-llvm -o %t
// DEFINE: -test-lower-to-arm-sme -test-lower-to-llvm -o %t
// DEFINE: %{run} = %mcr_aarch64_cmd %t \
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
// DEFINE: -e %{entry_point} -entry-point-result=void \
Expand Down
Original file line number Diff line number Diff line change
@@ -1,10 +1,6 @@
// DEFINE: %{entry_point} = test_outerproduct_no_accumulator_2x2xf64
// DEFINE: %{compile} = mlir-opt %s \
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za" \
// DEFINE: -convert-vector-to-arm-sme -convert-arith-to-arm-sme \
// DEFINE: -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
// DEFINE: -test-lower-to-llvm -o %t
// DEFINE: -test-lower-to-arm-sme -test-lower-to-llvm -o %t
// DEFINE: %{run} = %mcr_aarch64_cmd %t \
// DEFINE: -march=aarch64 -mattr=+sve,+sme-f64f64 \
// DEFINE: -e %{entry_point} -entry-point-result=void \
Expand Down
Original file line number Diff line number Diff line change
@@ -1,11 +1,5 @@
// DEFINE: %{entry} = main
// DEFINE: %{compile} = mlir-opt %s \
// DEFINE: -convert-vector-to-arm-sme -convert-arith-to-arm-sme \
// DEFINE: -arm-sme-outer-product-fusion \
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za only-if-required-by-ops" \
// DEFINE: -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
// DEFINE: -test-lower-to-llvm
// DEFINE: %{compile} = mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm
// DEFINE: %{run} = %mcr_aarch64_cmd \
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
// DEFINE: -e %{entry} -entry-point-result=void \
Expand Down
Original file line number Diff line number Diff line change
@@ -1,9 +1,5 @@
// DEFINE: %{entry_point} = entry
// DEFINE: %{compile} = mlir-opt %s \
// DEFINE: -convert-vector-to-arm-sme -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za only-if-required-by-ops" \
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
// DEFINE: -test-lower-to-llvm
// DEFINE: %{compile} = mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm
// DEFINE: %{run} = %mcr_aarch64_cmd \
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
// DEFINE: -e %{entry_point} -entry-point-result=void \
Expand Down
Original file line number Diff line number Diff line change
@@ -1,10 +1,5 @@
// DEFINE: %{entry_point} = entry
// DEFINE: %{compile} = mlir-opt %s \
// DEFINE: -convert-vector-to-arm-sme -convert-arith-to-arm-sme \
// DEFINE: -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za only-if-required-by-ops" \
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
// DEFINE: -test-lower-to-llvm
// DEFINE: %{compile} = mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm
// DEFINE: %{run} = %mcr_aarch64_cmd \
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
// DEFINE: -e %{entry_point} -entry-point-result=void \
Expand Down
Original file line number Diff line number Diff line change
@@ -1,9 +1,5 @@
// DEFINE: %{entry_point} = entry
// DEFINE: %{compile} = mlir-opt %s \
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za" \
// DEFINE: -convert-vector-to-arm-sme -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
// DEFINE: -test-lower-to-llvm
// DEFINE: %{compile} = mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm
// DEFINE: %{run} = %mcr_aarch64_cmd \
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
// DEFINE: -e %{entry_point} -entry-point-result=void \
Expand Down
Original file line number Diff line number Diff line change
@@ -1,8 +1,4 @@
// RUN: mlir-opt %s -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za" \
// RUN: -convert-vector-to-arm-sme -convert-arith-to-arm-sme \
// RUN: -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
// RUN: -convert-arm-sme-to-llvm -cse -canonicalize \
// RUN: -test-lower-to-llvm | \
// RUN: mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm | \
// RUN: %mcr_aarch64_cmd \
// RUN: -march=aarch64 -mattr=+sve,+sme \
// RUN: -e entry -entry-point-result=i32 \
Expand Down
Original file line number Diff line number Diff line change
@@ -1,9 +1,5 @@
// DEFINE: %{entry_point} = za0_d_f64
// DEFINE: %{compile} = mlir-opt %s \
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za" \
// DEFINE: -convert-vector-to-arm-sme -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
// DEFINE: -test-lower-to-llvm
// DEFINE: %{compile} = mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm
// DEFINE: %{run} = %mcr_aarch64_cmd \
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
// DEFINE: -e %{entry_point} -entry-point-result=i32 \
Expand Down
Original file line number Diff line number Diff line change
@@ -1,8 +1,5 @@
// DEFINE: %{entry_point} = entry
// DEFINE: %{compile} = mlir-opt %s -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za" \
// DEFINE: -convert-vector-to-arm-sme -convert-arith-to-arm-sme \
// DEFINE: -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
// DEFINE: -convert-arm-sme-to-llvm -test-lower-to-llvm
// DEFINE: %{compile} = mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm
// DEFINE: %{run} = %mcr_aarch64_cmd \
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
// DEFINE: -e %{entry_point} -entry-point-result=i32 \
Expand Down
16 changes: 16 additions & 0 deletions mlir/test/lib/Dialect/ArmSME/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
# Exclude tests from libMLIR.so
add_mlir_library(MLIRArmSMETestPasses
TestLowerToArmSME.cpp

EXCLUDE_FROM_LIBMLIR

LINK_LIBS PUBLIC
MLIRArithToArmSME
MLIRArmSMEToLLVM
MLIRArmSMEToSCF
MLIRIR
MLIRPass
MLIRTransforms
MLIRVectorToArmSME
MLIRVectorToSCF
)
99 changes: 99 additions & 0 deletions mlir/test/lib/Dialect/ArmSME/TestLowerToArmSME.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,99 @@
//===- TestLowerToArmSME.cpp - Test lowering to ArmSME as a sink pass -----===//
//
// 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 file implements a pass for testing the lowering to ArmSME as a
// generally usable sink pass.
//
//===----------------------------------------------------------------------===//

#include "mlir/Conversion/ArithToArmSME/ArithToArmSME.h"
#include "mlir/Conversion/ArmSMEToLLVM/ArmSMEToLLVM.h"
#include "mlir/Conversion/ArmSMEToSCF/ArmSMEToSCF.h"
#include "mlir/Conversion/VectorToArmSME/VectorToArmSME.h"
#include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
#include "mlir/Dialect/ArmSME/Transforms/Passes.h"
#include "mlir/Dialect/ArmSVE/Transforms/Passes.h"
#include "mlir/IR/DialectRegistry.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Pass/PassOptions.h"
#include "mlir/Transforms/Passes.h"

using namespace mlir;

namespace {
struct TestLowerToArmSMEOptions
: public PassPipelineOptions<TestLowerToArmSMEOptions> {
PassOptions::Option<bool> fuseOuterProducts{
*this, "fuse-outer-products",
llvm::cl::desc("Fuse outer product operations via "
"'-arm-sme-outer-product-fusion' pass"),
llvm::cl::init(true)};
};

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you leave a comment with the full pipeline in a form that can be passed mlir-opt? I'm pretty sure in the future we'll want to try different orders/passes, so I think that'd be helpful.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

-dump-pass-pipeline can be used for that? This is what it looks like

mlir-opt input.mlir -test-lower-to-arm-sme -dump-pass-pipeline

Pass Manager with 13 passes:
builtin.module(arm-sme-vector-legalization,canonicalize{  max-iterations=10 max-num-rewrites=-1 region-simplify=true test-convergence=false top-down=true},cse,convert-arith-to-arm-sme,convert-vector-to-arm-sme,func.func(arm-sme-outer-product-fusion),convert-arm-sme-to-scf,convert-vector-to-scf{full-unroll=true lower-tensors=false target-rank=1},func.func(allocate-arm-sme-tiles),func.func(enable-arm-streaming{only-if-required-by-ops=true streaming-mode=streaming-locally za-mode=new-za}),convert-arm-sme-to-llvm,canonicalize{  max-iterations=10 max-num-rewrites=-1 region-simplify=true test-convergence=false top-down=true},cse)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think that's still more noisy than I'd like, I'd prefer the simple form (e.g. -convert-vector-to-arm-sme -convert-arith-to-arm-sme ...).

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the comment will be formatted by clang-format and not easy to copy/paste. It'll also have to be kept in sync with the pipeline.

Copy link
Member

@MacDue MacDue Feb 22, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My concern is this change makes it harder experiment with the passes for an ArmSME test. You either have to use -dump-pass-pipeline which is in a pretty inconvenient form, or edit the passes in the pipeline and recompile.

I normally copy passes from a test (or from the LLVM lit logs) then play around with the order or adding/removing them, and this obscures that.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

copy/pasting:

  // -arm-sme-vector-legalization -canonicalize -cse -convert-arith-to-arm-sme
  // -convert-vector-to-arm-sme -arm-sme-outer-product-fusion
  // -convert-arm-sme-to-scf -convert-vector-to-scf=full-unroll
  // -allocate-arm-sme-tiles
  // -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za
  // only-if-required-by-ops" -convert-arm-sme-to-llvm -canonicalize -cse

seems less convenient than -dump-pass-pipeline to me. There's a bit of noise mostly from canonicalization options but it's not the end of the world.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Feel free to land this and we can see how things go 🙂

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think that's still more noisy than I'd like, I'd prefer the simple form (e.g. -convert-vector-to-arm-sme -convert-arith-to-arm-sme ...).

This seems hardly like a reason to encode this in the codebase IMO, this isn't really maintainable.

More importantly the textual pipeline is the only real correct form to describe a pipeline: the short one is only a convenience for running single passes but ambiguous when forming a more complex pipeline.
The option parsing for example can be very confusing when the same pass is attempted to be executed multiple time with different options.

If you would like to skip the default value for the options that aren't touched, this is likely doable: patch welcome for -dump-pass-pipeline here!

void buildTestLowerToArmSME(OpPassManager &pm,
const TestLowerToArmSMEOptions &options) {
// Legalize vector operations so they can be converted to ArmSME.
pm.addPass(arm_sme::createVectorLegalizationPass());

// Sprinkle some cleanups.
pm.addPass(createCanonicalizerPass());
pm.addPass(createCSEPass());

// Passes that convert operations on vectors to ArmSME operations.

// Convert Arith to ArmSME.
pm.addPass(createArithToArmSMEConversionPass());
// Convert Vector to ArmSME.
pm.addPass(createConvertVectorToArmSMEPass());

// Fuse outer products.
if (options.fuseOuterProducts)
pm.addPass(arm_sme::createOuterProductFusionPass());

// Convert operations on high-level vectors to loops.

// Convert ArmSME to SCF.
pm.addPass(createConvertArmSMEToSCFPass());

// Convert Vector to SCF (with full unroll enabled).
pm.addPass(createConvertVectorToSCFPass(
VectorTransferToSCFOptions().enableFullUnroll()));

// Allocate tiles for ArmSME operations.
//
// Later passes may create further ArmSME ops that implement the
// ArmSMETileOpInterface, but tiles are allocated for root operations,
// all of which should now exist.
pm.addPass(arm_sme::createTileAllocationPass());

// Enable streaming-mode and ZA.
pm.addPass(arm_sme::createEnableArmStreamingPass(
arm_sme::ArmStreamingMode::StreamingLocally, arm_sme::ArmZaMode::NewZA,
/*onlyIfRequiredByOps=*/true));

// Convert ArmSME to LLVM.
pm.addPass(createConvertArmSMEToLLVMPass());

// Sprinkle some cleanups.
pm.addPass(createCanonicalizerPass());
pm.addPass(createCSEPass());
}
} // namespace

namespace mlir {
namespace test {
void registerTestLowerToArmSME() {
PassPipelineRegistration<TestLowerToArmSMEOptions>(
"test-lower-to-arm-sme",
"An example pipeline to lower operations on vectors (arith, vector) to "
"LLVM via ArmSME.",
buildTestLowerToArmSME);
}
} // namespace test
} // namespace mlir
1 change: 1 addition & 0 deletions mlir/test/lib/Dialect/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
add_subdirectory(Affine)
add_subdirectory(Arith)
add_subdirectory(ArmSME)
add_subdirectory(Bufferization)
add_subdirectory(ControlFlow)
add_subdirectory(DLTI)
Expand Down
1 change: 1 addition & 0 deletions mlir/tools/mlir-opt/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ if(MLIR_INCLUDE_TESTS)
MLIRTestFuncToLLVM
MLIRAffineTransformsTestPasses
MLIRArithTestPasses
MLIRArmSMETestPasses
MLIRBufferizationTestPasses
MLIRControlFlowTestPasses
MLIRDLTITestPasses
Expand Down
2 changes: 2 additions & 0 deletions mlir/tools/mlir-opt/mlir-opt.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,7 @@ void registerTestLoopFusion();
void registerTestCFGLoopInfoPass();
void registerTestLoopMappingPass();
void registerTestLoopUnrollingPass();
void registerTestLowerToArmSME();
void registerTestLowerToLLVM();
void registerTestMakeIsolatedFromAbovePass();
void registerTestMatchReductionPass();
Expand Down Expand Up @@ -233,6 +234,7 @@ void registerTestPasses() {
mlir::test::registerTestCFGLoopInfoPass();
mlir::test::registerTestLoopMappingPass();
mlir::test::registerTestLoopUnrollingPass();
mlir::test::registerTestLowerToArmSME();
mlir::test::registerTestLowerToLLVM();
mlir::test::registerTestMakeIsolatedFromAbovePass();
mlir::test::registerTestMatchReductionPass();
Expand Down