AMDGPUOps.cpp

Go to the documentation of this file.

//===- AMDGPUOps.cpp - MLIR AMDGPU dialect operations ---------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file implements the AMDGPU dialect operations, their verifiers, and
// their canonicalizations.
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
 
#include "mlir/Dialect/AMDGPU/Utils/MemorySpaceUtils.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/TypeUtilities.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallVector.h"
 
#include <algorithm>
#include <cstdint>
#include <limits>
#include <optional>
 
using namespace mlir;
using namespace mlir::amdgpu;
 
/// Verifies that the number of indices matches the rank of the indexed memref,
/// emitting an op error mentioning `indexName` on mismatch.
template <typename OpTy>
static LogicalResult verifyIndexCount(OpTy op, StringRef indexName,
                                      MemRefType memrefType,
                                      int64_t numIndices) {
  int64_t rank = memrefType.getRank();
  if (rank != numIndices)
    return op.emitOpError("expected ")
           << rank << " " << indexName << " indices, got " << numIndices;
  return success();
}
 
//===----------------------------------------------------------------------===//
// 8-bit float ops
//===----------------------------------------------------------------------===//
LogicalResult PackedTrunc2xFp8Op::verify() {
  if (getExisting() && getExisting().getType() != getResult().getType())
    return emitOpError("existing values must have same type as result");
  return success();
}
 
LogicalResult PackedStochRoundFp8Op::verify() {
  if (getExisting() && getExisting().getType() != getResult().getType())
    return emitOpError("existing values must have same type as result");
  return success();
}
 
//===----------------------------------------------------------------------===//
// mxfp float ops
//===----------------------------------------------------------------------===//
LogicalResult PackedScaledTruncOp::verify() {
  if (getExisting() && getExisting().getType() != getResult().getType())
    return emitOpError("existing values must have same type as result");
  return success();
}
 
//===----------------------------------------------------------------------===//
// FatRawBufferCastOp
//===----------------------------------------------------------------------===//
 
/// Convert the type `source` to one with the same sizes and strides - and
/// offset, unless `stripOffset` is true, in which case the offset is reset to
/// 0, if the offset should be reset but the layout of `source` isn't either the
/// identity layout or a strided layout, this function fails.
static FailureOr<MemRefType> getFatRawBufferTypeLike(MemRefType source,
                                                     bool resetOffset) {
  MLIRContext *ctx = source.getContext();
  MemRefType::Builder mb(source);
  mb.setMemorySpace(
      amdgpu::AddressSpaceAttr::get(ctx, amdgpu::AddressSpace::FatRawBuffer));
  MemRefLayoutAttrInterface layout = source.getLayout();
  if (resetOffset && !layout.isIdentity()) {
    auto stridedLayout = dyn_cast<StridedLayoutAttr>(layout);
    if (!stridedLayout)
      return failure();
    MemRefLayoutAttrInterface newLayout =
        StridedLayoutAttr::get(ctx, 0, stridedLayout.getStrides());
    // Special case: if resetting the offset causes the strided layout to become
    // the identity layout, then reset to the identity layout.
    // TODO: this'll get a lot simpler when we have the contiguous layout.
    SmallVector<int64_t> stridesIfIdentity;
    if (source.hasStaticShape()) {
      stridesIfIdentity = computeSuffixProduct(source.getShape());
    } else if (source.getRank() <= 1) {
      stridesIfIdentity = SmallVector<int64_t>(source.getRank(), 1);
    }
    if (stridesIfIdentity == stridedLayout.getStrides()) {
      newLayout = AffineMapAttr::get(
          AffineMap::getMultiDimIdentityMap(source.getRank(), ctx));
    }
    mb.setLayout(newLayout);
  }
  return (MemRefType)(mb);
}
 
LogicalResult FatRawBufferCastOp::inferReturnTypes(
    MLIRContext *context, std::optional<Location> location, ValueRange operands,
    DictionaryAttr attributes, PropertyRef properties, RegionRange regions,
    SmallVectorImpl<Type> &inferredReturnTypes) {
  Adaptor adaptor(operands, attributes, properties, regions);
  auto sourceType =
      dyn_cast_if_present<MemRefType>(adaptor.getSource().getType());
  if (!sourceType)
    return failure();
  FailureOr<MemRefType> resultType =
      getFatRawBufferTypeLike(sourceType, adaptor.getResetOffset());
  if (failed(resultType))
    return failure();
  inferredReturnTypes = SmallVector<Type>{*resultType};
  return success();
}
 
FailureOr<OpFoldResult> FatRawBufferCastOp::reifyDimOfResult(OpBuilder &builder,
                                                             int resultIndex,
                                                             int dim) {
  assert(resultIndex == 0 && "FatRawBufferCastOp has a single result");
  return memref::getMixedSize(builder, getLoc(), getSource(), dim);
}
 
LogicalResult FatRawBufferCastOp::verify() {
  FailureOr<MemRefType> expectedResultType =
      getFatRawBufferTypeLike(getSource().getType(), getResetOffset());
  if (failed(expectedResultType))
    return emitOpError("source type ")
           << getSource().getType() << " can't have its offset reset";
  if (getResult().getType() != *expectedResultType)
    return emitOpError("expected result type to be ")
           << *expectedResultType << " but got " << getResult().getType();
  return success();
}
 
//===----------------------------------------------------------------------===//
// RawBuffer*Op
//===----------------------------------------------------------------------===//
template <typename T>
static LogicalResult verifyRawBufferOp(T &op) {
  MemRefType bufferType = llvm::cast<MemRefType>(op.getMemref().getType());
  bool isGlobal =
      isGlobalMemorySpace(bufferType.getMemorySpace(), /*allowFlat=*/true);
 
  if (!isGlobal)
    return op.emitOpError(
        "buffer ops must operate on a memref in global memory");
  if (!bufferType.hasRank())
    return op.emitOpError(
        "cannot meaningfully buffer_store to an unranked memref");
  return verifyIndexCount(op, "buffer", bufferType, op.getIndices().size());
}
 
LogicalResult RawBufferLoadOp::verify() { return verifyRawBufferOp(*this); }
 
LogicalResult RawBufferStoreOp::verify() { return verifyRawBufferOp(*this); }
 
LogicalResult RawBufferAtomicFaddOp::verify() {
  return verifyRawBufferOp(*this);
}
 
LogicalResult RawBufferAtomicFmaxOp::verify() {
  return verifyRawBufferOp(*this);
}
 
LogicalResult RawBufferAtomicSmaxOp::verify() {
  return verifyRawBufferOp(*this);
}
 
LogicalResult RawBufferAtomicUminOp::verify() {
  return verifyRawBufferOp(*this);
}
 
LogicalResult RawBufferAtomicCmpswapOp::verify() {
  return verifyRawBufferOp(*this);
}
 
static std::optional<uint32_t> getConstantUint32(Value v) {
  APInt cst;
  if (!v.getType().isInteger(32))
    return std::nullopt;
  if (matchPattern(v, m_ConstantInt(&cst)))
    return cst.getZExtValue();
  return std::nullopt;
}
 
