MLIR 22.0.0git
AMDGPUDialect.cpp
Go to the documentation of this file.
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-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This file implements the AMDGPU dialect and its operations.
10//
11//===----------------------------------------------------------------------===//
12
14
22#include "mlir/IR/Builders.h"
24#include "mlir/IR/Diagnostics.h"
26#include "mlir/IR/Matchers.h"
31#include "llvm/ADT/DenseMap.h"
32#include "llvm/ADT/SmallVector.h"
33#include "llvm/ADT/TypeSwitch.h"
34
35#include <algorithm>
36#include <cstdint>
37#include <limits>
38#include <optional>
39
40using namespace mlir;
41using namespace mlir::amdgpu;
42
43#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.cpp.inc"
44
45namespace {
46struct AMDGPUInlinerInterface final : DialectInlinerInterface {
47 using DialectInlinerInterface::DialectInlinerInterface;
48 bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
49 return true;
50 }
51};
52} // namespace
53
54void AMDGPUDialect::initialize() {
55 addOperations<
56#define GET_OP_LIST
57#include "mlir/Dialect/AMDGPU/IR/AMDGPU.cpp.inc"
58 >();
59 addTypes<
60#define GET_TYPEDEF_LIST
61#include "mlir/Dialect/AMDGPU/IR/AMDGPUTypes.cpp.inc"
62 >();
63 addAttributes<
64#define GET_ATTRDEF_LIST
65#include "mlir/Dialect/AMDGPU/IR/AMDGPUAttributes.cpp.inc"
66 >();
67 addInterfaces<AMDGPUInlinerInterface>();
68}
69
70//===----------------------------------------------------------------------===//
71// 8-bit float ops
72//===----------------------------------------------------------------------===//
73LogicalResult PackedTrunc2xFp8Op::verify() {
74 if (getExisting() && getExisting().getType() != getResult().getType())
75 return emitOpError("existing values must have same type as result");
76 return success();
77}
78
79LogicalResult PackedStochRoundFp8Op::verify() {
80 if (getExisting() && getExisting().getType() != getResult().getType())
81 return emitOpError("existing values must have same type as result");
82 return success();
83}
84
85//===----------------------------------------------------------------------===//
86// mxfp float ops
87//===----------------------------------------------------------------------===//
88LogicalResult PackedScaledTruncOp::verify() {
89 if (getExisting() && getExisting().getType() != getResult().getType())
90 return emitOpError("existing values must have same type as result");
91 return success();
92}
93
94//===----------------------------------------------------------------------===//
95// FatRawBufferCastOp
96//===----------------------------------------------------------------------===//
97
98/// Convert the type `source` to one with the same sizes and strides - and
99/// offset, unless `stripOffset` is true, in which case the offset is reset to
100/// 0, if the offset should be reset but the layout of `source` isn't either the
101/// identity layout or a strided layout, this function fails.
102static FailureOr<MemRefType> getFatRawBufferTypeLike(MemRefType source,
103 bool resetOffset) {
104 MLIRContext *ctx = source.getContext();
105 MemRefType::Builder mb(source);
107 amdgpu::AddressSpaceAttr::get(ctx, amdgpu::AddressSpace::FatRawBuffer));
108 MemRefLayoutAttrInterface layout = source.getLayout();
109 if (resetOffset && !layout.isIdentity()) {
110 auto stridedLayout = dyn_cast<StridedLayoutAttr>(layout);
111 if (!stridedLayout)
112 return failure();
113 MemRefLayoutAttrInterface newLayout =
114 StridedLayoutAttr::get(ctx, 0, stridedLayout.getStrides());
115 // Special case: if resetting the offset causes the strided layout to become
116 // the identity layout, then reset to the identity layout.
117 // TODO: this'll get a lot simpler when we have the contiguous layout.
118 SmallVector<int64_t> stridesIfIdentity;
119 if (source.hasStaticShape()) {
120 stridesIfIdentity = computeSuffixProduct(source.getShape());
121 } else if (source.getRank() <= 1) {
122 stridesIfIdentity = SmallVector<int64_t>(source.getRank(), 1);
123 }
124 if (stridesIfIdentity == stridedLayout.getStrides()) {
125 newLayout = AffineMapAttr::get(
126 AffineMap::getMultiDimIdentityMap(source.getRank(), ctx));
127 }
128 mb.setLayout(newLayout);
129 }
130 return (MemRefType)(mb);
131}
132
133LogicalResult FatRawBufferCastOp::inferReturnTypes(
134 MLIRContext *context, std::optional<Location> location, ValueRange operands,
135 DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions,
136 SmallVectorImpl<Type> &inferredReturnTypes) {
137 Adaptor adaptor(operands, attributes, properties, regions);
138 auto sourceType =
139 dyn_cast_if_present<MemRefType>(adaptor.getSource().getType());
140 if (!sourceType)
141 return failure();
142 FailureOr<MemRefType> resultType =
143 getFatRawBufferTypeLike(sourceType, adaptor.getResetOffset());
144 if (failed(resultType))
145 return failure();
146 inferredReturnTypes = SmallVector<Type>{*resultType};
147 return success();
148}
149
150FailureOr<OpFoldResult> FatRawBufferCastOp::reifyDimOfResult(OpBuilder &builder,
151 int resultIndex,
152 int dim) {
153 assert(resultIndex == 0 && "FatRawBufferCastOp has a single result");
154 Value source = getSource();
155 auto sourceType = cast<MemRefType>(source.getType());
156 if (sourceType.isDynamicDim(dim))
157 return OpFoldResult(
158 builder.createOrFold<memref::DimOp>(getLoc(), source, dim));
159 return OpFoldResult(builder.getIndexAttr(sourceType.getDimSize(dim)));
160}
161
162LogicalResult FatRawBufferCastOp::verify() {
163 FailureOr<MemRefType> expectedResultType =
164 getFatRawBufferTypeLike(getSource().getType(), getResetOffset());
165 if (failed(expectedResultType))
166 return emitOpError("source type ")
167 << getSource().getType() << " can't have its offset reset";
168 if (getResult().getType() != *expectedResultType)
169 return emitOpError("expected result type to be ")
170 << *expectedResultType << " but got " << getResult().getType();
171 return success();
172}
173
174static bool hasGlobalMemorySpace(Attribute memorySpace) {
175 if (!memorySpace)
176 return true;
177 if (auto intMemorySpace = dyn_cast<IntegerAttr>(memorySpace))
178 return intMemorySpace.getInt() == 0 || intMemorySpace.getInt() == 1;
179 if (auto gpuMemorySpace = dyn_cast<gpu::AddressSpaceAttr>(memorySpace))
180 return gpuMemorySpace.getValue() == gpu::AddressSpace::Global;
181 return false;
182}
183
184static bool hasWorkgroupMemorySpace(Attribute memorySpace) {
185 if (!memorySpace)
186 return false;
187 if (auto intMemorySpace = dyn_cast<IntegerAttr>(memorySpace))
188 return intMemorySpace.getInt() == 3;
189 if (auto gpuMemorySpace = dyn_cast<gpu::AddressSpaceAttr>(memorySpace))
190 return gpuMemorySpace.getValue() == gpu::AddressSpace::Workgroup;
191 return false;
192}
193
194static bool hasFatRawBufferMemorySpace(Attribute memorySpace) {
195 if (!memorySpace)
196 return false;
197 if (auto intMemorySpace = dyn_cast<IntegerAttr>(memorySpace))
198 return intMemorySpace.getInt() == 7;
199 if (auto gpuMemorySpace = dyn_cast<amdgpu::AddressSpaceAttr>(memorySpace))
200 return gpuMemorySpace.getValue() == amdgpu::AddressSpace::FatRawBuffer;
201 return false;
202}
203
204//===----------------------------------------------------------------------===//
205// RawBuffer*Op
206//===----------------------------------------------------------------------===//
207template <typename T>
208static LogicalResult verifyRawBufferOp(T &op) {
209 MemRefType bufferType = llvm::cast<MemRefType>(op.getMemref().getType());
210 bool isGlobal = hasGlobalMemorySpace(bufferType.getMemorySpace());
211
212 if (!isGlobal)
213 return op.emitOpError(
214 "Buffer ops must operate on a memref in global memory");
215 if (!bufferType.hasRank())
216 return op.emitOpError(
217 "Cannot meaningfully buffer_store to an unranked memref");
218 if (static_cast<int64_t>(op.getIndices().size()) != bufferType.getRank())
219 return op.emitOpError("Expected " + Twine(bufferType.getRank()) +
220 " indices to memref");
221 return success();
222}
223
224LogicalResult RawBufferLoadOp::verify() { return verifyRawBufferOp(*this); }
225
226LogicalResult RawBufferStoreOp::verify() { return verifyRawBufferOp(*this); }
227
228LogicalResult RawBufferAtomicFaddOp::verify() {
229 return verifyRawBufferOp(*this);
230}
231
232LogicalResult RawBufferAtomicFmaxOp::verify() {
233 return verifyRawBufferOp(*this);
234}
235
236LogicalResult RawBufferAtomicSmaxOp::verify() {
237 return verifyRawBufferOp(*this);
238}
239
240LogicalResult RawBufferAtomicUminOp::verify() {
241 return verifyRawBufferOp(*this);
242}
243
244LogicalResult RawBufferAtomicCmpswapOp::verify() {
245 return verifyRawBufferOp(*this);
246}
247
248static std::optional<uint32_t> getConstantUint32(Value v) {
249 APInt cst;
250 if (!v.getType().isInteger(32))
251 return std::nullopt;
252 if (matchPattern(v, m_ConstantInt(&cst)))
253 return cst.getZExtValue();
254 return std::nullopt;
255}
256
257template <typename OpType>
258static bool staticallyOutOfBounds(OpType op) {
259 if (!op.getBoundsCheck())
260 return false;
261 MemRefType bufferType = op.getMemref().getType();
262 if (!bufferType.hasStaticShape())
263 return false;
264 int64_t offset;
265 SmallVector<int64_t> strides;
266 if (failed(bufferType.getStridesAndOffset(strides, offset)))
267 return false;
268 int64_t result = offset + op.getIndexOffset().value_or(0);
269 if (op.getSgprOffset()) {
270 std::optional<uint32_t> sgprOffset = getConstantUint32(op.getSgprOffset());
271 if (!sgprOffset)
272 return false;
273 result += *sgprOffset;
274 }
275 if (strides.size() != op.getIndices().size())
276 return false;
277 int64_t indexVal = 0;
278 for (auto pair : llvm::zip(strides, op.getIndices())) {
279 int64_t stride = std::get<0>(pair);
280 Value idx = std::get<1>(pair);
281 std::optional<uint32_t> idxVal = getConstantUint32(idx);
282 if (!idxVal)
283 return false;
284 indexVal += stride * *idxVal;
285 }
286 result += indexVal;
287 if (result > std::numeric_limits<uint32_t>::max())
288 // Overflow means don't drop
289 return false;
290 return result >= bufferType.getNumElements();
291}
292
293namespace {
294template <typename OpType>
295struct RemoveStaticallyOobBufferLoads final : public OpRewritePattern<OpType> {
296 using OpRewritePattern<OpType>::OpRewritePattern;
297
298 LogicalResult matchAndRewrite(OpType op, PatternRewriter &rw) const override {
299 if (!staticallyOutOfBounds(op))
300 return failure();
301 Type loadType = op.getResult().getType();
302 rw.replaceOpWithNewOp<arith::ConstantOp>(op, loadType,
303 rw.getZeroAttr(loadType));
304 return success();
305 }
306};
307
308template <typename OpType>
309struct RemoveStaticallyOobBufferWrites final : public OpRewritePattern<OpType> {
310 using OpRewritePattern<OpType>::OpRewritePattern;
311
312 LogicalResult matchAndRewrite(OpType op, PatternRewriter &rw) const override {
313 if (!staticallyOutOfBounds(op))
314 return failure();
315
316 rw.eraseOp(op);
317 return success();
318 }
319};
320} // end namespace
321
322void RawBufferLoadOp::getCanonicalizationPatterns(RewritePatternSet &results,
323 MLIRContext *context) {
324 results.add<RemoveStaticallyOobBufferLoads<RawBufferLoadOp>>(context);
325}
326
327void RawBufferStoreOp::getCanonicalizationPatterns(RewritePatternSet &results,
328 MLIRContext *context) {
329 results.add<RemoveStaticallyOobBufferWrites<RawBufferStoreOp>>(context);
330}
331
332void RawBufferAtomicFaddOp::getCanonicalizationPatterns(
333 RewritePatternSet &results, MLIRContext *context) {
334 results.add<RemoveStaticallyOobBufferWrites<RawBufferAtomicFaddOp>>(context);
335}
336
337void RawBufferAtomicFmaxOp::getCanonicalizationPatterns(
338 RewritePatternSet &results, MLIRContext *context) {
339 results.add<RemoveStaticallyOobBufferWrites<RawBufferAtomicFmaxOp>>(context);
340}
341
342void RawBufferAtomicSmaxOp::getCanonicalizationPatterns(
343 RewritePatternSet &results, MLIRContext *context) {
344 results.add<RemoveStaticallyOobBufferWrites<RawBufferAtomicSmaxOp>>(context);
345}
346
347void RawBufferAtomicUminOp::getCanonicalizationPatterns(
348 RewritePatternSet &results, MLIRContext *context) {
349 results.add<RemoveStaticallyOobBufferWrites<RawBufferAtomicUminOp>>(context);
350}
351
352void RawBufferAtomicCmpswapOp::getCanonicalizationPatterns(
353 RewritePatternSet &results, MLIRContext *context) {
354 results.add<RemoveStaticallyOobBufferLoads<RawBufferAtomicCmpswapOp>>(
355 context);
356}
357
358//===----------------------------------------------------------------------===//
359// ScaledExtPackedMatrixOp
360//===----------------------------------------------------------------------===//
361LogicalResult ScaledExtPackedMatrixOp::verify() {
362 int blockSize = getBlockSize();
363 assert(llvm::is_contained({16, 32}, blockSize) && "invalid block size");
364
365 int firstScaleByte = getFirstScaleByte();
366 int firstScaleLane = getFirstScaleLane();
367 auto sourceType = cast<VectorType>(getSource().getType());
368 Type elementType = sourceType.getElementType();
369 auto floatType = cast<FloatType>(elementType);
370 unsigned bitWidth = floatType.getWidth();
371
372 assert(llvm::is_contained(llvm::ArrayRef<unsigned>{4, 6, 8}, bitWidth));
373
374 const bool is_fp8 = bitWidth == 8;
375 const bool is_block_16 = blockSize == 16;
376
377 if (!is_fp8) {
378 if (is_block_16) {
379 if (!llvm::is_contained({0, 1}, firstScaleByte)) {
380 return emitOpError("blockSize of 16 can only have firstScaleByte be 0 "
381 "or 1 for f4 and f6.");
382 }
383 } else {
384 if (!llvm::is_contained({0, 2}, firstScaleByte)) {
385 return emitOpError("blockSize of 32 can only have firstScaleByte be 0 "
386 "or 2 for f4 and f6.");
387 }
388 }
389 } else {
390 if (is_block_16) {
391 bool is_valid = ((firstScaleLane == 0) && (firstScaleByte == 0)) ||
392 ((firstScaleLane == 16) && (firstScaleByte == 2));
393 if (!is_valid) {
394 return emitOpError("blockSize of 16 can only have (firstScaleLane, "
395 "firstScaleByte) be (0, 0) or (16, 2) for f8.");
396 }
397 }
398 }
399
400 return success();
401}
402
403//===----------------------------------------------------------------------===//
404// WMMAOp
405//===----------------------------------------------------------------------===//
406
408 IntegerAttr &m, IntegerAttr &n,
409 IntegerAttr &k) {
410 SmallVector<int64_t, 3> dimensions;
411 if (parser.parseDimensionList(dimensions, false, false))
412 return failure();
413 if (dimensions.size() != 3)
414 return parser.emitError(parser.getCurrentLocation())
415 << "expected 3 dimensions in MNK dimension list";
416
417 m = parser.getBuilder().getI32IntegerAttr(dimensions[0]);
418 n = parser.getBuilder().getI32IntegerAttr(dimensions[1]);
419 k = parser.getBuilder().getI32IntegerAttr(dimensions[2]);
420 return success();
421}
422
423LogicalResult WMMAOp::verify() {
424 auto sourceAType = cast<VectorType>(getSourceA().getType());
425 auto sourceBType = cast<VectorType>(getSourceB().getType());
426 auto destType = cast<VectorType>(getDestC().getType());
427
428 Type sourceAElemType = sourceAType.getElementType();
429 Type sourceBElemType = sourceBType.getElementType();
430 if (sourceAType.getNumElements() != sourceBType.getNumElements()) {
431 return emitOpError("source vectors have different lengths: ")
432 << sourceAType << " vs. " << sourceBType;
433 }
434
435 bool isDestFloat = destType.getElementType().isFloat();
436 bool isSrcFloat = sourceAElemType.isFloat();
437
438 if (isDestFloat && !isSrcFloat)
439 return emitOpError("expected float sources with float destination");
440 if (!isDestFloat && isSrcFloat)
441 return emitOpError("expected int sources with int destination");
442
443 if (!sourceAElemType.isFloat(8) && sourceAElemType != sourceBElemType) {
444 return emitOpError(
445 "source element types must match (except for fp8/bf8) but have ")
446 << sourceAType << " and " << sourceBType;
447 }
448
449 if (isSrcFloat) {
450 if (getClamp())
451 return emitOpError("clamp flag is not supported for float types");
452 if (getUnsignedA() || getUnsignedB())
453 return emitOpError("unsigned flags are not supported for float types");
454 }
455 return success();
456}
457
458//===----------------------------------------------------------------------===//
459// ScaledWMMAOp
460//===----------------------------------------------------------------------===//
461
462LogicalResult ScaledWMMAOp::verify() {
463 // Helper functions for type classification.
464 auto isF8 = llvm::IsaPred<Float8E4M3FNType, Float8E5M2Type>;
465 auto isF6 = llvm::IsaPred<Float6E2M3FNType, Float6E3M2FNType>;
466 auto isF4 = llvm::IsaPred<Float4E2M1FNType>;
467 auto isScaleF8 = llvm::IsaPred<Float8E8M0FNUType, Float8E4M3FNType>;
468 auto isE8M0 = llvm::IsaPred<Float8E8M0FNUType>;
469 auto isE4M3 = llvm::IsaPred<Float8E4M3FNType>;
470
471 auto sourceAType = cast<VectorType>(getSourceA().getType());
472 auto sourceBType = cast<VectorType>(getSourceB().getType());
473 auto destType = cast<VectorType>(getDestC().getType());
474
475 // Validate source element types are small floats (fp4/fp6/fp8).
476 Type aElemType = sourceAType.getElementType();
477 Type bElemType = sourceBType.getElementType();
478
479 // Validate vector lengths based on dimensions.
480 int64_t m = getM();
481 int64_t aLen = sourceAType.getNumElements();
482 int64_t bLen = sourceBType.getNumElements();
483 int64_t expectedOutLen = (m == 16) ? 8 : 16;
484
485 if (destType.getNumElements() != expectedOutLen)
486 return emitOpError("expected output vector of length ")
487 << expectedOutLen << " but got " << destType.getNumElements();
488
489 if (m == 16) {
490 // For 16×16×128: both A and B must be 64 elements.
491 if (aLen != 64)
492 return emitOpError(
493 "for 16x16x128, sourceA must have 64 elements but got ")
494 << aLen;
495 if (bLen != 64)
496 return emitOpError(
497 "for 16x16x128, sourceB must have 64 elements but got ")
498 << bLen;
499 } else { // m == 32
500 // For 32×16×128: only fp4 is supported, A is 128, B is 64.
501 if (!isF4(aElemType) && !isF4(bElemType))
502 return emitOpError("32x16x128 only supports fp4 element types");
503
504 if (aLen != 128)
505 return emitOpError(
506 "for 32x16x128, sourceA must have 128 elements but got ")
507 << aLen;
508 if (bLen != 64)
509 return emitOpError(
510 "for 32x16x128, sourceB must have 64 elements but got ")
511 << bLen;
512
513 // For 32x16x128, matrix A uses all 32 lanes so a_first_scale_lane must be
514 // 0.
515 if (getAFirstScaleLane() != 0)
516 return emitOpError("for 32x16x128, a_first_scale_lane must be 0");
517 }
518
519 // Validate scale types and their compatibility with matrix element types.
520 auto scaleAType = cast<VectorType>(getScaleA().getType());
521 auto scaleBType = cast<VectorType>(getScaleB().getType());
522 Type scaleAElemType = scaleAType.getElementType();
523 Type scaleBElemType = scaleBType.getElementType();
524
525 // Validate scale element types are valid scale f8 types (E8M0FNU or E4M3FN).
526 if (!isScaleF8(scaleAElemType) || !isScaleF8(scaleBElemType))
527 return emitOpError(
528 "scale operands must have f8 element types (E8M0FNU or E4M3FN)");
529
530 // Any matrices A/B (fp8|fp6|fp4) with E8M0 scales for matrix A/B are valid.
531 if (isE8M0(scaleAElemType) && isE8M0(scaleBElemType))
532 return success();
533
534 // Matrix A (F8|F6) x Matrix B (F4) with Scale A (E8M0), Scale B (E5M3|E4M3).
535 if ((isF8(aElemType) || isF6(aElemType)) && isE8M0(scaleAElemType) &&
536 isF4(bElemType) && isE4M3(scaleBElemType))
537 return success();
538
539 // Matrix A (F4) x Matrix B (F8|F6) with Scale A (E5M3|E4M3), Scale B (E8M0).
540 if (isF4(aElemType) && isE4M3(scaleAElemType) &&
541 (isF8(bElemType) || isF6(bElemType)) && isE8M0(scaleBElemType))
542 return success();
543
544 // Matrix A (F4) x Matrix B (F4) with Scale A (E4M3), Scale B (E4M3).
545 if (isF4(aElemType) && isF4(bElemType) && isE4M3(scaleAElemType) &&
546 isE4M3(scaleBElemType))
547 return success();
548
549 // No valid combination matched.
550 return emitOpError("invalid combination of matrix and scale types: ")
551 << "sourceA=" << aElemType << ", scaleA=" << scaleAElemType
552 << ", sourceB=" << bElemType << ", scaleB=" << scaleBElemType;
553}
554
555//===----------------------------------------------------------------------===//
556// MFMAOp
557//===----------------------------------------------------------------------===//
558LogicalResult MFMAOp::verify() {
559 constexpr uint32_t waveSize = 64;
561
562 Type sourceType = getSourceA().getType();
563 Type destType = getDestC().getType();
564
565 Type sourceElem = sourceType, destElem = destType;
566 uint32_t sourceLen = 1, destLen = 1;
567 if (auto sourceVector = dyn_cast<VectorType>(sourceType)) {
568 sourceLen = sourceVector.getNumElements();
569 sourceElem = sourceVector.getElementType();
570 }
571 if (auto destVector = dyn_cast<VectorType>(destType)) {
572 destLen = destVector.getNumElements();
573 destElem = destVector.getElementType();
574 }
575
576 Type sourceBType = getSourceB().getType();
577 if (sourceElem.isFloat(8) || sourceElem.isFloat(6) || sourceElem.isFloat(4)) {
578 int64_t sourceBLen = 1;
579 Type sourceBElem = sourceBType;
580 if (auto sourceBVector = llvm::dyn_cast<VectorType>(sourceBType)) {
581 sourceBLen = sourceBVector.getNumElements();
582 sourceBElem = sourceBVector.getElementType();
583 }
584 if (!sourceBElem.isFloat(8) && !sourceBElem.isFloat(6) &&
585 !sourceBElem.isFloat(4))
586 return emitOpError("expected both source operands to have small-float "
587 "elements if one does");
588 if (sourceLen != sourceBLen)
589 return emitOpError(
590 "expected both small-float source vectors to have the same length");
591 } else {
592 if (sourceType != sourceBType)
593 return emitOpError("expected both non-small-float source operand types "
594 "to match exactly");
595 }
596 // Normalize the wider integer types the compiler expects to i8.
597 if (sourceElem.isInteger(32)) {
598 sourceLen *= 4;
599 sourceElem = b.getI8Type();
600 }
601 if (sourceElem.isInteger(64)) {
602 sourceLen *= 8;
603 sourceElem = b.getI8Type();
604 }
605
606 int64_t numSourceElems = (getM() * getK() * getBlocks()) / waveSize;
607 if (sourceLen != numSourceElems)
608 return emitOpError("expected " + Twine(numSourceElems) +
609 " source values for this operation but got " +
610 Twine(sourceLen));
611
612 int64_t numDestElems = (getM() * getN() * getBlocks()) / waveSize;
613 if (destLen != numDestElems)
614 return emitOpError("expected " + Twine(numDestElems) +
615 " result values for this operation but got " +
616 Twine(destLen));
617
618 if (destElem.isF64() && getBlgp() != MFMAPermB::none)
619 return emitOpError(
620 "double-precision ops do not support permuting lanes of B");
621 if (destElem.isF64() && getCbsz() != 0)
622 return emitOpError(
623 "double-precision ops do not support permuting lanes of A");
624 if (getAbid() >= (1u << getCbsz()))
625 return emitOpError(
626 "block ID for permuting A (abid) must be below 2 ** cbsz");
627
628 if ((getNegateA() || getNegateB() || getNegateC()) && !destElem.isF64())
629 return emitOpError(
630 "negation flags only available for double-precision operations");
631
632 return success();
633}
634
635//===----------------------------------------------------------------------===//
636// SparseMFMAOp
637//===----------------------------------------------------------------------===//
638
639LogicalResult SparseMFMAOp::verify() {
640 constexpr uint32_t waveSize = 64;
641
642 auto sparseType = cast<VectorType>(getSourceA().getType());
643 auto denseType = cast<VectorType>(getSourceB().getType());
644 auto destType = cast<VectorType>(getDestC().getType());
645
646 Type sparseElem = sparseType.getElementType();
647 Type denseElem = denseType.getElementType();
648 int64_t sparseLen = sparseType.getNumElements();
649 int64_t denseLen = denseType.getNumElements();
650 int64_t destLen = destType.getNumElements();
651
652 if (denseLen != 2 * sparseLen)
653 return emitOpError("expected dense source operand to have exactly double "
654 "the number of elements of the sparse source operand");
655
656 // Check that source element types are compatible.
657 // For fp8/bf8 mixed operations, element types can differ (e.g., fp8 * bf8).
658 // For other types, element types must match exactly.
659 bool bothFloat8 = sparseElem.isFloat(8) && denseElem.isFloat(8);
660 if (!bothFloat8 && sparseElem != denseElem)
661 return emitOpError(
662 "expected source operands to have the same element type");
663
664 // When CBSZ == 0, ABID selects the index set within the sparse index VGPR.
665 // When CBSZ != 0, the first index set is always used (ABID ignored).
666 bool is8BitSource = sparseElem.isFloat(8) || sparseElem.isInteger(8);
667 // 8-bit source: ABID selects one of two 16-bit index sets.
668 if (getCbsz() == 0 && is8BitSource && getAbid() > 1)
669 return emitOpError("ABID must be 0 or 1 for 8-bit source data");
670 // 16-bit source: ABID selects one of four 8-bit index sets (0-3 all valid).
671 if (getCbsz() == 0 && !is8BitSource && getAbid() > 3)
672 return emitOpError("ABID must be between 0 and 3 for 16-bit source data");
673
674 // Validate sparseIdx type matches source element type.
675 auto sparseIdxType = cast<VectorType>(getSparseIdx().getType());
676 if (is8BitSource) {
677 // 8-bit source data requires vector<2xi16> sparse indices.
678 if (sparseIdxType.getNumElements() != 2 ||
679 !sparseIdxType.getElementType().isInteger(16))
680 return emitOpError("expected vector<2xi16> sparse indices for 8-bit "
681 "source data, but got ")
682 << getSparseIdx().getType();
683 } else {
684 // 16-bit source data requires vector<4xi8> sparse indices.
685 if (sparseIdxType.getNumElements() != 4 ||
686 !sparseIdxType.getElementType().isInteger(8))
687 return emitOpError("expected vector<4xi8> sparse indices for 16-bit "
688 "source data, but got ")
689 << getSparseIdx().getType();
690 }
691
692 int64_t expectedSourceElems = (getM() * getK()) / waveSize;
693 if (denseLen != expectedSourceElems)
694 return emitOpError("expected " + Twine(expectedSourceElems) +
695 " source values for this operation but got " +
696 Twine(denseLen));
697
698 int64_t expectedDestElems = (getM() * getN()) / waveSize;
699 if (destLen != expectedDestElems)
700 return emitOpError("expected " + Twine(expectedDestElems) +
701 " result values for this operation but got " +
702 Twine(destLen));
703
704 return success();
705}
706
707//===----------------------------------------------------------------------===//
708// DPPOp
709//===----------------------------------------------------------------------===//
710LogicalResult DPPOp::verify() {
711 Type srcType = getSrc().getType();
712 if (srcType.getIntOrFloatBitWidth() > 64) {
713 return emitOpError("integer and floating point types larger than 64 bits "
714 "are not supported");
715 }
716
717 DPPPerm kind = getKind();
718 Attribute permArgument = getPermArgument().value_or(Attribute{});
719
720 switch (kind) {
721
722 case DPPPerm::quad_perm: {
723 auto quadPermAttr = dyn_cast_or_null<ArrayAttr>(permArgument);
724 if (!quadPermAttr || quadPermAttr.size() != 4) {
725 return emitOpError("quad_perm attribute must have exactly 4 elements");
726 }
727 for (auto elem : quadPermAttr.getAsRange<IntegerAttr>()) {
728 int32_t num = elem.getInt();
729 if (num < 0 || num > 3) {
730 return emitOpError(
731 "Each element of quad_perm must be in the range [0, 3]");
732 }
733 }
734 } break;
735
736 case DPPPerm::row_shl:
737 case DPPPerm::row_shr:
738 case DPPPerm::row_ror: {
739 if (!permArgument) {
740 return emitOpError("Attribute '" + Twine(stringifyDPPPerm(kind)) +
741 "' value not specified");
742 }
743 if (auto intAttr = dyn_cast<IntegerAttr>(permArgument)) {
744 uint32_t attrValue = intAttr.getInt();
745 if (attrValue < 1 || attrValue > 15) {
746 return emitOpError("Attribute value must be between 1 and 15");
747 }
748 }
749 } break;
750
751 case DPPPerm::wave_shl:
752 case DPPPerm::wave_shr:
753 case DPPPerm::wave_rol:
754 case DPPPerm::wave_ror:
755 case DPPPerm::row_mirror:
756 case DPPPerm::row_half_mirror:
757 case DPPPerm::row_bcast_15:
758 case DPPPerm::row_bcast_31: {
759 if (permArgument && !isa<UnitAttr>(permArgument)) {
760 return emitOpError("Expected unit attribute for permArgument, but found "
761 "non-trivial argument");
762 }
763 break;
764 }
765 }
766 return success();
767}
768
769//===----------------------------------------------------------------------===//
770// PermlaneSwapOp
771//===----------------------------------------------------------------------===//
772LogicalResult PermlaneSwapOp::verify() {
773 unsigned rowLength = getRowLength();
774
775 if (rowLength != 16 && rowLength != 32)
776 return emitOpError("row_length attribute must either be 16 or 32.");
777
778 return success();
779}
780
781//===----------------------------------------------------------------------===//
782// MemoryCounterWaitOp
783//===----------------------------------------------------------------------===//
784
785namespace {
786/// Fuse adjacent memory counter wait ops, taking the minimum value of the
787/// counters.
788struct FuseMemoryCounterWaitOp final : OpRewritePattern<MemoryCounterWaitOp> {
789 using Base::Base;
790
791 LogicalResult matchAndRewrite(MemoryCounterWaitOp op,
792 PatternRewriter &rewriter) const override {
793 auto next = dyn_cast<MemoryCounterWaitOp>(op->getNextNode());
794 if (!next)
795 return failure();
796
797 auto setters = {&MemoryCounterWaitOp::setLoad,
798 &MemoryCounterWaitOp::setStore, &MemoryCounterWaitOp::setDs,
799 &MemoryCounterWaitOp::setExp,
800 &MemoryCounterWaitOp::setTensor};
801 auto lhsVals = {op.getLoad(), op.getStore(), op.getDs(), op.getExp(),
802 op.getTensor()};
803 auto rhsVals = {next.getLoad(), next.getStore(), next.getDs(),
804 next.getExp(), next.getTensor()};
805 rewriter.modifyOpInPlace(op, [&] {
806 for (auto [setter, lhs, rhs] :
807 llvm::zip_equal(setters, lhsVals, rhsVals)) {
808 if (lhs && rhs) {
809 (op.*setter)(std::min(*lhs, *rhs));
810 } else if (lhs) {
811 (op.*setter)(*lhs);
812 } else if (rhs) {
813 (op.*setter)(*rhs);
814 }
815 }
816 });
817 rewriter.eraseOp(next);
818 return success();
819 }
820};
821} // namespace
822
823void MemoryCounterWaitOp::getCanonicalizationPatterns(
824 RewritePatternSet &results, MLIRContext *context) {
825 results.add<FuseMemoryCounterWaitOp>(context);
826}
827
828//===----------------------------------------------------------------------===//
829// GatherToLDSOp
830//===----------------------------------------------------------------------===//
831
832LogicalResult GatherToLDSOp::verify() {
833 MemRefType srcType = cast<MemRefType>(getSrc().getType());
834 MemRefType dstType = cast<MemRefType>(getDst().getType());
835
836 if (dstType.getRank() > 0 && !dstType.areTrailingDimsContiguous(1))
837 return emitOpError("destination type inner most dim must be contiguous");
838
839 auto elemType = srcType.getElementType();
840 // Check $src and $dst element types are the same.
841 if (elemType != dstType.getElementType())
842 return emitOpError("source and destination element types must match");
843
844 // copy type sizes should be 1, 2, 4, 12 or 16 bytes.
845 auto transferType = getTransferType();
846 int transferSize;
847 if (auto vectorTransfer = dyn_cast<VectorType>(transferType)) {
848 transferSize = vectorTransfer.getNumElements() *
849 vectorTransfer.getElementTypeBitWidth();
850 } else {
851 transferSize = transferType.getIntOrFloatBitWidth();
852 }
853 if (!llvm::is_contained({8, 16, 32, 96, 128}, transferSize))
854 return emitOpError(
855 "Transfering type size must be 8, 16, 32, 96 or 128 bits");
856
857 if (!hasGlobalMemorySpace(srcType.getMemorySpace()) &&
858 !hasFatRawBufferMemorySpace(srcType.getMemorySpace()))
859 return emitOpError(
860 "source memory address space must be global or fat raw buffer");
861
862 if (!hasWorkgroupMemorySpace(dstType.getMemorySpace()))
863 return emitOpError("destination memory address space must be Workgroup");
864
865 return success();
866}
867
868namespace {
869/// If the source/target of a GatherToLDSOp is a CastOp that only removes static
870/// information or changes layout, the cast can be skipped.
871struct FoldGatherToLDSOfCast final : OpRewritePattern<GatherToLDSOp> {
873
874 LogicalResult matchAndRewrite(GatherToLDSOp gatherOp,
875 PatternRewriter &rewriter) const override {
876 bool modified = false;
877 auto foldCast = [&](OpOperand &operand) {
878 if (auto castOp = operand.get().getDefiningOp<memref::CastOp>()) {
879 if (memref::CastOp::canFoldIntoConsumerOp(castOp)) {
880 rewriter.modifyOpInPlace(gatherOp,
881 [&] { operand.assign(castOp.getSource()); });
882 modified = true;
883 }
884 }
885 };
886
887 foldCast(gatherOp.getSrcMutable());
888 foldCast(gatherOp.getDstMutable());
889
890 return success(modified);
891 }
892};
893} // namespace
894
895void GatherToLDSOp::getCanonicalizationPatterns(RewritePatternSet &results,
896 MLIRContext *context) {
897 results.add<FoldGatherToLDSOfCast>(context);
898}
899
900//===----------------------------------------------------------------------===//
901// TransposeLoadOp
902//===----------------------------------------------------------------------===//
903
904LogicalResult TransposeLoadOp::verify() {
905 MemRefType srcType = cast<MemRefType>(getSrc().getType());
906
907 if (!hasWorkgroupMemorySpace(srcType.getMemorySpace()))
908 return emitOpError("source memory address space must be Workgroup");
909
910 auto transferType = cast<VectorType>(getType());
911 size_t numElements = transferType.getNumElements();
912 size_t elementTypeSize =
913 transferType.getElementType().getIntOrFloatBitWidth();
914
915 // ElementSize -> NumElements
916 const llvm::SmallDenseMap<size_t, size_t> kValidLoadSizeMap = {
917 {4, 16},
918 {6, 16},
919 {8, 8},
920 {16, 4},
921 };
922
923 auto validNumElems = kValidLoadSizeMap.find(elementTypeSize);
924 if (validNumElems == kValidLoadSizeMap.end())
925 return emitOpError("Unsupported element type size for transpose load: ")
926 << elementTypeSize << " bits";
927
928 if (numElements != validNumElems->second)
929 return emitOpError(
930 "Transferring type size mismatch: expected num of elements: ")
931 << validNumElems->second;
932
933 return success();
934}
935
936//===----------------------------------------------------------------------===//
937// MakeDmaBaseOp
938//===----------------------------------------------------------------------===//
939
940template <typename BaseOp>
941static LogicalResult verifyBase(BaseOp op) {
942 auto ldsType = cast<MemRefType>(op.getLds().getType());
943 auto globalType = cast<MemRefType>(op.getGlobal().getType());
944 if (!hasWorkgroupMemorySpace(ldsType.getMemorySpace()))
945 return op.emitOpError(
946 "lds memref must have workgroup address space attribute.");
947 if (!hasGlobalMemorySpace(globalType.getMemorySpace()))
948 return op.emitOpError(
949 "global memref must have global address space attribute.");
950
951 Type elementType = ldsType.getElementType();
952 unsigned width = elementType.getIntOrFloatBitWidth();
953
954 if (!llvm::is_contained({8u, 16u, 32u, 64u}, width))
955 return op.emitOpError(
956 "element type must be 1, 2, 4, or 8 bytes long but type was ")
957 << width << " bits long.";
958 return success();
959}
960
961LogicalResult MakeDmaBaseOp::verify() { return verifyBase(*this); }
962
963//===----------------------------------------------------------------------===//
964// MakeGatherDmaBaseOp
965//===----------------------------------------------------------------------===//
966
967LogicalResult
968TDMGatherBaseType::verify(function_ref<InFlightDiagnostic()> emitError,
969 Type elementType, Type indexType) {
970 unsigned width = elementType.getIntOrFloatBitWidth();
971 if (!llvm::is_contained({8u, 16u, 32u, 64u}, width))
972 return emitError()
973 << "element type must be 1, 2, 4, or 8 bytes wide but type "
974 << elementType << " is " << width / 8 << " bytes wide.";
975 MLIRContext *ctx = elementType.getContext();
976 Type i16 = IntegerType::get(ctx, 32);
977 Type i32 = IntegerType::get(ctx, 16);
978 if (!llvm::is_contained({i16, i32}, indexType))
979 return emitError() << "index type must be i16 or i32 but index type is "
980 << indexType << ".";
981 return success();
982}
983
984LogicalResult MakeGatherDmaBaseOp::verify() { return verifyBase(*this); }
985
986//===----------------------------------------------------------------------===//
987// MakeDmaDescriptorOp
988//===----------------------------------------------------------------------===//
989
990template <typename DescriptorOp>
991static LogicalResult verifyDescriptorOp(DescriptorOp op) {
992 ArrayRef<int64_t> globalStaticStrides = op.getGlobalStaticStrides();
993
994 if (globalStaticStrides.empty())
995 return op.emitOpError("strides must not be empty.");
996 if (globalStaticStrides.back() != 1)
997 return op.emitOpError("strides for the innermost dimension must be 1.");
998
999 ArrayRef<int64_t> globalStaticSizes = op.getGlobalStaticSizes();
1000 size_t rank = globalStaticSizes.size();
1001 if (rank > 5)
1002 return op.emitOpError("tensor and tile must be at most of rank 5.");
1003 if (rank != globalStaticStrides.size())
1004 return op.emitOpError("strides and sizes must have same rank.");
1005
1006 ArrayRef<int64_t> sharedStaticSizes = op.getSharedStaticSizes();
1007 if (rank != sharedStaticSizes.size())
1008 return op.emitOpError("tensor must have same rank as tile.");
1009
1010 unsigned elementTypeWidth = op.getElementTypeWidth();
1011 if (!llvm::is_contained({8u, 16u, 32u, 64u}, elementTypeWidth))
1012 return op.emitOpError(
1013 "element type width must be 1, 2, 4 or 8 bytes, but was ")
1014 << elementTypeWidth << " bits long";
1015
1016 if (Value atomicBarrierAddress = op.getAtomicBarrierAddress()) {
1017 auto atomicBarrierAddressType =
1018 cast<MemRefType>(atomicBarrierAddress.getType());
1019 bool barrierInLDS =
1020 hasWorkgroupMemorySpace(atomicBarrierAddressType.getMemorySpace());
1021 if (!barrierInLDS)
1022 return op.emitOpError("atomic barrier address must be in LDS.");
1023 }
1024
1025 if (op.getEarlyTimeout() && !op.getWorkgroupMask())
1026 return op.emitOpError(
1027 "early timeout does not apply when workgroup_mask is not set.");
1028 return success();
1029}
1030
1031template <typename DescriptorOp, typename FoldAdaptor>
1032static OpFoldResult foldDescriptorOp(DescriptorOp op, FoldAdaptor adaptor) {
1033 SmallVector<OpFoldResult> mixedGlobalSizes(op.getMixedGlobalSizes());
1034 SmallVector<OpFoldResult> mixedGlobalStrides(op.getMixedGlobalStrides());
1035 SmallVector<OpFoldResult> mixedSharedSizes(op.getMixedSharedSizes());
1036
1037 if (failed(foldDynamicIndexList(mixedGlobalSizes, /*onlyNonNegative=*/true,
1038 /*onlyNonZero=*/true)) &&
1039 failed(foldDynamicIndexList(mixedGlobalStrides, /*onlyNonNegative=*/true,
1040 /*onlyNonZero=*/true)) &&
1041 failed(foldDynamicIndexList(mixedSharedSizes, /*onlyNonNegative=*/true,
1042 /*onlyNonZero=*/true)))
1043 return nullptr;
1044
1045 SmallVector<Value> dynamicGlobalSizes, dynamicGlobalStrides,
1046 dynamicSharedSizes;
1047 SmallVector<int64_t> staticGlobalSizes, staticGlobalStrides,
1048 staticSharedSizes;
1049
1050 dispatchIndexOpFoldResults(mixedGlobalSizes, dynamicGlobalSizes,
1051 staticGlobalSizes);
1052 op.setGlobalStaticSizes(staticGlobalSizes);
1053 op.getGlobalDynamicSizesMutable().assign(dynamicGlobalSizes);
1054
1055 dispatchIndexOpFoldResults(mixedGlobalStrides, dynamicGlobalStrides,
1056 staticGlobalStrides);
1057 op.setGlobalStaticStrides(staticGlobalStrides);
1058 op.getGlobalDynamicStridesMutable().assign(dynamicGlobalStrides);
1059
1060 dispatchIndexOpFoldResults(mixedSharedSizes, dynamicSharedSizes,
1061 staticSharedSizes);
1062 op.setSharedStaticSizes(staticSharedSizes);
1063 op.getSharedDynamicSizesMutable().assign(dynamicSharedSizes);
1064 return op.getResult();
1065}
1066
1067LogicalResult MakeDmaDescriptorOp::verify() {
1068 return verifyDescriptorOp(*this);
1069}
1070
1071OpFoldResult MakeDmaDescriptorOp::fold(FoldAdaptor adaptor) {
1072 return foldDescriptorOp(*this, adaptor);
1073}
1074
1075//===----------------------------------------------------------------------===//
1076// MakeGatherDmaDescriptorOp
1077//===----------------------------------------------------------------------===//
1078
1079LogicalResult MakeGatherDmaDescriptorOp::verify() {
1080 ArrayRef<int64_t> globalStaticSizes = getGlobalStaticSizes();
1081 size_t rank = globalStaticSizes.size();
1082 if (rank > 2)
1083 return emitOpError(
1084 "tensor and tile must be at most of rank two in gather mode.");
1086 Type elementType = cast<VectorType>(indices.getType()).getElementType();
1087 if (elementType != getBase().getType().getIndexType())
1088 return emitOpError("indices' element type must match base's element type.");
1089
1090 return verifyDescriptorOp(*this);
1091}
1092
1093OpFoldResult MakeGatherDmaDescriptorOp::fold(FoldAdaptor adaptor) {
1094 return foldDescriptorOp(*this, adaptor);
1095}
1096
1097//===----------------------------------------------------------------------===//
1098// ScaledMFMAOp
1099//===----------------------------------------------------------------------===//
1100
1101namespace {
1102/// Check if the scales input is used in other scaled mfma's while they exist.
1103/// If theyre unused then pack the scales.
1104struct PackScales final : OpRewritePattern<ScaledMFMAOp> {
1106
1107 LogicalResult matchAndRewrite(ScaledMFMAOp op,
1108 PatternRewriter &rewriter) const override {
1109 Location loc = op.getLoc();
1110 auto setOpsel = [&op](unsigned idx, int64_t val) {
1111 switch (idx) {
1112 case 3:
1113 op.setScalesIdxA(val);
1114 break;
1115 case 4:
1116 op.setScalesIdxB(val);
1117 break;
1118 default:
1119 break;
1120 }
1121 };
1122
1123 // For every scale operand of this ScaledMFMAOp, if the scale is produced by
1124 // the extraction of a single scale from some vector, then attempt to
1125 // extract 4 values from that vector instead.
1126 //
1127 // Example: (f8 here means f8E8M0FNU)
1128 // %unit = vector.extract %ScaleSrc[offsets] : f8 from vector<...>
1129 // %scale = vector.insert %unit, ... : f8 into vector<4xf8>
1130 // amdgpu.scaled_mfma(%scale[0] * ...
1131 //
1132 // rewrite to:
1133 //
1134 // %reshaped = vector.shape_cast %ScaleSrc : vector<...> to vector<?xf8>
1135 // %scale = vector.extract %reshaped[?] : vector<4xf8> from vector<?xf8>
1136 // amdgpu.scaled_mfma(%scale[0-3] * ...
1137 //
1138 // This creates duplicate shape_casts for every use but these will be
1139 // removed in CSE.
1140 for (auto opIdx : std::array<int64_t, 2>({3, 4})) {
1141 auto insertOp = op.getOperand(opIdx).getDefiningOp<vector::InsertOp>();
1142 if (!insertOp) {
1143 return rewriter.notifyMatchFailure(op,
1144 "defining op not a vector.insert");
1145 }
1146 // If the extracted value is not a single scalar, then it has been packed.
1147 if (isa<VectorType>(insertOp.getValueToStore().getType())) {
1148 return rewriter.notifyMatchFailure(
1149 op, "scaled mfma operand already packed");
1150 }
1151
1152 auto extractOp =
1153 insertOp.getValueToStore().getDefiningOp<vector::ExtractOp>();
1154 if (!extractOp) {
1155 return rewriter.notifyMatchFailure(op,
1156 "defining op not a vector.extract");
1157 }
1158
1159 Value scaleSrc = extractOp.getOperand(0);
1160 auto scaleSrcType = dyn_cast<VectorType>(scaleSrc.getType());
1161 if (!scaleSrcType) {
1162 return rewriter.notifyMatchFailure(op, "not a vector type");
1163 }
1164
1165 // We do not handle dynamic dims yet, assume that the input is padded to
1166 // a static shape now.
1167 if (!scaleSrcType.hasStaticShape()) {
1168 return rewriter.notifyMatchFailure(op,
1169 "dynamic dims not yet supported");
1170 }
1171
1172 int64_t numElements = scaleSrcType.getNumElements();
1173 if (numElements <= 4) {
1174 return rewriter.notifyMatchFailure(
1175 op, "no packing if # of scales less than four");
1176 }
1177
1178 // Find a linearized idx using the size and offsets of the extract op.
1179 auto extractedPos = llvm::to_vector_of<int64_t>(
1180 llvm::reverse(extractOp.getStaticPosition()));
1181 ArrayRef<int64_t> scaleSrcShape = scaleSrcType.getShape();
1182 int64_t scaleSrcRank = scaleSrcType.getRank();
1183 SmallVector<int64_t> extractSizes(scaleSrcRank, 1);
1184 for (int64_t i = 1; i < scaleSrcRank; ++i) {
1185 extractSizes[i] = extractSizes[i - 1] * scaleSrcShape[scaleSrcRank - i];
1186 }
1187 int64_t idx = linearize(extractedPos, extractSizes);
1188
1189 // All n scales (where n is the total number of scales) must now be
1190 // extracted in chunks of 4 elements. This is done by dividing the
1191 // original vector of scales into groups of 4 elements
1192 // at offsets 0, 4, ..., m (where m = n/4). All extractions of a
1193 // scale at a particular index are now replaced with an extraction
1194 // of the entire group of 4 elements to which that index belongs.
1195 //
1196 // If the number of scales happens to be indivisible by 4, extract
1197 // the remaining n - m scales in a chunk of 4 elements starting at
1198 // offset n - 4.
1199 int64_t offset = idx - (idx % 4);
1200 int64_t opsel = idx - offset;
1201 int64_t size = 4l;
1202 // Accomdate remaining elements in the case of non-4-divisible vectors.
1203 if (numElements - offset < size) {
1204 opsel = size - (numElements - idx);
1205 offset = numElements - 4l;
1206 }
1207 Type scaleSrcElemType = scaleSrcType.getElementType();
1208 auto newSrcType =
1209 VectorType::get(ArrayRef{numElements}, scaleSrcElemType);
1210 Value newScaleSrc =
1211 vector::ShapeCastOp::create(rewriter, loc, newSrcType, scaleSrc);
1212 auto extract = vector::ExtractStridedSliceOp::create(
1213 rewriter, loc, newScaleSrc, ArrayRef{offset}, ArrayRef{size},
1214 ArrayRef{int64_t(1)});
1215 rewriter.modifyOpInPlace(op, [&] {
1216 op->setOperand(opIdx, extract);
1217 setOpsel(opIdx, opsel);
1218 });
1219 }
1220 return success();
1221 }
1222};
1223} // namespace
1224
1225void ScaledMFMAOp::getCanonicalizationPatterns(RewritePatternSet &results,
1226 MLIRContext *context) {
1227 results.add<PackScales>(context);
1228}
1229
1230#include "mlir/Dialect/AMDGPU/IR/AMDGPUEnums.cpp.inc"
1231
1232#define GET_ATTRDEF_CLASSES
1233#include "mlir/Dialect/AMDGPU/IR/AMDGPUAttributes.cpp.inc"
1234
1235#define GET_TYPEDEF_CLASSES
1236#include "mlir/Dialect/AMDGPU/IR/AMDGPUTypes.cpp.inc"
1237
1238#define GET_OP_CLASSES
1239#include "mlir/Dialect/AMDGPU/IR/AMDGPU.cpp.inc"
static LogicalResult verifyDescriptorOp(DescriptorOp op)
static LogicalResult verifyRawBufferOp(T &op)
static OpFoldResult foldDescriptorOp(DescriptorOp op, FoldAdaptor adaptor)
static bool hasGlobalMemorySpace(Attribute memorySpace)
static bool hasWorkgroupMemorySpace(Attribute memorySpace)
static FailureOr< MemRefType > getFatRawBufferTypeLike(MemRefType source, bool resetOffset)
Convert the type source to one with the same sizes and strides - and offset, unless stripOffset is tr...
static bool hasFatRawBufferMemorySpace(Attribute memorySpace)
static LogicalResult verifyBase(BaseOp op)
static bool staticallyOutOfBounds(OpType op)
static std::optional< uint32_t > getConstantUint32(Value v)
return success()
p<< " : "<< getMemRefType()<< ", "<< getType();}static LogicalResult verifyVectorMemoryOp(Operation *op, MemRefType memrefType, VectorType vectorType) { if(memrefType.getElementType() !=vectorType.getElementType()) return op-> emitOpError("requires memref and vector types of the same elemental type")
Given a list of lists of parsed operands, populates uniqueOperands with unique operands.
static Value getBase(Value v)
Looks through known "view-like" ops to find the base memref.
lhs
static bool isLegalToInline(InlinerInterface &interface, Region *src, Region *insertRegion, bool shouldCloneInlinedRegion, IRMapping &valueMapping)
Utility to check that all of the operations within 'src' can be inlined.
b
Return true if permutation is a valid permutation of the outer_dims_perm (case OuterOrInnerPerm::Oute...
b getContext())
static AffineMap getMultiDimIdentityMap(unsigned numDims, MLIRContext *context)
Returns an AffineMap with 'numDims' identity result dim exprs.
virtual Builder & getBuilder() const =0
Return a builder which provides useful access to MLIRContext, global objects like types and attribute...
virtual InFlightDiagnostic emitError(SMLoc loc, const Twine &message={})=0
Emit a diagnostic at the specified location and return failure.
virtual ParseResult parseDimensionList(SmallVectorImpl< int64_t > &dimensions, bool allowDynamic=true, bool withTrailingX=true)=0
Parse a dimension list of a tensor or memref type.
virtual SMLoc getCurrentLocation()=0
Get the location of the next token and store it into the argument.
Attributes are known-constant values of operations.
Definition Attributes.h:25
This class is a general helper class for creating context-global objects like types,...
Definition Builders.h:51
IntegerAttr getIndexAttr(int64_t value)
Definition Builders.cpp:108
IntegerAttr getI32IntegerAttr(int32_t value)
Definition Builders.cpp:200
TypedAttr getZeroAttr(Type type)
Definition Builders.cpp:324
This class represents a diagnostic that is inflight and set to be reported.
MLIRContext is the top-level object for a collection of MLIR operations.
Definition MLIRContext.h:63
This is a builder type that keeps local references to arguments.
Builder & setMemorySpace(Attribute newMemorySpace)
Builder & setLayout(MemRefLayoutAttrInterface newLayout)
The OpAsmParser has methods for interacting with the asm parser: parsing things from it,...
This class helps build Operations.
Definition Builders.h:207
void createOrFold(SmallVectorImpl< Value > &results, Location location, Args &&...args)
Create an operation of specific op type at the current insertion point, and immediately try to fold i...
Definition Builders.h:526
This class represents a single result from folding an operation.
Simple wrapper around a void* in order to express generically how to pass in op properties through AP...
This class provides an abstraction over the different types of ranges over Regions.
Definition Region.h:346
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments.
virtual void eraseOp(Operation *op)
This method erases an operation that is known to have no uses.
std::enable_if_t<!std::is_convertible< CallbackT, Twine >::value, LogicalResult > notifyMatchFailure(Location loc, CallbackT &&reasonCallback)
Used to notify the listener that the IR failed to be rewritten because of a match failure,...
void modifyOpInPlace(Operation *root, CallableT &&callable)
This method is a utility wrapper around an in-place modification of an operation.
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition Types.h:74
bool isFloat() const
Return true if this is an float type (with the specified width).
Definition Types.cpp:45
MLIRContext * getContext() const
Return the MLIRContext in which this type was uniqued.
Definition Types.cpp:35
bool isInteger() const
Return true if this is an integer type (with the specified width).
Definition Types.cpp:56
unsigned getIntOrFloatBitWidth() const
Return the bit width of an integer or a float type, assert failure on other types.
Definition Types.cpp:122
This class provides an abstraction over the different types of ranges over Values.
Definition ValueRange.h:387
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition Value.h:96
Type getType() const
Return the type of this value.
Definition Value.h:105
ParseResult parseMNKDimensionList(OpAsmParser &parser, IntegerAttr &m, IntegerAttr &n, IntegerAttr &k)
Parser for the custom<MNKDimensionList> custom assembly format used by WMMAOp.
Operation::operand_range getIndices(Operation *op)
Get the indices that the given load/store operation is operating on.
Definition Utils.cpp:18
detail::InFlightRemark failed(Location loc, RemarkOpts opts)
Report an optimization remark that failed.
Definition Remarks.h:573
uint64_t getN(LevelType lt)
Definition Enums.h:442
uint64_t getM(LevelType lt)
Definition Enums.h:443
SmallVector< unsigned > getBlockSize(AffineMap dimToLvl)
Given the dimToLvl map, returns the block sizes in a vector.
Include the generated interface declarations.
bool matchPattern(Value value, const Pattern &pattern)
Entry point for matching a pattern over a Value.
Definition Matchers.h:490
detail::constant_int_value_binder m_ConstantInt(IntegerAttr::ValueType *bind_value)
Matches a constant holding a scalar/vector/tensor integer (splat) and writes the integer value to bin...
Definition Matchers.h:527
Type getType(OpFoldResult ofr)
Returns the int type of the integer in ofr.
Definition Utils.cpp:304
InFlightDiagnostic emitError(Location loc)
Utility method to emit an error message using this location.
void dispatchIndexOpFoldResults(ArrayRef< OpFoldResult > ofrs, SmallVectorImpl< Value > &dynamicVec, SmallVectorImpl< int64_t > &staticVec)
Helper function to dispatch multiple OpFoldResults according to the behavior of dispatchIndexOpFoldRe...
SmallVector< int64_t > computeSuffixProduct(ArrayRef< int64_t > sizes)
Given a set of sizes, return the suffix product.
int64_t linearize(ArrayRef< int64_t > offsets, ArrayRef< int64_t > basis)
Return the linearized index of 'offsets' w.r.t.
llvm::function_ref< Fn > function_ref
Definition LLVM.h:152
LogicalResult foldDynamicIndexList(SmallVectorImpl< OpFoldResult > &ofrs, bool onlyNonNegative=false, bool onlyNonZero=false)
Returns "success" when any of the elements in ofrs is a constant value.
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
OpRewritePattern(MLIRContext *context, PatternBenefit benefit=1, ArrayRef< StringRef > generatedNames={})
Patterns must specify the root operation name they match against, and can also specify the benefit of...