889 lines · cpp
1//===- AMDGPUDialect.cpp - MLIR AMDGPU dialect 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 AMDGPU dialect and its operations.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"14 15#include "mlir/Dialect/Arith/IR/Arith.h"16#include "mlir/Dialect/GPU/IR/GPUDialect.h"17#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"18#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"19#include "mlir/Dialect/Utils/IndexingUtils.h"20#include "mlir/Dialect/Vector/IR/VectorOps.h"21#include "mlir/IR/Builders.h"22#include "mlir/IR/BuiltinTypes.h"23#include "mlir/IR/Diagnostics.h"24#include "mlir/IR/DialectImplementation.h"25#include "mlir/IR/Matchers.h"26#include "mlir/IR/OpImplementation.h"27#include "mlir/IR/PatternMatch.h"28#include "mlir/IR/TypeUtilities.h"29#include "mlir/Transforms/InliningUtils.h"30#include "llvm/ADT/DenseMap.h"31#include "llvm/ADT/SmallVector.h"32#include "llvm/ADT/TypeSwitch.h"33 34#include <algorithm>35#include <cstdint>36#include <limits>37#include <optional>38 39using namespace mlir;40using namespace mlir::amdgpu;41 42#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.cpp.inc"43 44namespace {45struct AMDGPUInlinerInterface final : DialectInlinerInterface {46 using DialectInlinerInterface::DialectInlinerInterface;47 bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {48 return true;49 }50};51} // namespace52 53void AMDGPUDialect::initialize() {54 addOperations<55#define GET_OP_LIST56#include "mlir/Dialect/AMDGPU/IR/AMDGPU.cpp.inc"57 >();58 addTypes<59#define GET_TYPEDEF_LIST60#include "mlir/Dialect/AMDGPU/IR/AMDGPUTypes.cpp.inc"61 >();62 addAttributes<63#define GET_ATTRDEF_LIST64#include "mlir/Dialect/AMDGPU/IR/AMDGPUAttributes.cpp.inc"65 >();66 addInterfaces<AMDGPUInlinerInterface>();67}68 69//===----------------------------------------------------------------------===//70// 8-bit float ops71//===----------------------------------------------------------------------===//72LogicalResult PackedTrunc2xFp8Op::verify() {73 if (getExisting() && getExisting().getType() != getResult().getType())74 return emitOpError("existing values must have same type as result");75 return success();76}77 78LogicalResult PackedStochRoundFp8Op::verify() {79 if (getExisting() && getExisting().getType() != getResult().getType())80 return emitOpError("existing values must have same type as result");81 return success();82}83 84//===----------------------------------------------------------------------===//85// mxfp float ops86//===----------------------------------------------------------------------===//87LogicalResult PackedScaledTruncOp::verify() {88 if (getExisting() && getExisting().getType() != getResult().getType())89 return emitOpError("existing values must have same type as result");90 return success();91}92 93//===----------------------------------------------------------------------===//94// FatRawBufferCastOp95//===----------------------------------------------------------------------===//96 97/// Convert the type `source` to one with the same sizes and strides - and98/// offset, unless `stripOffset` is true, in which case the offset is reset to99/// 0, if the offset should be reset but the layout of `source` isn't either the100/// identity layout or a strided layout, this function fails.101static FailureOr<MemRefType> getFatRawBufferTypeLike(MemRefType source,102 bool resetOffset) {103 MLIRContext *ctx = source.getContext();104 MemRefType::Builder mb(source);105 mb.setMemorySpace(106 amdgpu::AddressSpaceAttr::get(ctx, amdgpu::AddressSpace::FatRawBuffer));107 MemRefLayoutAttrInterface layout = source.getLayout();108 if (resetOffset && !layout.isIdentity()) {109 auto stridedLayout = dyn_cast<StridedLayoutAttr>(layout);110 if (!stridedLayout)111 return failure();112 MemRefLayoutAttrInterface newLayout =113 StridedLayoutAttr::get(ctx, 0, stridedLayout.getStrides());114 // Special case: if resetting the offset causes the strided layout to become115 // the identity layout, then reset to the identity layout.116 // TODO: this'll get a lot simpler when we have the contiguous layout.117 SmallVector<int64_t> stridesIfIdentity;118 if (source.hasStaticShape()) {119 stridesIfIdentity = computeSuffixProduct(source.getShape());120 } else if (source.getRank() <= 1) {121 stridesIfIdentity = SmallVector<int64_t>(source.getRank(), 1);122 }123 if (stridesIfIdentity == stridedLayout.getStrides()) {124 newLayout = AffineMapAttr::get(125 AffineMap::getMultiDimIdentityMap(source.getRank(), ctx));126 }127 mb.setLayout(newLayout);128 }129 return (MemRefType)(mb);130}131 132LogicalResult FatRawBufferCastOp::inferReturnTypes(133 MLIRContext *context, std::optional<Location> location, ValueRange operands,134 DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions,135 SmallVectorImpl<Type> &inferredReturnTypes) {136 Adaptor adaptor(operands, attributes, properties, regions);137 auto sourceType =138 dyn_cast_if_present<MemRefType>(adaptor.getSource().getType());139 if (!sourceType)140 return failure();141 FailureOr<MemRefType> resultType =142 getFatRawBufferTypeLike(sourceType, adaptor.getResetOffset());143 if (failed(resultType))144 return failure();145 inferredReturnTypes = SmallVector<Type>{*resultType};146 return success();147}148 149LogicalResult FatRawBufferCastOp::verify() {150 FailureOr<MemRefType> expectedResultType =151 getFatRawBufferTypeLike(getSource().getType(), getResetOffset());152 if (failed(expectedResultType))153 return emitOpError("source type ")154 << getSource().getType() << " can't have its offset reset";155 if (getResult().getType() != *expectedResultType)156 return emitOpError("expected result type to be ")157 << *expectedResultType << " but got " << getResult().getType();158 return success();159}160 161static bool hasGlobalMemorySpace(Attribute memorySpace) {162 if (!memorySpace)163 return true;164 if (auto intMemorySpace = dyn_cast<IntegerAttr>(memorySpace))165 return intMemorySpace.getInt() == 0 || intMemorySpace.getInt() == 1;166 if (auto gpuMemorySpace = dyn_cast<gpu::AddressSpaceAttr>(memorySpace))167 return gpuMemorySpace.getValue() == gpu::AddressSpace::Global;168 return false;169}170 171static bool hasWorkgroupMemorySpace(Attribute memorySpace) {172 if (!memorySpace)173 return false;174 if (auto intMemorySpace = dyn_cast<IntegerAttr>(memorySpace))175 return intMemorySpace.getInt() == 3;176 if (auto gpuMemorySpace = dyn_cast<gpu::AddressSpaceAttr>(memorySpace))177 return gpuMemorySpace.getValue() == gpu::AddressSpace::Workgroup;178 return false;179}180 181static bool hasFatRawBufferMemorySpace(Attribute memorySpace) {182 if (!memorySpace)183 return false;184 if (auto intMemorySpace = dyn_cast<IntegerAttr>(memorySpace))185 return intMemorySpace.getInt() == 7;186 if (auto gpuMemorySpace = dyn_cast<amdgpu::AddressSpaceAttr>(memorySpace))187 return gpuMemorySpace.getValue() == amdgpu::AddressSpace::FatRawBuffer;188 return false;189}190 191//===----------------------------------------------------------------------===//192// RawBuffer*Op193//===----------------------------------------------------------------------===//194template <typename T>195static LogicalResult verifyRawBufferOp(T &op) {196 MemRefType bufferType = llvm::cast<MemRefType>(op.getMemref().getType());197 bool isGlobal = hasGlobalMemorySpace(bufferType.getMemorySpace());198 199 if (!isGlobal)200 return op.emitOpError(201 "Buffer ops must operate on a memref in global memory");202 if (!bufferType.hasRank())203 return op.emitOpError(204 "Cannot meaningfully buffer_store to an unranked memref");205 if (static_cast<int64_t>(op.getIndices().size()) != bufferType.getRank())206 return op.emitOpError("Expected " + Twine(bufferType.getRank()) +207 " indices to memref");208 return success();209}210 211LogicalResult RawBufferLoadOp::verify() { return verifyRawBufferOp(*this); }212 213LogicalResult RawBufferStoreOp::verify() { return verifyRawBufferOp(*this); }214 215LogicalResult RawBufferAtomicFaddOp::verify() {216 return verifyRawBufferOp(*this);217}218 219LogicalResult RawBufferAtomicFmaxOp::verify() {220 return verifyRawBufferOp(*this);221}222 223LogicalResult RawBufferAtomicSmaxOp::verify() {224 return verifyRawBufferOp(*this);225}226 227LogicalResult RawBufferAtomicUminOp::verify() {228 return verifyRawBufferOp(*this);229}230 231LogicalResult RawBufferAtomicCmpswapOp::verify() {232 return verifyRawBufferOp(*this);233}234 235static std::optional<uint32_t> getConstantUint32(Value v) {236 APInt cst;237 if (!v.getType().isInteger(32))238 return std::nullopt;239 if (matchPattern(v, m_ConstantInt(&cst)))240 return cst.getZExtValue();241 return std::nullopt;242}243 244template <typename OpType>245static bool staticallyOutOfBounds(OpType op) {246 if (!op.getBoundsCheck())247 return false;248 MemRefType bufferType = op.getMemref().getType();249 if (!bufferType.hasStaticShape())250 return false;251 int64_t offset;252 SmallVector<int64_t> strides;253 if (failed(bufferType.getStridesAndOffset(strides, offset)))254 return false;255 int64_t result = offset + op.getIndexOffset().value_or(0);256 if (op.getSgprOffset()) {257 std::optional<uint32_t> sgprOffset = getConstantUint32(op.getSgprOffset());258 if (!sgprOffset)259 return false;260 result += *sgprOffset;261 }262 if (strides.size() != op.getIndices().size())263 return false;264 int64_t indexVal = 0;265 for (auto pair : llvm::zip(strides, op.getIndices())) {266 int64_t stride = std::get<0>(pair);267 Value idx = std::get<1>(pair);268 std::optional<uint32_t> idxVal = getConstantUint32(idx);269 if (!idxVal)270 return false;271 indexVal += stride * *idxVal;272 }273 result += indexVal;274 if (result > std::numeric_limits<uint32_t>::max())275 // Overflow means don't drop276 return false;277 return result >= bufferType.getNumElements();278}279 280namespace {281template <typename OpType>282struct RemoveStaticallyOobBufferLoads final : public OpRewritePattern<OpType> {283 using OpRewritePattern<OpType>::OpRewritePattern;284 285 LogicalResult matchAndRewrite(OpType op, PatternRewriter &rw) const override {286 if (!staticallyOutOfBounds(op))287 return failure();288 Type loadType = op.getResult().getType();289 rw.replaceOpWithNewOp<arith::ConstantOp>(op, loadType,290 rw.getZeroAttr(loadType));291 return success();292 }293};294 295template <typename OpType>296struct RemoveStaticallyOobBufferWrites final : public OpRewritePattern<OpType> {297 using OpRewritePattern<OpType>::OpRewritePattern;298 299 LogicalResult matchAndRewrite(OpType op, PatternRewriter &rw) const override {300 if (!staticallyOutOfBounds(op))301 return failure();302 303 rw.eraseOp(op);304 return success();305 }306};307} // end namespace308 309void RawBufferLoadOp::getCanonicalizationPatterns(RewritePatternSet &results,310 MLIRContext *context) {311 results.add<RemoveStaticallyOobBufferLoads<RawBufferLoadOp>>(context);312}313 314void RawBufferStoreOp::getCanonicalizationPatterns(RewritePatternSet &results,315 MLIRContext *context) {316 results.add<RemoveStaticallyOobBufferWrites<RawBufferStoreOp>>(context);317}318 319void RawBufferAtomicFaddOp::getCanonicalizationPatterns(320 RewritePatternSet &results, MLIRContext *context) {321 results.add<RemoveStaticallyOobBufferWrites<RawBufferAtomicFaddOp>>(context);322}323 324void RawBufferAtomicFmaxOp::getCanonicalizationPatterns(325 RewritePatternSet &results, MLIRContext *context) {326 results.add<RemoveStaticallyOobBufferWrites<RawBufferAtomicFmaxOp>>(context);327}328 329void RawBufferAtomicSmaxOp::getCanonicalizationPatterns(330 RewritePatternSet &results, MLIRContext *context) {331 results.add<RemoveStaticallyOobBufferWrites<RawBufferAtomicSmaxOp>>(context);332}333 334void RawBufferAtomicUminOp::getCanonicalizationPatterns(335 RewritePatternSet &results, MLIRContext *context) {336 results.add<RemoveStaticallyOobBufferWrites<RawBufferAtomicUminOp>>(context);337}338 339void RawBufferAtomicCmpswapOp::getCanonicalizationPatterns(340 RewritePatternSet &results, MLIRContext *context) {341 results.add<RemoveStaticallyOobBufferLoads<RawBufferAtomicCmpswapOp>>(342 context);343}344 345//===----------------------------------------------------------------------===//346// ScaledExtPacked816Op347//===----------------------------------------------------------------------===//348LogicalResult ScaledExtPacked816Op::verify() {349 int blockSize = getBlockSize();350 assert(llvm::is_contained({16, 32}, blockSize) && "invalid block size");351 352 int firstScaleByte = getFirstScaleByte();353 int firstScaleLane = getFirstScaleLane();354 auto sourceType = cast<VectorType>(getSource().getType());355 Type elementType = sourceType.getElementType();356 auto floatType = cast<FloatType>(elementType);357 unsigned bitWidth = floatType.getWidth();358 359 assert(llvm::is_contained(llvm::ArrayRef<unsigned>{4, 6, 8}, bitWidth));360 361 const bool is_fp8 = bitWidth == 8;362 const bool is_block_16 = blockSize == 16;363 364 if (!is_fp8) {365 if (is_block_16) {366 if (!llvm::is_contained({0, 1}, firstScaleByte)) {367 return emitOpError("blockSize of 16 can only have firstScaleByte be 0 "368 "or 1 for f4 and f6.");369 }370 } else {371 if (!llvm::is_contained({0, 2}, firstScaleByte)) {372 return emitOpError("blockSize of 32 can only have firstScaleByte be 0 "373 "or 2 for f4 and f6.");374 }375 }376 } else {377 if (is_block_16) {378 bool is_valid = ((firstScaleLane == 0) && (firstScaleByte == 0)) ||379 ((firstScaleLane == 1) && (firstScaleByte == 2));380 if (!is_valid) {381 return emitOpError("blockSize of 16 can only have (firstScaleLane, "382 "firstScaleByte) be (0, 0) or (1, 2) for f8.");383 }384 }385 }386 387 return success();388}389 390//===----------------------------------------------------------------------===//391// WMMAOp392//===----------------------------------------------------------------------===//393 394ParseResult mlir::amdgpu::parseMNKDimensionList(OpAsmParser &parser,395 IntegerAttr &m, IntegerAttr &n,396 IntegerAttr &k) {397 SmallVector<int64_t, 3> dimensions;398 if (parser.parseDimensionList(dimensions, false, false))399 return failure();400 if (dimensions.size() != 3)401 return parser.emitError(parser.getCurrentLocation())402 << "expected 3 dimensions in MNK dimension list";403 404 m = parser.getBuilder().getI32IntegerAttr(dimensions[0]);405 n = parser.getBuilder().getI32IntegerAttr(dimensions[1]);406 k = parser.getBuilder().getI32IntegerAttr(dimensions[2]);407 return success();408}409 410LogicalResult WMMAOp::verify() {411 auto sourceAType = cast<VectorType>(getSourceA().getType());412 auto sourceBType = cast<VectorType>(getSourceB().getType());413 auto destType = cast<VectorType>(getDestC().getType());414 415 Type sourceAElemType = sourceAType.getElementType();416 Type sourceBElemType = sourceBType.getElementType();417 if (sourceAType.getNumElements() != sourceBType.getNumElements()) {418 return emitOpError("source vectors have different lengths: ")419 << sourceAType << " vs. " << sourceBType;420 }421 422 bool isDestFloat = destType.getElementType().isFloat();423 bool isSrcFloat = sourceAElemType.isFloat();424 425 if (isDestFloat && !isSrcFloat)426 return emitOpError("expected float sources with float destination");427 if (!isDestFloat && isSrcFloat)428 return emitOpError("expected int sources with int destination");429 430 if (!sourceAElemType.isFloat(8) && sourceAElemType != sourceBElemType) {431 return emitOpError(432 "source element types must match (except for fp8/bf8) but have ")433 << sourceAType << " and " << sourceBType;434 }435 436 if (isSrcFloat) {437 if (getClamp())438 return emitOpError("clamp flag is not supported for float types");439 if (getUnsignedA() || getUnsignedB())440 return emitOpError("unsigned flags are not supported for float types");441 }442 return success();443}444 445//===----------------------------------------------------------------------===//446// MFMAOp447//===----------------------------------------------------------------------===//448LogicalResult MFMAOp::verify() {449 constexpr uint32_t waveSize = 64;450 Builder b(getContext());451 452 Type sourceType = getSourceA().getType();453 Type destType = getDestC().getType();454 455 Type sourceElem = sourceType, destElem = destType;456 uint32_t sourceLen = 1, destLen = 1;457 if (auto sourceVector = dyn_cast<VectorType>(sourceType)) {458 sourceLen = sourceVector.getNumElements();459 sourceElem = sourceVector.getElementType();460 }461 if (auto destVector = dyn_cast<VectorType>(destType)) {462 destLen = destVector.getNumElements();463 destElem = destVector.getElementType();464 }465 466 Type sourceBType = getSourceB().getType();467 if (sourceElem.isFloat(8) || sourceElem.isFloat(6) || sourceElem.isFloat(4)) {468 int64_t sourceBLen = 1;469 Type sourceBElem = sourceBType;470 if (auto sourceBVector = llvm::dyn_cast<VectorType>(sourceBType)) {471 sourceBLen = sourceBVector.getNumElements();472 sourceBElem = sourceBVector.getElementType();473 }474 if (!sourceBElem.isFloat(8) && !sourceBElem.isFloat(6) &&475 !sourceBElem.isFloat(4))476 return emitOpError("expected both source operands to have small-float "477 "elements if one does");478 if (sourceLen != sourceBLen)479 return emitOpError(480 "expected both small-float source vectors to have the same length");481 } else {482 if (sourceType != sourceBType)483 return emitOpError("expected both non-small-float source operand types "484 "to match exactly");485 }486 // Normalize the wider integer types the compiler expects to i8.487 if (sourceElem.isInteger(32)) {488 sourceLen *= 4;489 sourceElem = b.getI8Type();490 }491 if (sourceElem.isInteger(64)) {492 sourceLen *= 8;493 sourceElem = b.getI8Type();494 }495 496 int64_t numSourceElems = (getM() * getK() * getBlocks()) / waveSize;497 if (sourceLen != numSourceElems)498 return emitOpError("expected " + Twine(numSourceElems) +499 " source values for this operation but got " +500 Twine(sourceLen));501 502 int64_t numDestElems = (getM() * getN() * getBlocks()) / waveSize;503 if (destLen != numDestElems)504 return emitOpError("expected " + Twine(numDestElems) +505 " result values for this operation but got " +506 Twine(destLen));507 508 if (destElem.isF64() && getBlgp() != MFMAPermB::none)509 return emitOpError(510 "double-precision ops do not support permuting lanes of B");511 if (destElem.isF64() && getCbsz() != 0)512 return emitOpError(513 "double-precision ops do not support permuting lanes of A");514 if (getAbid() >= (1u << getCbsz()))515 return emitOpError(516 "block ID for permuting A (abid) must be below 2 ** cbsz");517 518 if ((getNegateA() || getNegateB() || getNegateC()) && !destElem.isF64())519 return emitOpError(520 "negation flags only available for double-precision operations");521 522 return success();523}524 525//===----------------------------------------------------------------------===//526// DPPOp527//===----------------------------------------------------------------------===//528LogicalResult DPPOp::verify() {529 Type srcType = getSrc().getType();530 if (srcType.getIntOrFloatBitWidth() > 64) {531 return emitOpError("integer and floating point types larger than 64 bits "532 "are not supported");533 }534 535 DPPPerm kind = getKind();536 Attribute permArgument = getPermArgument().value_or(Attribute{});537 538 switch (kind) {539 540 case DPPPerm::quad_perm: {541 auto quadPermAttr = dyn_cast_or_null<ArrayAttr>(permArgument);542 if (!quadPermAttr || quadPermAttr.size() != 4) {543 return emitOpError("quad_perm attribute must have exactly 4 elements");544 }545 for (auto elem : quadPermAttr.getAsRange<IntegerAttr>()) {546 int32_t num = elem.getInt();547 if (num < 0 || num > 3) {548 return emitOpError(549 "Each element of quad_perm must be in the range [0, 3]");550 }551 }552 } break;553 554 case DPPPerm::row_shl:555 case DPPPerm::row_shr:556 case DPPPerm::row_ror: {557 if (!permArgument) {558 return emitOpError("Attribute '" + Twine(stringifyDPPPerm(kind)) +559 "' value not specified");560 }561 if (auto intAttr = dyn_cast<IntegerAttr>(permArgument)) {562 uint32_t attrValue = intAttr.getInt();563 if (attrValue < 1 || attrValue > 15) {564 return emitOpError("Attribute value must be between 1 and 15");565 }566 }567 } break;568 569 case DPPPerm::wave_shl:570 case DPPPerm::wave_shr:571 case DPPPerm::wave_rol:572 case DPPPerm::wave_ror:573 case DPPPerm::row_mirror:574 case DPPPerm::row_half_mirror:575 case DPPPerm::row_bcast_15:576 case DPPPerm::row_bcast_31: {577 if (permArgument && !isa<UnitAttr>(permArgument)) {578 return emitOpError("Expected unit attribute for permArgument, but found "579 "non-trivial argument");580 }581 break;582 }583 }584 return success();585}586 587//===----------------------------------------------------------------------===//588// PermlaneSwapOp589//===----------------------------------------------------------------------===//590LogicalResult PermlaneSwapOp::verify() {591 unsigned rowLength = getRowLength();592 593 if (rowLength != 16 && rowLength != 32)594 return emitOpError("row_length attribute must either be 16 or 32.");595 596 return success();597}598 599//===----------------------------------------------------------------------===//600// GatherToLDSOp601//===----------------------------------------------------------------------===//602 603LogicalResult GatherToLDSOp::verify() {604 MemRefType srcType = cast<MemRefType>(getSrc().getType());605 MemRefType dstType = cast<MemRefType>(getDst().getType());606 607 if (!dstType.areTrailingDimsContiguous(1))608 return emitOpError("destination type inner most dim must be contiguous");609 610 auto elemType = srcType.getElementType();611 // Check $src and $dst element types are the same.612 if (elemType != dstType.getElementType())613 return emitOpError("source and destination element types must match");614 615 // copy type sizes should be 1, 2, 4, 12 or 16 bytes.616 auto transferType = getTransferType();617 int transferSize;618 if (auto vectorTransfer = dyn_cast<VectorType>(transferType)) {619 transferSize = vectorTransfer.getNumElements() *620 vectorTransfer.getElementTypeBitWidth();621 } else {622 transferSize = transferType.getIntOrFloatBitWidth();623 }624 if (!llvm::is_contained({8, 16, 32, 96, 128}, transferSize))625 return emitOpError(626 "Transfering type size must be 8, 16, 32, 96 or 128 bits");627 628 if (!hasGlobalMemorySpace(srcType.getMemorySpace()) &&629 !hasFatRawBufferMemorySpace(srcType.getMemorySpace()))630 return emitOpError(631 "source memory address space must be global or fat raw buffer");632 633 if (!hasWorkgroupMemorySpace(dstType.getMemorySpace()))634 return emitOpError("destination memory address space must be Workgroup");635 636 return success();637}638 639namespace {640/// If the source/target of a GatherToLDSOp is a CastOp that only removes static641/// information or changes layout, the cast can be skipped.642struct FoldGatherToLDSOfCast final : OpRewritePattern<GatherToLDSOp> {643 using OpRewritePattern::OpRewritePattern;644 645 LogicalResult matchAndRewrite(GatherToLDSOp gatherOp,646 PatternRewriter &rewriter) const override {647 bool modified = false;648 auto foldCast = [&](OpOperand &operand) {649 if (auto castOp = operand.get().getDefiningOp<memref::CastOp>()) {650 if (memref::CastOp::canFoldIntoConsumerOp(castOp)) {651 rewriter.modifyOpInPlace(gatherOp,652 [&] { operand.assign(castOp.getSource()); });653 modified = true;654 }655 }656 };657 658 foldCast(gatherOp.getSrcMutable());659 foldCast(gatherOp.getDstMutable());660 661 return success(modified);662 }663};664} // namespace665 666void GatherToLDSOp::getCanonicalizationPatterns(RewritePatternSet &results,667 MLIRContext *context) {668 results.add<FoldGatherToLDSOfCast>(context);669}670 671//===----------------------------------------------------------------------===//672// TransposeLoadOp673//===----------------------------------------------------------------------===//674 675LogicalResult TransposeLoadOp::verify() {676 MemRefType srcType = cast<MemRefType>(getSrc().getType());677 678 if (!hasWorkgroupMemorySpace(srcType.getMemorySpace()))679 return emitOpError("source memory address space must be Workgroup");680 681 auto transferType = cast<VectorType>(getType());682 size_t numElements = transferType.getNumElements();683 size_t elementTypeSize =684 transferType.getElementType().getIntOrFloatBitWidth();685 686 // ElementSize -> NumElements687 const llvm::SmallDenseMap<size_t, size_t> kValidLoadSizeMap = {688 {4, 16},689 {6, 16},690 {8, 8},691 {16, 4},692 };693 694 auto validNumElems = kValidLoadSizeMap.find(elementTypeSize);695 if (validNumElems == kValidLoadSizeMap.end()) {696 return emitOpError("Unsupported element type size for transpose load: ")697 << elementTypeSize << " bits";698 }699 if (numElements != validNumElems->second) {700 return emitOpError(701 "Transferring type size mismatch: expected num of elements: ")702 << validNumElems->second;703 }704 705 return success();706}707 708//===----------------------------------------------------------------------===//709// MakeDmaDescriptorOp710//===----------------------------------------------------------------------===//711 712LogicalResult MakeDmaDescriptorOp::verify() {713 ArrayRef<int64_t> globalStaticStrides = getGlobalStaticStrides();714 715 if (globalStaticStrides.empty()) {716 return emitOpError("strides must not be empty.");717 }718 if (globalStaticStrides.back() != 1) {719 return emitOpError("strides for the innermost dimension must be 1.");720 }721 722 ArrayRef<int64_t> globalStaticSizes = getGlobalStaticSizes();723 size_t rank = globalStaticSizes.size();724 if (rank != globalStaticStrides.size()) {725 return emitOpError("strides and sizes must have same rank.");726 }727 728 ArrayRef<int64_t> sharedStaticSizes = getSharedStaticSizes();729 if (rank != sharedStaticSizes.size()) {730 return emitOpError("tensor must have same rank as tile.");731 }732 733 if (Value atomicBarrierAddress = getAtomicBarrierAddress()) {734 MemRefType atomicBarrierAddressType =735 cast<MemRefType>(atomicBarrierAddress.getType());736 bool barrierInLDS =737 hasWorkgroupMemorySpace(atomicBarrierAddressType.getMemorySpace());738 if (!barrierInLDS) {739 return emitOpError("atomic barrier address must be in LDS.");740 }741 }742 743 return success();744}745 746//===----------------------------------------------------------------------===//747// ScaledMFMAOp748//===----------------------------------------------------------------------===//749 750namespace {751/// Check if the scales input is used in other scaled mfma's while they exist.752/// If theyre unused then pack the scales.753struct PackScales final : OpRewritePattern<ScaledMFMAOp> {754 using OpRewritePattern::OpRewritePattern;755 756 LogicalResult matchAndRewrite(ScaledMFMAOp op,757 PatternRewriter &rewriter) const override {758 Location loc = op.getLoc();759 auto setOpsel = [&op](unsigned idx, int64_t val) {760 switch (idx) {761 case 3:762 op.setScalesIdxA(val);763 break;764 case 4:765 op.setScalesIdxB(val);766 break;767 default:768 break;769 }770 };771 772 // For every scale operand of this ScaledMFMAOp, if the scale is produced by773 // the extraction of a single scale from some vector, then attempt to774 // extract 4 values from that vector instead.775 //776 // Example: (f8 here means f8E8M0FNU)777 // %unit = vector.extract %ScaleSrc[offsets] : f8 from vector<...>778 // %scale = vector.insert %unit, ... : f8 into vector<4xf8>779 // amdgpu.scaled_mfma(%scale[0] * ...780 //781 // rewrite to:782 //783 // %reshaped = vector.shape_cast %ScaleSrc : vector<...> to vector<?xf8>784 // %scale = vector.extract %reshaped[?] : vector<4xf8> from vector<?xf8>785 // amdgpu.scaled_mfma(%scale[0-3] * ...786 //787 // This creates duplicate shape_casts for every use but these will be788 // removed in CSE.789 for (auto opIdx : std::array<int64_t, 2>({3, 4})) {790 auto insertOp = op.getOperand(opIdx).getDefiningOp<vector::InsertOp>();791 if (!insertOp) {792 return rewriter.notifyMatchFailure(op,793 "defining op not a vector.insert");794 }795 // If the extracted value is not a single scalar, then it has been packed.796 if (isa<VectorType>(insertOp.getValueToStore().getType())) {797 return rewriter.notifyMatchFailure(798 op, "scaled mfma operand already packed");799 }800 801 auto extractOp =802 insertOp.getValueToStore().getDefiningOp<vector::ExtractOp>();803 if (!extractOp) {804 return rewriter.notifyMatchFailure(op,805 "defining op not a vector.extract");806 }807 808 Value scaleSrc = extractOp.getOperand(0);809 auto scaleSrcType = dyn_cast<VectorType>(scaleSrc.getType());810 if (!scaleSrcType) {811 return rewriter.notifyMatchFailure(op, "not a vector type");812 }813 814 // We do not handle dynamic dims yet, assume that the input is padded to815 // a static shape now.816 if (!scaleSrcType.hasStaticShape()) {817 return rewriter.notifyMatchFailure(op,818 "dynamic dims not yet supported");819 }820 821 int64_t numElements = scaleSrcType.getNumElements();822 if (numElements <= 4) {823 return rewriter.notifyMatchFailure(824 op, "no packing if # of scales less than four");825 }826 827 // Find a linearized idx using the size and offsets of the extract op.828 auto extractedPos = llvm::to_vector_of<int64_t>(829 llvm::reverse(extractOp.getStaticPosition()));830 ArrayRef<int64_t> scaleSrcShape = scaleSrcType.getShape();831 int64_t scaleSrcRank = scaleSrcType.getRank();832 SmallVector<int64_t> extractSizes(scaleSrcRank, 1);833 for (int64_t i = 1; i < scaleSrcRank; ++i) {834 extractSizes[i] = extractSizes[i - 1] * scaleSrcShape[scaleSrcRank - i];835 }836 int64_t idx = linearize(extractedPos, extractSizes);837 838 // All n scales (where n is the total number of scales) must now be839 // extracted in chunks of 4 elements. This is done by dividing the840 // original vector of scales into groups of 4 elements841 // at offsets 0, 4, ..., m (where m = n/4). All extractions of a842 // scale at a particular index are now replaced with an extraction843 // of the entire group of 4 elements to which that index belongs.844 //845 // If the number of scales happens to be indivisible by 4, extract846 // the remaining n - m scales in a chunk of 4 elements starting at847 // offset n - 4.848 int64_t offset = idx - (idx % 4);849 int64_t opsel = idx - offset;850 int64_t size = 4l;851 // Accomdate remaining elements in the case of non-4-divisible vectors.852 if (numElements - offset < size) {853 opsel = size - (numElements - idx);854 offset = numElements - 4l;855 }856 Type scaleSrcElemType = scaleSrcType.getElementType();857 auto newSrcType =858 VectorType::get(ArrayRef{numElements}, scaleSrcElemType);859 Value newScaleSrc =860 vector::ShapeCastOp::create(rewriter, loc, newSrcType, scaleSrc);861 auto extract = vector::ExtractStridedSliceOp::create(862 rewriter, loc, newScaleSrc, ArrayRef{offset}, ArrayRef{size},863 ArrayRef{int64_t(1)});864 rewriter.modifyOpInPlace(op, [&] {865 op->setOperand(opIdx, extract);866 setOpsel(opIdx, opsel);867 });868 }869 return success();870 }871};872} // namespace873 874void ScaledMFMAOp::getCanonicalizationPatterns(RewritePatternSet &results,875 MLIRContext *context) {876 results.add<PackScales>(context);877}878 879#include "mlir/Dialect/AMDGPU/IR/AMDGPUEnums.cpp.inc"880 881#define GET_ATTRDEF_CLASSES882#include "mlir/Dialect/AMDGPU/IR/AMDGPUAttributes.cpp.inc"883 884#define GET_TYPEDEF_CLASSES885#include "mlir/Dialect/AMDGPU/IR/AMDGPUTypes.cpp.inc"886 887#define GET_OP_CLASSES888#include "mlir/Dialect/AMDGPU/IR/AMDGPU.cpp.inc"889