319 lines · cpp
1//===- AMXDialect.cpp - MLIR AMX ops implementation -----------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements the AMX dialect and its operations.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/AMX/AMXDialect.h"14#include "mlir/Conversion/LLVMCommon/Pattern.h"15#include "mlir/Dialect/LLVMIR/LLVMDialect.h"16#include "mlir/Dialect/LLVMIR/LLVMTypes.h"17#include "mlir/IR/Builders.h"18#include "mlir/IR/DialectImplementation.h"19#include "mlir/IR/OpImplementation.h"20#include "mlir/IR/TypeUtilities.h"21 22#include "llvm/ADT/TypeSwitch.h"23 24using namespace mlir;25 26#include "mlir/Dialect/AMX/AMXInterfaces.cpp.inc"27 28#include "mlir/Dialect/AMX/AMXDialect.cpp.inc"29 30void amx::AMXDialect::initialize() {31 addTypes<32#define GET_TYPEDEF_LIST33#include "mlir/Dialect/AMX/AMXTypes.cpp.inc"34 >();35 36 addOperations<37#define GET_OP_LIST38#include "mlir/Dialect/AMX/AMX.cpp.inc"39 >();40}41 42/// Verify that AMX supports the implied tile shape.43static LogicalResult verifyTileSize(Operation *op, amx::TileType tp) {44 const unsigned kMaxRows = 16;45 const unsigned kBitsPerRow = 64 * 8;46 unsigned col = tp.getDimSize(1) * tp.getElementType().getIntOrFloatBitWidth();47 if (tp.getDimSize(0) > kMaxRows)48 return op->emitOpError("bad row height: ") << tp.getDimSize(0);49 if (col > kBitsPerRow || col & 0x1f)50 return op->emitOpError("bad column width: ") << (col >> 3);51 return success();52}53 54/// Verify that AMX supports the multiplication.55static LogicalResult verifyMultShape(Operation *op, amx::TileType atp,56 amx::TileType btp, amx::TileType ctp,57 unsigned scale) {58 unsigned am = atp.getDimSize(0), ak = atp.getDimSize(1) >> scale;59 unsigned bk = btp.getDimSize(0), bn = btp.getDimSize(1) >> scale;60 unsigned cm = ctp.getDimSize(0), cn = ctp.getDimSize(1);61 if (cm != am || cn != bn || ak != bk)62 return op->emitOpError("bad mult shape: ")63 << cm << " x " << cn << " x " << ak;64 return success();65}66 67/// Maps the 2-dim vector shape to the two 16-bit tile sizes. The first68/// dimension directly translates into the number of rows of the tiles.69/// The second dimensions needs to be scaled by the number of bytes.70static SmallVector<Value> getTileSizes(Location loc, amx::TileType tType,71 RewriterBase &rewriter) {72 Type llvmInt16Type = rewriter.getIntegerType(16);73 unsigned width = tType.getElementType().getIntOrFloatBitWidth();74 assert(llvm::isPowerOf2_64(width) && width >= 8);75 unsigned bytes = width >> 3;76 auto mattr = rewriter.getI16IntegerAttr(tType.getDimSize(0));77 auto nattr = rewriter.getI16IntegerAttr(tType.getDimSize(1) * bytes);78 return SmallVector<Value>{79 LLVM::ConstantOp::create(rewriter, loc, llvmInt16Type, mattr),80 LLVM::ConstantOp::create(rewriter, loc, llvmInt16Type, nattr)};81}82 83/// Returns stride expressed in number of bytes for the given `elementStride`84/// stride encoded in number of elements of the type `mType`.85static Value computeStrideInBytes(Location loc, MemRefType mType,86 Value elementStride, RewriterBase &rewriter) {87 Type llvmInt64Type = rewriter.getIntegerType(64);88 unsigned bytes = mType.getElementType().getIntOrFloatBitWidth() / 8;89 auto attr = rewriter.getI64IntegerAttr(bytes);90 Value scale = LLVM::ConstantOp::create(rewriter, loc, llvmInt64Type, attr);91 return LLVM::MulOp::create(rewriter, loc, llvmInt64Type, scale, elementStride)92 .getResult();93}94 95/// Maps the 2-dim memref shape to the 64-bit stride. Note that the buffer96/// shape may "envelop" the actual tile shape, and may be dynamically sized.97static Value inferStride(Location loc, MemRefType mType, Value base,98 RewriterBase &rewriter) {99 assert(mType.getRank() >= 2 && "Invalid shape for AMX strides");100 int64_t preLast = mType.getRank() - 2;101 Type llvmInt64Type = rewriter.getIntegerType(64);102 unsigned width = mType.getElementType().getIntOrFloatBitWidth();103 assert(llvm::isPowerOf2_64(width) && width >= 8);104 unsigned bytes = width >> 3;105 auto [strides, offset] = mType.getStridesAndOffset();106 if (strides[preLast] == ShapedType::kDynamic) {107 // Dynamic stride needs code to compute the stride at runtime.108 MemRefDescriptor memrefDescriptor(base);109 return computeStrideInBytes(110 loc, mType, memrefDescriptor.stride(rewriter, loc, preLast), rewriter);111 }112 // Use direct constant for static stride.113 auto attr = rewriter.getI64IntegerAttr(strides[preLast] * bytes);114 return LLVM::ConstantOp::create(rewriter, loc, llvmInt64Type, attr)115 .getResult();116}117 118LogicalResult amx::TileZeroOp::verify() {119 return verifyTileSize(*this, getTileType());120}121 122SmallVector<Value>123amx::TileZeroOp::getIntrinsicOperands(ArrayRef<Value> operands,124 const LLVMTypeConverter &typeConverter,125 RewriterBase &rewriter) {126 return getTileSizes(getLoc(), getTileType(), rewriter);127}128 129template <typename OpTy,130 typename = std::enable_if_t<std::is_same_v<OpTy, amx::TileLoadOp> ||131 std::is_same_v<OpTy, amx::TileStoreOp>>>132static LogicalResult tileTransferVerifier(OpTy op) {133 MemRefType memrefTy = op.getMemRefType();134 unsigned rank = memrefTy.getRank();135 if (op.getIndices().size() != rank)136 return op.emitOpError("requires ") << rank << " indices";137 138 if (failed(verifyTileSize(op, op.getTileType())))139 return failure();140 141 // Validate basic buffer properties when the stride is implicit.142 if (!op.getStride()) {143 if (rank < 2)144 return op.emitOpError("requires at least 2D memref");145 SmallVector<int64_t> strides;146 int64_t offset;147 if (failed(memrefTy.getStridesAndOffset(strides, offset)) ||148 strides.back() != 1)149 return op.emitOpError("requires memref with unit innermost stride");150 }151 152 return success();153}154 155void amx::TileLoadOp::build(OpBuilder &builder, OperationState &state, Type res,156 Value base, ValueRange indices) {157 build(builder, state, res, base, indices, /*stride=*/nullptr);158}159 160LogicalResult amx::TileLoadOp::verify() { return tileTransferVerifier(*this); }161 162SmallVector<Value>163amx::TileLoadOp::getIntrinsicOperands(ArrayRef<Value> operands,164 const LLVMTypeConverter &typeConverter,165 RewriterBase &rewriter) {166 auto loc = getLoc();167 Adaptor adaptor(operands, *this);168 169 SmallVector<Value> intrinsicOperands;170 intrinsicOperands.append(getTileSizes(loc, getTileType(), rewriter));171 intrinsicOperands.push_back(172 LLVM::getStridedElementPtr(rewriter, loc, typeConverter, getMemRefType(),173 adaptor.getBase(), adaptor.getIndices()));174 if (Value stride = adaptor.getStride())175 intrinsicOperands.push_back(176 computeStrideInBytes(loc, getMemRefType(), stride, rewriter));177 else178 intrinsicOperands.push_back(179 inferStride(loc, getMemRefType(), adaptor.getBase(), rewriter));180 181 return intrinsicOperands;182}183 184void amx::TileStoreOp::build(OpBuilder &builder, OperationState &state,185 Value base, ValueRange indices, Value val) {186 build(builder, state, base, indices, val, /*stride=*/nullptr);187}188 189LogicalResult amx::TileStoreOp::verify() { return tileTransferVerifier(*this); }190 191SmallVector<Value>192amx::TileStoreOp::getIntrinsicOperands(ArrayRef<Value> operands,193 const LLVMTypeConverter &typeConverter,194 RewriterBase &rewriter) {195 auto loc = getLoc();196 Adaptor adaptor(operands, *this);197 198 SmallVector<Value> intrinsicOperands;199 intrinsicOperands.append(getTileSizes(loc, getTileType(), rewriter));200 intrinsicOperands.push_back(201 LLVM::getStridedElementPtr(rewriter, loc, typeConverter, getMemRefType(),202 adaptor.getBase(), adaptor.getIndices()));203 if (Value stride = adaptor.getStride())204 intrinsicOperands.push_back(205 computeStrideInBytes(loc, getMemRefType(), stride, rewriter));206 else207 intrinsicOperands.push_back(208 inferStride(loc, getMemRefType(), adaptor.getBase(), rewriter));209 intrinsicOperands.push_back(adaptor.getVal());210 211 return intrinsicOperands;212}213 214LogicalResult amx::TileMulFOp::verify() {215 amx::TileType aType = getLhsTileType();216 amx::TileType bType = getRhsTileType();217 amx::TileType cType = getTileType();218 if (failed(verifyTileSize(*this, aType)) ||219 failed(verifyTileSize(*this, bType)) ||220 failed(verifyTileSize(*this, cType)) ||221 failed(verifyMultShape(*this, aType, bType, cType, 1)))222 return failure();223 Type ta = aType.getElementType();224 Type tb = bType.getElementType();225 Type tc = cType.getElementType();226 if ((!ta.isBF16() && !ta.isF16()) || (ta != tb) || !tc.isF32())227 return emitOpError("unsupported type combination");228 return success();229}230 231SmallVector<Value>232amx::TileMulFOp::getIntrinsicOperands(ArrayRef<Value> operands,233 const LLVMTypeConverter &typeConverter,234 RewriterBase &rewriter) {235 auto loc = getLoc();236 Adaptor adaptor(operands, *this);237 238 amx::TileType aType = getLhsTileType();239 amx::TileType bType = getRhsTileType();240 SmallVector<Value> tsza = getTileSizes(loc, aType, rewriter);241 SmallVector<Value> tszb = getTileSizes(loc, bType, rewriter);242 243 SmallVector<Value> intrinsicOperands = {tsza[0], tszb[1],244 tsza[1], adaptor.getAcc(),245 adaptor.getLhs(), adaptor.getRhs()};246 247 return intrinsicOperands;248}249 250LogicalResult amx::TileMulIOp::verify() {251 amx::TileType aType = getLhsTileType();252 amx::TileType bType = getRhsTileType();253 amx::TileType cType = getTileType();254 if (failed(verifyTileSize(*this, aType)) ||255 failed(verifyTileSize(*this, bType)) ||256 failed(verifyTileSize(*this, cType)) ||257 failed(verifyMultShape(*this, aType, bType, cType, 2)))258 return failure();259 Type ta = aType.getElementType();260 Type tb = bType.getElementType();261 Type tc = cType.getElementType();262 if (!ta.isInteger(8) || !tb.isInteger(8) || !tc.isInteger(32))263 return emitOpError("unsupported type combination");264 return success();265}266 267SmallVector<Value>268amx::TileMulIOp::getIntrinsicOperands(ArrayRef<Value> operands,269 const LLVMTypeConverter &typeConverter,270 RewriterBase &rewriter) {271 auto loc = getLoc();272 Adaptor adaptor(operands, *this);273 274 amx::TileType aType = getLhsTileType();275 amx::TileType bType = getRhsTileType();276 SmallVector<Value> tsza = getTileSizes(loc, aType, rewriter);277 SmallVector<Value> tszb = getTileSizes(loc, bType, rewriter);278 279 SmallVector<Value> intrinsicOperands = {tsza[0], tszb[1],280 tsza[1], adaptor.getAcc(),281 adaptor.getLhs(), adaptor.getRhs()};282 283 return intrinsicOperands;284}285 286Type amx::TileType::parse(AsmParser &parser) {287 if (parser.parseLess())288 return nullptr;289 290 SmallVector<int64_t, 2> shape;291 if (parser.parseDimensionList(shape, false, true))292 return nullptr;293 294 Type elementType;295 if (parser.parseType(elementType))296 return nullptr;297 298 if (parser.parseGreater())299 return nullptr;300 301 return TileType::getChecked(302 [&] { return parser.emitError(parser.getNameLoc()); }, shape,303 elementType);304}305 306void amx::TileType::print(AsmPrinter &os) const {307 os << "<";308 os.printDimensionList(getShape());309 os << 'x';310 os.printType(getElementType());311 os << '>';312}313 314#define GET_OP_CLASSES315#include "mlir/Dialect/AMX/AMX.cpp.inc"316 317#define GET_TYPEDEF_CLASSES318#include "mlir/Dialect/AMX/AMXTypes.cpp.inc"319