template <typename OpType>
static bool staticallyOutOfBounds(OpType op) {
  if (!op.getBoundsCheck())
    return false;
  MemRefType bufferType = op.getMemref().getType();
  if (!bufferType.hasStaticShape())
    return false;
  int64_t offset;
  SmallVector<int64_t> strides;
  if (failed(bufferType.getStridesAndOffset(strides, offset)))
    return false;
  int64_t result = offset + op.getIndexOffset().value_or(0);
  if (op.getSgprOffset()) {
    std::optional<uint32_t> sgprOffset = getConstantUint32(op.getSgprOffset());
    if (!sgprOffset)
      return false;
    result += *sgprOffset;
  }
  if (strides.size() != op.getIndices().size())
    return false;
  int64_t indexVal = 0;
  for (auto pair : llvm::zip(strides, op.getIndices())) {
    int64_t stride = std::get<0>(pair);
    Value idx = std::get<1>(pair);
    std::optional<uint32_t> idxVal = getConstantUint32(idx);
    if (!idxVal)
      return false;
    indexVal += stride * *idxVal;
  }
  result += indexVal;
  if (result > std::numeric_limits<uint32_t>::max())
    // Overflow means don't drop
    return false;
  return result >= bufferType.getNumElements();
}
 
namespace {
template <typename OpType>
struct RemoveStaticallyOobBufferLoads final : public OpRewritePattern<OpType> {
  using OpRewritePattern<OpType>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(OpType op, PatternRewriter &rw) const override {
    if (!staticallyOutOfBounds(op))
      return failure();
    Type loadType = op.getResult().getType();
    rw.replaceOpWithNewOp<arith::ConstantOp>(op, loadType,
                                             rw.getZeroAttr(loadType));
    return success();
  }
};
 
template <typename OpType>
struct RemoveStaticallyOobBufferWrites final : public OpRewritePattern<OpType> {
  using OpRewritePattern<OpType>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(OpType op, PatternRewriter &rw) const override {
    if (!staticallyOutOfBounds(op))
      return failure();
 
    rw.eraseOp(op);
    return success();
  }
};
} // end namespace
 
void RawBufferLoadOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                  MLIRContext *context) {
  results.add<RemoveStaticallyOobBufferLoads<RawBufferLoadOp>>(context);
}
 
void RawBufferStoreOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                   MLIRContext *context) {
  results.add<RemoveStaticallyOobBufferWrites<RawBufferStoreOp>>(context);
}
 
void RawBufferAtomicFaddOp::getCanonicalizationPatterns(
    RewritePatternSet &results, MLIRContext *context) {
  results.add<RemoveStaticallyOobBufferLoads<RawBufferAtomicFaddOp>>(context);
}
 
void RawBufferAtomicFmaxOp::getCanonicalizationPatterns(
    RewritePatternSet &results, MLIRContext *context) {
  results.add<RemoveStaticallyOobBufferLoads<RawBufferAtomicFmaxOp>>(context);
}
 
void RawBufferAtomicSmaxOp::getCanonicalizationPatterns(
    RewritePatternSet &results, MLIRContext *context) {
  results.add<RemoveStaticallyOobBufferLoads<RawBufferAtomicSmaxOp>>(context);
}
 
void RawBufferAtomicUminOp::getCanonicalizationPatterns(
    RewritePatternSet &results, MLIRContext *context) {
  results.add<RemoveStaticallyOobBufferLoads<RawBufferAtomicUminOp>>(context);
}
 
void RawBufferAtomicCmpswapOp::getCanonicalizationPatterns(
    RewritePatternSet &results, MLIRContext *context) {
  results.add<RemoveStaticallyOobBufferLoads<RawBufferAtomicCmpswapOp>>(
      context);
}
 
