Google GitGoogle Git
Sign in
llvm/llvm-project/cd9864892553166300a6ac8fec6d2445d30fa513/./mlir/lib/Dialect/ArmSVE/Transforms/LowerContractToSVEPatterns.cpp
blob: 6cb2a56aebdd12d212756ad24ddf500542cc1a92 [file]
//===- LowerContractToSVEPatterns.cpp - Contract to I8MM/BF16 ---*- 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
//
//===----------------------------------------------------------------------===//
//
// This file implements lowering patterns from vector.contract to operations
// that map to instructions from the SVE FEAT_I8MM and FEAT_BF16 extensions.
//
// TODO: There may be opportunities to unify this with a similar pattern
// for Neon. See:
// https://github.com/llvm/llvm-project/issues/145559
// LowerContractToNeonPatterns.cpp
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/ArmSVE/IR/ArmSVEDialect.h"
#include "mlir/Dialect/ArmSVE/Transforms/Transforms.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/UB/IR/UBOps.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/PatternMatch.h"
#include <cassert>
#include <numeric>
#define DEBUG_TYPE "lower-contract-to-arm-sve-i8mm"
using namespace mlir;
namespace {
// Get the operand of a `vector.contract`. This function is intended to abstract
// away from the particular way a value is extended before feeding it into the
// `vector.contract` - via zero-extend or an explicit or implicit sign-extend
// (for implicit sign-extension see `vector.contract` documentation).
//
// The template parameter `Op` indicates the extension operation (explicit or
// implicit) for which we are checking.
//
// Return success only for extensions from `i8` to `i32`.
template <typename Op>
std::optional<Value> getExtOperand(Value v) {
static_assert(llvm::is_one_of<Op, arith::ExtSIOp, arith::ExtUIOp>::value,
"Must be instantiated with either sign- or zero- extension op");
// If the operand is not defined by an explicit extend operation of the
// accepted operation type allow for an implicit sign-extension.
auto extOp = v.getDefiningOp<Op>();
if (!extOp) {
if constexpr (std::is_same<Op, arith::ExtSIOp>::value) {
auto vTy = cast<VectorType>(v.getType());
if (!vTy.getElementType().isSignlessInteger(8))
return {};
return v;
}
return {};
}
// If the operand is defined by an explicit extend operation of the accepted
// operation type, check it's extended from `i8` to `i32`.
auto inOp = extOp.getIn();
auto inTy = dyn_cast<VectorType>(inOp.getType());
if (!inTy || !inTy.getElementType().isSignlessInteger(8))
return {};
auto outTy = dyn_cast<VectorType>(extOp.getType());
if (!outTy || !outTy.getElementType().isSignlessInteger(32))
return {};
return inOp;
}
/// This class encapsulates the algorithm and parametrisation (in terms of types
/// and dimensions) of lowering a `vector.contract` to "primitive" matrix
/// multiplication operations of the SVE dialect (here "primitive" would mean
/// corresponding to a single target instruction).
///
/// Supported are lowering to FEAT_I8MM `smmla`, `ummla`, and `usmmla`, and to
/// FEAT_BF16 `bfmmla`. All the transformations are very similar to each other
/// for concreteness the description below is given for `smmla`.
///
/// The lowering triggers for a contraction operation that performs a matrix
/// multiply of two 8-bit integer matrix tiles with logical dimensions
/// <Mx8> and <8x[N]> for the left-hand side (LHS) and the right-hand side
/// (RHS), respectively, added to a 32-bit integer accumulator operand (ACC)
/// with dimensions <Mx[N]>, yielding a <Mx[N]> 32-bit integer result (OUT).
///
/// The operands' shapes are such that the operands can be evenly split into
/// sub-tiles with dimensions as expected by the targeted FEAT_I8MM
/// instructions. The intent is that M and N are chosen (by higher level
/// transforms) in such a way as to maximise register usage. The main use case
/// we envision as of now is MMT4D, thus the RHS operand is expected
/// pre-transposed.
///
/// The matrix multiplication is performed by unrolling the usual tiled matrix
/// multiplication algorithm using sub-tiles with dimensions <2x8> for the
/// LHS, <8x[2]> for the RHS, and <2x[2]> for the result and the input
/// accumulator.
///
/// One way to illustrate the operation is as follows:
///
/// RHS<8x[N]>: <8x[2]> <8x[2]> ... <8x[2]>
/// +-----------------------------
/// LHS<Mx8>: <2x8> | <2x[2]> <2x[2]> ... <2x[2]>
/// <2x8> | <2x[2]> <2x[2]> ... <2x[2]>
/// ... | ... ... ... ...
/// <2x8> | <2x[2]> <2x[2]> ... <2x[2]>
///
/// The RHS operand is unpacked into N/2 values, each representing a sequence
/// of VSCALE number of sub-tiles with dimensions <8x2>.
/// The LHS operand is initially unpacked into M/2 values, each representing a
/// sub-tile with dimensions <2x8>, and then each such sub-tile is replicated
/// VSCALE times. Multiplying thus replicated LHS sub-tile by the corresponding
/// RHS sub-tile correctly computes an entire result sub-tile.
/// The 2x2 sub-tiles of the ACC and OUT have rows that are not adjacent
/// (in memory or when imposing a row-major layout on the 2D vector value).
/// Reading the ACC is implemented as reading two consecutive rows and
/// interleaving the by pairs to obtain a vector having length twice the length
/// of an ACC row. This vector now is a sequence of one-dimensional tiles with
/// the exact layout needed by the `smmla`/`bfmmla`/etc instructions, which
/// tiles are extracted one by one. For illustration, if we have an 2x4 ACC tile
/// a0 a1 b0 b1
/// a2 a3 b2 b3
/// we read the two rows as separate values and then interleave by pairs
/// to obtain
/// a0 a1 a2 a3 b0 b1 b2 b3
/// from which we extract `a0 a1 a2 a3` and `b0 b1 b2 b3`.
///
/// Writing the OUT tile is done by the reverse of the above procedure,
/// concatenate two "flattened" sub-tiles into
/// c0 c1 c2 c3 d0 d1 d2 d3
/// deinterleave by pairs to obtain as separate values
/// c0 c1 d0 d1
/// c2 c3 d2 d3
/// which are then inserted into the final result.
///
/// Multiplication of a signed LHS by an unsigned LHS is performed by
/// swapping the order of the operands and emitting an `usmmla` (since there
/// isn't an `summla` instruction). Therefore each ACC sub-tile needs
/// to be transposed before the addition and the sum, an OUT sub-tile,
/// needs to be transposed before insertion into the final result.
/// This is done very elegantly by a modification of the above to
/// interleave/deinterleave not by pairs, but by individual elements, e.g.
/// after ordinary interleave we obtain
/// a0 a2 a1 a3 b0 b2 b1 b3
/// which is exactly the desired layout of having each individual 2x2 tile
/// transposed.
///
/// All of the above readily applies to FEAT_BF16 `bfmmla` with the
/// difference that the shapes of the LHS, RHS are <Mx4>, <4x[M]>, and
/// respectively, that is the "K" dimension is fixed to 4, instead of 8 (like
/// for the integer case).
class VectorContractRewriter {
protected:
// Designate the operation (resp. instruction) used to do sub-tile matrix
// multiplications.
enum class MMLA {
Nop,
SignedInt, // smmla
UnsignedInt, // ummla
MixedInt, // usmmla
Bfloat // bfmmla
};
// Lower-level operation to be emitted.
MMLA mmlaOp = MMLA::Nop;
// Indicate if the operands for the ArmSVE dialect operation need to be
// swapped. Currently this is needed in order to emulate an "summla"
// operation.
bool swapOperands = false;
// The operand tiles. These are not necessarily the operends of
// `vector.contract`, for example they could be operands to `arith.extsi`
// that is in turn fed into `vector.contract`.
Value lhs;
Value rhs;
Value acc;
// Conventional names for matrix dimensions.
int64_t m = 0;
int64_t n = 0;
int64_t k = 0;
// Create the matrix mulitply and accumulate operation according to
// `mmlaOp`.
Value createMMLA(PatternRewriter &rewriter, Location loc, Value acc,
Value lhs, Value rhs);
// Check general preconditions for applying the transformation, common to the
// integer and the bfloat16 case.
LogicalResult match(vector::ContractionOp op, PatternRewriter &rewriter);
public:
VectorContractRewriter() = default;
// Do the actuall rewrite. This member function is shared by both integer and
// bfloat16 rewrites.
Value lower(vector::ContractionOp op, PatternRewriter &rewriter);
};
Value VectorContractRewriter::createMMLA(PatternRewriter &rewriter,
Location loc, Value acc, Value lhs,
Value rhs) {
Type resTy = acc.getType();
if (swapOperands)
std::swap(lhs, rhs);
switch (mmlaOp) {
case MMLA::SignedInt:
return arm_sve::SmmlaOp::create(rewriter, loc, resTy, acc, lhs, rhs);
case MMLA::UnsignedInt:
return arm_sve::UmmlaOp::create(rewriter, loc, resTy, acc, lhs, rhs);
case MMLA::MixedInt:
return arm_sve::UsmmlaOp::create(rewriter, loc, resTy, acc, lhs, rhs);
case MMLA::Bfloat:
return arm_sve::BfmmlaOp::create(rewriter, loc, resTy, acc, lhs, rhs);
default:
llvm_unreachable("Uninitialized operation kind");
}
}
LogicalResult VectorContractRewriter::match(vector::ContractionOp op,
PatternRewriter &rewriter) {
// Check iterator types for matrix multiplication.
auto itTypes = op.getIteratorTypesArray();
if (itTypes.size() != 3 || itTypes[0] != vector::IteratorType::parallel ||
itTypes[1] != vector::IteratorType::parallel ||
itTypes[2] != vector::IteratorType::reduction)
return rewriter.notifyMatchFailure(
op, "iterator types do not correspond to matrix multiplication");
// Check permutation maps. For now only accept
// lhs: (d0, d1, d2) -> (d0, d2)
// rhs: (d0, d1, d2) -> (d1, d2)
// acc: (d0, d1, d2) -> (d0, d1)
// This corresponds to matrix multiplication with transposed RHS.
if (op.getIndexingMapsArray()[0] !=
AffineMap::getMultiDimMapWithTargets(3, ArrayRef{0u, 2u},
op.getContext()) ||
op.getIndexingMapsArray()[1] !=
AffineMap::getMultiDimMapWithTargets(3, ArrayRef{1u, 2u},
op.getContext()) ||
op.getIndexingMapsArray()[2] != AffineMap::getMultiDimMapWithTargets(
3, ArrayRef{0u, 1u}, op.getContext()))
return rewriter.notifyMatchFailure(op, "non-matching permutation maps");
// Check the combining kind is addition.
if (op.getKind() != vector::CombiningKind::ADD)
return rewriter.notifyMatchFailure(op, "combining kind is not an addition");
return success();
}
Value VectorContractRewriter::lower(vector::ContractionOp op,
PatternRewriter &rewriter) {
// Initialize some helper types.
Type operandEltType = cast<VectorType>(lhs.getType()).getElementType();
Type resultEltType = cast<VectorType>(op.getResultType()).getElementType();
const int64_t numOperandSubTileElts =
128 / operandEltType.getIntOrFloatBitWidth();
assert(resultEltType.getIntOrFloatBitWidth() == 32 &&
"Only implemented for i32 or f32 output");
const int64_t numResultSubTileElts = 4;
// Single-dimensional vector types for the operands of the ArmSVE dialect
// op.
auto flatLhsType =
VectorType::get(/*shape=*/numOperandSubTileElts, operandEltType,
/*scalableDims=*/{true});
auto flatRhsType =
VectorType::get(/*shape=*/numOperandSubTileElts, operandEltType,
/*scalableDims=*/{true});
auto flatAccType =
VectorType::get(/*shape=*/numResultSubTileElts, resultEltType,
/*scalableDims=*/{true});
// Single-dimension vector type for the entire RHS tile.
auto flatRhsTileType = VectorType::get(/*shape=*/k * n, operandEltType,
/*scalableDims=*/{true});
// Vector type having the same number of elements as a row in the
// accumulator/output tile and the same element type.
auto accRowTy = VectorType::get(/*shape=*/n, resultEltType,
/*scalableDims=*/{true});
// Vector type having twice the number of elements as a row in the
// accumulator/output tile the same element type.
auto accRowX2Ty = VectorType::get(/*shape=*/2 * n, resultEltType,
/*scalableDims=*/{true});
// Vector type having half the number of elements as a row in the
// accumulator/output tile and an integer element type with twice the bit
// width.
auto accRow64Ty = VectorType::get(/*shape=*/n / 2, rewriter.getI64Type(),
/*scalableDims=*/{true});
// Vector type having the same the number of elements as a row in the
// accumulator/output tile and an integer element type with twice the bit
// width.
auto accRowX264Ty = VectorType::get(/*shape=*/n, rewriter.getI64Type(),
/*scalableDims=*/{true});
Location loc = op.getLoc();
// Extract LHS sub-tiles with logical shape <2xK>.
SmallVector<Value> lhsTile;
for (int64_t i = 0; i < m; i += 2) {
// Extract two consecutive rows of the LHS tile.
auto r0 =
vector::ExtractOp::create(rewriter, loc, lhs, ArrayRef<int64_t>{i});
auto r1 =
vector::ExtractOp::create(rewriter, loc, lhs, ArrayRef<int64_t>{i + 1});
// Concatenate to obtain a 2 x K x <input-type> flattened sub-tile.
SmallVector<int64_t> shuffleIdx(2 * k);
std::iota(shuffleIdx.begin(), shuffleIdx.end(), 0);
auto t = vector::ShuffleOp::create(rewriter, loc, r0, r1, shuffleIdx);
// Turn it into a scalable vector.
auto s = vector::ScalableInsertOp::create(
rewriter, loc, t, ub::PoisonOp::create(rewriter, loc, flatLhsType), 0);
// Replicate the sub-tile VSCALE times to fill the entire vector.
auto r = arm_sve::DupQLaneOp::create(rewriter, loc, s, 0);
lhsTile.push_back(r);
}
// "Flatten" the RHS tile from <[N]xK> to <[N*K]>.
auto rhs = vector::ShapeCastOp::create(rewriter, this->rhs.getLoc(),
flatRhsTileType, this->rhs);
// Extract the RHS sub-tiles with logical shape <Kx[2]>.
SmallVector<Value> rhsTile;
for (int64_t j = 0; j < n; j += 2)
rhsTile.push_back(vector::ScalableExtractOp::create(
rewriter, loc, flatRhsType, rhs, j * k));
// Extract and pack the ACC sub-tiles.
SmallVector<Value> accTile;
for (int64_t i = 0; i < m; i += 2) {
// Extract two consecutive rows of the accumulator tile.
auto r0 = vector::ExtractOp::create(rewriter, loc, op.getAcc(),
ArrayRef<int64_t>{i});
auto r1 = vector::ExtractOp::create(rewriter, loc, op.getAcc(),
ArrayRef<int64_t>{i + 1});
Value accTileVec;
if (swapOperands) {
// We are performing the operation with swapped LHS and RHS we need to
// transpose each individual 2x2 tile of the accumulator and (later) the
// final result.
accTileVec = vector::InterleaveOp::create(rewriter, loc, r0, r1);
} else {
// Bitcast accumulator rows to double-width integer elements, so
// subsequent interleave/deinterleave work on pairs of elements.
auto r0I64 = vector::BitCastOp::create(rewriter, loc, accRow64Ty, r0);
auto r1I64 = vector::BitCastOp::create(rewriter, loc, accRow64Ty, r1);
// Interleave the rows, effectively flattening each 2x2 tile into 4
// consecutive elements.
auto intrI64 = vector::InterleaveOp::create(rewriter, loc, r0I64, r1I64);
// Bitcast back to original element type.
accTileVec =
vector::BitCastOp::create(rewriter, loc, accRowX2Ty, intrI64);
}
// Extract ACC sub-tiles.
for (int64_t j = 0; j < n; j += 2)
accTile.push_back(vector::ScalableExtractOp::create(
rewriter, loc, flatAccType, accTileVec, j * 2));
}
// Emit sub-tile matrix multiplications.
SmallVector<Value> outTile;
for (int64_t i = 0; i < m / 2; ++i)
for (int64_t j = 0; j < n / 2; ++j) {
Value mmla = createMMLA(rewriter, loc, accTile[i * n / 2 + j], lhsTile[i],
rhsTile[j]);
outTile.push_back(mmla);
}
// Unpack the OUT sub-tiles and insert into the result.
Value result = ub::PoisonOp::create(rewriter, loc, op.getResultType());
for (int64_t i = 0; i < m / 2; ++i) {
// Collect a number of sub-tiles in a row.
Value row = ub::PoisonOp::create(rewriter, loc, accRowX2Ty);
for (int64_t j = 0; j < n / 2; ++j)
row = vector::ScalableInsertOp::create(
rewriter, loc, outTile[i * n / 2 + j], row, j * 4);
// Unpack the row to obtain two rows of the output. If we have the out
// sub-tiles transposed we obtain two consecutive output rows by
// separating even and odd elements, i.e. a simple deinterleave.
// Otherwise, the interleave is by pairs.
Value out0, out1;
if (swapOperands) {
auto tmp = vector::DeinterleaveOp::create(rewriter, loc, row);
out0 = tmp.getRes1();
out1 = tmp.getRes2();
} else {
// Deinterleave by pairs.
auto row64 = vector::BitCastOp::create(rewriter, loc, accRowX264Ty, row);
auto deintr64 = vector::DeinterleaveOp::create(rewriter, loc, row64);
// Bitcast back into original element type and insert into the result.
out0 = vector::BitCastOp::create(rewriter, loc, accRowTy,
deintr64.getRes1());
out1 = vector::BitCastOp::create(rewriter, loc, accRowTy,
deintr64.getRes2());
}
result = vector::InsertOp::create(rewriter, loc, out0, result, i * 2);
result = vector::InsertOp::create(rewriter, loc, out1, result, i * 2 + 1);
}
return result;
}
class VectorContractRewriterI8MM : public VectorContractRewriter {
public:
// Check the specific preconditions for the integer case. Initialise
// parametrisation types and dimensions.
LogicalResult matchAndInit(vector::ContractionOp op,
PatternRewriter &rewriter) {
if (failed(match(op, rewriter)))
return failure();
VectorType lhsType = op.getLhsType();
VectorType rhsType = op.getRhsType();
m = lhsType.getDimSize(0);
n = rhsType.getDimSize(0);
k = rhsType.getDimSize(1);
// Check the operands have the expected shape:
// * for LHS: fixed vector MxK
// * for RHS: scalable vector [N]xK
// * K == 8
// * M and N even and at least 2
if (lhsType.isScalable() || !rhsType.getScalableDims()[0] ||
rhsType.getScalableDims()[1] || lhsType.getDimSize(1) != k || k != 8 ||
m < 2 || m % 2 != 0 || n < 2 || n % 2 != 0 ||
!rhsType.getScalableDims()[0])
return rewriter.notifyMatchFailure(op, "non-matching operand shape");
// Check the output is a vector of i32 elements.
auto outTy = dyn_cast<VectorType>(op.getResultType());
if (!outTy || outTy.getElementType() != rewriter.getI32Type())
return rewriter.notifyMatchFailure(op,
"output type is not a vector of i32");
// Check inputs are sign-/zero- extensions from i8 to i32. Get the values
// before the extension. All four signed/unsigned combinations for input
// operands are supported, but they are lowered to different operations.
// Determine which is the appropriate operation to lower to.
mmlaOp = MMLA::SignedInt;
swapOperands = false;
auto maybeLhs = getExtOperand<arith::ExtSIOp>(op.getLhs());
if (!maybeLhs) {
mmlaOp = MMLA::UnsignedInt;
maybeLhs = getExtOperand<arith::ExtUIOp>(op.getLhs());
}
if (!maybeLhs)
return rewriter.notifyMatchFailure(
op, "LHS is not a sign- or zero- extended i8");
auto maybeRhs = getExtOperand<arith::ExtSIOp>(op.getRhs());
if (maybeRhs) {
if (mmlaOp == MMLA::UnsignedInt)
mmlaOp = MMLA::MixedInt;
} else {
if (mmlaOp == MMLA::SignedInt) {
mmlaOp = MMLA::MixedInt;
swapOperands = true;
}
maybeRhs = getExtOperand<arith::ExtUIOp>(op.getRhs());
}
if (!maybeRhs)
return rewriter.notifyMatchFailure(
op, "RHS is not a sign- or zero- extended i8");
// Initialise algorithm parameters.
lhs = *maybeLhs;
rhs = *maybeRhs;
acc = op.getAcc();
return success();
}
};
class VectorContractRewriterBfloat : public VectorContractRewriter {
public:
// Check the specific preconditions for the bfloat16 case. Initialise
// parametrisation types and dimensions.
LogicalResult matchAndInit(vector::ContractionOp op,
PatternRewriter &rewriter) {
if (failed(match(op, rewriter)))
return failure();
VectorType lhsType = op.getLhsType();
VectorType rhsType = op.getRhsType();
m = lhsType.getDimSize(0);
n = rhsType.getDimSize(0);
k = rhsType.getDimSize(1);
// Check the operands have the expected shape:
// * for LHS: fixed vector MxK
// * for RHS: scalable vector [N]xK
// * K == 4
// * M and N even and at least 2
if (lhsType.isScalable() || !rhsType.getScalableDims()[0] ||
rhsType.getScalableDims()[1] || lhsType.getDimSize(1) != k || k != 4 ||
m < 2 || m % 2 != 0 || n < 2 || n % 2 != 0 ||
!rhsType.getScalableDims()[0])
return rewriter.notifyMatchFailure(op, "non-matching operand shape");
// Check the output is a vector of Float32 elements.
auto outTy = dyn_cast<VectorType>(op.getResultType());
if (!outTy || outTy.getElementType() != rewriter.getF32Type())
return rewriter.notifyMatchFailure(op,
"output type is not a vector of f32");
// Check the inputs are vectors of BFloat16 elements.
if (lhsType.getElementType() != rewriter.getBF16Type())
return rewriter.notifyMatchFailure(op,
"input type is not a vector of bf16");
// Initialise algorithm parameters.
mmlaOp = MMLA::Bfloat;
swapOperands = false;
lhs = op.getLhs();
rhs = op.getRhs();
acc = op.getAcc();
return success();
}
};
class LowerContractionToSVEI8MMPattern
: public OpRewritePattern<vector::ContractionOp> {
public:
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(vector::ContractionOp op,
PatternRewriter &rewriter) const override {
// Match i8xi8 -> i32 matrix multiply and accumulate.
VectorContractRewriterI8MM vcr;
if (failed(vcr.matchAndInit(op, rewriter)))
return failure();
Value result = vcr.lower(op, rewriter);
rewriter.replaceOp(op, result);
return success();
}
};
class LowerContractionToSVEBFMMLAPattern
: public OpRewritePattern<vector::ContractionOp> {
public:
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(vector::ContractionOp op,
PatternRewriter &rewriter) const override {
// Match bf16xbf16 -> f32 matrix multiply and accumulate.
VectorContractRewriterBfloat vcr;
if (failed(vcr.matchAndInit(op, rewriter)))
return failure();
Value result = vcr.lower(op, rewriter);
rewriter.replaceOp(op, result);
return success();
}
};
} // namespace
void mlir::populateLowerContractionToSVEI8MMPatterns(
RewritePatternSet &patterns) {
MLIRContext *context = patterns.getContext();
patterns.add<LowerContractionToSVEI8MMPattern>(context, /*benefit=*/2);
}
void mlir::populateLowerContractionToSVEBFMMLAPatterns(
RewritePatternSet &patterns) {
MLIRContext *context = patterns.getContext();
patterns.add<LowerContractionToSVEBFMMLAPattern>(context, /*benefit=*/2);
}