b9877ec82f
Adds an optional stride argument to `amx.tile_load` and `amx.tile_store` operations. The stride argument aligns ops closer to the hardware intrinsics. However, stride remains optional to preserve current op behavior. Explicit stride allows greater flexibility in terms of the base buffer shapes and allows different read and write memory patterns. When stride is not provided, it is inferred from the buffer shape as before. Operations documentation is expanded to make ops easier to use.
319 lines
12 KiB
C++
319 lines
12 KiB
C++
//===- AMXDialect.cpp - MLIR AMX ops implementation -----------------------===//
|
|
//
|
|
// 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 the AMX dialect and its operations.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/AMX/AMXDialect.h"
|
|
#include "mlir/Conversion/LLVMCommon/Pattern.h"
|
|
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
|
|
#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
|
|
#include "mlir/IR/Builders.h"
|
|
#include "mlir/IR/DialectImplementation.h"
|
|
#include "mlir/IR/OpImplementation.h"
|
|
#include "mlir/IR/TypeUtilities.h"
|
|
|
|
#include "llvm/ADT/TypeSwitch.h"
|
|
|
|
using namespace mlir;
|
|
|
|
#include "mlir/Dialect/AMX/AMXInterfaces.cpp.inc"
|
|
|
|
#include "mlir/Dialect/AMX/AMXDialect.cpp.inc"
|
|
|
|
void amx::AMXDialect::initialize() {
|
|
addTypes<
|
|
#define GET_TYPEDEF_LIST
|
|
#include "mlir/Dialect/AMX/AMXTypes.cpp.inc"
|
|
>();
|
|
|
|
addOperations<
|
|
#define GET_OP_LIST
|
|
#include "mlir/Dialect/AMX/AMX.cpp.inc"
|
|
>();
|
|
}
|
|
|
|
/// Verify that AMX supports the implied tile shape.
|
|
static LogicalResult verifyTileSize(Operation *op, amx::TileType tp) {
|
|
const unsigned kMaxRows = 16;
|
|
const unsigned kBitsPerRow = 64 * 8;
|
|
unsigned col = tp.getDimSize(1) * tp.getElementType().getIntOrFloatBitWidth();
|
|
if (tp.getDimSize(0) > kMaxRows)
|
|
return op->emitOpError("bad row height: ") << tp.getDimSize(0);
|
|
if (col > kBitsPerRow || col & 0x1f)
|
|
return op->emitOpError("bad column width: ") << (col >> 3);
|
|
return success();
|
|
}
|
|
|
|
/// Verify that AMX supports the multiplication.
|
|
static LogicalResult verifyMultShape(Operation *op, amx::TileType atp,
|
|
amx::TileType btp, amx::TileType ctp,
|
|
unsigned scale) {
|
|
unsigned am = atp.getDimSize(0), ak = atp.getDimSize(1) >> scale;
|
|
unsigned bk = btp.getDimSize(0), bn = btp.getDimSize(1) >> scale;
|
|
unsigned cm = ctp.getDimSize(0), cn = ctp.getDimSize(1);
|
|
if (cm != am || cn != bn || ak != bk)
|
|
return op->emitOpError("bad mult shape: ")
|
|
<< cm << " x " << cn << " x " << ak;
|
|
return success();
|
|
}
|
|
|
|
/// Maps the 2-dim vector shape to the two 16-bit tile sizes. The first
|
|
/// dimension directly translates into the number of rows of the tiles.
|
|
/// The second dimensions needs to be scaled by the number of bytes.
|
|
static SmallVector<Value> getTileSizes(Location loc, amx::TileType tType,
|
|
RewriterBase &rewriter) {
|
|
Type llvmInt16Type = rewriter.getIntegerType(16);
|
|
unsigned width = tType.getElementType().getIntOrFloatBitWidth();
|
|
assert(llvm::isPowerOf2_64(width) && width >= 8);
|
|
unsigned bytes = width >> 3;
|
|
auto mattr = rewriter.getI16IntegerAttr(tType.getDimSize(0));
|
|
auto nattr = rewriter.getI16IntegerAttr(tType.getDimSize(1) * bytes);
|
|
return SmallVector<Value>{
|
|
LLVM::ConstantOp::create(rewriter, loc, llvmInt16Type, mattr),
|
|
LLVM::ConstantOp::create(rewriter, loc, llvmInt16Type, nattr)};
|
|
}
|
|
|
|
/// Returns stride expressed in number of bytes for the given `elementStride`
|
|
/// stride encoded in number of elements of the type `mType`.
|
|
static Value computeStrideInBytes(Location loc, MemRefType mType,
|
|
Value elementStride, RewriterBase &rewriter) {
|
|
Type llvmInt64Type = rewriter.getIntegerType(64);
|
|
unsigned bytes = mType.getElementType().getIntOrFloatBitWidth() / 8;
|
|
auto attr = rewriter.getI64IntegerAttr(bytes);
|
|
Value scale = LLVM::ConstantOp::create(rewriter, loc, llvmInt64Type, attr);
|
|
return LLVM::MulOp::create(rewriter, loc, llvmInt64Type, scale, elementStride)
|
|
.getResult();
|
|
}
|
|
|
|
/// Maps the 2-dim memref shape to the 64-bit stride. Note that the buffer
|
|
/// shape may "envelop" the actual tile shape, and may be dynamically sized.
|
|
static Value inferStride(Location loc, MemRefType mType, Value base,
|
|
RewriterBase &rewriter) {
|
|
assert(mType.getRank() >= 2 && "Invalid shape for AMX strides");
|
|
int64_t preLast = mType.getRank() - 2;
|
|
Type llvmInt64Type = rewriter.getIntegerType(64);
|
|
unsigned width = mType.getElementType().getIntOrFloatBitWidth();
|
|
assert(llvm::isPowerOf2_64(width) && width >= 8);
|
|
unsigned bytes = width >> 3;
|
|
auto [strides, offset] = mType.getStridesAndOffset();
|
|
if (strides[preLast] == ShapedType::kDynamic) {
|
|
// Dynamic stride needs code to compute the stride at runtime.
|
|
MemRefDescriptor memrefDescriptor(base);
|
|
return computeStrideInBytes(
|
|
loc, mType, memrefDescriptor.stride(rewriter, loc, preLast), rewriter);
|
|
}
|
|
// Use direct constant for static stride.
|
|
auto attr = rewriter.getI64IntegerAttr(strides[preLast] * bytes);
|
|
return LLVM::ConstantOp::create(rewriter, loc, llvmInt64Type, attr)
|
|
.getResult();
|
|
}
|
|
|
|
LogicalResult amx::TileZeroOp::verify() {
|
|
return verifyTileSize(*this, getTileType());
|
|
}
|
|
|
|
SmallVector<Value>
|
|
amx::TileZeroOp::getIntrinsicOperands(ArrayRef<Value> operands,
|
|
const LLVMTypeConverter &typeConverter,
|
|
RewriterBase &rewriter) {
|
|
return getTileSizes(getLoc(), getTileType(), rewriter);
|
|
}
|
|
|
|
template <typename OpTy,
|
|
typename = std::enable_if_t<std::is_same_v<OpTy, amx::TileLoadOp> ||
|
|
std::is_same_v<OpTy, amx::TileStoreOp>>>
|
|
static LogicalResult tileTransferVerifier(OpTy op) {
|
|
MemRefType memrefTy = op.getMemRefType();
|
|
unsigned rank = memrefTy.getRank();
|
|
if (op.getIndices().size() != rank)
|
|
return op.emitOpError("requires ") << rank << " indices";
|
|
|
|
if (failed(verifyTileSize(op, op.getTileType())))
|
|
return failure();
|
|
|
|
// Validate basic buffer properties when the stride is implicit.
|
|
if (!op.getStride()) {
|
|
if (rank < 2)
|
|
return op.emitOpError("requires at least 2D memref");
|
|
SmallVector<int64_t> strides;
|
|
int64_t offset;
|
|
if (failed(memrefTy.getStridesAndOffset(strides, offset)) ||
|
|
strides.back() != 1)
|
|
return op.emitOpError("requires memref with unit innermost stride");
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
void amx::TileLoadOp::build(OpBuilder &builder, OperationState &state, Type res,
|
|
Value base, ValueRange indices) {
|
|
build(builder, state, res, base, indices, /*stride=*/nullptr);
|
|
}
|
|
|
|
LogicalResult amx::TileLoadOp::verify() { return tileTransferVerifier(*this); }
|
|
|
|
SmallVector<Value>
|
|
amx::TileLoadOp::getIntrinsicOperands(ArrayRef<Value> operands,
|
|
const LLVMTypeConverter &typeConverter,
|
|
RewriterBase &rewriter) {
|
|
auto loc = getLoc();
|
|
Adaptor adaptor(operands, *this);
|
|
|
|
SmallVector<Value> intrinsicOperands;
|
|
intrinsicOperands.append(getTileSizes(loc, getTileType(), rewriter));
|
|
intrinsicOperands.push_back(
|
|
LLVM::getStridedElementPtr(rewriter, loc, typeConverter, getMemRefType(),
|
|
adaptor.getBase(), adaptor.getIndices()));
|
|
if (Value stride = adaptor.getStride())
|
|
intrinsicOperands.push_back(
|
|
computeStrideInBytes(loc, getMemRefType(), stride, rewriter));
|
|
else
|
|
intrinsicOperands.push_back(
|
|
inferStride(loc, getMemRefType(), adaptor.getBase(), rewriter));
|
|
|
|
return intrinsicOperands;
|
|
}
|
|
|
|
void amx::TileStoreOp::build(OpBuilder &builder, OperationState &state,
|
|
Value base, ValueRange indices, Value val) {
|
|
build(builder, state, base, indices, val, /*stride=*/nullptr);
|
|
}
|
|
|
|
LogicalResult amx::TileStoreOp::verify() { return tileTransferVerifier(*this); }
|
|
|
|
SmallVector<Value>
|
|
amx::TileStoreOp::getIntrinsicOperands(ArrayRef<Value> operands,
|
|
const LLVMTypeConverter &typeConverter,
|
|
RewriterBase &rewriter) {
|
|
auto loc = getLoc();
|
|
Adaptor adaptor(operands, *this);
|
|
|
|
SmallVector<Value> intrinsicOperands;
|
|
intrinsicOperands.append(getTileSizes(loc, getTileType(), rewriter));
|
|
intrinsicOperands.push_back(
|
|
LLVM::getStridedElementPtr(rewriter, loc, typeConverter, getMemRefType(),
|
|
adaptor.getBase(), adaptor.getIndices()));
|
|
if (Value stride = adaptor.getStride())
|
|
intrinsicOperands.push_back(
|
|
computeStrideInBytes(loc, getMemRefType(), stride, rewriter));
|
|
else
|
|
intrinsicOperands.push_back(
|
|
inferStride(loc, getMemRefType(), adaptor.getBase(), rewriter));
|
|
intrinsicOperands.push_back(adaptor.getVal());
|
|
|
|
return intrinsicOperands;
|
|
}
|
|
|
|
LogicalResult amx::TileMulFOp::verify() {
|
|
amx::TileType aType = getLhsTileType();
|
|
amx::TileType bType = getRhsTileType();
|
|
amx::TileType cType = getTileType();
|
|
if (failed(verifyTileSize(*this, aType)) ||
|
|
failed(verifyTileSize(*this, bType)) ||
|
|
failed(verifyTileSize(*this, cType)) ||
|
|
failed(verifyMultShape(*this, aType, bType, cType, 1)))
|
|
return failure();
|
|
Type ta = aType.getElementType();
|
|
Type tb = bType.getElementType();
|
|
Type tc = cType.getElementType();
|
|
if ((!ta.isBF16() && !ta.isF16()) || (ta != tb) || !tc.isF32())
|
|
return emitOpError("unsupported type combination");
|
|
return success();
|
|
}
|
|
|
|
SmallVector<Value>
|
|
amx::TileMulFOp::getIntrinsicOperands(ArrayRef<Value> operands,
|
|
const LLVMTypeConverter &typeConverter,
|
|
RewriterBase &rewriter) {
|
|
auto loc = getLoc();
|
|
Adaptor adaptor(operands, *this);
|
|
|
|
amx::TileType aType = getLhsTileType();
|
|
amx::TileType bType = getRhsTileType();
|
|
SmallVector<Value> tsza = getTileSizes(loc, aType, rewriter);
|
|
SmallVector<Value> tszb = getTileSizes(loc, bType, rewriter);
|
|
|
|
SmallVector<Value> intrinsicOperands = {tsza[0], tszb[1],
|
|
tsza[1], adaptor.getAcc(),
|
|
adaptor.getLhs(), adaptor.getRhs()};
|
|
|
|
return intrinsicOperands;
|
|
}
|
|
|
|
LogicalResult amx::TileMulIOp::verify() {
|
|
amx::TileType aType = getLhsTileType();
|
|
amx::TileType bType = getRhsTileType();
|
|
amx::TileType cType = getTileType();
|
|
if (failed(verifyTileSize(*this, aType)) ||
|
|
failed(verifyTileSize(*this, bType)) ||
|
|
failed(verifyTileSize(*this, cType)) ||
|
|
failed(verifyMultShape(*this, aType, bType, cType, 2)))
|
|
return failure();
|
|
Type ta = aType.getElementType();
|
|
Type tb = bType.getElementType();
|
|
Type tc = cType.getElementType();
|
|
if (!ta.isInteger(8) || !tb.isInteger(8) || !tc.isInteger(32))
|
|
return emitOpError("unsupported type combination");
|
|
return success();
|
|
}
|
|
|
|
SmallVector<Value>
|
|
amx::TileMulIOp::getIntrinsicOperands(ArrayRef<Value> operands,
|
|
const LLVMTypeConverter &typeConverter,
|
|
RewriterBase &rewriter) {
|
|
auto loc = getLoc();
|
|
Adaptor adaptor(operands, *this);
|
|
|
|
amx::TileType aType = getLhsTileType();
|
|
amx::TileType bType = getRhsTileType();
|
|
SmallVector<Value> tsza = getTileSizes(loc, aType, rewriter);
|
|
SmallVector<Value> tszb = getTileSizes(loc, bType, rewriter);
|
|
|
|
SmallVector<Value> intrinsicOperands = {tsza[0], tszb[1],
|
|
tsza[1], adaptor.getAcc(),
|
|
adaptor.getLhs(), adaptor.getRhs()};
|
|
|
|
return intrinsicOperands;
|
|
}
|
|
|
|
Type amx::TileType::parse(AsmParser &parser) {
|
|
if (parser.parseLess())
|
|
return nullptr;
|
|
|
|
SmallVector<int64_t, 2> shape;
|
|
if (parser.parseDimensionList(shape, false, true))
|
|
return nullptr;
|
|
|
|
Type elementType;
|
|
if (parser.parseType(elementType))
|
|
return nullptr;
|
|
|
|
if (parser.parseGreater())
|
|
return nullptr;
|
|
|
|
return TileType::getChecked(
|
|
[&] { return parser.emitError(parser.getNameLoc()); }, shape,
|
|
elementType);
|
|
}
|
|
|
|
void amx::TileType::print(AsmPrinter &os) const {
|
|
os << "<";
|
|
os.printDimensionList(getShape());
|
|
os << 'x';
|
|
os.printType(getElementType());
|
|
os << '>';
|
|
}
|
|
|
|
#define GET_OP_CLASSES
|
|
#include "mlir/Dialect/AMX/AMX.cpp.inc"
|
|
|
|
#define GET_TYPEDEF_CLASSES
|
|
#include "mlir/Dialect/AMX/AMXTypes.cpp.inc"
|