Skip to content

Commit b39f566

Browse files
c-rhodesc-rhodes
authored
[mlir][ArmSME] Add test-lower-to-arm-sme pipeline (#81732)
The ArmSME compilation pipeline has evolved significantly and is now sufficiently complex enough that it warrants a proper lowering pipeline that encapsulates the various passes and orderings. Currently the pipeline is loosely defined in our integration tests, but these have diverged and are not using the same passes or ordering everywhere. This patch introduces a test-lower-to-arm-sme pipeline mirroring test-lower-to-llvm that provides some sanity when running e2e examples and can be used a reference for targeting ArmSME in MLIR. All the integration tests are updated to use this pipeline. The intention is to productize the pipeline once it becomes more mature.
1 parent a622b21 commit b39f566

23 files changed

+141
-103
lines changed

‎mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/fill-2d.mlir

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,8 @@
11
// RUN: mlir-opt %s \
2-
// RUN: -transform-interpreter \
3-
// RUN: -test-transform-dialect-erase-schedule \
2+
// RUN: -transform-interpreter -test-transform-dialect-erase-schedule \
43
// RUN: -lower-vector-mask \
54
// RUN: -one-shot-bufferize="bufferize-function-boundaries" \
6-
// RUN: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za" \
7-
// RUN: -convert-vector-to-arm-sme -convert-arith-to-arm-sme \
8-
// RUN: -allocate-arm-sme-tiles -convert-arm-sme-to-scf \
9-
// RUN: -convert-arm-sme-to-llvm -cse -canonicalize \
10-
// RUN: -test-lower-to-llvm | \
5+
// RUN: -test-lower-to-arm-sme -test-lower-to-llvm | \
116
// RUN: %mcr_aarch64_cmd \
127
// RUN: -e=entry -entry-point-result=void \
138
// RUN: -march=aarch64 -mattr="+sve,+sme" \

‎mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul-transpose-a.mlir

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,7 @@
11
// RUN: mlir-opt %s \
22
// RUN: -transform-interpreter -test-transform-dialect-erase-schedule \
3-
// RUN: -one-shot-bufferize="bufferize-function-boundaries" -canonicalize \
4-
// RUN: -convert-vector-to-arm-sme -allocate-arm-sme-tiles -convert-arm-sme-to-scf \
5-
// RUN: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za only-if-required-by-ops" \
6-
// RUN: -convert-vector-to-scf -cse -arm-sve-legalize-vector-storage \
7-
// RUN: -convert-arm-sme-to-llvm \
8-
// RUN: -convert-vector-to-llvm=enable-arm-sve \
9-
// RUN: -cse -canonicalize -test-lower-to-llvm | \
3+
// RUN: -one-shot-bufferize="bufferize-function-boundaries" \
4+
// RUN: -test-lower-to-arm-sme -test-lower-to-llvm | \
105
// RUN: %mcr_aarch64_cmd \
116
// RUN: -e=main -entry-point-result=void \
127
// RUN: -march=aarch64 -mattr="+sve,+sme" \

‎mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul.mlir

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,6 @@
11
// RUN: mlir-opt %s \
22
// RUN: -transform-interpreter -test-transform-dialect-erase-schedule \
3-
// RUN: -canonicalize \
4-
// RUN: -convert-vector-to-arm-sme -allocate-arm-sme-tiles -convert-arm-sme-to-scf \
5-
// RUN: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za only-if-required-by-ops" \
6-
// RUN: -convert-vector-to-scf -cse -arm-sve-legalize-vector-storage \
7-
// RUN: -convert-arm-sme-to-llvm \
8-
// RUN: -convert-vector-to-llvm=enable-arm-sve \
9-
// RUN: -cse -canonicalize -test-lower-to-llvm | \
3+
// RUN: -test-lower-to-arm-sme -test-lower-to-llvm | \
104
// RUN: %mcr_aarch64_cmd \
115
// RUN: -e=main -entry-point-result=void \
126
// RUN: -march=aarch64 -mattr="+sve,+sme" \

‎mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/multi-tile-matmul.mlir

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,7 @@
11
// RUN: mlir-opt %s \
22
// RUN: -transform-interpreter -test-transform-dialect-erase-schedule \
33
// RUN: -one-shot-bufferize="bufferize-function-boundaries" -canonicalize \
4-
// RUN: -arm-sme-vector-legalization -canonicalize -cse \
5-
// RUN: -convert-vector-to-arm-sme -allocate-arm-sme-tiles -convert-arm-sme-to-scf \
6-
// RUN: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za only-if-required-by-ops" \
7-
// RUN: -convert-vector-to-scf=full-unroll -convert-arm-sme-to-llvm \
8-
// RUN: -test-lower-to-llvm | \
4+
// RUN: -test-lower-to-arm-sme -test-lower-to-llvm | \
95
// RUN: %mcr_aarch64_cmd \
106
// RUN: -e=main -entry-point-result=void \
117
// RUN: -march=aarch64 -mattr="+sve,+sme" \

‎mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/use-too-many-tiles.mlir

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,5 @@
11
// RUN: mlir-opt %s \
2-
// RUN: -convert-vector-to-arm-sme -convert-arith-to-arm-sme \
3-
// RUN: -allocate-arm-sme-tiles -convert-arm-sme-to-scf \
4-
// RUN: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za only-if-required-by-ops" \
5-
// RUN: -convert-vector-to-scf -cse -arm-sve-legalize-vector-storage \
6-
// RUN: -convert-arm-sme-to-llvm -convert-vector-to-llvm=enable-arm-sve -cse \
7-
// RUN: -canonicalize -test-lower-to-llvm -verify-diagnostics | \
2+
// RUN: -test-lower-to-arm-sme -test-lower-to-llvm -verify-diagnostics | \
83
// RUN: %mcr_aarch64_cmd \
94
// RUN: -e=main -entry-point-result=void \
105
// RUN: -march=aarch64 -mattr="+sve,+sme" \

‎mlir/test/Integration/Dialect/Vector/CPU/ArmSME/load-store-128-bit-tile.mlir

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,5 @@
11
// DEFINE: %{entry_point} = test_load_store_zaq0
2-
// DEFINE: %{compile} = mlir-opt %s \
3-
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za" \
4-
// DEFINE: -convert-vector-to-arm-sme -convert-arm-sme-to-scf \
5-
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
6-
// DEFINE: -allocate-arm-sme-tiles -test-lower-to-llvm
2+
// DEFINE: %{compile} = mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm
73
// DEFINE: %{run} = %mcr_aarch64_cmd \
84
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
95
// DEFINE: -e %{entry_point} -entry-point-result=void \

‎mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-load-vertical.mlir

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,5 @@
11
// DEFINE: %{entry_point} = entry
2-
// DEFINE: %{compile} = mlir-opt %s \
3-
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za" \
4-
// DEFINE: -convert-vector-to-arm-sme -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
5-
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
6-
// DEFINE: -test-lower-to-llvm
2+
// DEFINE: %{compile} = mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm
73
// DEFINE: %{run} = %mcr_aarch64_cmd \
84
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
95
// DEFINE: -e %{entry_point} -entry-point-result=void \

‎mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-multi-tile-transpose.mlir

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,4 @@
1-
// RUN: mlir-opt %s -arm-sme-vector-legalization -cse -canonicalize \
2-
// RUN: -convert-vector-to-arm-sme -allocate-arm-sme-tiles -convert-arm-sme-to-scf \
3-
// RUN: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za only-if-required-by-ops" \
4-
// RUN: -convert-vector-to-scf -cse -arm-sve-legalize-vector-storage \
5-
// RUN: -convert-arm-sme-to-llvm \
6-
// RUN: -convert-vector-to-llvm=enable-arm-sve \
7-
// RUN: -cse -canonicalize -test-lower-to-llvm | \
1+
// RUN: mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm | \
82
// RUN: %mcr_aarch64_cmd \
93
// RUN: -e=main -entry-point-result=void \
104
// RUN: -march=aarch64 -mattr="+sve,+sme" \

‎mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-outerproduct-f16f16f32.mlir

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,7 @@
1+
// DEFINE: %{opts} =
12
// DEFINE: %{entry} = main
2-
// DEFINE: %{fusion_opts} = -arm-sme-outer-product-fusion
33
// DEFINE: %{compile} = mlir-opt %s \
4-
// DEFINE: -convert-vector-to-arm-sme -convert-arith-to-arm-sme %{fusion_opts} \
5-
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za only-if-required-by-ops" \
6-
// DEFINE: -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
7-
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
8-
// DEFINE: -test-lower-to-llvm -o %t
4+
// DEFINE: -test-lower-to-arm-sme=%{opts} -test-lower-to-llvm -o %t
95
// DEFINE: %{run} = %mcr_aarch64_cmd %t \
106
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
117
// DEFINE: -e %{entry} -entry-point-result=void \
@@ -18,7 +14,7 @@
1814
// Check result is the same when outerproducts are not combined into widening
1915
// variant.
2016

21-
// REDEFINE: %{fusion_opts} =
17+
// REDEFINE: %{opts} = fuse-outer-products=false
2218
// RUN: %{run} | FileCheck %s
2319

2420
func.func @main() {

‎mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-outerproduct-f32.mlir

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,6 @@
11
// DEFINE: %{entry_point} = test_outerproduct_no_accumulator_4x4xf32
22
// DEFINE: %{compile} = mlir-opt %s \
3-
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za" \
4-
// DEFINE: -convert-vector-to-arm-sme -convert-arith-to-arm-sme \
5-
// DEFINE: -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
6-
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
7-
// DEFINE: -test-lower-to-llvm -o %t
3+
// DEFINE: -test-lower-to-arm-sme -test-lower-to-llvm -o %t
84
// DEFINE: %{run} = %mcr_aarch64_cmd %t \
95
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
106
// DEFINE: -e %{entry_point} -entry-point-result=void \

‎mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-outerproduct-f64.mlir

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,6 @@
11
// DEFINE: %{entry_point} = test_outerproduct_no_accumulator_2x2xf64
22
// DEFINE: %{compile} = mlir-opt %s \
3-
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za" \
4-
// DEFINE: -convert-vector-to-arm-sme -convert-arith-to-arm-sme \
5-
// DEFINE: -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
6-
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
7-
// DEFINE: -test-lower-to-llvm -o %t
3+
// DEFINE: -test-lower-to-arm-sme -test-lower-to-llvm -o %t
84
// DEFINE: %{run} = %mcr_aarch64_cmd %t \
95
// DEFINE: -march=aarch64 -mattr=+sve,+sme-f64f64 \
106
// DEFINE: -e %{entry_point} -entry-point-result=void \

‎mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-outerproduct-i8i8i32.mlir

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,5 @@
11
// DEFINE: %{entry} = main
2-
// DEFINE: %{compile} = mlir-opt %s \
3-
// DEFINE: -convert-vector-to-arm-sme -convert-arith-to-arm-sme \
4-
// DEFINE: -arm-sme-outer-product-fusion \
5-
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za only-if-required-by-ops" \
6-
// DEFINE: -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
7-
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
8-
// DEFINE: -test-lower-to-llvm
2+
// DEFINE: %{compile} = mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm
93
// DEFINE: %{run} = %mcr_aarch64_cmd \
104
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
115
// DEFINE: -e %{entry} -entry-point-result=void \

‎mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-transfer-read-2d.mlir

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,5 @@
11
// DEFINE: %{entry_point} = entry
2-
// DEFINE: %{compile} = mlir-opt %s \
3-
// DEFINE: -convert-vector-to-arm-sme -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
4-
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za only-if-required-by-ops" \
5-
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
6-
// DEFINE: -test-lower-to-llvm
2+
// DEFINE: %{compile} = mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm
73
// DEFINE: %{run} = %mcr_aarch64_cmd \
84
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
95
// DEFINE: -e %{entry_point} -entry-point-result=void \

‎mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-transfer-write-2d.mlir

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,5 @@
11
// DEFINE: %{entry_point} = entry
2-
// DEFINE: %{compile} = mlir-opt %s \
3-
// DEFINE: -convert-vector-to-arm-sme -convert-arith-to-arm-sme \
4-
// DEFINE: -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
5-
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za only-if-required-by-ops" \
6-
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
7-
// DEFINE: -test-lower-to-llvm
2+
// DEFINE: %{compile} = mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm
83
// DEFINE: %{run} = %mcr_aarch64_cmd \
94
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
105
// DEFINE: -e %{entry_point} -entry-point-result=void \

‎mlir/test/Integration/Dialect/Vector/CPU/ArmSME/test-transpose.mlir

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,5 @@
11
// DEFINE: %{entry_point} = entry
2-
// DEFINE: %{compile} = mlir-opt %s \
3-
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za" \
4-
// DEFINE: -convert-vector-to-arm-sme -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
5-
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
6-
// DEFINE: -test-lower-to-llvm
2+
// DEFINE: %{compile} = mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm
73
// DEFINE: %{run} = %mcr_aarch64_cmd \
84
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
95
// DEFINE: -e %{entry_point} -entry-point-result=void \

‎mlir/test/Integration/Dialect/Vector/CPU/ArmSME/tile_fill.mlir

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,4 @@
1-
// RUN: mlir-opt %s -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za" \
2-
// RUN: -convert-vector-to-arm-sme -convert-arith-to-arm-sme \
3-
// RUN: -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
4-
// RUN: -convert-arm-sme-to-llvm -cse -canonicalize \
5-
// RUN: -test-lower-to-llvm | \
1+
// RUN: mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm | \
62
// RUN: %mcr_aarch64_cmd \
73
// RUN: -march=aarch64 -mattr=+sve,+sme \
84
// RUN: -e entry -entry-point-result=i32 \

‎mlir/test/Integration/Dialect/Vector/CPU/ArmSME/vector-load-store.mlir

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,5 @@
11
// DEFINE: %{entry_point} = za0_d_f64
2-
// DEFINE: %{compile} = mlir-opt %s \
3-
// DEFINE: -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za" \
4-
// DEFINE: -convert-vector-to-arm-sme -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
5-
// DEFINE: -convert-arm-sme-to-llvm -cse -canonicalize \
6-
// DEFINE: -test-lower-to-llvm
2+
// DEFINE: %{compile} = mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm
73
// DEFINE: %{run} = %mcr_aarch64_cmd \
84
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
95
// DEFINE: -e %{entry_point} -entry-point-result=i32 \

‎mlir/test/Integration/Dialect/Vector/CPU/ArmSME/vector-ops.mlir

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,5 @@
11
// DEFINE: %{entry_point} = entry
2-
// DEFINE: %{compile} = mlir-opt %s -enable-arm-streaming="streaming-mode=streaming-locally za-mode=new-za" \
3-
// DEFINE: -convert-vector-to-arm-sme -convert-arith-to-arm-sme \
4-
// DEFINE: -convert-arm-sme-to-scf -allocate-arm-sme-tiles \
5-
// DEFINE: -convert-arm-sme-to-llvm -test-lower-to-llvm
2+
// DEFINE: %{compile} = mlir-opt %s -test-lower-to-arm-sme -test-lower-to-llvm
63
// DEFINE: %{run} = %mcr_aarch64_cmd \
74
// DEFINE: -march=aarch64 -mattr=+sve,+sme \
85
// DEFINE: -e %{entry_point} -entry-point-result=i32 \

‎mlir/test/lib/Dialect/ArmSME/CMakeLists.txt

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
# Exclude tests from libMLIR.so
2+
add_mlir_library(MLIRArmSMETestPasses
3+
TestLowerToArmSME.cpp
4+
5+
EXCLUDE_FROM_LIBMLIR
6+
7+
LINK_LIBS PUBLIC
8+
MLIRArithToArmSME
9+
MLIRArmSMEToLLVM
10+
MLIRArmSMEToSCF
11+
MLIRIR
12+
MLIRPass
13+
MLIRTransforms
14+
MLIRVectorToArmSME
15+
MLIRVectorToSCF
16+
)

‎mlir/test/lib/Dialect/ArmSME/TestLowerToArmSME.cpp

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
//===- TestLowerToArmSME.cpp - Test lowering to ArmSME as a sink pass -----===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// This file implements a pass for testing the lowering to ArmSME as a
10+
// generally usable sink pass.
11+
//
12+
//===----------------------------------------------------------------------===//
13+
14+
#include "mlir/Conversion/ArithToArmSME/ArithToArmSME.h"
15+
#include "mlir/Conversion/ArmSMEToLLVM/ArmSMEToLLVM.h"
16+
#include "mlir/Conversion/ArmSMEToSCF/ArmSMEToSCF.h"
17+
#include "mlir/Conversion/VectorToArmSME/VectorToArmSME.h"
18+
#include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
19+
#include "mlir/Dialect/ArmSME/Transforms/Passes.h"
20+
#include "mlir/Dialect/ArmSVE/Transforms/Passes.h"
21+
#include "mlir/IR/DialectRegistry.h"
22+
#include "mlir/Pass/Pass.h"
23+
#include "mlir/Pass/PassManager.h"
24+
#include "mlir/Pass/PassOptions.h"
25+
#include "mlir/Transforms/Passes.h"
26+
27+
using namespace mlir;
28+
29+
namespace {
30+
struct TestLowerToArmSMEOptions
31+
: public PassPipelineOptions<TestLowerToArmSMEOptions> {
32+
PassOptions::Option<bool> fuseOuterProducts{
33+
*this, "fuse-outer-products",
34+
llvm::cl::desc("Fuse outer product operations via "
35+
"'-arm-sme-outer-product-fusion' pass"),
36+
llvm::cl::init(true)};
37+
};
38+
39+
void buildTestLowerToArmSME(OpPassManager &pm,
40+
const TestLowerToArmSMEOptions &options) {
41+
// Legalize vector operations so they can be converted to ArmSME.
42+
pm.addPass(arm_sme::createVectorLegalizationPass());
43+
44+
// Sprinkle some cleanups.
45+
pm.addPass(createCanonicalizerPass());
46+
pm.addPass(createCSEPass());
47+
48+
// Passes that convert operations on vectors to ArmSME operations.
49+
50+
// Convert Arith to ArmSME.
51+
pm.addPass(createArithToArmSMEConversionPass());
52+
// Convert Vector to ArmSME.
53+
pm.addPass(createConvertVectorToArmSMEPass());
54+
55+
// Fuse outer products.
56+
if (options.fuseOuterProducts)
57+
pm.addPass(arm_sme::createOuterProductFusionPass());
58+
59+
// Convert operations on high-level vectors to loops.
60+
61+
// Convert ArmSME to SCF.
62+
pm.addPass(createConvertArmSMEToSCFPass());
63+
64+
// Convert Vector to SCF (with full unroll enabled).
65+
pm.addPass(createConvertVectorToSCFPass(
66+
VectorTransferToSCFOptions().enableFullUnroll()));
67+
68+
// Allocate tiles for ArmSME operations.
69+
//
70+
// Later passes may create further ArmSME ops that implement the
71+
// ArmSMETileOpInterface, but tiles are allocated for root operations,
72+
// all of which should now exist.
73+
pm.addPass(arm_sme::createTileAllocationPass());
74+
75+
// Enable streaming-mode and ZA.
76+
pm.addPass(arm_sme::createEnableArmStreamingPass(
77+
arm_sme::ArmStreamingMode::StreamingLocally, arm_sme::ArmZaMode::NewZA,
78+
/*onlyIfRequiredByOps=*/true));
79+
80+
// Convert ArmSME to LLVM.
81+
pm.addPass(createConvertArmSMEToLLVMPass());
82+
83+
// Sprinkle some cleanups.
84+
pm.addPass(createCanonicalizerPass());
85+
pm.addPass(createCSEPass());
86+
}
87+
} // namespace
88+
89+
namespace mlir {
90+
namespace test {
91+
void registerTestLowerToArmSME() {
92+
PassPipelineRegistration<TestLowerToArmSMEOptions>(
93+
"test-lower-to-arm-sme",
94+
"An example pipeline to lower operations on vectors (arith, vector) to "
95+
"LLVM via ArmSME.",
96+
buildTestLowerToArmSME);
97+
}
98+
} // namespace test
99+
} // namespace mlir

‎mlir/test/lib/Dialect/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
add_subdirectory(Affine)
22
add_subdirectory(Arith)
3+
add_subdirectory(ArmSME)
34
add_subdirectory(Bufferization)
45
add_subdirectory(ControlFlow)
56
add_subdirectory(DLTI)

‎mlir/tools/mlir-opt/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ if(MLIR_INCLUDE_TESTS)
1717
MLIRTestFuncToLLVM
1818
MLIRAffineTransformsTestPasses
1919
MLIRArithTestPasses
20+
MLIRArmSMETestPasses
2021
MLIRBufferizationTestPasses
2122
MLIRControlFlowTestPasses
2223
MLIRDLTITestPasses

‎mlir/tools/mlir-opt/mlir-opt.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,7 @@ void registerTestLoopFusion();
109109
void registerTestCFGLoopInfoPass();
110110
void registerTestLoopMappingPass();
111111
void registerTestLoopUnrollingPass();
112+
void registerTestLowerToArmSME();
112113
void registerTestLowerToLLVM();
113114
void registerTestMakeIsolatedFromAbovePass();
114115
void registerTestMatchReductionPass();
@@ -233,6 +234,7 @@ void registerTestPasses() {
233234
mlir::test::registerTestCFGLoopInfoPass();
234235
mlir::test::registerTestLoopMappingPass();
235236
mlir::test::registerTestLoopUnrollingPass();
237+
mlir::test::registerTestLowerToArmSME();
236238
mlir::test::registerTestLowerToLLVM();
237239
mlir::test::registerTestMakeIsolatedFromAbovePass();
238240
mlir::test::registerTestMatchReductionPass();

0 commit comments

Comments
 (0)