[Mlir-commits] [mlir] [mlir][ArmSME] Add initial SME vector legalization pass (PR #79152)
Cullen Rhodes
llvmlistbot at llvm.org
Tue Jan 30 03:33:10 PST 2024
================
@@ -0,0 +1,380 @@
+//===- VectorLegalization.cpp - Legalize vectors for lowering to ArmSME ---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass legalizes vector operations so they can be lowered to ArmSME.
+// Currently, this only implements the decomposition of vector operations that
+// use vector sizes larger than an SME tile, into multiple SME-sized operations.
+//
+// Note: In the context of this pass 'tile' always refers to an SME tile.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/ArmSME/IR/ArmSME.h"
+#include "mlir/Dialect/ArmSME/Transforms/Passes.h"
+#include "mlir/Dialect/ArmSME/Utils/Utils.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Func/Transforms/OneToNFuncConversions.h"
+#include "mlir/Dialect/SCF/Transforms/Patterns.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
+#include "mlir/Transforms/OneToNTypeConversion.h"
+
+#define DEBUG_TYPE "arm-sme-vector-legalization"
+
+namespace mlir::arm_sme {
+#define GEN_PASS_DEF_VECTORLEGALIZATION
+#include "mlir/Dialect/ArmSME/Transforms/Passes.h.inc"
+} // namespace mlir::arm_sme
+
+using namespace mlir;
+using namespace mlir::arm_sme;
+
+namespace {
+
+// Common match failure reasons.
+static constexpr StringLiteral MATCH_FAILURE_NOT_SME_TILE_TYPE_MULTIPLE(
+ "op vector size is not multiple of SME tiles");
+static constexpr StringLiteral MATCH_FAILURE_UNSUPPORTED_MASK_OP(
+ "op mask is unsupported for legalization/decomposition");
+static constexpr StringLiteral
+ MATCH_FAILURE_NON_PERMUTATION_MAP("op affine map is not a permutation");
+
+/// An SMESubTile represents a single SME-sized sub-tile from decomposing a
+/// larger vector type. The (`row`, `col`) are the position of the tile in the
+/// original vector type. For example for an [8]x[8] tile with four [4]x[4]
+/// sub-tiles, we would have:
+///
+/// 8 x vscale
+/// ┌─────────────┬─────────────┐
+/// │(0,0) │(0,4) │
+/// │ │ │
+/// ├─────────────┼─────────────┤ 8 x vscale
+/// │(4,0) │(4,4) │
+/// │ │ │
+/// └─────────────┴─────────────┘
+struct SMESubTile {
+ // Note: The units of (row, col) are vscale (as SME tiles are scalable).
+ int row{0};
+ int col{0};
+ // The SME tile type.
+ VectorType type;
+};
+
+/// Adds a constant elementwise scalable offset to `indices` (which are of equal
+/// length). For example, in the 2D case this would return:
+// { indices[0] + offset[0] * vscale, indices[1] + offset[1] * vscale }
+SmallVector<Value, 2> addConstantScalableOffset(OpBuilder &builder,
+ Location loc,
+ ValueRange indices,
+ ArrayRef<int> scalableOffsets) {
+ auto vscale = builder.create<vector::VectorScaleOp>(loc);
+ return llvm::map_to_vector(
+ llvm::zip_equal(indices, scalableOffsets), [&](auto pair) -> Value {
+ auto [index, base] = pair;
+ auto offset = builder.create<arith::MulIOp>(
+ loc, builder.create<arith::ConstantIndexOp>(loc, base), vscale);
+ return builder.create<arith::AddIOp>(loc, index, offset);
+ });
+}
+
+/// Adjusts `indices` (e.g. from a load/store) for a larger vector type to
+/// indices for one of the SME sub-tiles it will decompose into.
+///
+/// For example, if you were to decompose an 8x8 load into four 4x4 tiles, the
+/// indices for each tile would need to be adjusted as follows:
+///
+/// initial indices = [a,b], inital size = 8x8, target size = 4x4
+/// ┌─────────────┬─────────────┐
+/// │[a,b] │[a,b+4] │
+/// │ │ │
+/// ├─────────────┼─────────────┤
+/// │[a+4,b] │[a+4,b+4] │
+/// │ │ │
+/// └─────────────┴─────────────┘
+SmallVector<Value, 2> getSMESubTileIndices(OpBuilder &builder, Location loc,
+ ValueRange indices,
+ SMESubTile smeTile) {
+ return addConstantScalableOffset(builder, loc, indices,
+ {smeTile.row, smeTile.col});
+}
+
+/// Returns true if `mask` is generated by an operation that can be decomposed
+/// for SME. Currently, that is just no mask, or vector.create_mask.
+/// TODO: Add support for vector.constant_mask once required for SME.
+bool isSupportedMaskOp(Value mask) {
+ return !mask || mask.getDefiningOp<vector::CreateMaskOp>();
+}
+
+/// Extracts a mask for an SME sub-tile from the mask of a larger vector type.
+Value extractSMEMask(OpBuilder &builder, Location loc, Value mask,
+ SMESubTile smeTile) {
+ assert(isSupportedMaskOp(mask));
+ if (!mask)
+ return Value{};
+ auto createMask = mask.getDefiningOp<vector::CreateMaskOp>();
+ // The operands of `vector.create_mask` (from a 2D perspective) are the
+ // coordinates where the mask ends. So we subtract where this tile starts,
+ // from the mask operands to get the parameters for this sub-tile.
+ auto smeTileMaskDims = addConstantScalableOffset(
+ builder, loc, createMask.getOperands(), {-smeTile.row, -smeTile.col});
+ auto smeTileCreateMask = builder.create<vector::CreateMaskOp>(
+ loc, smeTile.type.clone(builder.getI1Type()), smeTileMaskDims);
+ return smeTileCreateMask.getResult();
+}
+
+/// Constructs an iterator that returns each SME tile (with coordinates)
+/// contained within a VectorType. For example, if decomposing an [8]x[8] into
+/// [4]x[4] tiles, the iterator would yield the tiles: (0, 0), (0, 4), (4, 0),
+/// (4, 4).
+auto decomposeToSMETiles(OpBuilder &builder, VectorType type,
+ VectorType smeTileType,
+ bool transposeIndices = false) {
+ assert(isMultipleOfSMETileVectorType(type) &&
+ "`type` not multiple of SME tiles");
+ return llvm::map_range(
+ StaticTileOffsetRange(type.getShape(), {smeTileType.getDimSize(0),
+ smeTileType.getDimSize(1)}),
+ [=](auto indices) {
+ int row = int(indices[0]);
+ int col = int(indices[1]);
+ if (transposeIndices)
+ std::swap(row, col);
+ return SMESubTile{row, col, smeTileType};
+ });
+}
+
+/// Returns the number of SME tiles that fit into the (2D-scalable) vector type
+/// `type`.
+int getNumberOfSMETilesForVectorType(VectorType type) {
+ assert(isMultipleOfSMETileVectorType(type) &&
+ "`type` not multiple of SME tiles");
+ int64_t vectorRows = type.getDimSize(0);
+ int64_t vectorCols = type.getDimSize(1);
+ auto elementType = type.getElementType();
+ unsigned minNumElts = getSMETileSliceMinNumElts(elementType);
+ return (vectorRows * vectorCols) / (minNumElts * minNumElts);
+}
+
+/// Legalize `vector.outerproduct` operations to fit within SME tiles by
+/// decomposing them into tile-sized operations.
+struct LegalizeVectorOuterProductOpsByDecomposition
+ : public OneToNOpConversionPattern<vector::OuterProductOp> {
+ using OneToNOpConversionPattern::OneToNOpConversionPattern;
+
+ LogicalResult
+ matchAndRewrite(vector::OuterProductOp outerProductOp, OpAdaptor adaptor,
+ OneToNPatternRewriter &rewriter) const override {
+ auto vectorType = outerProductOp.getResultVectorType();
+ if (!isMultipleOfSMETileVectorType(vectorType))
+ return rewriter.notifyMatchFailure(
+ outerProductOp, MATCH_FAILURE_NOT_SME_TILE_TYPE_MULTIPLE);
+
+ Value mask;
+ Operation *rootOp = outerProductOp;
+ auto loc = outerProductOp.getLoc();
+ if (outerProductOp.isMasked()) {
+ auto maskOp = outerProductOp.getMaskingOp();
+ mask = maskOp.getMask();
+ rootOp = maskOp;
+ }
+
+ if (!isSupportedMaskOp(mask))
+ return rewriter.notifyMatchFailure(outerProductOp,
+ MATCH_FAILURE_UNSUPPORTED_MASK_OP);
+
+ ValueRange accSMETiles = adaptor.getAcc();
+ auto smeTileType = getSMETileTypeForElement(vectorType.getElementType());
+ VectorType sliceType = VectorType::Builder(smeTileType).dropDim(0);
+
+ SmallVector<Value> resultSMETiles;
+ for (auto [index, smeTile] : llvm::enumerate(
+ decomposeToSMETiles(rewriter, vectorType, smeTileType))) {
+
+ auto smeMask = extractSMEMask(rewriter, loc, mask, smeTile);
+ auto lhs = rewriter.create<vector::ScalableExtractOp>(
+ loc, sliceType, outerProductOp.getLhs(), smeTile.row);
+ auto rhs = rewriter.create<vector::ScalableExtractOp>(
+ loc, sliceType, outerProductOp.getRhs(), smeTile.col);
+ auto smeOuterProduct = rewriter.create<vector::OuterProductOp>(
+ loc, smeTileType, lhs, rhs,
+ !accSMETiles.empty() ? accSMETiles[index] : Value{},
+ outerProductOp.getKind());
+
+ auto maskedOuterProduct =
+ vector::maskOperation(rewriter, smeOuterProduct, smeMask);
+ resultSMETiles.push_back(maskedOuterProduct->getResult(0));
+ }
+
+ rewriter.replaceOp(rootOp, resultSMETiles, adaptor.getResultMapping());
+ return success();
+ }
+};
+
+// Workaround for `vector.mask`. We want to match on `vector.outerproduct` (to
+// get the help of the type conversion), but doing so results in the type
+// conversion adding target materializations in the `vector.mask` region
+// (invalid). This pattern matches on `vector.mask` then calls into the
+// `vector.outerproduct` pattern to work around this issue.
+struct LegalizeMaskedVectorOuterProductOpsByDecomposition
+ : public OneToNOpConversionPattern<vector::MaskOp> {
+ using OneToNOpConversionPattern::OneToNOpConversionPattern;
+
+ LogicalResult
+ matchAndRewrite(vector::MaskOp maskOp, OpAdaptor adaptor,
+ OneToNPatternRewriter &rewriter) const override {
+ if (auto outerProductOp =
+ llvm::dyn_cast<vector::OuterProductOp>(maskOp.getMaskableOp())) {
+ LegalizeVectorOuterProductOpsByDecomposition pattern(*getTypeConverter(),
+ getContext());
+ return static_cast<RewritePattern &>(pattern).matchAndRewrite(
+ outerProductOp, rewriter);
+ }
+ return failure();
+ }
+};
+
+/// Legalize `vector.transfer_read` operations to fit within SME tiles by
+/// decomposing them into tile-sized operations.
+struct LegalizeTransferReadOpsByDecomposition
+ : public OneToNOpConversionPattern<vector::TransferReadOp> {
+ using OneToNOpConversionPattern::OneToNOpConversionPattern;
+
+ LogicalResult
+ matchAndRewrite(vector::TransferReadOp readOp, OpAdaptor adaptor,
+ OneToNPatternRewriter &rewriter) const override {
+ auto vectorType = readOp.getVectorType();
+ if (!isMultipleOfSMETileVectorType(vectorType))
+ return rewriter.notifyMatchFailure(
+ readOp, MATCH_FAILURE_NOT_SME_TILE_TYPE_MULTIPLE);
+
+ auto mask = readOp.getMask();
+ if (!isSupportedMaskOp(mask))
+ return rewriter.notifyMatchFailure(readOp,
+ MATCH_FAILURE_UNSUPPORTED_MASK_OP);
+
+ auto permutationMap = readOp.getPermutationMap();
+ if (!permutationMap.isPermutation())
+ return rewriter.notifyMatchFailure(readOp,
+ MATCH_FAILURE_NON_PERMUTATION_MAP);
+
+ // Note: For 2D vector types the only non-identity permutation is a simple
+ // tranpose [1, 0].
+ bool transposed = !permutationMap.isIdentity();
+
+ auto loc = readOp.getLoc();
+ auto smeTileType = getSMETileTypeForElement(vectorType.getElementType());
+
+ SmallVector<Value> resultSMETiles;
+ for (SMESubTile smeTile :
+ decomposeToSMETiles(rewriter, vectorType, smeTileType, transposed)) {
+ auto smeMask = extractSMEMask(rewriter, loc, mask, smeTile);
+ auto smeRead = rewriter.create<vector::TransferReadOp>(
+ loc, smeTileType, readOp.getSource(),
+ getSMESubTileIndices(rewriter, loc, readOp.getIndices(), smeTile),
+ readOp.getPermutationMapAttr(), readOp.getPadding(), smeMask,
+ readOp.getInBoundsAttr());
+ resultSMETiles.push_back(smeRead);
+ }
+
+ rewriter.replaceOp(readOp, resultSMETiles, adaptor.getResultMapping());
+ return success();
+ }
+};
+
+/// Legalize `vector.transfer_write` operations to fit within SME tiles by
+/// decomposing them into tile-sized operations.
+struct LegalizeTransferWriteOpsByDecomposition
+ : public OneToNOpConversionPattern<vector::TransferWriteOp> {
+ using OneToNOpConversionPattern::OneToNOpConversionPattern;
+
+ LogicalResult
+ matchAndRewrite(vector::TransferWriteOp writeOp, OpAdaptor adaptor,
+ OneToNPatternRewriter &rewriter) const override {
+ auto vectorType = writeOp.getVectorType();
+ if (!isMultipleOfSMETileVectorType(vectorType))
+ return rewriter.notifyMatchFailure(
+ writeOp, MATCH_FAILURE_NOT_SME_TILE_TYPE_MULTIPLE);
+
+ auto mask = writeOp.getMask();
+ if (!isSupportedMaskOp(mask))
+ return rewriter.notifyMatchFailure(writeOp,
+ MATCH_FAILURE_UNSUPPORTED_MASK_OP);
+
+ auto permutationMap = writeOp.getPermutationMap();
+ if (!permutationMap.isPermutation())
+ return rewriter.notifyMatchFailure(writeOp,
+ MATCH_FAILURE_NON_PERMUTATION_MAP);
+
+ // Note: For 2D vector types the only non-identity permutation is a simple
+ // tranpose [1, 0].
+ bool transposed = !permutationMap.isIdentity();
+
+ auto loc = writeOp.getLoc();
+ auto smeTileType = getSMETileTypeForElement(vectorType.getElementType());
+ auto inputSMETiles = adaptor.getVector();
+
+ Value destTensorOrMemref = writeOp.getSource();
+ for (auto [index, smeTile] : llvm::enumerate(decomposeToSMETiles(
+ rewriter, vectorType, smeTileType, transposed))) {
+ auto smeMask = extractSMEMask(rewriter, loc, mask, smeTile);
+ auto smeWrite = rewriter.create<vector::TransferWriteOp>(
+ loc, inputSMETiles[index], destTensorOrMemref,
+ getSMESubTileIndices(rewriter, loc, writeOp.getIndices(), smeTile),
+ writeOp.getPermutationMapAttr(), smeMask, writeOp.getInBoundsAttr());
+ if (writeOp.hasPureTensorSemantics())
+ destTensorOrMemref = smeWrite.getResult();
+ }
+
+ if (writeOp.hasPureTensorSemantics())
+ rewriter.replaceOp(writeOp, destTensorOrMemref);
+ else
+ rewriter.eraseOp(writeOp);
+
+ return success();
+ }
+};
+
+struct VectorLegalizationPass
+ : public arm_sme::impl::VectorLegalizationBase<VectorLegalizationPass> {
+ void runOnOperation() override {
+ auto *context = &getContext();
+ OneToNTypeConverter converter;
+ RewritePatternSet patterns(context);
+
+ converter.addConversion([](Type type) { return type; });
+ converter.addConversion(
+ [](VectorType vectorType,
+ SmallVectorImpl<Type> &types) -> std::optional<LogicalResult> {
+ if (!isMultipleOfSMETileVectorType(vectorType))
+ return std::nullopt;
+ auto smeTileTileCount = getNumberOfSMETilesForVectorType(vectorType);
----------------
c-rhodes wrote:
```suggestion
auto smeTileCount = getNumberOfSMETilesForVectorType(vectorType);
```
?
https://github.com/llvm/llvm-project/pull/79152
More information about the Mlir-commits
mailing list