//===----------------------------------------------------------------------===//
// ScaledExtPackedMatrixOp
//===----------------------------------------------------------------------===//
LogicalResult ScaledExtPackedMatrixOp::verify() {
  int blockSize = getBlockSize();
  assert(llvm::is_contained({16, 32}, blockSize) && "invalid block size");
 
  int firstScaleByte = getFirstScaleByte();
  int firstScaleLane = getFirstScaleLane();
  auto sourceType = cast<VectorType>(getSource().getType());
  Type elementType = sourceType.getElementType();
  auto floatType = cast<FloatType>(elementType);
  unsigned bitWidth = floatType.getWidth();
 
  assert(llvm::is_contained(llvm::ArrayRef<unsigned>{4, 6, 8}, bitWidth));
 
  const bool is_fp8 = bitWidth == 8;
  const bool is_block_16 = blockSize == 16;
 
  if (!is_fp8) {
    if (is_block_16) {
      if (!llvm::is_contained({0, 1}, firstScaleByte)) {
        return emitOpError("blockSize of 16 can only have firstScaleByte be 0 "
                           "or 1 for f4 and f6.");
      }
    } else {
      if (!llvm::is_contained({0, 2}, firstScaleByte)) {
        return emitOpError("blockSize of 32 can only have firstScaleByte be 0 "
                           "or 2 for f4 and f6.");
      }
    }
  } else {
    if (is_block_16) {
      bool is_valid = ((firstScaleLane == 0) && (firstScaleByte == 0)) ||
                      ((firstScaleLane == 16) && (firstScaleByte == 2));
      if (!is_valid) {
        return emitOpError("blockSize of 16 can only have (firstScaleLane, "
                           "firstScaleByte) be (0, 0) or (16, 2) for f8.");
      }
    }
  }
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// WMMAOp
//===----------------------------------------------------------------------===//
 
ParseResult mlir::amdgpu::parseMNKDimensionList(OpAsmParser &parser,
                                                IntegerAttr &m, IntegerAttr &n,
                                                IntegerAttr &k) {
  SmallVector<int64_t, 3> dimensions;
  if (parser.parseDimensionList(dimensions, false, false))
    return failure();
  if (dimensions.size() != 3)
    return parser.emitError(parser.getCurrentLocation())
           << "expected 3 dimensions in MNK dimension list";
 
  m = parser.getBuilder().getI32IntegerAttr(dimensions[0]);
  n = parser.getBuilder().getI32IntegerAttr(dimensions[1]);
  k = parser.getBuilder().getI32IntegerAttr(dimensions[2]);
  return success();
}
 
LogicalResult WMMAOp::verify() {
  auto sourceAType = cast<VectorType>(getSourceA().getType());
  auto sourceBType = cast<VectorType>(getSourceB().getType());
  auto destType = cast<VectorType>(getDestC().getType());
 
  Type sourceAElemType = sourceAType.getElementType();
  Type sourceBElemType = sourceBType.getElementType();
  if (sourceAType.getNumElements() != sourceBType.getNumElements()) {
    return emitOpError("source vectors have different lengths: ")
           << sourceAType << " vs. " << sourceBType;
  }
 
  bool isDestFloat = destType.getElementType().isFloat();
  bool isSrcFloat = sourceAElemType.isFloat();
 
  if (isDestFloat && !isSrcFloat)
    return emitOpError("expected float sources with float destination");
  if (!isDestFloat && isSrcFloat)
    return emitOpError("expected int sources with int destination");
 
  if (!sourceAElemType.isFloat(8) && sourceAElemType != sourceBElemType) {
    return emitOpError(
               "source element types must match (except for fp8/bf8) but have ")
           << sourceAType << " and " << sourceBType;
  }
 
  if (isSrcFloat) {
    if (getClamp())
      return emitOpError("clamp flag is not supported for float types");
    if (getUnsignedA() || getUnsignedB())
      return emitOpError("unsigned flags are not supported for float types");
  }
  return success();
}
 
//===----------------------------------------------------------------------===//
// ScaledWMMAOp
//===----------------------------------------------------------------------===//
 
LogicalResult ScaledWMMAOp::verify() {
  // Helper functions for type classification.
  auto isF8 = llvm::IsaPred<Float8E4M3FNType, Float8E5M2Type>;
  auto isF6 = llvm::IsaPred<Float6E2M3FNType, Float6E3M2FNType>;
  auto isF4 = llvm::IsaPred<Float4E2M1FNType>;
  auto isScaleF8 = llvm::IsaPred<Float8E8M0FNUType, Float8E4M3FNType>;
  auto isE8M0 = llvm::IsaPred<Float8E8M0FNUType>;
  auto isE4M3 = llvm::IsaPred<Float8E4M3FNType>;
 
  auto sourceAType = cast<VectorType>(getSourceA().getType());
  auto sourceBType = cast<VectorType>(getSourceB().getType());
  auto destType = cast<VectorType>(getDestC().getType());
 
  // Validate source element types are small floats (fp4/fp6/fp8).
  Type aElemType = sourceAType.getElementType();
  Type bElemType = sourceBType.getElementType();
 
  // Validate vector lengths based on dimensions.
  int64_t m = getM();
  int64_t aLen = sourceAType.getNumElements();
  int64_t bLen = sourceBType.getNumElements();
  int64_t expectedOutLen = (m == 16) ? 8 : 16;
 
  if (destType.getNumElements() != expectedOutLen)
    return emitOpError("expected output vector of length ")
           << expectedOutLen << " but got " << destType.getNumElements();
 
  if (m == 16) {
    // For 16×16×128: both A and B must be 64 elements.
    if (aLen != 64)
      return emitOpError(
                 "for 16x16x128, sourceA must have 64 elements but got ")
             << aLen;
    if (bLen != 64)
      return emitOpError(
                 "for 16x16x128, sourceB must have 64 elements but got ")
             << bLen;
  } else { // m == 32
    // For 32×16×128: only fp4 is supported, A is 128, B is 64.
    if (!isF4(aElemType) && !isF4(bElemType))
      return emitOpError("32x16x128 only supports fp4 element types");
 
    if (aLen != 128)
      return emitOpError(
                 "for 32x16x128, sourceA must have 128 elements but got ")
             << aLen;
    if (bLen != 64)
      return emitOpError(
                 "for 32x16x128, sourceB must have 64 elements but got ")
             << bLen;
 
    // For 32x16x128, matrix A uses all 32 lanes so a_first_scale_lane must be
    // 0.
    if (getAFirstScaleLane() != 0)
      return emitOpError("for 32x16x128, a_first_scale_lane must be 0");
  }
 
  // Validate scale types and their compatibility with matrix element types.
  auto scaleAType = cast<VectorType>(getScaleA().getType());
  auto scaleBType = cast<VectorType>(getScaleB().getType());
  Type scaleAElemType = scaleAType.getElementType();
  Type scaleBElemType = scaleBType.getElementType();
 
  // Validate scale element types are valid scale f8 types (E8M0FNU or E4M3FN).
  if (!isScaleF8(scaleAElemType) || !isScaleF8(scaleBElemType))
    return emitOpError(
        "scale operands must have f8 element types (E8M0FNU or E4M3FN)");
 
  // Any matrices A/B (fp8|fp6|fp4) with E8M0 scales for matrix A/B are valid.
  if (isE8M0(scaleAElemType) && isE8M0(scaleBElemType))
    return success();
 
  // Matrix A (F8|F6) x Matrix B (F4) with Scale A (E8M0), Scale B (E5M3|E4M3).
  if ((isF8(aElemType) || isF6(aElemType)) && isE8M0(scaleAElemType) &&
      isF4(bElemType) && isE4M3(scaleBElemType))
    return success();
 
  // Matrix A (F4) x Matrix B (F8|F6) with Scale A (E5M3|E4M3), Scale B (E8M0).
  if (isF4(aElemType) && isE4M3(scaleAElemType) &&
      (isF8(bElemType) || isF6(bElemType)) && isE8M0(scaleBElemType))
    return success();
 
  // Matrix A (F4) x Matrix B (F4) with Scale A (E4M3), Scale B (E4M3).
  if (isF4(aElemType) && isF4(bElemType) && isE4M3(scaleAElemType) &&
      isE4M3(scaleBElemType))
    return success();
 
  // No valid combination matched.
  return emitOpError("invalid combination of matrix and scale types: ")
         << "sourceA=" << aElemType << ", scaleA=" << scaleAElemType
         << ", sourceB=" << bElemType << ", scaleB=" << scaleBElemType;
}
 
//===----------------------------------------------------------------------===//
// MFMAOp
//===----------------------------------------------------------------------===//
LogicalResult MFMAOp::verify() {
  constexpr uint32_t waveSize = 64;
  Builder b(getContext());
 
  Type sourceType = getSourceA().getType();
  Type destType = getDestC().getType();
 
  Type sourceElem = sourceType, destElem = destType;
  uint32_t sourceLen = 1, destLen = 1;
  if (auto sourceVector = dyn_cast<VectorType>(sourceType)) {
    sourceLen = sourceVector.getNumElements();
    sourceElem = sourceVector.getElementType();
  }
  if (auto destVector = dyn_cast<VectorType>(destType)) {
    destLen = destVector.getNumElements();
    destElem = destVector.getElementType();
  }
 
  Type sourceBType = getSourceB().getType();
  if (sourceElem.isFloat(8) || sourceElem.isFloat(6) || sourceElem.isFloat(4)) {
    int64_t sourceBLen = 1;
    Type sourceBElem = sourceBType;
    if (auto sourceBVector = llvm::dyn_cast<VectorType>(sourceBType)) {
      sourceBLen = sourceBVector.getNumElements();
      sourceBElem = sourceBVector.getElementType();
    }
    if (!sourceBElem.isFloat(8) && !sourceBElem.isFloat(6) &&
        !sourceBElem.isFloat(4))
      return emitOpError("expected both source operands to have small-float "
                         "elements if one does");
    if (sourceLen != sourceBLen)
      return emitOpError(
          "expected both small-float source vectors to have the same length");
  } else {
    if (sourceType != sourceBType)
      return emitOpError("expected both non-small-float source operand types "
                         "to match exactly");
  }
  // Normalize the wider integer types the compiler expects to i8.
  if (sourceElem.isInteger(32)) {
    sourceLen *= 4;
    sourceElem = b.getI8Type();
  }
  if (sourceElem.isInteger(64)) {
    sourceLen *= 8;
    sourceElem = b.getI8Type();
  }
 
  int64_t numSourceElems = (getM() * getK() * getBlocks()) / waveSize;
  if (sourceLen != numSourceElems)
    return emitOpError("expected " + Twine(numSourceElems) +
                       " source values for this operation but got " +
                       Twine(sourceLen));
 
  int64_t numDestElems = (getM() * getN() * getBlocks()) / waveSize;
  if (destLen != numDestElems)
    return emitOpError("expected " + Twine(numDestElems) +
                       " result values for this operation but got " +
                       Twine(destLen));
 
  if (destElem.isF64() && getBlgp() != MFMAPermB::none)
    return emitOpError(
        "double-precision ops do not support permuting lanes of B");
  if (destElem.isF64() && getCbsz() != 0)
    return emitOpError(
        "double-precision ops do not support permuting lanes of A");
  if (getAbid() >= (1u << getCbsz()))
    return emitOpError(
        "block ID for permuting A (abid) must be below 2 ** cbsz");
 
  if ((getNegateA() || getNegateB() || getNegateC()) && !destElem.isF64())
    return emitOpError(
        "negation flags only available for double-precision operations");
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// SparseMFMAOp
//===----------------------------------------------------------------------===//
 
LogicalResult SparseMFMAOp::verify() {
  constexpr uint32_t waveSize = 64;
 
  auto sparseType = cast<VectorType>(getSourceA().getType());
  auto denseType = cast<VectorType>(getSourceB().getType());
  auto destType = cast<VectorType>(getDestC().getType());
 
  Type sparseElem = sparseType.getElementType();
  Type denseElem = denseType.getElementType();
  int64_t sparseLen = sparseType.getNumElements();
  int64_t denseLen = denseType.getNumElements();
  int64_t destLen = destType.getNumElements();
 
  if (denseLen != 2 * sparseLen)
    return emitOpError("expected dense source operand to have exactly double "
                       "the number of elements of the sparse source operand");
 
  // Check that source element types are compatible.
  // For fp8/bf8 mixed operations, element types can differ (e.g., fp8 * bf8).
  // For other types, element types must match exactly.
  bool bothFloat8 = sparseElem.isFloat(8) && denseElem.isFloat(8);
  if (!bothFloat8 && sparseElem != denseElem)
    return emitOpError(
        "expected source operands to have the same element type");
 
  // Classify the sparse MFMA variant. The three flavors differ in CBSZ/ABID
  // handling and in the sparse-index layout:
  //   - gfx942 16-bit:              max ABID = 3, sparse idx = vector<4xi8>
  //   - gfx950 16-bit / gfx942 8-bit: max ABID = 1, sparse idx = vector<2xi16>
  //   - gfx950 8-bit:  CBSZ/ABID ignored by hw,   sparse idx = i32
  uint32_t m = getM(), k = getK();
  bool is8BitSource = sparseElem.isFloat(8) || sparseElem.isInteger(8);
  bool is16BitGfx942 =
      !is8BitSource && ((m == 16 && k == 32) || (m == 32 && k == 16));
  bool is8BitGfx950 =
      is8BitSource && ((m == 16 && k == 128) || (m == 32 && k == 64));
 
  // CBSZ/ABID range check. On gfx950 8-bit the hardware always uses the first
  // set and ignores these fields, so require zeros in IR. Otherwise ABID is
  // only meaningful when CBSZ == 0 (when CBSZ != 0 the first set is always
  // used and ABID is irrelevant, so the verifier accepts any value).
  if (is8BitGfx950) {
    if (getCbsz() != 0)
      return emitOpError(
          "CBSZ must be 0 for this variant (field is ignored by hardware)");
    if (getAbid() != 0)
      return emitOpError(
          "ABID must be 0 for this variant (field is ignored by hardware)");
  } else if (getCbsz() == 0) {
    unsigned maxAbid = is16BitGfx942 ? 3u : 1u;
    if (getAbid() > maxAbid)
      return emitOpError("ABID must be in [0, ")
             << maxAbid << "] for this variant";
  }
 
  Type sparseIdxType = getSparseIdx().getType();
  if (is8BitGfx950) {
    if (!sparseIdxType.isInteger(32))
      return emitOpError("expected i32 sparse indices for this variant "
                         "(no internal set structure), but got ")
             << sparseIdxType;
  } else {
    unsigned expectedIdxElems = is16BitGfx942 ? 4 : 2;
    unsigned expectedIdxBits = is16BitGfx942 ? 8 : 16;
    auto vecType = dyn_cast<VectorType>(sparseIdxType);
    if (!vecType || vecType.getNumElements() != expectedIdxElems ||
        !vecType.getElementType().isInteger(expectedIdxBits))
      return emitOpError("expected vector<")
             << expectedIdxElems << "xi" << expectedIdxBits
             << "> sparse indices for this variant, but got " << sparseIdxType;
  }
 
  int64_t expectedSourceElems = (getM() * getK()) / waveSize;
  if (denseLen != expectedSourceElems)
    return emitOpError("expected " + Twine(expectedSourceElems) +
                       " source values for this operation but got " +
                       Twine(denseLen));
 
  int64_t expectedDestElems = (getM() * getN()) / waveSize;
  if (destLen != expectedDestElems)
    return emitOpError("expected " + Twine(expectedDestElems) +
                       " result values for this operation but got " +
                       Twine(destLen));
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// SparseWMMAOp
//===----------------------------------------------------------------------===//
 
LogicalResult SparseWMMAOp::verify() {
  auto sparseType = cast<VectorType>(getSourceA().getType());
  auto denseType = cast<VectorType>(getSourceB().getType());
  auto destType = cast<VectorType>(getDestC().getType());
 
  Type sparseElem = sparseType.getElementType();
  Type denseElem = denseType.getElementType();
  Type destElem = destType.getElementType();
  int64_t sparseLen = sparseType.getNumElements();
  int64_t denseLen = denseType.getNumElements();
  int64_t destLen = destType.getNumElements();
 
  uint32_t m = getM(), n = getN(), k = getK();
  if ((m != 16) || (n != 16))
    return emitOpError("expected MxN to be exactly 16x16");
 
  const bool isWavesize64 = getWave64();
  const bool isInt4Input = sparseElem.isInteger(4) && denseElem.isInteger(4);
  const bool isEqualLengthAllowed = isWavesize64 && isInt4Input && k == 32;
 
  if ((denseLen != 2 * sparseLen) && !isEqualLengthAllowed)
    return emitOpError("expected dense source operand to have exactly double "
                       "the number of elements of the sparse source operand");
 
  if (isEqualLengthAllowed && (denseLen != sparseLen))
    return emitOpError("expected dense source operand to have exactly the "
                       "same the number of elements");
 
  if (destElem.isInteger()) {
    if (!(sparseElem.isInteger() && denseElem.isInteger())) {
      return emitOpError("source operand and destination operands must all be "
                         "either integer or float types");
    }
  }
 
  if (destElem.isFloat()) {
    if (!(sparseElem.isFloat() && denseElem.isFloat())) {
      return emitOpError("source operand and destination operands must all be "
                         "either integer or float types");
    }
  }
 
  // Check that source element types are compatible.
  // For fp8/bf8 mixed operations, element types can differ (e.g., fp8 * bf8).
  // For other types, element types must match exactly.
  bool bothFloat8 = sparseElem.isFloat(8) && denseElem.isFloat(8);
  if (!bothFloat8 && sparseElem != denseElem)
    return emitOpError(
        "expected source operands to have the same element type");
 
  const int64_t waveSize = isWavesize64 ? 64 : 32;
 
  int64_t expectedSourceElems = (getM() * getK()) / waveSize;
  if (denseLen != expectedSourceElems)
    return emitOpError("expected " + Twine(expectedSourceElems) +
                       " source values for this operation but got " +
                       Twine(denseLen));
 
  int64_t expectedDestElems = (getM() * getN()) / waveSize;
  if (destLen != expectedDestElems)
    return emitOpError("expected " + Twine(expectedDestElems) +
                       " result values for this operation but got " +
                       Twine(destLen));
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// DotOp
//===----------------------------------------------------------------------===//
LogicalResult DotOp::verify() {
  Type aElem = cast<VectorType>(getSourceA().getType()).getElementType();
  Type bElem = cast<VectorType>(getSourceB().getType()).getElementType();
  Type dest = getDestC().getType();
 
  bool aIsFloat8 = aElem.isFloat(8);
  bool bIsFloat8 = bElem.isFloat(8);
  bool aIsInteger = isa<IntegerType>(aElem);
 
  bool bothFloat8 = aIsFloat8 && bIsFloat8;
  if (!bothFloat8 && aElem != bElem)
    return emitOpError(
        "expected source operands to have the same element type");
 
  if (aElem.isF16()) {
    if (!dest.isF32() && !dest.isF16())
      return emitOpError("expected f32 or f16 accumulator for f16 sources");
  } else if (aElem.isBF16()) {
    if (!dest.isF32() && !dest.isBF16())
      return emitOpError("expected f32 or bf16 accumulator for bf16 sources");
  } else if (aIsInteger) {
    if (!dest.isInteger(32))
      return emitOpError("expected i32 accumulator for integer sources");
  } else if (aIsFloat8) {
    if (!dest.isF32())
      return emitOpError("expected f32 accumulator for fp8 sources");
  }
 
  if ((getUnsignedA() || getUnsignedB()) && !aIsInteger)
    return emitOpError(
        "unsignedA/unsignedB are only valid for integer source types");
 
  if (aElem.isInteger(16) && getUnsignedA() != getUnsignedB())
    return emitOpError(
        "mixed-sign dot is not supported for 16-bit integer sources");
 
  if (getClamp()) {
    bool noClamp = (aElem.isF16() && dest.isF16()) ||
                   (aElem.isBF16() && dest.isBF16()) || aIsFloat8;
    if (noClamp)
      return emitOpError(
          "clamp is not supported for this (source, accumulator) combination");
  }
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// DPPOp
//===----------------------------------------------------------------------===//
LogicalResult DPPOp::verify() {
  DPPPerm kind = getKind();
  Attribute permArgument = getPermArgument().value_or(Attribute{});
 
  switch (kind) {
 
  case DPPPerm::quad_perm: {
    auto quadPermAttr = dyn_cast_or_null<ArrayAttr>(permArgument);
    if (!quadPermAttr || quadPermAttr.size() != 4) {
      return emitOpError("quad_perm attribute must have exactly 4 elements");
    }
    for (auto elem : quadPermAttr.getAsRange<IntegerAttr>()) {
      int32_t num = elem.getInt();
      if (num < 0 || num > 3) {
        return emitOpError(
            "Each element of quad_perm must be in the range [0, 3]");
      }
    }
  } break;
 
  case DPPPerm::row_shl:
  case DPPPerm::row_shr:
  case DPPPerm::row_ror: {
    if (!permArgument) {
      return emitOpError("Attribute '" + Twine(stringifyDPPPerm(kind)) +
                         "' value not specified");
    }
    if (auto intAttr = dyn_cast<IntegerAttr>(permArgument)) {
      uint32_t attrValue = intAttr.getInt();
      if (attrValue < 1 || attrValue > 15) {
        return emitOpError("Attribute value must be between 1 and 15");
      }
    }
  } break;
 
  case DPPPerm::wave_shl:
  case DPPPerm::wave_shr:
  case DPPPerm::wave_rol:
  case DPPPerm::wave_ror:
  case DPPPerm::row_mirror:
  case DPPPerm::row_half_mirror:
  case DPPPerm::row_bcast_15:
  case DPPPerm::row_bcast_31: {
    if (permArgument && !isa<UnitAttr>(permArgument)) {
      return emitOpError("Expected unit attribute for permArgument, but found "
                         "non-trivial argument");
    }
    break;
  }
  }
  return success();
}
 
//===----------------------------------------------------------------------===//
// PermlaneSwapOp
//===----------------------------------------------------------------------===//
LogicalResult PermlaneSwapOp::verify() {
  unsigned rowLength = getRowLength();
 
  if (rowLength != 16 && rowLength != 32)
    return emitOpError("row_length attribute must either be 16 or 32.");
 
  return success();
}
 
/// Remove amdgpu.lds_barrier after amdgpu.lds_barrier.
static LogicalResult eraseRedundantLDSBarrierOps(LDSBarrierOp op,
                                                 PatternRewriter &rewriter) {
  if (isa_and_nonnull<LDSBarrierOp>(op->getNextNode())) {
    rewriter.eraseOp(op);
    return success();
  }
  return failure();
}
 
void LDSBarrierOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
  results.add(eraseRedundantLDSBarrierOps);
}
 
//===----------------------------------------------------------------------===//
// MemoryCounterWaitOp
//===----------------------------------------------------------------------===//
 
namespace {
/// Fuse adjacent memory counter wait ops, taking the minimum value of the
/// counters.
struct FuseMemoryCounterWaitOp final : OpRewritePattern<MemoryCounterWaitOp> {
  using Base::Base;
 
  LogicalResult matchAndRewrite(MemoryCounterWaitOp op,
                                PatternRewriter &rewriter) const override {
    auto next = dyn_cast<MemoryCounterWaitOp>(op->getNextNode());
    if (!next)
      return failure();
 
    auto setters = {&MemoryCounterWaitOp::setLoad,
                    &MemoryCounterWaitOp::setStore, &MemoryCounterWaitOp::setDs,
                    &MemoryCounterWaitOp::setExp,
                    &MemoryCounterWaitOp::setTensor};
    auto lhsVals = {op.getLoad(), op.getStore(), op.getDs(), op.getExp(),
                    op.getTensor()};
    auto rhsVals = {next.getLoad(), next.getStore(), next.getDs(),
                    next.getExp(), next.getTensor()};
    rewriter.modifyOpInPlace(op, [&] {
      for (auto [setter, lhs, rhs] :
           llvm::zip_equal(setters, lhsVals, rhsVals)) {
        if (lhs && rhs) {
          (op.*setter)(std::min(*lhs, *rhs));
        } else if (lhs) {
          (op.*setter)(*lhs);
        } else if (rhs) {
          (op.*setter)(*rhs);
        }
      }
    });
    rewriter.eraseOp(next);
    return success();
  }
};
} // namespace
 
void MemoryCounterWaitOp::getCanonicalizationPatterns(
    RewritePatternSet &results, MLIRContext *context) {
  results.add<FuseMemoryCounterWaitOp>(context);
}
 
//===----------------------------------------------------------------------===//
// GatherToLDSOp
//===----------------------------------------------------------------------===//
 
LogicalResult GatherToLDSOp::verify() {
  MemRefType srcType = cast<MemRefType>(getSrc().getType());
  MemRefType dstType = cast<MemRefType>(getDst().getType());
 
  if (failed(
          verifyIndexCount(*this, "source", srcType, getSrcIndices().size())) ||
      failed(verifyIndexCount(*this, "destination", dstType,
                              getDstIndices().size())))
    return failure();
 
  if (dstType.getRank() > 0 && !dstType.areTrailingDimsContiguous(1))
    return emitOpError("destination type inner most dim must be contiguous");
 
  auto elemType = srcType.getElementType();
  // Check $src and $dst element types are the same.
  if (elemType != dstType.getElementType())
    return emitOpError("source and destination element types must match");
 
  // copy type sizes should be 1, 2, 4, 12 or 16 bytes.
  auto transferType = getTransferType();
  int transferSize;
  if (auto vectorTransfer = dyn_cast<VectorType>(transferType)) {
    transferSize = vectorTransfer.getNumElements() *
                   vectorTransfer.getElementTypeBitWidth();
  } else {
    transferSize = transferType.getIntOrFloatBitWidth();
  }
  if (!llvm::is_contained({8, 16, 32, 96, 128}, transferSize))
    return emitOpError(
        "Transfering type size must be 8, 16, 32, 96 or 128 bits");
 
  if (!isGlobalMemorySpace(srcType.getMemorySpace(), /*allowFlat=*/true) &&
      !isFatRawBufferMemorySpace(srcType.getMemorySpace()))
    return emitOpError(
        "source memory address space must be global or fat raw buffer");
 
  if (!isWorkgroupMemorySpace(dstType.getMemorySpace()))
    return emitOpError("destination memory address space must be Workgroup");
 
  return success();
}
 
namespace {
/// If the source/target of a GatherToLDSOp is a CastOp that only removes static
/// information or changes layout, the cast can be skipped.
struct FoldGatherToLDSOfCast final : OpRewritePattern<GatherToLDSOp> {
  using OpRewritePattern::OpRewritePattern;
 
  LogicalResult matchAndRewrite(GatherToLDSOp gatherOp,
                                PatternRewriter &rewriter) const override {
    bool modified = false;
    auto foldCast = [&](OpOperand &operand) {
      if (auto castOp = operand.get().getDefiningOp<memref::CastOp>()) {
        if (memref::CastOp::canFoldIntoConsumerOp(castOp)) {
          rewriter.modifyOpInPlace(gatherOp,
                                   [&] { operand.assign(castOp.getSource()); });
          modified = true;
        }
      }
    };
 
    foldCast(gatherOp.getSrcMutable());
    foldCast(gatherOp.getDstMutable());
 
    return success(modified);
  }
};
} // namespace
 
void GatherToLDSOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                MLIRContext *context) {
  results.add<FoldGatherToLDSOfCast>(context);
}
 
//===----------------------------------------------------------------------===//
// GlobalLoadAsyncToLDSOp
//===----------------------------------------------------------------------===//
 
LogicalResult GlobalLoadAsyncToLDSOp::verify() {
  MemRefType srcType = cast<MemRefType>(getSrc().getType());
  MemRefType dstType = cast<MemRefType>(getDst().getType());
 
  if (failed(
          verifyIndexCount(*this, "source", srcType, getSrcIndices().size())) ||
      failed(verifyIndexCount(*this, "destination", dstType,
                              getDstIndices().size())))
    return failure();
 
  if (srcType.getElementType() != dstType.getElementType())
    return emitOpError("source and destination element types must match");
 
  Type transferType = getTransferType();
  int transferSize;
  if (auto vectorTransfer = dyn_cast<VectorType>(transferType)) {
    transferSize = vectorTransfer.getNumElements() *
                   vectorTransfer.getElementTypeBitWidth();
  } else {
    transferSize = transferType.getIntOrFloatBitWidth();
  }
  if (!llvm::is_contained({8, 32, 64, 128}, transferSize))
    return emitOpError("transfer type size must be 8, 32, 64, or 128 bits");
 
  if (!isGlobalMemorySpace(srcType.getMemorySpace(), /*allowFlat=*/false))
    return emitOpError("source memory address space must be global");
 
  if (!isWorkgroupMemorySpace(dstType.getMemorySpace()))
    return emitOpError("destination memory address space must be Workgroup");
 
  return success();
}
 
static LogicalResult
foldGlobalLoadAsyncToLDSConstantMask(GlobalLoadAsyncToLDSOp op,
                                     PatternRewriter &rewriter) {
  Value mask = op.getMask();
  if (!mask)
    return failure();
 
  APInt maskValue;
  if (!matchPattern(mask, m_ConstantInt(&maskValue)))
    return failure();
 
  if (maskValue.isZero()) {
    rewriter.eraseOp(op);
    return success();
  }
 
  rewriter.modifyOpInPlace(op, [&]() { op.getMaskMutable().clear(); });
  return success();
}
 
void GlobalLoadAsyncToLDSOp::getCanonicalizationPatterns(
    RewritePatternSet &results, MLIRContext *context) {
  results.add(foldGlobalLoadAsyncToLDSConstantMask);
}
 
//===----------------------------------------------------------------------===//
// TransposeLoadOp
//===----------------------------------------------------------------------===//
 
LogicalResult TransposeLoadOp::verify() {
  MemRefType srcType = cast<MemRefType>(getSrc().getType());
 
  if (failed(
          verifyIndexCount(*this, "source", srcType, getSrcIndices().size())))
    return failure();
 
  if (!isWorkgroupMemorySpace(srcType.getMemorySpace()))
    return emitOpError("source memory address space must be Workgroup");
 
  auto transferType = cast<VectorType>(getType());
  size_t numElements = transferType.getNumElements();
  size_t elementTypeSize =
      transferType.getElementType().getIntOrFloatBitWidth();
 
  auto emitNumElementsError = [&](StringRef expected) {
    return emitOpError(
               "Transferring type size mismatch: expected num of elements: ")
           << expected;
  };
 
  switch (elementTypeSize) {
  case 4:
  case 6:
    if (numElements != 16)
      return emitNumElementsError("16");
    break;
  case 8:
    if (numElements != 8)
      return emitNumElementsError("8");
    break;
  case 16:
    if (numElements != 4 && numElements != 8)
      return emitNumElementsError("4 or 8");
    break;
  default:
    return emitOpError("Unsupported element type size for transpose load: ")
           << elementTypeSize << " bits";
  }
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// GlobalTransposeLoadOp
//===----------------------------------------------------------------------===//
 
LogicalResult GlobalTransposeLoadOp::verify() {
  MemRefType srcType = cast<MemRefType>(getSrc().getType());
 
  if (failed(
          verifyIndexCount(*this, "source", srcType, getSrcIndices().size())))
    return failure();
 
  if (!isGlobalMemorySpace(srcType.getMemorySpace(), /*allowFlat=*/false))
    return emitOpError("source memory address space must be Global");
 
  auto resultType = cast<VectorType>(getType());
  size_t numElements = resultType.getNumElements();
  size_t elementTypeSize = resultType.getElementType().getIntOrFloatBitWidth();
 
  // ElementSize -> NumElements. Chipset gating (gfx1200 vs gfx1250) is
  // enforced in the lowering.
  static const llvm::SmallDenseMap<size_t, size_t> kValidLoadSizeMap = {
      {4, 16}, // global_load_tr4_b64  (gfx1250+)
      {6, 16}, // global_load_tr6_b96  (gfx1250+)
      {8, 8},  // global_load_tr_b64   (gfx1200+)
      {16, 8}, // global_load_tr_b128  (gfx1200+)
  };
 
  auto validNumElems = kValidLoadSizeMap.find(elementTypeSize);
  if (validNumElems == kValidLoadSizeMap.end())
    return emitOpError(
               "unsupported element type size for global transpose load: ")
           << elementTypeSize << " bits";
 
  if (numElements != validNumElems->second)
    return emitOpError(
               "transferring type size mismatch: expected num of elements: ")
           << validNumElems->second;
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// MakeDmaBaseOp
//===----------------------------------------------------------------------===//
 
template <typename BaseOp>
static LogicalResult verifyBase(BaseOp op) {
  auto ldsType = cast<MemRefType>(op.getLds().getType());
  auto globalType = cast<MemRefType>(op.getGlobal().getType());
  if (failed(verifyIndexCount(op, "global", globalType,
                              op.getGlobalIndices().size())) ||
      failed(verifyIndexCount(op, "lds", ldsType, op.getLdsIndices().size())))
    return failure();
 
  if (!isWorkgroupMemorySpace(ldsType.getMemorySpace()))
    return op.emitOpError(
        "lds memref must have workgroup address space attribute.");
  if (!isGlobalMemorySpace(globalType.getMemorySpace(), /*allowFlat=*/false))
    return op.emitOpError(
        "global memref must have global address space attribute.");
 
  Type elementType = ldsType.getElementType();
  unsigned width = elementType.getIntOrFloatBitWidth();
 
  if (!llvm::is_contained({8u, 16u, 32u, 64u}, width))
    return op.emitOpError(
               "element type must be 1, 2, 4, or 8 bytes long but type was ")
           << width << " bits long.";
  return success();
}
 
LogicalResult MakeDmaBaseOp::verify() { return verifyBase(*this); }
 
//===----------------------------------------------------------------------===//
// MakeGatherDmaBaseOp
//===----------------------------------------------------------------------===//
 
LogicalResult
TDMGatherBaseType::verify(function_ref<InFlightDiagnostic()> emitError,
                          Type elementType, Type indexType) {
  unsigned width = elementType.getIntOrFloatBitWidth();
  if (!llvm::is_contained({8u, 16u, 32u, 64u}, width))
    return emitError()
           << "element type must be 1, 2, 4, or 8 bytes wide but type "
           << elementType << " is " << width / 8 << " bytes wide.";
  MLIRContext *ctx = elementType.getContext();
  Type i16 = IntegerType::get(ctx, 32);
  Type i32 = IntegerType::get(ctx, 16);
  if (!llvm::is_contained({i16, i32}, indexType))
    return emitError() << "index type must be i16 or i32 but index type is "
                       << indexType << ".";
  return success();
}
 
LogicalResult MakeGatherDmaBaseOp::verify() { return verifyBase(*this); }
 
//===----------------------------------------------------------------------===//
// MakeDmaDescriptorOp
//===----------------------------------------------------------------------===//
 
template <typename DescriptorOp>
static LogicalResult verifyDescriptorOp(DescriptorOp op) {
  ArrayRef<int64_t> globalStaticStrides = op.getGlobalStaticStrides();
 
  if (globalStaticStrides.empty())
    return op.emitOpError("strides must not be empty.");
  if (globalStaticStrides.back() != 1)
    return op.emitOpError("strides for the innermost dimension must be 1.");
 
  ArrayRef<int64_t> globalStaticSizes = op.getGlobalStaticSizes();
  size_t rank = globalStaticSizes.size();
  if (rank > 5)
    return op.emitOpError("tensor and tile must be at most of rank 5.");
  if (rank != globalStaticStrides.size())
    return op.emitOpError("strides and sizes must have same rank.");
 
  ArrayRef<int64_t> sharedStaticSizes = op.getSharedStaticSizes();
  if (rank != sharedStaticSizes.size())
    return op.emitOpError("tensor must have same rank as tile.");
 
  unsigned elementTypeWidth = op.getElementTypeWidth();
  if (!llvm::is_contained({8u, 16u, 32u, 64u}, elementTypeWidth))
    return op.emitOpError(
               "element type width must be 1, 2, 4 or 8 bytes, but was ")
           << elementTypeWidth << " bits long";
 
  if (!op.getAtomicBarrierAddress() && !op.getAtomicBarrierIndices().empty())
    return op.emitOpError(
        "atomic barrier indices require an atomic barrier address");
 
  if (Value atomicBarrierAddress = op.getAtomicBarrierAddress()) {
    auto atomicBarrierAddressType =
        cast<MemRefType>(atomicBarrierAddress.getType());
    if (failed(verifyIndexCount(op, "atomic barrier", atomicBarrierAddressType,
                                op.getAtomicBarrierIndices().size())))
      return failure();
 
    bool barrierInLDS =
        isWorkgroupMemorySpace(atomicBarrierAddressType.getMemorySpace());
    if (!barrierInLDS)
      return op.emitOpError("atomic barrier address must be in LDS.");
  }
 
  if (op.getEarlyTimeout() && !op.getWorkgroupMask())
    return op.emitOpError(
        "early timeout does not apply when workgroup_mask is not set.");
  return success();
}
 
template <typename DescriptorOp, typename FoldAdaptor>
static OpFoldResult foldDescriptorOp(DescriptorOp op, FoldAdaptor adaptor) {
  SmallVector<OpFoldResult> mixedGlobalSizes(op.getMixedGlobalSizes());
  SmallVector<OpFoldResult> mixedGlobalStrides(op.getMixedGlobalStrides());
  SmallVector<OpFoldResult> mixedSharedSizes(op.getMixedSharedSizes());
 
  if (failed(foldDynamicIndexList(mixedGlobalSizes, /*onlyNonNegative=*/true,
                                  /*onlyNonZero=*/true)) &&
      failed(foldDynamicIndexList(mixedGlobalStrides, /*onlyNonNegative=*/true,
                                  /*onlyNonZero=*/true)) &&
      failed(foldDynamicIndexList(mixedSharedSizes, /*onlyNonNegative=*/true,
                                  /*onlyNonZero=*/true)))
    return nullptr;
 
  SmallVector<Value> dynamicGlobalSizes, dynamicGlobalStrides,
      dynamicSharedSizes;
  SmallVector<int64_t> staticGlobalSizes, staticGlobalStrides,
      staticSharedSizes;
 
  dispatchIndexOpFoldResults(mixedGlobalSizes, dynamicGlobalSizes,
                             staticGlobalSizes);
  op.setGlobalStaticSizes(staticGlobalSizes);
  op.getGlobalDynamicSizesMutable().assign(dynamicGlobalSizes);
 
  dispatchIndexOpFoldResults(mixedGlobalStrides, dynamicGlobalStrides,
                             staticGlobalStrides);
  op.setGlobalStaticStrides(staticGlobalStrides);
  op.getGlobalDynamicStridesMutable().assign(dynamicGlobalStrides);
 
  dispatchIndexOpFoldResults(mixedSharedSizes, dynamicSharedSizes,
                             staticSharedSizes);
  op.setSharedStaticSizes(staticSharedSizes);
  op.getSharedDynamicSizesMutable().assign(dynamicSharedSizes);
  return op.getResult();
}
 
LogicalResult MakeDmaDescriptorOp::verify() {
  return verifyDescriptorOp(*this);
}
 
OpFoldResult MakeDmaDescriptorOp::fold(FoldAdaptor adaptor) {
  return foldDescriptorOp(*this, adaptor);
}
 
//===----------------------------------------------------------------------===//
// MakeGatherDmaDescriptorOp
//===----------------------------------------------------------------------===//
 
LogicalResult MakeGatherDmaDescriptorOp::verify() {
  ArrayRef<int64_t> globalStaticSizes = getGlobalStaticSizes();
  size_t rank = globalStaticSizes.size();
  if (rank > 2)
    return emitOpError(
        "tensor and tile must be at most of rank two in gather mode.");
  Value indices = getIndices();
  Type elementType = cast<VectorType>(indices.getType()).getElementType();
  if (elementType != getBase().getType().getIndexType())
    return emitOpError("indices' element type must match base's element type.");
 
  return verifyDescriptorOp(*this);
}
 
OpFoldResult MakeGatherDmaDescriptorOp::fold(FoldAdaptor adaptor) {
  return foldDescriptorOp(*this, adaptor);
}
 
//===----------------------------------------------------------------------===//
// ScaledMFMAOp
//===----------------------------------------------------------------------===//
 
namespace {
/// Check if the scales input is used in other scaled mfma's while they exist.
/// If theyre unused then pack the scales.
struct PackScales final : OpRewritePattern<ScaledMFMAOp> {
  using OpRewritePattern::OpRewritePattern;
 
  LogicalResult matchAndRewrite(ScaledMFMAOp op,
                                PatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    auto setOpsel = [&op](unsigned idx, int64_t val) {
      switch (idx) {
      case 3:
        op.setScalesIdxA(val);
        break;
      case 4:
        op.setScalesIdxB(val);
        break;
      default:
        break;
      }
    };
 
    // For every scale operand of this ScaledMFMAOp, if the scale is produced by
    // the extraction of a single scale from some vector, then attempt to
    // extract 4 values from that vector instead.
    //
    // Example: (f8 here means f8E8M0FNU)
    // %unit = vector.extract %ScaleSrc[offsets] : f8 from vector<...>
    // %scale = vector.insert %unit, ... : f8 into vector<4xf8>
    // amdgpu.scaled_mfma(%scale[0] * ...
    //
    // rewrite to:
    //
    // %reshaped = vector.shape_cast %ScaleSrc : vector<...> to vector<?xf8>
    // %scale = vector.extract %reshaped[?] : vector<4xf8> from vector<?xf8>
    // amdgpu.scaled_mfma(%scale[0-3] * ...
    //
    // This creates duplicate shape_casts for every use but these will be
    // removed in CSE.
    for (auto opIdx : std::array<int64_t, 2>({3, 4})) {
      auto insertOp = op.getOperand(opIdx).getDefiningOp<vector::InsertOp>();
      if (!insertOp) {
        return rewriter.notifyMatchFailure(op,
                                           "defining op not a vector.insert");
      }
      // If the extracted value is not a single scalar, then it has been packed.
      if (isa<VectorType>(insertOp.getValueToStore().getType())) {
        return rewriter.notifyMatchFailure(
            op, "scaled mfma operand already packed");
      }
 
      auto extractOp =
          insertOp.getValueToStore().getDefiningOp<vector::ExtractOp>();
      if (!extractOp) {
        return rewriter.notifyMatchFailure(op,
                                           "defining op not a vector.extract");
      }
 
      Value scaleSrc = extractOp.getOperand(0);
      auto scaleSrcType = dyn_cast<VectorType>(scaleSrc.getType());
      if (!scaleSrcType) {
        return rewriter.notifyMatchFailure(op, "not a vector type");
      }
 
      // We do not handle dynamic dims yet, assume that the input is padded to
      // a static shape now.
      if (!scaleSrcType.hasStaticShape()) {
        return rewriter.notifyMatchFailure(op,
                                           "dynamic dims not yet supported");
      }
 
      int64_t numElements = scaleSrcType.getNumElements();
      if (numElements < 4) {
        return rewriter.notifyMatchFailure(
            op, "do not pack if # of scales less than four");
      }
 
      // Find a linearized idx using the size and offsets of the extract op.
      auto extractedPos = llvm::to_vector_of<int64_t>(
          llvm::reverse(extractOp.getStaticPosition()));
      ArrayRef<int64_t> scaleSrcShape = scaleSrcType.getShape();
      int64_t scaleSrcRank = scaleSrcType.getRank();
      SmallVector<int64_t> extractSizes(scaleSrcRank, 1);
      for (int64_t i = 1; i < scaleSrcRank; ++i) {
        extractSizes[i] = extractSizes[i - 1] * scaleSrcShape[scaleSrcRank - i];
      }
      int64_t idx = linearize(extractedPos, extractSizes);
 
      // All n scales (where n is the total number of scales) must now be
      // extracted in chunks of 4 elements. This is done by dividing the
      // original vector of scales into groups of 4 elements
      // at offsets 0, 4, ..., m (where m = n/4). All extractions of a
      // scale at a particular index are now replaced with an extraction
      // of the entire group of 4 elements to which that index belongs.
      //
      // If the number of scales happens to be indivisible by 4, extract
      // the remaining n - m scales in a chunk of 4 elements starting at
      // offset n - 4.
      int64_t offset = idx - (idx % 4);
      int64_t opsel = idx - offset;
      int64_t size = 4l;
      // Accomdate remaining elements in the case of non-4-divisible vectors.
      if (numElements - offset < size) {
        opsel = size - (numElements - idx);
        offset = numElements - 4l;
      }
      Type scaleSrcElemType = scaleSrcType.getElementType();
      auto newSrcType =
          VectorType::get(ArrayRef{numElements}, scaleSrcElemType);
      Value newScaleSrc =
          vector::ShapeCastOp::create(rewriter, loc, newSrcType, scaleSrc);
      auto extract = vector::ExtractStridedSliceOp::create(
          rewriter, loc, newScaleSrc, ArrayRef{offset}, ArrayRef{size},
          ArrayRef{int64_t(1)});
      rewriter.modifyOpInPlace(op, [&] {
        op->setOperand(opIdx, extract);
        setOpsel(opIdx, opsel);
      });
    }
    return success();
  }
};
} // namespace
 
void ScaledMFMAOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
  results.add<PackScales>(context);
}
 
//===----------------------------------------------------------------------===//
// In-LDS Barrier Operations (gfx1250+)
//===----------------------------------------------------------------------===//
 
template <typename T>
static LogicalResult verifyDsBarrierOpCommon(T &op) {
  MemRefType memrefType = llvm::cast<MemRefType>(op.getBase().getType());
  if (failed(
          verifyIndexCount(op, "barrier", memrefType, op.getIndices().size())))
    return failure();
 
  if (!isWorkgroupMemorySpace(memrefType.getMemorySpace()))
    return op.emitOpError("barrier must be in workgroup (LDS) memory");
 
  return success();
}
 
LogicalResult DsBarrierInitOp::verify() {
  return verifyDsBarrierOpCommon(*this);
}
 
LogicalResult DsBarrierPollStateOp::verify() {
  return verifyDsBarrierOpCommon(*this);
}
 
LogicalResult DsAsyncBarrierArriveOp::verify() {
  return verifyDsBarrierOpCommon(*this);
}
 
LogicalResult DsBarrierArriveOp::verify() {
  return verifyDsBarrierOpCommon(*this);
}
 
//===----------------------------------------------------------------------===//
// GlobalPrefetchOp
//===----------------------------------------------------------------------===//
 
LogicalResult GlobalPrefetchOp::verify() {
  auto src = cast<MemRefType>(getSrc().getType());
 
  if (failed(verifyIndexCount(*this, "source", src, getIndices().size())))
    return failure();
 
  Attribute memSpace = src.getMemorySpace();
  if (!memSpace)
    return this->emitOpError("the source must have address space attribute");
  if (!isGlobalMemorySpace(memSpace, /*allowFlat=*/false))
    return this->emitOpError("the source must reside in global address space");
 
  const LoadTemporalHint temporalHint = getTemporalHint();
  const Scope scope = getCacheScope();
  const bool isSpeculative = getSpeculative();
 
  // See GFX1250 SPG for a detail explanation
  if (isSpeculative && scope == Scope::WGP)
    return this->emitOpError(
        "does not support speculative prefetch in WGP scope");
 
  // Note that temporal hints are shared between load, store,
  // prefetch, etc. instructions. However, some instructions
  // operate only with a subset of hints according to the ISA
  // documentation. In case of global prefetch, non-temporal (NT)
  // and last-use (LU) hints are not used. The extra bits of encoding
  // are used to encode speculative or non-speculative instruction behavior
  if (llvm::is_contained({LoadTemporalHint::NT, LoadTemporalHint::LU},
                         temporalHint))
    return this->emitOpError("does not support NT and LU modes");
 
  if (llvm::is_contained({LoadTemporalHint::NT_RT, LoadTemporalHint::RT_NT,
                          LoadTemporalHint::NT_HT},
                         temporalHint) &&
      !isSpeculative) {
    return this->emitOpError("operates only in the speculative mode");
  }
  return success();
}
 
#define GET_OP_CLASSES
#include "mlir/Dialect/AMDGPU/IR/AMDGPU.cpp.inc"