[Mlir-commits] [mlir] [mlir][SVE] Add an e2e test for vector.contract (PR #69845)
Andrzej Warzyński
llvmlistbot at llvm.org
Tue Oct 24 07:22:32 PDT 2023
https://github.com/banach-space updated https://github.com/llvm/llvm-project/pull/69845
>From 2ba4626eafcc45244b46579bd3e9f130f57c4fa8 Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Tue, 10 Oct 2023 18:53:05 +0000
Subject: [PATCH 1/3] [mlir][ArmSVE] Add `-arm-sve-legalize-vector-storage`
pass
This patch adds a pass that ensures that loads, stores, and allocations
of SVE vector types will be legal in the LLVM backend. It does this at
the memref level, so this pass must be applied before lowering all the
way to LLVM.
This pass currently fixes two issues.
It is only legal to load/store predicate types equal to (or greater
than) a full predicate register, which in MLIR is `vector<[16]xi1>`.
Smaller predicate types (`vector<[1|2|4|8]xi1>`) must be converted
to/from a full predicate type (referred to as a `svbool`) before and
after storing and loading respectively. This pass does this by widening
allocations and inserting conversion intrinsics.
For example:
```mlir
%alloca = memref.alloca() : memref<vector<[4]xi1>>
%mask = vector.constant_mask [4] : vector<[4]xi1>
memref.store %mask, %alloca[] : memref<vector<[4]xi1>>
%reload = memref.load %alloca[] : memref<vector<[4]xi1>>
```
Becomes:
```mlir
%alloca = memref.alloca() {alignment = 1 : i64} : memref<vector<[16]xi1>>
%mask = vector.constant_mask [4] : vector<[4]xi1>
%svbool = arm_sve.convert_to_svbool %mask : vector<[4]xi1>
memref.store %svbool, %alloca[] : memref<vector<[16]xi1>>
%reload_svbool = memref.load %alloca[] : memref<vector<[16]xi1>>
%reload = arm_sve.convert_from_svbool %reload_svbool : vector<[4]xi1>
```
The storage for SVE vector types only needs to have an alignment that
matches the element type (for example 4 byte alignment for `f32`s).
However, the LLVM backend currently defaults to aligning to `base size`
x `element size` bytes. For non-legal vector types like
`vector<[8]xf32>` this results in 8 x 4 = 32-byte alignment, but the
backend only supports up to 16-byte alignment for SVE vectors on the
stack. Explicitly setting a smaller alignment prevents this issue.
---
.../mlir/Dialect/ArmSVE/CMakeLists.txt | 1 +
.../Dialect/ArmSVE/Transforms/CMakeLists.txt | 5 +
.../mlir/Dialect/ArmSVE/Transforms/Passes.h | 33 ++
.../mlir/Dialect/ArmSVE/Transforms/Passes.td | 67 ++++
mlir/include/mlir/InitAllPasses.h | 2 +
.../Dialect/ArmSVE/Transforms/CMakeLists.txt | 2 +
.../Transforms/LegalizeVectorStorage.cpp | 310 ++++++++++++++++++
.../ArmSVE/legalize-vector-storage.mlir | 160 +++++++++
.../ArmSVE/arrays-of-scalable-vectors.mlir | 121 +++++++
9 files changed, 701 insertions(+)
create mode 100644 mlir/include/mlir/Dialect/ArmSVE/Transforms/CMakeLists.txt
create mode 100644 mlir/include/mlir/Dialect/ArmSVE/Transforms/Passes.h
create mode 100644 mlir/include/mlir/Dialect/ArmSVE/Transforms/Passes.td
create mode 100644 mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp
create mode 100644 mlir/test/Dialect/ArmSVE/legalize-vector-storage.mlir
create mode 100644 mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/arrays-of-scalable-vectors.mlir
diff --git a/mlir/include/mlir/Dialect/ArmSVE/CMakeLists.txt b/mlir/include/mlir/Dialect/ArmSVE/CMakeLists.txt
index f33061b2d87cffc..9f57627c321fb0c 100644
--- a/mlir/include/mlir/Dialect/ArmSVE/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/ArmSVE/CMakeLists.txt
@@ -1 +1,2 @@
add_subdirectory(IR)
+add_subdirectory(Transforms)
diff --git a/mlir/include/mlir/Dialect/ArmSVE/Transforms/CMakeLists.txt b/mlir/include/mlir/Dialect/ArmSVE/Transforms/CMakeLists.txt
new file mode 100644
index 000000000000000..7226642daf86172
--- /dev/null
+++ b/mlir/include/mlir/Dialect/ArmSVE/Transforms/CMakeLists.txt
@@ -0,0 +1,5 @@
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name ArmSVE)
+add_public_tablegen_target(MLIRArmSVEPassIncGen)
+
+add_mlir_doc(Passes ArmSVEPasses ./ -gen-pass-doc)
diff --git a/mlir/include/mlir/Dialect/ArmSVE/Transforms/Passes.h b/mlir/include/mlir/Dialect/ArmSVE/Transforms/Passes.h
new file mode 100644
index 000000000000000..317fb9021b3c577
--- /dev/null
+++ b/mlir/include/mlir/Dialect/ArmSVE/Transforms/Passes.h
@@ -0,0 +1,33 @@
+//===- Passes.h - Pass Entrypoints ------------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_ARMSVE_TRANSFORMS_PASSES_H
+#define MLIR_DIALECT_ARMSVE_TRANSFORMS_PASSES_H
+
+#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir::arm_sve {
+
+#define GEN_PASS_DECL
+#include "mlir/Dialect/ArmSVE/Transforms/Passes.h.inc"
+
+/// Pass to legalize the types of mask stores.
+std::unique_ptr<Pass> createLegalizeVectorStoragePass();
+
+//===----------------------------------------------------------------------===//
+// Registration
+//===----------------------------------------------------------------------===//
+
+/// Generate the code for registering passes.
+#define GEN_PASS_REGISTRATION
+#include "mlir/Dialect/ArmSVE/Transforms/Passes.h.inc"
+
+} // namespace mlir::arm_sve
+
+#endif // MLIR_DIALECT_ARMSVE_TRANSFORMS_PASSES_H
diff --git a/mlir/include/mlir/Dialect/ArmSVE/Transforms/Passes.td b/mlir/include/mlir/Dialect/ArmSVE/Transforms/Passes.td
new file mode 100644
index 000000000000000..35c49607181da0c
--- /dev/null
+++ b/mlir/include/mlir/Dialect/ArmSVE/Transforms/Passes.td
@@ -0,0 +1,67 @@
+//===-- Passes.td - ArmSVE pass definition file ------------*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_ARMSVE_TRANSFORMS_PASSES_TD
+#define MLIR_DIALECT_ARMSVE_TRANSFORMS_PASSES_TD
+
+include "mlir/Pass/PassBase.td"
+
+def LegalizeVectorStorage
+ : Pass<"arm-sve-legalize-vector-storage", "mlir::func::FuncOp"> {
+ let summary = "Ensures stores of SVE vector types will be legal";
+ let description = [{
+ This pass ensures that loads, stores, and allocations of SVE vector types
+ will be legal in the LLVM backend. It does this at the memref level, so this
+ pass must be applied before lowering all the way to LLVM.
+
+ This pass currently fixes two issues.
+
+ ## Loading and storing predicate types
+
+ It is only legal to load/store predicate types equal to (or greater than) a
+ full predicate register, which in MLIR is `vector<[16]xi1>`. Smaller
+ predicate types (`vector<[1|2|4|8]xi1>`) must be converted to/from a full
+ predicate type (referred to as a `svbool`) before and after storing and
+ loading respectively. This pass does this by widening allocations and
+ inserting conversion intrinsics.
+
+ For example:
+
+ ```mlir
+ %alloca = memref.alloca() : memref<vector<[4]xi1>>
+ %mask = vector.constant_mask [4] : vector<[4]xi1>
+ memref.store %mask, %alloca[] : memref<vector<[4]xi1>>
+ %reload = memref.load %alloca[] : memref<vector<[4]xi1>>
+ ```
+ Becomes:
+ ```mlir
+ %alloca = memref.alloca() {alignment = 1 : i64} : memref<vector<[16]xi1>>
+ %mask = vector.constant_mask [4] : vector<[4]xi1>
+ %svbool = arm_sve.convert_to_svbool %mask : vector<[4]xi1>
+ memref.store %svbool, %alloca[] : memref<vector<[16]xi1>>
+ %reload_svbool = memref.load %alloca[] : memref<vector<[16]xi1>>
+ %reload = arm_sve.convert_from_svbool %reload_svbool : vector<[4]xi1>
+ ```
+
+ ## Relax alignments for SVE vector allocas
+
+ The storage for SVE vector types only needs to have an alignment that
+ matches the element type (for example 4 byte alignment for `f32`s). However,
+ the LLVM backend currently defaults to aligning to `base size` x
+ `element size` bytes. For non-legal vector types like `vector<[8]xf32>` this
+ results in 8 x 4 = 32-byte alignment, but the backend only supports up to
+ 16-byte alignment for SVE vectors on the stack. Explicitly setting a smaller
+ alignment prevents this issue.
+ }];
+ let constructor = "mlir::arm_sve::createLegalizeVectorStoragePass()";
+ let dependentDialects = ["func::FuncDialect",
+ "memref::MemRefDialect", "vector::VectorDialect",
+ "arm_sve::ArmSVEDialect"];
+}
+
+#endif // MLIR_DIALECT_ARMSVE_TRANSFORMS_PASSES_TD
diff --git a/mlir/include/mlir/InitAllPasses.h b/mlir/include/mlir/InitAllPasses.h
index 5489a13a8040bdb..7301905954f56d8 100644
--- a/mlir/include/mlir/InitAllPasses.h
+++ b/mlir/include/mlir/InitAllPasses.h
@@ -19,6 +19,7 @@
#include "mlir/Dialect/Affine/Passes.h"
#include "mlir/Dialect/Arith/Transforms/Passes.h"
#include "mlir/Dialect/ArmSME/Transforms/Passes.h"
+#include "mlir/Dialect/ArmSVE/Transforms/Passes.h"
#include "mlir/Dialect/Async/Passes.h"
#include "mlir/Dialect/Bufferization/Pipelines/Passes.h"
#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
@@ -82,6 +83,7 @@ inline void registerAllPasses() {
transform::registerTransformPasses();
vector::registerVectorPasses();
arm_sme::registerArmSMEPasses();
+ arm_sve::registerArmSVEPasses();
// Dialect pipelines
bufferization::registerBufferizationPipelines();
diff --git a/mlir/lib/Dialect/ArmSVE/Transforms/CMakeLists.txt b/mlir/lib/Dialect/ArmSVE/Transforms/CMakeLists.txt
index 2f1c43fae240d51..a70c489a51fea9a 100644
--- a/mlir/lib/Dialect/ArmSVE/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/ArmSVE/Transforms/CMakeLists.txt
@@ -1,8 +1,10 @@
add_mlir_dialect_library(MLIRArmSVETransforms
LegalizeForLLVMExport.cpp
+ LegalizeVectorStorage.cpp
DEPENDS
MLIRArmSVEConversionsIncGen
+ MLIRArmSVEPassIncGen
LINK_LIBS PUBLIC
MLIRArmSVEDialect
diff --git a/mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp b/mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp
new file mode 100644
index 000000000000000..610eb38089c4c88
--- /dev/null
+++ b/mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp
@@ -0,0 +1,310 @@
+//===- LegalizeVectorStorage.cpp - Ensures SVE loads/stores are legal -----===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/ArmSVE/IR/ArmSVEDialect.h"
+#include "mlir/Dialect/ArmSVE/Transforms/Passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::arm_sve {
+#define GEN_PASS_DEF_LEGALIZEVECTORSTORAGE
+#include "mlir/Dialect/ArmSVE/Transforms/Passes.h.inc"
+} // namespace mlir::arm_sve
+
+using namespace mlir;
+using namespace mlir::arm_sve;
+
+constexpr StringLiteral kPassLabel("__arm_sve_legalize_vector_storage__");
+
+namespace {
+
+/// A (legal) SVE predicate mask that has a logical size, i.e. the number of
+/// bits match the number of lanes it masks (such as vector<[4]xi1>), but is too
+/// small to be stored to memory.
+bool isLogicalSVEPredicateType(VectorType type) {
+ return type.getRank() > 0 && type.getElementType().isInteger(1) &&
+ type.getScalableDims().back() && type.getShape().back() < 16 &&
+ llvm::isPowerOf2_32(type.getShape().back()) &&
+ !llvm::is_contained(type.getScalableDims().drop_back(), true);
+}
+
+VectorType widenScalableMaskTypeToSvbool(VectorType type) {
+ assert(isLogicalSVEPredicateType(type));
+ return VectorType::Builder(type).setDim(type.getRank() - 1, 16);
+}
+
+template <typename TOp, typename TLegalizerCallback>
+void replaceOpWithLegalizedOp(PatternRewriter &rewriter, TOp op,
+ TLegalizerCallback callback) {
+ // Clone the previous op to preserve any properties/attributes.
+ auto newOp = op.clone();
+ rewriter.insert(newOp);
+ rewriter.replaceOp(op, callback(newOp));
+}
+
+template <typename TOp, typename TLegalizerCallback>
+void replaceOpWithUnrealizedConversion(PatternRewriter &rewriter, TOp op,
+ TLegalizerCallback callback) {
+ replaceOpWithLegalizedOp(rewriter, op, [&](TOp newOp) {
+ // Mark our `unrealized_conversion_casts` with a pass label.
+ return rewriter.create<UnrealizedConversionCastOp>(
+ op.getLoc(), TypeRange{op.getResult().getType()},
+ ValueRange{callback(newOp)},
+ NamedAttribute(rewriter.getStringAttr(kPassLabel),
+ rewriter.getUnitAttr()));
+ });
+}
+
+/// Extracts the legal memref value from the `unrealized_conversion_casts` added
+/// by this pass.
+static FailureOr<Value> getLegalMemRef(Value illegalMemref) {
+ Operation *definingOp = illegalMemref.getDefiningOp();
+ if (!definingOp || !definingOp->hasAttr(kPassLabel))
+ return failure();
+ auto unrealizedConversion =
+ llvm::cast<UnrealizedConversionCastOp>(definingOp);
+ return unrealizedConversion.getOperand(0);
+}
+
+/// The default alignment of an alloca may request overaligned sizes for SVE
+/// types, which will fail during stack frame allocation. This rewrite
+/// explicitly adds a reasonable alignment to allocas of scalable types.
+struct RelaxScalableVectorAllocaAlignment
+ : public OpRewritePattern<memref::AllocaOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(memref::AllocaOp allocaOp,
+ PatternRewriter &rewriter) const override {
+ auto elementType = allocaOp.getType().getElementType();
+ auto vectorType = llvm::dyn_cast<VectorType>(elementType);
+ if (!vectorType || !vectorType.isScalable() || allocaOp.getAlignment())
+ return failure();
+
+ unsigned elementByteSize =
+ vectorType.getElementType().getIntOrFloatBitWidth() / 8;
+
+ unsigned aligment = std::max(1u, elementByteSize);
+ allocaOp.setAlignment(aligment);
+
+ return success();
+ }
+};
+
+/// Replaces allocations of SVE predicates smaller than an svbool with a wider
+/// allocation and a tagged unrealized conversion.
+///
+/// Example
+/// ```
+/// %alloca = memref.alloca() : memref<vector<[4]xi1>>
+/// ```
+/// is rewritten into:
+/// ```
+/// %widened = memref.alloca() {alignment = 1 : i64} : memref<vector<[16]xi1>>
+/// %alloca = builtin.unrealized_conversion_cast %widened
+/// : memref<vector<[16]xi1>> to memref<vector<[4]xi1>>
+/// {__arm_sve_legalize_vector_storage__}
+/// ```
+template <typename AllocLikeOp>
+struct LegalizeAllocLikeOpConversion : public OpRewritePattern<AllocLikeOp> {
+ using OpRewritePattern<AllocLikeOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(AllocLikeOp allocLikeOp,
+ PatternRewriter &rewriter) const override {
+ auto vectorType =
+ llvm::dyn_cast<VectorType>(allocLikeOp.getType().getElementType());
+
+ if (!vectorType || !isLogicalSVEPredicateType(vectorType))
+ return failure();
+
+ // Replace this alloc-like op of an SVE mask with one of a (storable)
+ // svbool_t mask. A temporary unrealized_conversion_cast is added to the old
+ // type to allow local rewrites.
+ replaceOpWithUnrealizedConversion(
+ rewriter, allocLikeOp, [&](AllocLikeOp newAllocLikeOp) {
+ newAllocLikeOp.getResult().setType(
+ llvm::cast<MemRefType>(newAllocLikeOp.getType().cloneWith(
+ {}, widenScalableMaskTypeToSvbool(vectorType))));
+ return newAllocLikeOp;
+ });
+
+ return success();
+ }
+};
+
+/// Replaces vector.type_casts of unrealized conversions to illegal memref types
+/// with legal type casts, followed by unrealized conversions.
+///
+/// Example:
+/// ```
+/// %alloca = builtin.unrealized_conversion_cast %widened
+/// : memref<vector<[16]xi1>> to memref<vector<[8]xi1>>
+/// {__arm_sve_legalize_vector_storage__}
+/// %cast = vector.type_cast %alloca
+/// : memref<vector<3x[8]xi1>> to memref<3xvector<[8]xi1>>
+/// ```
+/// is rewritten into:
+/// ```
+/// %widened_cast = vector.type_cast %widened
+/// : memref<vector<3x[16]xi1>> to memref<3xvector<[16]xi1>>
+/// %cast = builtin.unrealized_conversion_cast %widened_cast
+/// : memref<3xvector<[16]xi1>> to memref<3xvector<[8]xi1>>
+/// {__arm_sve_legalize_vector_storage__}
+/// ```
+struct LegalizeVectorTypeCastConversion
+ : public OpRewritePattern<vector::TypeCastOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(vector::TypeCastOp typeCastOp,
+ PatternRewriter &rewriter) const override {
+ auto resultType = typeCastOp.getResultMemRefType();
+ auto vectorType = llvm::dyn_cast<VectorType>(resultType.getElementType());
+
+ if (!vectorType || !isLogicalSVEPredicateType(vectorType))
+ return failure();
+
+ auto legalMemref = getLegalMemRef(typeCastOp.getMemref());
+ if (failed(legalMemref))
+ return failure();
+
+ // Replace this vector.type_cast with one of a (storable) svbool_t mask.
+ replaceOpWithUnrealizedConversion(
+ rewriter, typeCastOp, [&](vector::TypeCastOp newTypeCast) {
+ newTypeCast.setOperand(*legalMemref);
+ newTypeCast.getResult().setType(
+ llvm::cast<MemRefType>(newTypeCast.getType().cloneWith(
+ {}, widenScalableMaskTypeToSvbool(vectorType))));
+ return newTypeCast;
+ });
+
+ return success();
+ }
+};
+
+/// Replaces stores to unrealized conversions to illegal memref types with
+/// `arm_sve.convert_to_svbool`s followed by (legal) wider stores.
+///
+/// Example:
+/// ```
+/// memref.store %mask, %alloca[] : memref<vector<[8]xi1>>
+/// ```
+/// is rewritten into:
+/// ```
+/// %svbool = arm_sve.convert_to_svbool %mask : vector<[8]xi1>
+/// memref.store %svbool, %widened[] : memref<vector<[16]xi1>>
+/// ```
+struct LegalizeMemrefStoreConversion
+ : public OpRewritePattern<memref::StoreOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(memref::StoreOp storeOp,
+ PatternRewriter &rewriter) const override {
+ auto loc = storeOp.getLoc();
+
+ Value valueToStore = storeOp.getValueToStore();
+ auto vectorType = llvm::dyn_cast<VectorType>(valueToStore.getType());
+
+ if (!vectorType || !isLogicalSVEPredicateType(vectorType))
+ return failure();
+
+ auto legalMemref = getLegalMemRef(storeOp.getMemref());
+ if (failed(legalMemref))
+ return failure();
+
+ auto legalMaskType = widenScalableMaskTypeToSvbool(
+ llvm::cast<VectorType>(valueToStore.getType()));
+ auto convertToSvbool = rewriter.create<arm_sve::ConvertToSvboolOp>(
+ loc, legalMaskType, valueToStore);
+ // Replace this store with a conversion to a storable svbool_t mask,
+ // followed by a wider store.
+ replaceOpWithLegalizedOp(rewriter, storeOp,
+ [&](memref::StoreOp newStoreOp) {
+ newStoreOp.setOperand(0, convertToSvbool);
+ newStoreOp.setOperand(1, *legalMemref);
+ return newStoreOp;
+ });
+
+ return success();
+ }
+};
+
+/// Replaces loads from unrealized conversions to illegal memref types with
+/// (legal) wider loads, followed by `arm_sve.convert_from_svbool`s.
+///
+/// Example:
+/// ```
+/// %reload = memref.load %alloca[] : memref<vector<[4]xi1>>
+/// ```
+/// is rewritten into:
+/// ```
+/// %svbool = memref.load %widened[] : memref<vector<[16]xi1>>
+/// %reload = arm_sve.convert_from_svbool %reload : vector<[4]xi1>
+/// ```
+struct LegalizeMemrefLoadConversion : public OpRewritePattern<memref::LoadOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(memref::LoadOp loadOp,
+ PatternRewriter &rewriter) const override {
+ auto loc = loadOp.getLoc();
+
+ Value loadedMask = loadOp.getResult();
+ auto vectorType = llvm::dyn_cast<VectorType>(loadedMask.getType());
+
+ if (!vectorType || !isLogicalSVEPredicateType(vectorType))
+ return failure();
+
+ auto legalMemref = getLegalMemRef(loadOp.getMemref());
+ if (failed(legalMemref))
+ return failure();
+
+ auto legalMaskType = widenScalableMaskTypeToSvbool(vectorType);
+ // Replace this load with a legal load of an svbool_t type, followed by a
+ // conversion back to the original type.
+ replaceOpWithLegalizedOp(rewriter, loadOp, [&](memref::LoadOp newLoadOp) {
+ newLoadOp.setMemRef(*legalMemref);
+ newLoadOp.getResult().setType(legalMaskType);
+ return rewriter.create<arm_sve::ConvertFromSvboolOp>(
+ loc, loadedMask.getType(), newLoadOp);
+ });
+
+ return success();
+ }
+};
+
+struct LegalizeVectorStorage
+ : public arm_sve::impl::LegalizeVectorStorageBase<LegalizeVectorStorage> {
+
+ void runOnOperation() override {
+ RewritePatternSet patterns(&getContext());
+ patterns.add<RelaxScalableVectorAllocaAlignment,
+ LegalizeAllocLikeOpConversion<memref::AllocaOp>,
+ LegalizeAllocLikeOpConversion<memref::AllocOp>,
+ LegalizeVectorTypeCastConversion,
+ LegalizeMemrefStoreConversion, LegalizeMemrefLoadConversion>(
+ patterns.getContext());
+ if (failed(applyPatternsAndFoldGreedily(getOperation(),
+ std::move(patterns)))) {
+ signalPassFailure();
+ }
+ ConversionTarget target(getContext());
+ target.addDynamicallyLegalOp<UnrealizedConversionCastOp>(
+ [](UnrealizedConversionCastOp unrealizedConversion) {
+ return !unrealizedConversion->hasAttr(kPassLabel);
+ });
+ // This detects if we failed to completely legalize the IR.
+ if (failed(applyPartialConversion(getOperation(), target, {})))
+ signalPassFailure();
+ }
+};
+
+} // namespace
+
+std::unique_ptr<Pass> mlir::arm_sve::createLegalizeVectorStoragePass() {
+ return std::make_unique<LegalizeVectorStorage>();
+}
diff --git a/mlir/test/Dialect/ArmSVE/legalize-vector-storage.mlir b/mlir/test/Dialect/ArmSVE/legalize-vector-storage.mlir
new file mode 100644
index 000000000000000..fda8e2e0fab9618
--- /dev/null
+++ b/mlir/test/Dialect/ArmSVE/legalize-vector-storage.mlir
@@ -0,0 +1,160 @@
+// RUN: mlir-opt %s -allow-unregistered-dialect -arm-sve-legalize-vector-storage -split-input-file -verify-diagnostics | FileCheck %s
+
+/// This tests the basic functionality of the -arm-sve-legalize-vector-storage pass.
+
+// -----
+
+// CHECK-LABEL: @store_and_reload_sve_predicate_nxv1i1(
+// CHECK-SAME: %[[MASK:.*]]: vector<[1]xi1>)
+func.func @store_and_reload_sve_predicate_nxv1i1(%mask: vector<[1]xi1>) -> vector<[1]xi1> {
+ // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 1 : i64} : memref<vector<[16]xi1>>
+ %alloca = memref.alloca() : memref<vector<[1]xi1>>
+ // CHECK-NEXT: %[[SVBOOL:.*]] = arm_sve.convert_to_svbool %[[MASK]] : vector<[1]xi1>
+ // CHECK-NEXT: memref.store %[[SVBOOL]], %[[ALLOCA]][] : memref<vector<[16]xi1>>
+ memref.store %mask, %alloca[] : memref<vector<[1]xi1>>
+ // CHECK-NEXT: %[[RELOAD:.*]] = memref.load %[[ALLOCA]][] : memref<vector<[16]xi1>>
+ // CHECK-NEXT: %[[MASK:.*]] = arm_sve.convert_from_svbool %[[RELOAD]] : vector<[1]xi1>
+ %reload = memref.load %alloca[] : memref<vector<[1]xi1>>
+ // CHECK-NEXT: return %[[MASK]] : vector<[1]xi1>
+ return %reload : vector<[1]xi1>
+}
+
+// -----
+
+// CHECK-LABEL: @store_and_reload_sve_predicate_nxv2i1(
+// CHECK-SAME: %[[MASK:.*]]: vector<[2]xi1>)
+func.func @store_and_reload_sve_predicate_nxv2i1(%mask: vector<[2]xi1>) -> vector<[2]xi1> {
+ // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 1 : i64} : memref<vector<[16]xi1>>
+ %alloca = memref.alloca() : memref<vector<[2]xi1>>
+ // CHECK-NEXT: %[[SVBOOL:.*]] = arm_sve.convert_to_svbool %[[MASK]] : vector<[2]xi1>
+ // CHECK-NEXT: memref.store %[[SVBOOL]], %[[ALLOCA]][] : memref<vector<[16]xi1>>
+ memref.store %mask, %alloca[] : memref<vector<[2]xi1>>
+ // CHECK-NEXT: %[[RELOAD:.*]] = memref.load %[[ALLOCA]][] : memref<vector<[16]xi1>>
+ // CHECK-NEXT: %[[MASK:.*]] = arm_sve.convert_from_svbool %[[RELOAD]] : vector<[2]xi1>
+ %reload = memref.load %alloca[] : memref<vector<[2]xi1>>
+ // CHECK-NEXT: return %[[MASK]] : vector<[2]xi1>
+ return %reload : vector<[2]xi1>
+}
+
+// -----
+
+// CHECK-LABEL: @store_and_reload_sve_predicate_nxv4i1(
+// CHECK-SAME: %[[MASK:.*]]: vector<[4]xi1>)
+func.func @store_and_reload_sve_predicate_nxv4i1(%mask: vector<[4]xi1>) -> vector<[4]xi1> {
+ // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 1 : i64} : memref<vector<[16]xi1>>
+ %alloca = memref.alloca() : memref<vector<[4]xi1>>
+ // CHECK-NEXT: %[[SVBOOL:.*]] = arm_sve.convert_to_svbool %[[MASK]] : vector<[4]xi1>
+ // CHECK-NEXT: memref.store %[[SVBOOL]], %[[ALLOCA]][] : memref<vector<[16]xi1>>
+ memref.store %mask, %alloca[] : memref<vector<[4]xi1>>
+ // CHECK-NEXT: %[[RELOAD:.*]] = memref.load %[[ALLOCA]][] : memref<vector<[16]xi1>>
+ // CHECK-NEXT: %[[MASK:.*]] = arm_sve.convert_from_svbool %[[RELOAD]] : vector<[4]xi1>
+ %reload = memref.load %alloca[] : memref<vector<[4]xi1>>
+ // CHECK-NEXT: return %[[MASK]] : vector<[4]xi1>
+ return %reload : vector<[4]xi1>
+}
+
+// -----
+
+// CHECK-LABEL: @store_and_reload_sve_predicate_nxv8i1(
+// CHECK-SAME: %[[MASK:.*]]: vector<[8]xi1>)
+func.func @store_and_reload_sve_predicate_nxv8i1(%mask: vector<[8]xi1>) -> vector<[8]xi1> {
+ // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 1 : i64} : memref<vector<[16]xi1>>
+ %alloca = memref.alloca() : memref<vector<[8]xi1>>
+ // CHECK-NEXT: %[[SVBOOL:.*]] = arm_sve.convert_to_svbool %[[MASK]] : vector<[8]xi1>
+ // CHECK-NEXT: memref.store %[[SVBOOL]], %[[ALLOCA]][] : memref<vector<[16]xi1>>
+ memref.store %mask, %alloca[] : memref<vector<[8]xi1>>
+ // CHECK-NEXT: %[[RELOAD:.*]] = memref.load %[[ALLOCA]][] : memref<vector<[16]xi1>>
+ // CHECK-NEXT: %[[MASK:.*]] = arm_sve.convert_from_svbool %[[RELOAD]] : vector<[8]xi1>
+ %reload = memref.load %alloca[] : memref<vector<[8]xi1>>
+ // CHECK-NEXT: return %[[MASK]] : vector<[8]xi1>
+ return %reload : vector<[8]xi1>
+}
+
+// -----
+
+// CHECK-LABEL: @store_2d_mask_and_reload_slice(
+// CHECK-SAME: %[[MASK:.*]]: vector<3x[8]xi1>)
+func.func @store_2d_mask_and_reload_slice(%mask: vector<3x[8]xi1>) -> vector<[8]xi1> {
+ // CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
+ %c0 = arith.constant 0 : index
+ // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 1 : i64} : memref<vector<3x[16]xi1>>
+ %alloca = memref.alloca() : memref<vector<3x[8]xi1>>
+ // CHECK-NEXT: %[[SVBOOL:.*]] = arm_sve.convert_to_svbool %[[MASK]] : vector<3x[8]xi1>
+ // CHECK-NEXT: memref.store %[[SVBOOL]], %[[ALLOCA]][] : memref<vector<3x[16]xi1>>
+ memref.store %mask, %alloca[] : memref<vector<3x[8]xi1>>
+ // CHECK-NEXT: %[[UNPACK:.*]] = vector.type_cast %[[ALLOCA]] : memref<vector<3x[16]xi1>> to memref<3xvector<[16]xi1>>
+ %unpack = vector.type_cast %alloca : memref<vector<3x[8]xi1>> to memref<3xvector<[8]xi1>>
+ // CHECK-NEXT: %[[RELOAD:.*]] = memref.load %[[UNPACK]][%[[C0]]] : memref<3xvector<[16]xi1>>
+ // CHECK-NEXT: %[[SLICE:.*]] = arm_sve.convert_from_svbool %[[RELOAD]] : vector<[8]xi1>
+ %slice = memref.load %unpack[%c0] : memref<3xvector<[8]xi1>>
+ // CHECK-NEXT: return %[[SLICE]] : vector<[8]xi1>
+ return %slice : vector<[8]xi1>
+}
+
+// -----
+
+// CHECK-LABEL: @set_sve_alloca_alignment
+func.func @set_sve_alloca_alignment() {
+ // CHECK-COUNT-6: alignment = 1
+ %a1 = memref.alloca() : memref<vector<[32]xi8>>
+ %a2 = memref.alloca() : memref<vector<[16]xi8>>
+ %a3 = memref.alloca() : memref<vector<[8]xi8>>
+ %a4 = memref.alloca() : memref<vector<[4]xi8>>
+ %a5 = memref.alloca() : memref<vector<[2]xi8>>
+ %a6 = memref.alloca() : memref<vector<[1]xi8>>
+
+ // CHECK-COUNT-6: alignment = 2
+ %b1 = memref.alloca() : memref<vector<[32]xi16>>
+ %b2 = memref.alloca() : memref<vector<[16]xi16>>
+ %b3 = memref.alloca() : memref<vector<[8]xi16>>
+ %b4 = memref.alloca() : memref<vector<[4]xi16>>
+ %b5 = memref.alloca() : memref<vector<[2]xi16>>
+ %b6 = memref.alloca() : memref<vector<[1]xi16>>
+
+ // CHECK-COUNT-6: alignment = 4
+ %c1 = memref.alloca() : memref<vector<[32]xi32>>
+ %c2 = memref.alloca() : memref<vector<[16]xi32>>
+ %c3 = memref.alloca() : memref<vector<[8]xi32>>
+ %c4 = memref.alloca() : memref<vector<[4]xi32>>
+ %c5 = memref.alloca() : memref<vector<[2]xi32>>
+ %c6 = memref.alloca() : memref<vector<[1]xi32>>
+
+ // CHECK-COUNT-6: alignment = 8
+ %d1 = memref.alloca() : memref<vector<[32]xi64>>
+ %d2 = memref.alloca() : memref<vector<[16]xi64>>
+ %d3 = memref.alloca() : memref<vector<[8]xi64>>
+ %d4 = memref.alloca() : memref<vector<[4]xi64>>
+ %d5 = memref.alloca() : memref<vector<[2]xi64>>
+ %d6 = memref.alloca() : memref<vector<[1]xi64>>
+
+ // CHECK-COUNT-6: alignment = 4
+ %e1 = memref.alloca() : memref<vector<[32]xf32>>
+ %e2 = memref.alloca() : memref<vector<[16]xf32>>
+ %e3 = memref.alloca() : memref<vector<[8]xf32>>
+ %e4 = memref.alloca() : memref<vector<[4]xf32>>
+ %e5 = memref.alloca() : memref<vector<[2]xf32>>
+ %e6 = memref.alloca() : memref<vector<[1]xf32>>
+
+ // CHECK-COUNT-6: alignment = 8
+ %f1 = memref.alloca() : memref<vector<[32]xf64>>
+ %f2 = memref.alloca() : memref<vector<[16]xf64>>
+ %f3 = memref.alloca() : memref<vector<[8]xf64>>
+ %f4 = memref.alloca() : memref<vector<[4]xf64>>
+ %f5 = memref.alloca() : memref<vector<[2]xf64>>
+ %f6 = memref.alloca() : memref<vector<[1]xf64>>
+
+ "prevent.dce"(
+ %a1, %a2, %a3, %a4, %a5, %a6,
+ %b1, %b2, %b3, %b4, %b5, %b6,
+ %c1, %c2, %c3, %c4, %c5, %c6,
+ %d1, %d2, %d3, %d4, %d5, %d6,
+ %e1, %e2, %e3, %e4, %e5, %e6,
+ %f1, %f2, %f3, %f4, %f5, %f6)
+ : (memref<vector<[32]xi8>>, memref<vector<[16]xi8>>, memref<vector<[8]xi8>>, memref<vector<[4]xi8>>, memref<vector<[2]xi8>>, memref<vector<[1]xi8>>,
+ memref<vector<[32]xi16>>, memref<vector<[16]xi16>>, memref<vector<[8]xi16>>, memref<vector<[4]xi16>>, memref<vector<[2]xi16>>, memref<vector<[1]xi16>>,
+ memref<vector<[32]xi32>>, memref<vector<[16]xi32>>, memref<vector<[8]xi32>>, memref<vector<[4]xi32>>, memref<vector<[2]xi32>>, memref<vector<[1]xi32>>,
+ memref<vector<[32]xi64>>, memref<vector<[16]xi64>>, memref<vector<[8]xi64>>, memref<vector<[4]xi64>>, memref<vector<[2]xi64>>, memref<vector<[1]xi64>>,
+ memref<vector<[32]xf32>>, memref<vector<[16]xf32>>, memref<vector<[8]xf32>>, memref<vector<[4]xf32>>, memref<vector<[2]xf32>>, memref<vector<[1]xf32>>,
+ memref<vector<[32]xf64>>, memref<vector<[16]xf64>>, memref<vector<[8]xf64>>, memref<vector<[4]xf64>>, memref<vector<[2]xf64>>, memref<vector<[1]xf64>>) -> ()
+ return
+}
diff --git a/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/arrays-of-scalable-vectors.mlir b/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/arrays-of-scalable-vectors.mlir
new file mode 100644
index 000000000000000..260cd1eabe7dae1
--- /dev/null
+++ b/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/arrays-of-scalable-vectors.mlir
@@ -0,0 +1,121 @@
+// RUN: mlir-opt %s -convert-vector-to-scf -arm-sve-legalize-vector-storage -convert-vector-to-llvm="enable-arm-sve" -test-lower-to-llvm | \
+// RUN: %mcr_aarch64_cmd -e=entry -entry-point-result=void --march=aarch64 --mattr="+sve" -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+/// This tests basic functionality of arrays of scalable vectors, which in MLIR
+/// are vectors with a single trailing scalable dimension. This test requires
+/// the -arm-sve-legalize-vector-storage pass to ensure the loads/stores done
+/// here are be legal for the LLVM backend.
+
+func.func @read_and_print_2d_vector(%memref: memref<3x?xf32>) {
+ %cst = arith.constant 0.000000e+00 : f32
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %dim = memref.dim %memref, %c1 : memref<3x?xf32>
+ %mask = vector.create_mask %c2, %dim : vector<3x[8]xi1>
+ %vector = vector.transfer_read %memref[%c0,%c0], %cst, %mask {in_bounds = [true, true]} : memref<3x?xf32>, vector<3x[8]xf32>
+
+ /// TODO: Support vector.print for arrays of scalable vectors.
+ %row0 = vector.extract %vector[0] : vector<[8]xf32> from vector<3x[8]xf32>
+ %row1 = vector.extract %vector[1] : vector<[8]xf32> from vector<3x[8]xf32>
+ %row2 = vector.extract %vector[2] : vector<[8]xf32> from vector<3x[8]xf32>
+
+ /// Print each of the vectors. This only checks the first eight elements (which
+ /// works for all vscale >= 1).
+
+ // CHECK-LABEL: TEST 1
+ vector.print str "TEST 1 (print and read 2D arrays of scalable vectors)"
+ // CHECK: ( 8, 8, 8, 8, 8, 8, 8, 8
+ vector.print %row0 : vector<[8]xf32>
+ // CHECK: ( 8, 8, 8, 8, 8, 8, 8, 8
+ vector.print %row1 : vector<[8]xf32>
+ /// This last row is all zero due to our mask.
+ // CHECK: ( 0, 0, 0, 0, 0, 0, 0, 0
+ vector.print %row2 : vector<[8]xf32>
+
+ return
+}
+
+func.func @print_1x2xVSCALExf32(%vector: vector<1x2x[4]xf32>) {
+ /// TODO: Support vector.print for arrays of scalable vectors.
+ %slice0 = vector.extract %vector[0, 1] : vector<[4]xf32> from vector<1x2x[4]xf32>
+ %slice1 = vector.extract %vector[0, 1] : vector<[4]xf32> from vector<1x2x[4]xf32>
+ vector.print %slice0 : vector<[4]xf32>
+ vector.print %slice1 : vector<[4]xf32>
+ return
+}
+
+func.func @add_arrays_of_scalable_vectors(%a: memref<1x2x?xf32>, %b: memref<1x2x?xf32>) {
+ %c0 = arith.constant 0 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 2 : index
+ %cst = arith.constant 0.000000e+00 : f32
+ %dim_a = memref.dim %a, %c2 : memref<1x2x?xf32>
+ %dim_b = memref.dim %b, %c2 : memref<1x2x?xf32>
+ %mask_a = vector.create_mask %c2, %c3, %dim_a : vector<1x2x[4]xi1>
+ %mask_b = vector.create_mask %c2, %c3, %dim_b : vector<1x2x[4]xi1>
+
+ vector.print str "TEST 2 (reading and adding two 3D arrays of scalable vectors)"
+
+ /// Print each of the vectors. This only checks the first four elements (which
+ /// works for all vscale >= 1).
+
+ // CHECK-LABEL: Vector A
+ // CHECK-NEXT: ( 5, 5, 5, 5
+ // CHECK-NEXT: ( 5, 5, 5, 5
+ vector.print str "\nVector A"
+ %vector_a = vector.transfer_read %a[%c0, %c0, %c0], %cst, %mask_a {in_bounds = [true, true, true]} : memref<1x2x?xf32>, vector<1x2x[4]xf32>
+ func.call @print_1x2xVSCALExf32(%vector_a) : (vector<1x2x[4]xf32>) -> ()
+
+ vector.print str "\nVector B"
+ // CHECK-LABEL: Vector B
+ // CHECK-NEXT: ( 4, 4, 4, 4
+ // CHECK-NEXT: ( 4, 4, 4, 4
+ %vector_b = vector.transfer_read %b[%c0, %c0, %c0], %cst, %mask_b {in_bounds = [true, true, true]} : memref<1x2x?xf32>, vector<1x2x[4]xf32>
+ func.call @print_1x2xVSCALExf32(%vector_b) : (vector<1x2x[4]xf32>) -> ()
+
+ %sum = arith.addf %vector_a, %vector_b : vector<1x2x[4]xf32>
+ // CHECK-LABEL: Sum
+ // CHECK-NEXT: ( 9, 9, 9, 9
+ // CHECK-NEXT: ( 9, 9, 9, 9
+ vector.print str "\nSum"
+ func.call @print_1x2xVSCALExf32(%sum) : (vector<1x2x[4]xf32>) -> ()
+
+ return
+}
+
+func.func @entry() {
+ %vscale = vector.vscale
+
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %f32_8 = arith.constant 8.0 : f32
+ %f32_5 = arith.constant 5.0 : f32
+ %f32_4 = arith.constant 4.0 : f32
+
+ vector.print str "\n====================\n"
+
+ %test_1_memref_size = arith.muli %vscale, %c8 : index
+ %test_1_memref = memref.alloca(%test_1_memref_size) : memref<3x?xf32>
+
+ linalg.fill ins(%f32_8 : f32) outs(%test_1_memref :memref<3x?xf32>)
+
+ func.call @read_and_print_2d_vector(%test_1_memref) : (memref<3x?xf32>) -> ()
+
+ vector.print str "\n====================\n"
+
+ %test_2_memref_size = arith.muli %vscale, %c4 : index
+ %test_2_memref_a = memref.alloca(%test_2_memref_size) : memref<1x2x?xf32>
+ %test_2_memref_b = memref.alloca(%test_2_memref_size) : memref<1x2x?xf32>
+
+ linalg.fill ins(%f32_5 : f32) outs(%test_2_memref_a :memref<1x2x?xf32>)
+ linalg.fill ins(%f32_4 : f32) outs(%test_2_memref_b :memref<1x2x?xf32>)
+
+ func.call @add_arrays_of_scalable_vectors(
+ %test_2_memref_a, %test_2_memref_b) : (memref<1x2x?xf32>, memref<1x2x?xf32>) -> ()
+
+ vector.print str "\n====================\n"
+
+ return
+}
>From 72f19279dcd0a95eb21dc7f72e4d5850569361f8 Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Mon, 23 Oct 2023 12:59:03 +0000
Subject: [PATCH 2/3] Always align SVE vectors to 16 bytes and predicates to 2
bytes
---
.../Transforms/LegalizeVectorStorage.cpp | 10 ++++-----
.../ArmSVE/legalize-vector-storage.mlir | 22 +++++++++----------
2 files changed, 15 insertions(+), 17 deletions(-)
diff --git a/mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp b/mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp
index 610eb38089c4c88..e7edcf6ec789235 100644
--- a/mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp
+++ b/mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp
@@ -82,15 +82,13 @@ struct RelaxScalableVectorAllocaAlignment
LogicalResult matchAndRewrite(memref::AllocaOp allocaOp,
PatternRewriter &rewriter) const override {
- auto elementType = allocaOp.getType().getElementType();
- auto vectorType = llvm::dyn_cast<VectorType>(elementType);
+ auto memrefElementType = allocaOp.getType().getElementType();
+ auto vectorType = llvm::dyn_cast<VectorType>(memrefElementType);
if (!vectorType || !vectorType.isScalable() || allocaOp.getAlignment())
return failure();
- unsigned elementByteSize =
- vectorType.getElementType().getIntOrFloatBitWidth() / 8;
-
- unsigned aligment = std::max(1u, elementByteSize);
+ // Set alignment based on the defaults for SVE vectors and predicates.
+ unsigned aligment = vectorType.getElementType().isInteger(1) ? 2 : 16;
allocaOp.setAlignment(aligment);
return success();
diff --git a/mlir/test/Dialect/ArmSVE/legalize-vector-storage.mlir b/mlir/test/Dialect/ArmSVE/legalize-vector-storage.mlir
index fda8e2e0fab9618..61879c48712f4d2 100644
--- a/mlir/test/Dialect/ArmSVE/legalize-vector-storage.mlir
+++ b/mlir/test/Dialect/ArmSVE/legalize-vector-storage.mlir
@@ -7,7 +7,7 @@
// CHECK-LABEL: @store_and_reload_sve_predicate_nxv1i1(
// CHECK-SAME: %[[MASK:.*]]: vector<[1]xi1>)
func.func @store_and_reload_sve_predicate_nxv1i1(%mask: vector<[1]xi1>) -> vector<[1]xi1> {
- // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 1 : i64} : memref<vector<[16]xi1>>
+ // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 2 : i64} : memref<vector<[16]xi1>>
%alloca = memref.alloca() : memref<vector<[1]xi1>>
// CHECK-NEXT: %[[SVBOOL:.*]] = arm_sve.convert_to_svbool %[[MASK]] : vector<[1]xi1>
// CHECK-NEXT: memref.store %[[SVBOOL]], %[[ALLOCA]][] : memref<vector<[16]xi1>>
@@ -24,7 +24,7 @@ func.func @store_and_reload_sve_predicate_nxv1i1(%mask: vector<[1]xi1>) -> vecto
// CHECK-LABEL: @store_and_reload_sve_predicate_nxv2i1(
// CHECK-SAME: %[[MASK:.*]]: vector<[2]xi1>)
func.func @store_and_reload_sve_predicate_nxv2i1(%mask: vector<[2]xi1>) -> vector<[2]xi1> {
- // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 1 : i64} : memref<vector<[16]xi1>>
+ // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 2 : i64} : memref<vector<[16]xi1>>
%alloca = memref.alloca() : memref<vector<[2]xi1>>
// CHECK-NEXT: %[[SVBOOL:.*]] = arm_sve.convert_to_svbool %[[MASK]] : vector<[2]xi1>
// CHECK-NEXT: memref.store %[[SVBOOL]], %[[ALLOCA]][] : memref<vector<[16]xi1>>
@@ -41,7 +41,7 @@ func.func @store_and_reload_sve_predicate_nxv2i1(%mask: vector<[2]xi1>) -> vecto
// CHECK-LABEL: @store_and_reload_sve_predicate_nxv4i1(
// CHECK-SAME: %[[MASK:.*]]: vector<[4]xi1>)
func.func @store_and_reload_sve_predicate_nxv4i1(%mask: vector<[4]xi1>) -> vector<[4]xi1> {
- // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 1 : i64} : memref<vector<[16]xi1>>
+ // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 2 : i64} : memref<vector<[16]xi1>>
%alloca = memref.alloca() : memref<vector<[4]xi1>>
// CHECK-NEXT: %[[SVBOOL:.*]] = arm_sve.convert_to_svbool %[[MASK]] : vector<[4]xi1>
// CHECK-NEXT: memref.store %[[SVBOOL]], %[[ALLOCA]][] : memref<vector<[16]xi1>>
@@ -58,7 +58,7 @@ func.func @store_and_reload_sve_predicate_nxv4i1(%mask: vector<[4]xi1>) -> vecto
// CHECK-LABEL: @store_and_reload_sve_predicate_nxv8i1(
// CHECK-SAME: %[[MASK:.*]]: vector<[8]xi1>)
func.func @store_and_reload_sve_predicate_nxv8i1(%mask: vector<[8]xi1>) -> vector<[8]xi1> {
- // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 1 : i64} : memref<vector<[16]xi1>>
+ // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 2 : i64} : memref<vector<[16]xi1>>
%alloca = memref.alloca() : memref<vector<[8]xi1>>
// CHECK-NEXT: %[[SVBOOL:.*]] = arm_sve.convert_to_svbool %[[MASK]] : vector<[8]xi1>
// CHECK-NEXT: memref.store %[[SVBOOL]], %[[ALLOCA]][] : memref<vector<[16]xi1>>
@@ -77,7 +77,7 @@ func.func @store_and_reload_sve_predicate_nxv8i1(%mask: vector<[8]xi1>) -> vecto
func.func @store_2d_mask_and_reload_slice(%mask: vector<3x[8]xi1>) -> vector<[8]xi1> {
// CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
%c0 = arith.constant 0 : index
- // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 1 : i64} : memref<vector<3x[16]xi1>>
+ // CHECK-NEXT: %[[ALLOCA:.*]] = memref.alloca() {alignment = 2 : i64} : memref<vector<3x[16]xi1>>
%alloca = memref.alloca() : memref<vector<3x[8]xi1>>
// CHECK-NEXT: %[[SVBOOL:.*]] = arm_sve.convert_to_svbool %[[MASK]] : vector<3x[8]xi1>
// CHECK-NEXT: memref.store %[[SVBOOL]], %[[ALLOCA]][] : memref<vector<3x[16]xi1>>
@@ -95,7 +95,7 @@ func.func @store_2d_mask_and_reload_slice(%mask: vector<3x[8]xi1>) -> vector<[8]
// CHECK-LABEL: @set_sve_alloca_alignment
func.func @set_sve_alloca_alignment() {
- // CHECK-COUNT-6: alignment = 1
+ // CHECK-COUNT-6: alignment = 16
%a1 = memref.alloca() : memref<vector<[32]xi8>>
%a2 = memref.alloca() : memref<vector<[16]xi8>>
%a3 = memref.alloca() : memref<vector<[8]xi8>>
@@ -103,7 +103,7 @@ func.func @set_sve_alloca_alignment() {
%a5 = memref.alloca() : memref<vector<[2]xi8>>
%a6 = memref.alloca() : memref<vector<[1]xi8>>
- // CHECK-COUNT-6: alignment = 2
+ // CHECK-COUNT-6: alignment = 16
%b1 = memref.alloca() : memref<vector<[32]xi16>>
%b2 = memref.alloca() : memref<vector<[16]xi16>>
%b3 = memref.alloca() : memref<vector<[8]xi16>>
@@ -111,7 +111,7 @@ func.func @set_sve_alloca_alignment() {
%b5 = memref.alloca() : memref<vector<[2]xi16>>
%b6 = memref.alloca() : memref<vector<[1]xi16>>
- // CHECK-COUNT-6: alignment = 4
+ // CHECK-COUNT-6: alignment = 16
%c1 = memref.alloca() : memref<vector<[32]xi32>>
%c2 = memref.alloca() : memref<vector<[16]xi32>>
%c3 = memref.alloca() : memref<vector<[8]xi32>>
@@ -119,7 +119,7 @@ func.func @set_sve_alloca_alignment() {
%c5 = memref.alloca() : memref<vector<[2]xi32>>
%c6 = memref.alloca() : memref<vector<[1]xi32>>
- // CHECK-COUNT-6: alignment = 8
+ // CHECK-COUNT-6: alignment = 16
%d1 = memref.alloca() : memref<vector<[32]xi64>>
%d2 = memref.alloca() : memref<vector<[16]xi64>>
%d3 = memref.alloca() : memref<vector<[8]xi64>>
@@ -127,7 +127,7 @@ func.func @set_sve_alloca_alignment() {
%d5 = memref.alloca() : memref<vector<[2]xi64>>
%d6 = memref.alloca() : memref<vector<[1]xi64>>
- // CHECK-COUNT-6: alignment = 4
+ // CHECK-COUNT-6: alignment = 16
%e1 = memref.alloca() : memref<vector<[32]xf32>>
%e2 = memref.alloca() : memref<vector<[16]xf32>>
%e3 = memref.alloca() : memref<vector<[8]xf32>>
@@ -135,7 +135,7 @@ func.func @set_sve_alloca_alignment() {
%e5 = memref.alloca() : memref<vector<[2]xf32>>
%e6 = memref.alloca() : memref<vector<[1]xf32>>
- // CHECK-COUNT-6: alignment = 8
+ // CHECK-COUNT-6: alignment = 16
%f1 = memref.alloca() : memref<vector<[32]xf64>>
%f2 = memref.alloca() : memref<vector<[16]xf64>>
%f3 = memref.alloca() : memref<vector<[8]xf64>>
>From d5985c7b521dd30f06da24d060f9f1909db1d2ea Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Fri, 20 Oct 2023 16:13:49 +0000
Subject: [PATCH 3/3] [mlir][SVE] Add an e2e test for vector.contract
Adds an end-to-end test for `vector.contract` that targets SVE (i.e.
scalable vectors). Note that this requires lifting the restriction on
`vector.outerproduct` (to which `vector.contract` is lowered to) that
would deem the following as invalid by the Op verifier (*):
```
vector.outerproduct %27, %28, %26 {kind = #vector.kind<add>} : vector<3xf32>, vector<[2]xf32>
```
This is indeed valid as the end-to-end test demonstrates (at least when
compiling for SVE).
Depends on #68794
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 9 +-
...r-contract-to-outerproduct-transforms.mlir | 59 ++++++--
.../Vector/vector-scalable-outerproduct.mlir | 5 +-
.../Vector/CPU/ArmSVE/test-contraction.mlir | 137 ++++++++++++++++++
4 files changed, 191 insertions(+), 19 deletions(-)
create mode 100644 mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/test-contraction.mlir
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 8de132daf3c6a5d..d77476c10908395 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -3067,9 +3067,12 @@ LogicalResult OuterProductOp::verify() {
return emitOpError("expected #1 operand dim to match result dim #1");
if (vRHS.getDimSize(0) != vRES.getDimSize(1))
return emitOpError("expected #2 operand dim to match result dim #2");
- if (vRHS.isScalable() != vLHS.isScalable())
- return emitOpError("expected either all or none of vector operands #1 "
- "and #2 to be scalable");
+ if (vLHS.isScalable() && !vRHS.isScalable()) {
+ // This restriction reflects what's currently supported in terms of
+ // scalable vectors. However, we could relax this if there's a use case.
+ return emitOpError(
+ "expected either both or only #2 operand dim to be scalable");
+ }
} else {
// An AXPY operation.
if (vRES.getRank() != 1)
diff --git a/mlir/test/Dialect/Vector/vector-contract-to-outerproduct-transforms.mlir b/mlir/test/Dialect/Vector/vector-contract-to-outerproduct-transforms.mlir
index 1e92fcff64dea57..a5cd9b8f39173b3 100644
--- a/mlir/test/Dialect/Vector/vector-contract-to-outerproduct-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-contract-to-outerproduct-transforms.mlir
@@ -79,21 +79,21 @@ func.func @masked_extract_contract2_scalable_parallel_dim(%arg0: vector<[2]x3xf3
}
// CHECK-LABEL: func.func @masked_extract_contract4(
-// CHECK-SAME: %[[VAL_0:.*]]: vector<3x5xf32>,
-// CHECK-SAME: %[[VAL_1:.*]]: vector<5x7xf32>,
-// CHECK-SAME: %[[VAL_2:.*]]: vector<3x7xf32>,
-// CHECK-SAME: %[[VAL_3:.*]]: vector<3x7x5xi1>) -> vector<3x7xf32> {
-// CHECK: %[[VAL_5:.*]] = vector.transpose %[[VAL_3]], [2, 0, 1] : vector<3x7x5xi1> to vector<5x3x7xi1>
-// CHECK: %[[VAL_8:.*]] = vector.extract %[[VAL_5]][0] : vector<3x7xi1> from vector<5x3x7xi1>
-// CHECK: %[[VAL_9:.*]] = vector.mask %[[VAL_8]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
-// CHECK: %[[VAL_12:.*]] = vector.extract %[[VAL_5]][1] : vector<3x7xi1> from vector<5x3x7xi1>
-// CHECK: %[[VAL_13:.*]] = vector.mask %[[VAL_12]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
-// CHECK: %[[VAL_16:.*]] = vector.extract %[[VAL_5]][2] : vector<3x7xi1> from vector<5x3x7xi1>
-// CHECK: %[[VAL_17:.*]] = vector.mask %[[VAL_16]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
-// CHECK: %[[VAL_20:.*]] = vector.extract %[[VAL_5]][3] : vector<3x7xi1> from vector<5x3x7xi1>
-// CHECK: %[[VAL_21:.*]] = vector.mask %[[VAL_20]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
-// CHECK: %[[VAL_24:.*]] = vector.extract %[[VAL_5]][4] : vector<3x7xi1> from vector<5x3x7xi1>
-// CHECK: %[[VAL_25:.*]] = vector.mask %[[VAL_24]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
+// CHECK-SAME: %{{.*}}: vector<3x5xf32>,
+// CHECK-SAME: %{{.*}}: vector<5x7xf32>,
+// CHECK-SAME: %{{.*}}: vector<3x7xf32>,
+// CHECK-SAME: %[[IN_MASK:.*]]: vector<3x7x5xi1>) -> vector<3x7xf32> {
+// CHECK: %[[T_MASK:.*]] = vector.transpose %[[IN_MASK]], [2, 0, 1] : vector<3x7x5xi1> to vector<5x3x7xi1>
+// CHECK: %[[T_MASK_R0:.*]] = vector.extract %[[T_MASK]][0] : vector<3x7xi1> from vector<5x3x7xi1>
+// CHECK: %{{.*}} = vector.mask %[[T_MASK_R0]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
+// CHECK: %[[T_MASK_R1:.*]] = vector.extract %[[T_MASK]][1] : vector<3x7xi1> from vector<5x3x7xi1>
+// CHECK: %{{.*}} = vector.mask %[[T_MASK_R1]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
+// CHECK: %[[T_MASK_R2:.*]] = vector.extract %[[T_MASK]][2] : vector<3x7xi1> from vector<5x3x7xi1>
+// CHECK: %{{.*}} = vector.mask %[[T_MASK_R2]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
+// CHECK: %[[T_MASK_R3:.*]] = vector.extract %[[T_MASK]][3] : vector<3x7xi1> from vector<5x3x7xi1>
+// CHECK: %{{.*}} = vector.mask %[[T_MASK_R3]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
+// CHECK: %[[T_MASK_R4:.*]] = vector.extract %[[T_MASK]][4] : vector<3x7xi1> from vector<5x3x7xi1>
+// CHECK: %{{.*}} = vector.mask %[[T_MASK_R4]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<7xf32> } : vector<3x7xi1> -> vector<3x7xf32>
func.func @masked_extract_contract4(%arg0: vector<3x5xf32>,
%arg1: vector<5x7xf32>,
@@ -104,6 +104,35 @@ func.func @masked_extract_contract4(%arg0: vector<3x5xf32>,
return %0 : vector<3x7xf32>
}
+// CHECK-LABEL: func.func @masked_extract_contract4_scalable_J_dim(
+// CHECK-SAME: %{{.*}}: vector<3x5xf32>,
+// CHECK-SAME: %{{.*}}: vector<5x[7]xf32>,
+// CHECK-SAME: %{{.*}}: vector<3x[7]xf32>,
+// CHECK-SAME: %[[IN_MASK:.*]]: vector<3x[7]x5xi1>) -> vector<3x[7]xf32> {
+// CHECK: %[[T_MASK:.*]] = vector.transpose %[[IN_MASK]], [2, 0, 1] : vector<3x[7]x5xi1> to vector<5x3x[7]xi1>
+// CHECK: %[[T_MASK_R0:.*]] = vector.extract %[[T_MASK]][0] : vector<3x[7]xi1> from vector<5x3x[7]xi1>
+// CHECK: %{{.*}} = vector.mask %[[T_MASK_R0]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<[7]xf32> } : vector<3x[7]xi1> -> vector<3x[7]xf32>
+// CHECK: %[[T_MASK_R1:.*]] = vector.extract %[[T_MASK]][1] : vector<3x[7]xi1> from vector<5x3x[7]xi1>
+// CHECK: %[[VAL_13:.*]] = vector.mask %[[T_MASK_R1]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<[7]xf32> } : vector<3x[7]xi1> -> vector<3x[7]xf32>
+// CHECK: %[[T_MASK_R2:.*]] = vector.extract %[[T_MASK]][2] : vector<3x[7]xi1> from vector<5x3x[7]xi1>
+// CHECK: %{{.*}} = vector.mask %[[T_MASK_R2]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<[7]xf32> } : vector<3x[7]xi1> -> vector<3x[7]xf32>
+// CHECK: %[[T_MASK_R3:.*]] = vector.extract %[[T_MASK]][3] : vector<3x[7]xi1> from vector<5x3x[7]xi1>
+// CHECK: %{{.*}} = vector.mask %[[T_MASK_R3]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<[7]xf32> } : vector<3x[7]xi1> -> vector<3x[7]xf32>
+// CHECK: %[[T_MASK_R4:.*]] = vector.extract %[[T_MASK]][4] : vector<3x[7]xi1> from vector<5x3x[7]xi1>
+// CHECK: %{{.*}} = vector.mask %[[T_MASK_R4]] { vector.outerproduct %{{.*}} {kind = #vector.kind<add>} : vector<3xf32>, vector<[7]xf32> } : vector<3x[7]xi1> -> vector<3x[7]xf32>
+
+// Note that only the J dimension is scalable in this example. In theory, all
+// dimensions could be be scalable, but there is no target yet for which this
+// would make sense.
+func.func @masked_extract_contract4_scalable_J_dim(%arg0: vector<3x5xf32>,
+ %arg1: vector<5x[7]xf32>,
+ %arg2: vector<3x[7]xf32>,
+ %m : vector<3x[7]x5xi1>) -> vector<3x[7]xf32> {
+ %0 = vector.mask %m { vector.contract #matmat_trait %arg0, %arg1, %arg2
+ : vector<3x5xf32>, vector<5x[7]xf32> into vector<3x[7]xf32> } : vector<3x[7]x5xi1> -> vector<3x[7]xf32>
+ return %0 : vector<3x[7]xf32>
+}
+
// CHECK-LABEL: func @matmul
// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: vector<2x4xf32>,
// CHECK-SAME: %[[B:[a-zA-Z0-9]*]]: vector<4x3xf32>,
diff --git a/mlir/test/Dialect/Vector/vector-scalable-outerproduct.mlir b/mlir/test/Dialect/Vector/vector-scalable-outerproduct.mlir
index 7d9923e036660c9..3b4e24da92aaacc 100644
--- a/mlir/test/Dialect/Vector/vector-scalable-outerproduct.mlir
+++ b/mlir/test/Dialect/Vector/vector-scalable-outerproduct.mlir
@@ -21,9 +21,12 @@ func.func @invalid_outerproduct(%src : memref<?xf32>) {
%0 = vector.load %src[%idx] : memref<?xf32>, vector<[4]xf32>
%1 = vector.load %src[%idx] : memref<?xf32>, vector<4xf32>
- // expected-error @+1 {{expected either all or none of vector operands #1 and #2 to be scalable}}
+ // expected-error @+1 {{expected either both or only #2 operand dim to be scalable}}
%op = vector.outerproduct %0, %1 : vector<[4]xf32>, vector<4xf32>
+
+ return
}
+
// -----
func.func @invalid_outerproduct1(%src : memref<?xf32>) {
diff --git a/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/test-contraction.mlir b/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/test-contraction.mlir
new file mode 100644
index 000000000000000..12187dfd7787155
--- /dev/null
+++ b/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/test-contraction.mlir
@@ -0,0 +1,137 @@
+// DEFINE: %{compile} = mlir-opt %s -test-transform-dialect-interpreter -test-transform-dialect-erase-schedule\
+// DEFINE: -cse -canonicalize -convert-vector-to-scf -arm-sve-legalize-vector-storage\
+// DEFINE: -convert-vector-to-llvm="enable-arm-sve" -test-lower-to-llvm
+// DEFINE: %{entry} =
+// DEFINE: %{run} = %mcr_aarch64_cmd -e=%{entry} -entry-point-result=void --march=aarch64 --mattr="+sve" -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext
+
+// REDEFINE: %{entry} = entry_i32
+// RUN: %{compile} | %{run} | FileCheck %s --check-prefix=I32
+
+// REDEFINE: %{entry} = entry_f32
+// RUN: %{compile} | %{run} | FileCheck %s --check-prefix=F32
+
+#matmat_accesses = [
+ affine_map<(i, j, k) -> (i, k)>,
+ affine_map<(i, j, k) -> (k, j)>,
+ affine_map<(i, j, k) -> (i, j)>
+]
+#matmat_trait = {
+ indexing_maps = #matmat_accesses,
+ iterator_types = ["parallel", "parallel", "reduction"]
+}
+
+func.func @entry_i32() {
+ %vscale = vector.vscale
+
+ %c0 = arith.constant 0 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ %c5 = arith.constant 5 : index
+ %n_rows = arith.muli %vscale, %c2 : index
+
+ %cst = arith.constant 0: i32
+ %i32_123 = arith.constant 123 : i32
+ %i32_314 = arith.constant 314 : i32
+
+ // Allocate and initialize matrix A
+ %A_alloc = memref.alloca() : memref<3x5xi32>
+ linalg.fill ins(%i32_123 : i32) outs(%A_alloc :memref<3x5xi32>)
+ %mask_a = vector.create_mask %c3, %c5 : vector<3x5xi1>
+ %vector_a = vector.transfer_read %A_alloc[%c0, %c0], %cst, %mask_a {in_bounds = [true, true]} : memref<3x5xi32>, vector<3x5xi32>
+
+ // Allocate and initialize matrix B
+ %B_alloc = memref.alloca(%n_rows) : memref<5x?xi32>
+ linalg.fill ins(%i32_123 : i32) outs(%B_alloc :memref<5x?xi32>)
+ %mask_b = vector.create_mask %c5, %n_rows : vector<5x[2]xi1>
+ %vector_b = vector.transfer_read %B_alloc[%c0, %c0], %cst, %mask_b {in_bounds = [true, true]} : memref<5x?xi32>, vector<5x[2]xi32>
+
+ // Allocate and initialize matrix C
+ %C_alloc = memref.alloca(%n_rows) : memref<3x?xi32>
+ linalg.fill ins(%i32_314 : i32) outs(%C_alloc :memref<3x?xi32>)
+ %mask_c = vector.create_mask %c3, %n_rows : vector<3x[2]xi1>
+ %vector_c = vector.transfer_read %C_alloc[%c0, %c0], %cst, %mask_c {in_bounds = [true, true]} : memref<3x?xi32>, vector<3x[2]xi32>
+
+ // Matmul
+ %m = vector.create_mask %c3, %n_rows, %c5 : vector<3x[2]x5xi1>
+ %0 = vector.mask %m { vector.contract #matmat_trait %vector_a, %vector_b, %vector_c
+ : vector<3x5xi32>, vector<5x[2]xi32> into vector<3x[2]xi32> } : vector<3x[2]x5xi1> -> vector<3x[2]xi32>
+
+ // Print the output
+ %slice1 = vector.extract %0[0] : vector<[2]xi32> from vector<3x[2]xi32>
+ // I32: ( 75959, 75959, 75959, 75959
+ vector.print %slice1 : vector<[2]xi32>
+ %slice2 = vector.extract %0[1] : vector<[2]xi32> from vector<3x[2]xi32>
+ // I32-NEXT: ( 75959, 75959, 75959, 75959
+ vector.print %slice2 : vector<[2]xi32>
+ %slice3 = vector.extract %0[2] : vector<[2]xi32> from vector<3x[2]xi32>
+ // I32-NEXT: ( 75959, 75959, 75959, 75959
+ vector.print %slice3 : vector<[2]xi32>
+
+ // CHECK: SVE: END OF TEST OUTPUT
+ vector.print str "SVE: END OF TEST OUTPUT"
+
+ return
+}
+
+func.func @entry_f32() {
+ %vscale = vector.vscale
+
+ %c0 = arith.constant 0 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ %c5 = arith.constant 5 : index
+ %n_rows = arith.muli %vscale, %c2 : index
+
+ %cst = arith.constant 0.0: f32
+ %f32_123 = arith.constant 1.23 : f32
+ %f32_314 = arith.constant 3.14 : f32
+
+ // Allocate and initialize matrix A
+ %A_alloc = memref.alloca() : memref<3x5xf32>
+ linalg.fill ins(%f32_123 : f32) outs(%A_alloc :memref<3x5xf32>)
+ %mask_a = vector.create_mask %c3, %c5 : vector<3x5xi1>
+ %vector_a = vector.transfer_read %A_alloc[%c0, %c0], %cst, %mask_a {in_bounds = [true, true]} : memref<3x5xf32>, vector<3x5xf32>
+
+ // Allocate and initialize matrix B
+ %B_alloc = memref.alloca(%n_rows) : memref<5x?xf32>
+ linalg.fill ins(%f32_123 : f32) outs(%B_alloc :memref<5x?xf32>)
+ %mask_b = vector.create_mask %c5, %n_rows : vector<5x[2]xi1>
+ %vector_b = vector.transfer_read %B_alloc[%c0, %c0], %cst, %mask_b {in_bounds = [true, true]} : memref<5x?xf32>, vector<5x[2]xf32>
+
+ // Allocate and initialize matrix C
+ %C_alloc = memref.alloca(%n_rows) : memref<3x?xf32>
+ linalg.fill ins(%f32_314 : f32) outs(%C_alloc :memref<3x?xf32>)
+ %mask_c = vector.create_mask %c3, %n_rows : vector<3x[2]xi1>
+ %vector_c = vector.transfer_read %C_alloc[%c0, %c0], %cst, %mask_c {in_bounds = [true, true]} : memref<3x?xf32>, vector<3x[2]xf32>
+
+ // Matmul
+ %m = vector.create_mask %c3, %n_rows, %c5 : vector<3x[2]x5xi1>
+ %0 = vector.mask %m { vector.contract #matmat_trait %vector_a, %vector_b, %vector_c
+ : vector<3x5xf32>, vector<5x[2]xf32> into vector<3x[2]xf32> } : vector<3x[2]x5xi1> -> vector<3x[2]xf32>
+
+ // Print the output
+ %slice1 = vector.extract %0[0] : vector<[2]xf32> from vector<3x[2]xf32>
+ // F32: ( 10.7045, 10.7045, 10.7045, 10.7045
+ vector.print %slice1 : vector<[2]xf32>
+ %slice2 = vector.extract %0[1] : vector<[2]xf32> from vector<3x[2]xf32>
+ // F32-NEXT: ( 10.7045, 10.7045, 10.7045, 10.7045
+ vector.print %slice2 : vector<[2]xf32>
+ %slice3 = vector.extract %0[2] : vector<[2]xf32> from vector<3x[2]xf32>
+ // F32-NEXT: ( 10.7045, 10.7045, 10.7045, 10.7045
+ vector.print %slice3 : vector<[2]xf32>
+
+ // CHECK: SVE: END OF TEST OUTPUT
+ vector.print str "SVE: END OF TEST OUTPUT"
+
+ return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%module_op: !transform.any_op):
+ %f = transform.structured.match ops{["func.func"]} in %module_op
+ : (!transform.any_op) -> !transform.any_op
+
+ transform.apply_patterns to %f {
+ transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"
+ } : !transform.any_op
+}
More information about the Mlir-commits
mailing list