//===- GPUDialect.cpp - MLIR Dialect for GPU Kernels implementation -------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file implements the GPU kernel-related dialect and its operations.
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/GPU/IR/GPUDialect.h"

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/DialectImplementation.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Interfaces/FunctionImplementation.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Transforms/InliningUtils.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/StringSaver.h"
#include <cassert>

using namespace mlir;
using namespace mlir::gpu;

#include "mlir/Dialect/GPU/IR/GPUOpsDialect.cpp.inc"

//===----------------------------------------------------------------------===//
// GPU Device Mapping Attributes
//===----------------------------------------------------------------------===//

int64_t GPUBlockMappingAttr::getMappingId() const {
  return static_cast<int64_t>(getBlock());
}

bool GPUBlockMappingAttr::isLinearMapping() const {
  return getMappingId() >= static_cast<int64_t>(MappingId::LinearDim0);
}

int64_t GPUBlockMappingAttr::getRelativeIndex() const {
  return isLinearMapping()
             ? getMappingId() - static_cast<int64_t>(MappingId::LinearDim0)
             : getMappingId();
}

int64_t GPUWarpgroupMappingAttr::getMappingId() const {
  return static_cast<int64_t>(getWarpgroup());
}

bool GPUWarpgroupMappingAttr::isLinearMapping() const {
  return getMappingId() >= static_cast<int64_t>(MappingId::LinearDim0);
}

int64_t GPUWarpgroupMappingAttr::getRelativeIndex() const {
  return isLinearMapping()
             ? getMappingId() - static_cast<int64_t>(MappingId::LinearDim0)
             : getMappingId();
}

int64_t GPUWarpMappingAttr::getMappingId() const {
  return static_cast<int64_t>(getWarp());
}

bool GPUWarpMappingAttr::isLinearMapping() const {
  return getMappingId() >= static_cast<int64_t>(MappingId::LinearDim0);
}

int64_t GPUWarpMappingAttr::getRelativeIndex() const {
  return isLinearMapping()
             ? getMappingId() - static_cast<int64_t>(MappingId::LinearDim0)
             : getMappingId();
}

int64_t GPUThreadMappingAttr::getMappingId() const {
  return static_cast<int64_t>(getThread());
}

bool GPUThreadMappingAttr::isLinearMapping() const {
  return getMappingId() >= static_cast<int64_t>(MappingId::LinearDim0);
}

int64_t GPUThreadMappingAttr::getRelativeIndex() const {
  return isLinearMapping()
             ? getMappingId() - static_cast<int64_t>(MappingId::LinearDim0)
             : getMappingId();
}

int64_t GPUMemorySpaceMappingAttr::getMappingId() const {
  return static_cast<int64_t>(getAddressSpace());
}

bool GPUMemorySpaceMappingAttr::isLinearMapping() const {
  llvm_unreachable("GPUMemorySpaceMappingAttr does not support linear mapping");
}

int64_t GPUMemorySpaceMappingAttr::getRelativeIndex() const {
  llvm_unreachable("GPUMemorySpaceMappingAttr does not support relative index");
}

//===----------------------------------------------------------------------===//
// MMAMatrixType
//===----------------------------------------------------------------------===//

MMAMatrixType MMAMatrixType::get(ArrayRef<int64_t> shape, Type elementType,
                                 StringRef operand) {
  return Base::get(elementType.getContext(), shape, elementType, operand);
}

MMAMatrixType
MMAMatrixType::getChecked(function_ref<InFlightDiagnostic()> emitError,
                          ArrayRef<int64_t> shape, Type elementType,
                          StringRef operand) {
  return Base::getChecked(emitError, elementType.getContext(), shape,
                          elementType, operand);
}

unsigned MMAMatrixType::getNumDims() const { return getImpl()->numDims; }

ArrayRef<int64_t> MMAMatrixType::getShape() const {
  return getImpl()->getShape();
}

Type MMAMatrixType::getElementType() const { return getImpl()->elementType; }

StringRef MMAMatrixType::getOperand() const { return getImpl()->getOperand(); }

bool MMAMatrixType::isValidElementType(Type elementType) {
  return elementType.isF16() || elementType.isF32() ||
         elementType.isUnsignedInteger(8) || elementType.isSignedInteger(8) ||
         elementType.isInteger(32);
}

LogicalResult
MMAMatrixType::verify(function_ref<InFlightDiagnostic()> emitError,
                      ArrayRef<int64_t> shape, Type elementType,
                      StringRef operand) {
  if (!operand.equals("AOp") && !operand.equals("BOp") &&
      !operand.equals("COp"))
    return emitError() << "operand expected to be one of AOp, BOp or COp";

  if (shape.size() != 2)
    return emitError() << "MMAMatrixType must have exactly two dimensions";

  if (!MMAMatrixType::isValidElementType(elementType))
    return emitError()
           << "MMAMatrixType elements must be SI8, UI8, I32, F16, or F32";

  return success();
}

//===----------------------------------------------------------------------===//
// GPUDialect
//===----------------------------------------------------------------------===//

bool GPUDialect::isWorkgroupMemoryAddressSpace(Attribute memorySpace) {
  if (!memorySpace)
    return false;
  if (auto gpuAttr = llvm::dyn_cast<gpu::AddressSpaceAttr>(memorySpace))
    return gpuAttr.getValue() == getWorkgroupAddressSpace();
  return false;
}

bool GPUDialect::hasWorkgroupMemoryAddressSpace(MemRefType type) {
  Attribute memorySpace = type.getMemorySpace();
  return isWorkgroupMemoryAddressSpace(memorySpace);
}

bool GPUDialect::isKernel(Operation *op) {
  UnitAttr isKernelAttr = op->getAttrOfType<UnitAttr>(getKernelFuncAttrName());
  return static_cast<bool>(isKernelAttr);
}

namespace {
/// This class defines the interface for handling inlining with gpu
/// operations.
struct GPUInlinerInterface : public DialectInlinerInterface {
  using DialectInlinerInterface::DialectInlinerInterface;

  /// All gpu dialect ops can be inlined.
  bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
    return true;
  }
};
} // namespace

void GPUDialect::initialize() {
  addTypes<AsyncTokenType>();
  addTypes<MMAMatrixType>();
  addTypes<SparseDnTensorHandleType>();
  addTypes<SparseSpMatHandleType>();
  addTypes<SparseSpGEMMOpHandleType>();
  addOperations<
#define GET_OP_LIST
#include "mlir/Dialect/GPU/IR/GPUOps.cpp.inc"
      >();
  addAttributes<
#define GET_ATTRDEF_LIST
#include "mlir/Dialect/GPU/IR/GPUOpsAttributes.cpp.inc"
      >();
  addInterfaces<GPUInlinerInterface>();
}

static std::string getSparseHandleKeyword(SparseHandleKind kind) {
  switch (kind) {
  case SparseHandleKind::DnTensor:
    return "sparse.dntensor_handle";
  case SparseHandleKind::SpMat:
    return "sparse.spmat_handle";
  case SparseHandleKind::SpGEMMOp:
    return "sparse.spgemmop_handle";
  }
  llvm_unreachable("unknown sparse handle kind");
  return "";
}

Type GPUDialect::parseType(DialectAsmParser &parser) const {
  // Parse the main keyword for the type.
  StringRef keyword;
  if (parser.parseKeyword(&keyword))
    return Type();
  MLIRContext *context = getContext();

  // Handle 'async token' types.
  if (keyword == "async.token")
    return AsyncTokenType::get(context);

  if (keyword == "mma_matrix") {
    SMLoc beginLoc = parser.getNameLoc();

    // Parse '<'.
    if (parser.parseLess())
      return nullptr;

    // Parse the size and elementType.
    SmallVector<int64_t> shape;
    Type elementType;
    if (parser.parseDimensionList(shape, /*allowDynamic=*/false) ||
        parser.parseType(elementType))
      return nullptr;

    // Parse ','
    if (parser.parseComma())
      return nullptr;

    // Parse operand.
    std::string operand;
    if (failed(parser.parseOptionalString(&operand)))
      return nullptr;

    // Parse '>'.
    if (parser.parseGreater())
      return nullptr;

    return MMAMatrixType::getChecked(mlir::detail::getDefaultDiagnosticEmitFn(
                                         parser.getEncodedSourceLoc(beginLoc)),
                                     shape, elementType, operand);
  }

  if (keyword == getSparseHandleKeyword(SparseHandleKind::DnTensor))
    return SparseDnTensorHandleType::get(context);
  if (keyword == getSparseHandleKeyword(SparseHandleKind::SpMat))
    return SparseSpMatHandleType::get(context);
  if (keyword == getSparseHandleKeyword(SparseHandleKind::SpGEMMOp))
    return SparseSpGEMMOpHandleType::get(context);

  parser.emitError(parser.getNameLoc(), "unknown gpu type: " + keyword);
  return Type();
}
// TODO: print refined type here. Notice that should be corresponding to the
// parser
void GPUDialect::printType(Type type, DialectAsmPrinter &os) const {
  TypeSwitch<Type>(type)
      .Case<AsyncTokenType>([&](Type) { os << "async.token"; })
      .Case<SparseDnTensorHandleType>([&](Type) {
        os << getSparseHandleKeyword(SparseHandleKind::DnTensor);
      })
      .Case<SparseSpMatHandleType>(
          [&](Type) { os << getSparseHandleKeyword(SparseHandleKind::SpMat); })
      .Case<SparseSpGEMMOpHandleType>([&](Type) {
        os << getSparseHandleKeyword(SparseHandleKind::SpGEMMOp);
      })
      .Case<MMAMatrixType>([&](MMAMatrixType fragTy) {
        os << "mma_matrix<";
        auto shape = fragTy.getShape();
        for (auto dim = shape.begin(), e = shape.end() - 1; dim != e; ++dim)
          os << *dim << 'x';
        os << shape.back() << 'x' << fragTy.getElementType();
        os << ", \"" << fragTy.getOperand() << "\"" << '>';
      })
      .Default([](Type) { llvm_unreachable("unexpected 'gpu' type kind"); });
}

LogicalResult GPUDialect::verifyOperationAttribute(Operation *op,
                                                   NamedAttribute attr) {
  if (!llvm::isa<UnitAttr>(attr.getValue()) ||
      attr.getName() != getContainerModuleAttrName())
    return success();

  auto module = dyn_cast<ModuleOp>(op);
  if (!module)
    return op->emitError("expected '")
           << getContainerModuleAttrName() << "' attribute to be attached to '"
           << ModuleOp::getOperationName() << '\'';

  auto walkResult = module.walk([&module](LaunchFuncOp launchOp) -> WalkResult {
    // Ignore launches that are nested more or less deep than functions in the
    // module we are currently checking.
    if (!launchOp->getParentOp() ||
        launchOp->getParentOp()->getParentOp() != module)
      return success();

    // Ignore launch ops with missing attributes here. The errors will be
    // reported by the verifiers of those ops.
    if (!launchOp->getAttrOfType<SymbolRefAttr>(
            LaunchFuncOp::getKernelAttrName(launchOp->getName())))
      return success();

    // Check that `launch_func` refers to a well-formed GPU kernel container.
    StringAttr kernelContainerName = launchOp.getKernelModuleName();
    Operation *kernelContainer = module.lookupSymbol(kernelContainerName);
    if (!kernelContainer)
      return launchOp.emitOpError()
             << "kernel container '" << kernelContainerName.getValue()
             << "' is undefined";

    // If the container is a GPU binary op return success.
    if (isa<BinaryOp>(kernelContainer))
      return success();

    auto kernelModule = dyn_cast<GPUModuleOp>(kernelContainer);
    if (!kernelModule)
      return launchOp.emitOpError()
             << "kernel module '" << kernelContainerName.getValue()
             << "' is undefined";

    // Check that `launch_func` refers to a well-formed kernel function.
    Operation *kernelFunc = module.lookupSymbol(launchOp.getKernelAttr());
    if (!kernelFunc)
      return launchOp.emitOpError("kernel function '")
             << launchOp.getKernel() << "' is undefined";
    auto kernelConvertedFunction = dyn_cast<FunctionOpInterface>(kernelFunc);
    if (!kernelConvertedFunction) {
      InFlightDiagnostic diag = launchOp.emitOpError()
                                << "referenced kernel '" << launchOp.getKernel()
                                << "' is not a function";
      diag.attachNote(kernelFunc->getLoc()) << "see the kernel definition here";
      return diag;
    }

    if (!kernelFunc->getAttrOfType<mlir::UnitAttr>(
            GPUDialect::getKernelFuncAttrName()))
      return launchOp.emitOpError("kernel function is missing the '")
             << GPUDialect::getKernelFuncAttrName() << "' attribute";

    // TODO: If the kernel isn't a GPU function (which happens during separate
    // compilation), do not check type correspondence as it would require the
    // verifier to be aware of the type conversion.
    auto kernelGPUFunction = dyn_cast<gpu::GPUFuncOp>(kernelFunc);
    if (!kernelGPUFunction)
      return success();

    unsigned actualNumArguments = launchOp.getNumKernelOperands();
    unsigned expectedNumArguments = kernelGPUFunction.getNumArguments();
    if (expectedNumArguments != actualNumArguments)
      return launchOp.emitOpError("got ")
             << actualNumArguments << " kernel operands but expected "
             << expectedNumArguments;

    auto functionType = kernelGPUFunction.getFunctionType();
    for (unsigned i = 0; i < expectedNumArguments; ++i) {
      if (launchOp.getKernelOperand(i).getType() != functionType.getInput(i)) {
        return launchOp.emitOpError("type of function argument ")
               << i << " does not match";
      }
    }

    return success();
  });

  return walkResult.wasInterrupted() ? failure() : success();
}

/// Parses an optional list of async operands with an optional leading keyword.
/// (`async`)? (`[` ssa-id-list `]`)?
///
/// This method is used by the tablegen assembly format for async ops as well.
static ParseResult parseAsyncDependencies(
    OpAsmParser &parser, Type &asyncTokenType,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &asyncDependencies) {
  auto loc = parser.getCurrentLocation();
  if (succeeded(parser.parseOptionalKeyword("async"))) {
    if (parser.getNumResults() == 0)
      return parser.emitError(loc, "needs to be named when marked 'async'");
    asyncTokenType = parser.getBuilder().getType<AsyncTokenType>();
  }
  return parser.parseOperandList(asyncDependencies,
                                 OpAsmParser::Delimiter::OptionalSquare);
}

/// Prints optional async dependencies with its leading keyword.
///   (`async`)? (`[` ssa-id-list `]`)?
// Used by the tablegen assembly format for several async ops.
static void printAsyncDependencies(OpAsmPrinter &printer, Operation *op,
                                   Type asyncTokenType,
                                   OperandRange asyncDependencies) {
  if (asyncTokenType)
    printer << "async";
  if (asyncDependencies.empty())
    return;
  if (asyncTokenType)
    printer << ' ';
  printer << '[';
  llvm::interleaveComma(asyncDependencies, printer);
  printer << ']';
}

// GPU Memory attributions functions shared by LaunchOp and GPUFuncOp.
/// Parses a GPU function memory attribution.
///
/// memory-attribution ::= (`workgroup` `(` ssa-id-and-type-list `)`)?
///                        (`private` `(` ssa-id-and-type-list `)`)?
///
/// Note that this function parses only one of the two similar parts, with the
/// keyword provided as argument.
static ParseResult
parseAttributions(OpAsmParser &parser, StringRef keyword,
                  SmallVectorImpl<OpAsmParser::Argument> &args) {
  // If we could not parse the keyword, just assume empty list and succeed.
  if (failed(parser.parseOptionalKeyword(keyword)))
    return success();

  return parser.parseArgumentList(args, OpAsmParser::Delimiter::Paren,
                                  /*allowType=*/true);
}

/// Prints a GPU function memory attribution.
static void printAttributions(OpAsmPrinter &p, StringRef keyword,
                              ArrayRef<BlockArgument> values) {
  if (values.empty())
    return;

  p << ' ' << keyword << '(';
  llvm::interleaveComma(
      values, p, [&p](BlockArgument v) { p << v << " : " << v.getType(); });
  p << ')';
}

/// Verifies a GPU function memory attribution.
static LogicalResult verifyAttributions(Operation *op,
                                        ArrayRef<BlockArgument> attributions,
                                        gpu::AddressSpace memorySpace) {
  for (Value v : attributions) {
    auto type = llvm::dyn_cast<MemRefType>(v.getType());
    if (!type)
      return op->emitOpError() << "expected memref type in attribution";

    // We can only verify the address space if it hasn't already been lowered
    // from the AddressSpaceAttr to a target-specific numeric value.
    auto addressSpace =
        llvm::dyn_cast_or_null<gpu::AddressSpaceAttr>(type.getMemorySpace());
    if (!addressSpace)
      continue;
    if (addressSpace.getValue() != memorySpace)
      return op->emitOpError()
             << "expected memory space " << stringifyAddressSpace(memorySpace)
             << " in attribution";
  }
  return success();
}

//===----------------------------------------------------------------------===//
// AllReduceOp
//===----------------------------------------------------------------------===//

static LogicalResult verifyReduceOpAndType(gpu::AllReduceOperation opName,
                                           Type resType) {
  using Kind = gpu::AllReduceOperation;
  if (llvm::is_contained(
          {Kind::MINNUMF, Kind::MAXNUMF, Kind::MINIMUMF, Kind::MAXIMUMF},
          opName)) {
    if (!isa<FloatType>(resType))
      return failure();
  }

  if (llvm::is_contained({Kind::MINSI, Kind::MINUI, Kind::MAXSI, Kind::MAXUI,
                          Kind::AND, Kind::OR, Kind::XOR},
                         opName)) {
    if (!isa<IntegerType>(resType))
      return failure();
  }

  return success();
}

LogicalResult gpu::AllReduceOp::verifyRegions() {
  if (getBody().empty() != getOp().has_value())
    return emitError("expected either an op attribute or a non-empty body");
  if (!getBody().empty()) {
    if (getBody().getNumArguments() != 2)
      return emitError("expected two region arguments");
    for (auto argument : getBody().getArguments()) {
      if (argument.getType() != getType())
        return emitError("incorrect region argument type");
    }
    unsigned yieldCount = 0;
    for (Block &block : getBody()) {
      if (auto yield = dyn_cast<gpu::YieldOp>(block.getTerminator())) {
        if (yield.getNumOperands() != 1)
          return emitError("expected one gpu.yield operand");
        if (yield.getOperand(0).getType() != getType())
          return emitError("incorrect gpu.yield type");
        ++yieldCount;
      }
    }
    if (yieldCount == 0)
      return emitError("expected gpu.yield op in region");
  } else {
    gpu::AllReduceOperation opName = *getOp();
    if (failed(verifyReduceOpAndType(opName, getType()))) {
      return emitError() << '`' << gpu::stringifyAllReduceOperation(opName)
                         << "` reduction operation is not compatible with type "
                         << getType();
    }
  }

  return success();
}

static bool canMakeGroupOpUniform(Operation *op) {
  auto launchOp = dyn_cast<gpu::LaunchOp>(op->getParentOp());
  if (!launchOp)
    return false;

  Region &body = launchOp.getBody();
  assert(!body.empty() && "Invalid region");

  // Only convert ops in gpu::launch entry block for now.
  return op->getBlock() == &body.front();
}

OpFoldResult gpu::AllReduceOp::fold(FoldAdaptor /*adaptor*/) {
  if (!getUniform() && canMakeGroupOpUniform(*this)) {
    setUniform(true);
    return getResult();
  }

  return nullptr;
}

// TODO: Support optional custom attributes (without dialect prefix).
static ParseResult parseAllReduceOperation(AsmParser &parser,
                                           AllReduceOperationAttr &attr) {
  StringRef enumStr;
  if (!parser.parseOptionalKeyword(&enumStr)) {
    std::optional<AllReduceOperation> op =
        gpu::symbolizeAllReduceOperation(enumStr);
    if (!op)
      return parser.emitError(parser.getCurrentLocation(), "invalid op kind");
    attr = AllReduceOperationAttr::get(parser.getContext(), *op);
  }
  return success();
}

static void printAllReduceOperation(AsmPrinter &printer, Operation *op,
                                    AllReduceOperationAttr attr) {
  if (attr)
    attr.print(printer);
}

//===----------------------------------------------------------------------===//
// SubgroupReduceOp
//===----------------------------------------------------------------------===//

LogicalResult gpu::SubgroupReduceOp::verify() {
  Type elemType = getType();
  if (auto vecTy = dyn_cast<VectorType>(elemType)) {
    if (vecTy.isScalable())
      return emitOpError() << "is not compatible with scalable vector types";

    elemType = vecTy.getElementType();
  }

  gpu::AllReduceOperation opName = getOp();
  if (failed(verifyReduceOpAndType(opName, elemType))) {
    return emitError() << '`' << gpu::stringifyAllReduceOperation(opName)
                       << "` reduction operation is not compatible with type "
                       << getType();
  }
  return success();
}

OpFoldResult gpu::SubgroupReduceOp::fold(FoldAdaptor /*adaptor*/) {
  if (!getUniform() && canMakeGroupOpUniform(*this)) {
    setUniform(true);
    return getResult();
  }

  return nullptr;
}

//===----------------------------------------------------------------------===//
// AsyncOpInterface
//===----------------------------------------------------------------------===//

void gpu::addAsyncDependency(Operation *op, Value token) {
  op->insertOperands(0, {token});
  if (!op->template hasTrait<OpTrait::AttrSizedOperandSegments>())
    return;
  auto attrName =
      OpTrait::AttrSizedOperandSegments<void>::getOperandSegmentSizeAttr();
  auto sizeAttr = op->template getAttrOfType<DenseI32ArrayAttr>(attrName);

  // Async dependencies is the only variadic operand.
  if (!sizeAttr)
    return;

  SmallVector<int32_t, 8> sizes(sizeAttr.asArrayRef());
  ++sizes.front();
  op->setAttr(attrName, Builder(op->getContext()).getDenseI32ArrayAttr(sizes));
}

//===----------------------------------------------------------------------===//
// LaunchOp
//===----------------------------------------------------------------------===//

void LaunchOp::build(OpBuilder &builder, OperationState &result,
                     Value gridSizeX, Value gridSizeY, Value gridSizeZ,
                     Value getBlockSizeX, Value getBlockSizeY,
                     Value getBlockSizeZ, Value dynamicSharedMemorySize,
                     Type asyncTokenType, ValueRange asyncDependencies,
                     TypeRange workgroupAttributions,
                     TypeRange privateAttributions, Value clusterSizeX,
                     Value clusterSizeY, Value clusterSizeZ) {
  // Add a WorkGroup attribution attribute. This attribute is required to
  // identify private attributions in the list of block argguments.
  result.addAttribute(getNumWorkgroupAttributionsAttrName(),
                      builder.getI64IntegerAttr(workgroupAttributions.size()));

  // Add Op operands.
  result.addOperands(asyncDependencies);
  if (asyncTokenType)
    result.types.push_back(builder.getType<AsyncTokenType>());

  // Add grid and block sizes as op operands, followed by the data operands.
  result.addOperands({gridSizeX, gridSizeY, gridSizeZ, getBlockSizeX,
                      getBlockSizeY, getBlockSizeZ});
  if (clusterSizeX)
    result.addOperands(clusterSizeX);
  if (clusterSizeY)
    result.addOperands(clusterSizeY);
  if (clusterSizeZ)
    result.addOperands(clusterSizeZ);
  if (dynamicSharedMemorySize)
    result.addOperands(dynamicSharedMemorySize);

  // Create a kernel body region with kNumConfigRegionAttributes + N memory
  // attributions, where the first kNumConfigRegionAttributes arguments have
  // `index` type and the rest have the same types as the data operands.
  Region *kernelRegion = result.addRegion();
  Block *body = new Block();
  // TODO: Allow passing in proper locations here.
  for (unsigned i = 0; i < kNumConfigRegionAttributes; ++i)
    body->addArgument(builder.getIndexType(), result.location);
  // Add WorkGroup & Private attributions to the region arguments.
  for (Type argTy : workgroupAttributions)
    body->addArgument(argTy, result.location);
  for (Type argTy : privateAttributions)
    body->addArgument(argTy, result.location);
  kernelRegion->push_back(body);
  // Fill OperandSegmentSize Attribute.
  SmallVector<int32_t, 11> segmentSizes(11, 1);
  segmentSizes.front() = asyncDependencies.size();
  segmentSizes.back() = dynamicSharedMemorySize ? 1 : 0;
  segmentSizes[7] = clusterSizeX ? 1 : 0;
  segmentSizes[8] = clusterSizeY ? 1 : 0;
  segmentSizes[9] = clusterSizeZ ? 1 : 0;
  result.addAttribute(getOperandSegmentSizeAttr(),
                      builder.getDenseI32ArrayAttr(segmentSizes));
}

KernelDim3 LaunchOp::getBlockIds() {
  assert(!getBody().empty() && "LaunchOp body must not be empty.");
  auto args = getBody().getArguments();
  return KernelDim3{args[0], args[1], args[2]};
}

KernelDim3 LaunchOp::getThreadIds() {
  assert(!getBody().empty() && "LaunchOp body must not be empty.");
  auto args = getBody().getArguments();
  return KernelDim3{args[3], args[4], args[5]};
}

KernelDim3 LaunchOp::getGridSize() {
  assert(!getBody().empty() && "LaunchOp body must not be empty.");
  auto args = getBody().getArguments();
  return KernelDim3{args[6], args[7], args[8]};
}

KernelDim3 LaunchOp::getBlockSize() {
  assert(!getBody().empty() && "LaunchOp body must not be empty.");
  auto args = getBody().getArguments();
  return KernelDim3{args[9], args[10], args[11]};
}

std::optional<KernelDim3> LaunchOp::getClusterIds() {
  assert(!getBody().empty() && "LaunchOp body must not be empty.");
  if (!hasClusterSize())
    return std::nullopt;
  auto args = getBody().getArguments();
  return KernelDim3{args[12], args[13], args[14]};
}

std::optional<KernelDim3> LaunchOp::getClusterSize() {
  assert(!getBody().empty() && "LaunchOp body must not be empty.");
  if (!hasClusterSize())
    return std::nullopt;
  auto args = getBody().getArguments();
  return KernelDim3{args[15], args[16], args[17]};
}

KernelDim3 LaunchOp::getGridSizeOperandValues() {
  auto operands = getOperands().drop_front(getAsyncDependencies().size());
  return KernelDim3{operands[0], operands[1], operands[2]};
}

KernelDim3 LaunchOp::getBlockSizeOperandValues() {
  auto operands = getOperands().drop_front(getAsyncDependencies().size());
  return KernelDim3{operands[3], operands[4], operands[5]};
}

std::optional<KernelDim3> LaunchOp::getClusterSizeOperandValues() {
  auto operands = getOperands().drop_front(getAsyncDependencies().size());
  if (!hasClusterSize())
    return std::nullopt;
  return KernelDim3{operands[6], operands[7], operands[8]};
}

LogicalResult LaunchOp::verify() {
  if (!(hasClusterSize()) &&
      (getClusterSizeX() || getClusterSizeY() || getClusterSizeZ()))
    return emitOpError() << "cluster size must be all present";
  return success();
}

LogicalResult LaunchOp::verifyRegions() {
  // Kernel launch takes kNumConfigOperands leading operands for grid/block
  // sizes and transforms them into kNumConfigRegionAttributes region arguments
  // for block/thread identifiers and grid/block sizes.
  if (!getBody().empty()) {
    if (getBody().getNumArguments() <
        kNumConfigRegionAttributes + getNumWorkgroupAttributions())
      return emitOpError("unexpected number of region arguments");
  }

  // Verify Attributions Address Spaces.
  if (failed(verifyAttributions(getOperation(), getWorkgroupAttributions(),
                                GPUDialect::getWorkgroupAddressSpace())) ||
      failed(verifyAttributions(getOperation(), getPrivateAttributions(),
                                GPUDialect::getPrivateAddressSpace())))
    return failure();

  // Block terminators without successors are expected to exit the kernel region
  // and must be `gpu.terminator`.
  for (Block &block : getBody()) {
    if (block.empty())
      continue;
    if (block.back().getNumSuccessors() != 0)
      continue;
    if (!isa<gpu::TerminatorOp>(&block.back())) {
      return block.back()
          .emitError()
          .append("expected '", gpu::TerminatorOp::getOperationName(),
                  "' or a terminator with successors")
          .attachNote(getLoc())
          .append("in '", LaunchOp::getOperationName(), "' body region");
    }
  }

  if (getNumResults() == 0 && getAsyncToken())
    return emitOpError("needs to be named when async keyword is specified");

  return success();
}

// Pretty-print the kernel grid/block size assignment as
//   (%iter-x, %iter-y, %iter-z) in
//   (%size-x = %ssa-use, %size-y = %ssa-use, %size-z = %ssa-use)
// where %size-* and %iter-* will correspond to the body region arguments.
static void printSizeAssignment(OpAsmPrinter &p, KernelDim3 size,
                                KernelDim3 operands, KernelDim3 ids) {
  p << '(' << ids.x << ", " << ids.y << ", " << ids.z << ") in (";
  p << size.x << " = " << operands.x << ", ";
  p << size.y << " = " << operands.y << ", ";
  p << size.z << " = " << operands.z << ')';
}

void LaunchOp::print(OpAsmPrinter &p) {
  if (getAsyncToken()) {
    p << " async";
    if (!getAsyncDependencies().empty())
      p << " [" << getAsyncDependencies() << ']';
  }
  // Print the launch configuration.
  if (hasClusterSize()) {
    p << ' ' << getClustersKeyword();
    printSizeAssignment(p, getClusterSize().value(),
                        getClusterSizeOperandValues().value(),
                        getClusterIds().value());
  }
  p << ' ' << getBlocksKeyword();
  printSizeAssignment(p, getGridSize(), getGridSizeOperandValues(),
                      getBlockIds());
  p << ' ' << getThreadsKeyword();
  printSizeAssignment(p, getBlockSize(), getBlockSizeOperandValues(),
                      getThreadIds());
  if (getDynamicSharedMemorySize())
    p << ' ' << getDynamicSharedMemorySizeKeyword() << ' '
      << getDynamicSharedMemorySize();

  printAttributions(p, getWorkgroupKeyword(), getWorkgroupAttributions());
  printAttributions(p, getPrivateKeyword(), getPrivateAttributions());

  p << ' ';

  p.printRegion(getBody(), /*printEntryBlockArgs=*/false);
  p.printOptionalAttrDict((*this)->getAttrs(), /*elidedAttrs=*/{
                              LaunchOp::getOperandSegmentSizeAttr(),
                              getNumWorkgroupAttributionsAttrName()});
}

// Parse the size assignment blocks for blocks and threads.  These have the form
//   (%region_arg, %region_arg, %region_arg) in
//   (%region_arg = %operand, %region_arg = %operand, %region_arg = %operand)
// where %region_arg are percent-identifiers for the region arguments to be
// introduced further (SSA defs), and %operand are percent-identifiers for the
// SSA value uses.
static ParseResult
parseSizeAssignment(OpAsmParser &parser,
                    MutableArrayRef<OpAsmParser::UnresolvedOperand> sizes,
                    MutableArrayRef<OpAsmParser::UnresolvedOperand> regionSizes,
                    MutableArrayRef<OpAsmParser::UnresolvedOperand> indices) {
  assert(indices.size() == 3 && "space for three indices expected");
  SmallVector<OpAsmParser::UnresolvedOperand, 3> args;
  if (parser.parseOperandList(args, OpAsmParser::Delimiter::Paren,
                              /*allowResultNumber=*/false) ||
      parser.parseKeyword("in") || parser.parseLParen())
    return failure();
  std::move(args.begin(), args.end(), indices.begin());

  for (int i = 0; i < 3; ++i) {
    if (i != 0 && parser.parseComma())
      return failure();
    if (parser.parseOperand(regionSizes[i], /*allowResultNumber=*/false) ||
        parser.parseEqual() || parser.parseOperand(sizes[i]))
      return failure();
  }

  return parser.parseRParen();
}

/// Parses a Launch operation.
/// operation ::= `gpu.launch` (`async` `[` ssa-id-list `]`)?
///       `clusters` `(` ssa-id-list `)` `in` ssa-reassignment (Optional)
///       `blocks` `(` ssa-id-list `)` `in` ssa-reassignment
///       `threads` `(` ssa-id-list `)` `in` ssa-reassignment
///       memory-attribution
///       region attr-dict?
/// ssa-reassignment ::= `(` ssa-id `=` ssa-use (`,` ssa-id `=` ssa-use)* `)`
ParseResult LaunchOp::parse(OpAsmParser &parser, OperationState &result) {
  // Sizes of the grid and block.
  SmallVector<OpAsmParser::UnresolvedOperand, LaunchOp::kNumConfigOperands>
      sizes(LaunchOp::kNumConfigOperands);

  // Actual (data) operands passed to the kernel.
  SmallVector<OpAsmParser::UnresolvedOperand, 4> dataOperands;

  // Region arguments to be created.
  SmallVector<OpAsmParser::UnresolvedOperand, 16> regionArgs(
      LaunchOp::kNumConfigRegionAttributes);

  // Parse optional async dependencies.
  SmallVector<OpAsmParser::UnresolvedOperand, 4> asyncDependencies;
  Type asyncTokenType;
  if (failed(
          parseAsyncDependencies(parser, asyncTokenType, asyncDependencies)) ||
      parser.resolveOperands(asyncDependencies, asyncTokenType,
                             result.operands))
    return failure();
  if (parser.getNumResults() > 0)
    result.types.push_back(asyncTokenType);

  bool hasCluster = false;
  if (succeeded(
          parser.parseOptionalKeyword(LaunchOp::getClustersKeyword().data()))) {
    hasCluster = true;
    sizes.resize(9);
    regionArgs.resize(18);
  }
  MutableArrayRef<OpAsmParser::UnresolvedOperand> sizesRef(sizes);
  MutableArrayRef<OpAsmParser::UnresolvedOperand> regionArgsRef(regionArgs);

  // Last three segment assigns the cluster size. In the region argument
  // list, this is last 6 arguments.
  if (hasCluster) {
    if (parseSizeAssignment(parser, sizesRef.drop_front(6),
                            regionArgsRef.slice(15, 3),
                            regionArgsRef.slice(12, 3)))
      return failure();
  }
  // Parse the size assignment segments: the first segment assigns grid sizes
  // and defines values for block identifiers; the second segment assigns block
  // sizes and defines values for thread identifiers.  In the region argument
  // list, identifiers precede sizes, and block-related values precede
  // thread-related values.
  if (parser.parseKeyword(LaunchOp::getBlocksKeyword().data()) ||
      parseSizeAssignment(parser, sizesRef.take_front(3),
                          regionArgsRef.slice(6, 3),
                          regionArgsRef.slice(0, 3)) ||
      parser.parseKeyword(LaunchOp::getThreadsKeyword().data()) ||
      parseSizeAssignment(parser, sizesRef.drop_front(3),
                          regionArgsRef.slice(9, 3),
                          regionArgsRef.slice(3, 3)) ||
      parser.resolveOperands(sizes, parser.getBuilder().getIndexType(),
                             result.operands))
    return failure();

  OpAsmParser::UnresolvedOperand dynamicSharedMemorySize;
  bool hasDynamicSharedMemorySize = false;
  if (!parser.parseOptionalKeyword(
          LaunchOp::getDynamicSharedMemorySizeKeyword())) {
    hasDynamicSharedMemorySize = true;
    if (parser.parseOperand(dynamicSharedMemorySize) ||
        parser.resolveOperand(dynamicSharedMemorySize,
                              parser.getBuilder().getI32Type(),
                              result.operands))
      return failure();
  }

  // Create the region arguments, it has kNumConfigRegionAttributes arguments
  // that correspond to block/thread identifiers and grid/block sizes, all
  // having `index` type, a variadic number of WorkGroup Attributions and
  // a variadic number of Private Attributions. The number of WorkGroup
  // Attributions is stored in the attr with name:
  // LaunchOp::getNumWorkgroupAttributionsAttrName().
  Type index = parser.getBuilder().getIndexType();
  SmallVector<Type, LaunchOp::kNumConfigRegionAttributes> dataTypes(
      LaunchOp::kNumConfigRegionAttributes + 6, index);

  SmallVector<OpAsmParser::Argument> regionArguments;
  for (auto ssaValueAndType : llvm::zip(regionArgs, dataTypes)) {
    OpAsmParser::Argument arg;
    arg.ssaName = std::get<0>(ssaValueAndType);
    arg.type = std::get<1>(ssaValueAndType);
    regionArguments.push_back(arg);
  }

  Builder &builder = parser.getBuilder();
  // Parse workgroup memory attributions.
  if (failed(parseAttributions(parser, LaunchOp::getWorkgroupKeyword(),
                               regionArguments)))
    return failure();

  // Store the number of operands we just parsed as the number of workgroup
  // memory attributions.
  unsigned numWorkgroupAttrs = regionArguments.size() -
                               LaunchOp::kNumConfigRegionAttributes -
                               (hasCluster ? 6 : 0);
  result.addAttribute(LaunchOp::getNumWorkgroupAttributionsAttrName(),
                      builder.getI64IntegerAttr(numWorkgroupAttrs));

  // Parse private memory attributions.
  if (failed(parseAttributions(parser, LaunchOp::getPrivateKeyword(),
                               regionArguments)))
    return failure();

  // Introduce the body region and parse it. The region has
  // kNumConfigRegionAttributes arguments that correspond to
  // block/thread identifiers and grid/block sizes, all having `index` type.
  Region *body = result.addRegion();
  if (parser.parseRegion(*body, regionArguments) ||
      parser.parseOptionalAttrDict(result.attributes))
    return failure();

  SmallVector<int32_t, 11> segmentSizes(11, 1);
  segmentSizes.front() = asyncDependencies.size();

  if (!hasCluster) {
    segmentSizes[7] = 0;
    segmentSizes[8] = 0;
    segmentSizes[9] = 0;
  }
  segmentSizes.back() = hasDynamicSharedMemorySize ? 1 : 0;
  result.addAttribute(LaunchOp::getOperandSegmentSizeAttr(),
                      parser.getBuilder().getDenseI32ArrayAttr(segmentSizes));
  return success();
}

/// Simplify the gpu.launch when the range of a thread or block ID is
/// trivially known to be one.
struct FoldLaunchArguments : public OpRewritePattern<LaunchOp> {
  using OpRewritePattern<LaunchOp>::OpRewritePattern;
  LogicalResult matchAndRewrite(LaunchOp op,
                                PatternRewriter &rewriter) const override {
    // If the range implies a single value for `id`, replace `id`'s uses by
    // zero.
    Value zero;
    bool simplified = false;
    auto constPropIdUses = [&](Value id, Value size) {
      // Check if size is trivially one.
      if (!matchPattern(size, m_One()))
        return;
      if (id.getUses().empty())
        return;
      if (!simplified) {
        // Create a zero value the first time.
        OpBuilder::InsertionGuard guard(rewriter);
        rewriter.setInsertionPointToStart(&op.getBody().front());
        zero =
            rewriter.create<arith::ConstantIndexOp>(op.getLoc(), /*value=*/0);
      }
      rewriter.replaceAllUsesWith(id, zero);
      simplified = true;
    };
    constPropIdUses(op.getBlockIds().x, op.getGridSizeX());
    constPropIdUses(op.getBlockIds().y, op.getGridSizeY());
    constPropIdUses(op.getBlockIds().z, op.getGridSizeZ());
    constPropIdUses(op.getThreadIds().x, op.getBlockSizeX());
    constPropIdUses(op.getThreadIds().y, op.getBlockSizeY());
    constPropIdUses(op.getThreadIds().z, op.getBlockSizeZ());

    return success(simplified);
  }
};

void LaunchOp::getCanonicalizationPatterns(RewritePatternSet &rewrites,
                                           MLIRContext *context) {
  rewrites.add<FoldLaunchArguments>(context);
}

/// Adds a new block argument that corresponds to buffers located in
/// workgroup memory.
BlockArgument LaunchOp::addWorkgroupAttribution(Type type, Location loc) {
  auto attrName = getNumWorkgroupAttributionsAttrName();
  auto attr = (*this)->getAttrOfType<IntegerAttr>(attrName);
  (*this)->setAttr(attrName,
                   IntegerAttr::get(attr.getType(), attr.getValue() + 1));
  return getBody().insertArgument(
      LaunchOp::getNumConfigRegionAttributes() + attr.getInt(), type, loc);
}

/// Adds a new block argument that corresponds to buffers located in
/// private memory.
BlockArgument LaunchOp::addPrivateAttribution(Type type, Location loc) {
  // Buffers on the private memory always come after buffers on the workgroup
  // memory.
  return getBody().addArgument(type, loc);
}

//===----------------------------------------------------------------------===//
// LaunchFuncOp
//===----------------------------------------------------------------------===//

void LaunchFuncOp::build(OpBuilder &builder, OperationState &result,
                         GPUFuncOp kernelFunc, KernelDim3 gridSize,
                         KernelDim3 getBlockSize, Value dynamicSharedMemorySize,
                         ValueRange kernelOperands, Type asyncTokenType,
                         ValueRange asyncDependencies,
                         std::optional<KernelDim3> clusterSize) {
  result.addOperands(asyncDependencies);
  if (asyncTokenType)
    result.types.push_back(builder.getType<AsyncTokenType>());

  // Add grid and block sizes as op operands, followed by the data operands.
  result.addOperands({gridSize.x, gridSize.y, gridSize.z, getBlockSize.x,
                      getBlockSize.y, getBlockSize.z});
  if (clusterSize.has_value())
    result.addOperands({clusterSize->x, clusterSize->y, clusterSize->z});
  if (dynamicSharedMemorySize)
    result.addOperands(dynamicSharedMemorySize);
  result.addOperands(kernelOperands);
  auto kernelModule = kernelFunc->getParentOfType<GPUModuleOp>();
  auto kernelSymbol =
      SymbolRefAttr::get(kernelModule.getNameAttr(),
                         {SymbolRefAttr::get(kernelFunc.getNameAttr())});

  Properties &prop = result.getOrAddProperties<Properties>();
  prop.kernel = kernelSymbol;
  size_t segmentSizesLen = std::size(prop.operandSegmentSizes);
  // Initialize the segment sizes to 1.
  for (auto &sz : prop.operandSegmentSizes)
    sz = 1;
  prop.operandSegmentSizes[0] = asyncDependencies.size();
  if (!clusterSize.has_value()) {
    prop.operandSegmentSizes[segmentSizesLen - 4] = 0;
    prop.operandSegmentSizes[segmentSizesLen - 5] = 0;
    prop.operandSegmentSizes[segmentSizesLen - 6] = 0;
  }
  prop.operandSegmentSizes[segmentSizesLen - 3] =
      dynamicSharedMemorySize ? 1 : 0;
  prop.operandSegmentSizes[segmentSizesLen - 2] =
      static_cast<int32_t>(kernelOperands.size());
  prop.operandSegmentSizes[segmentSizesLen - 1] = 0;
}

void LaunchFuncOp::build(OpBuilder &builder, OperationState &result,
                         SymbolRefAttr kernel, KernelDim3 gridSize,
                         KernelDim3 getBlockSize, Value dynamicSharedMemorySize,
                         ValueRange kernelOperands, Value asyncObject,
                         std::optional<KernelDim3> clusterSize) {
  // Add grid and block sizes as op operands, followed by the data operands.
  result.addOperands({gridSize.x, gridSize.y, gridSize.z, getBlockSize.x,
                      getBlockSize.y, getBlockSize.z});
  if (clusterSize.has_value())
    result.addOperands({clusterSize->x, clusterSize->y, clusterSize->z});
  if (dynamicSharedMemorySize)
    result.addOperands(dynamicSharedMemorySize);
  result.addOperands(kernelOperands);
  if (asyncObject)
    result.addOperands(asyncObject);
  Properties &prop = result.getOrAddProperties<Properties>();
  prop.kernel = kernel;
  size_t segmentSizesLen = std::size(prop.operandSegmentSizes);
  // Initialize the segment sizes to 1.
  for (auto &sz : prop.operandSegmentSizes)
    sz = 1;
  prop.operandSegmentSizes[0] = 0;
  if (!clusterSize.has_value()) {
    prop.operandSegmentSizes[segmentSizesLen - 4] = 0;
    prop.operandSegmentSizes[segmentSizesLen - 5] = 0;
    prop.operandSegmentSizes[segmentSizesLen - 6] = 0;
  }
  prop.operandSegmentSizes[segmentSizesLen - 3] =
      dynamicSharedMemorySize ? 1 : 0;
  prop.operandSegmentSizes[segmentSizesLen - 2] =
      static_cast<int32_t>(kernelOperands.size());
  prop.operandSegmentSizes[segmentSizesLen - 1] = asyncObject ? 1 : 0;
}

StringAttr LaunchFuncOp::getKernelModuleName() {
  return getKernel().getRootReference();
}

StringAttr LaunchFuncOp::getKernelName() {
  return getKernel().getLeafReference();
}

unsigned LaunchFuncOp::getNumKernelOperands() {
  return getKernelOperands().size();
}

Value LaunchFuncOp::getKernelOperand(unsigned i) {
  return getKernelOperands()[i];
}

KernelDim3 LaunchFuncOp::getGridSizeOperandValues() {
  auto operands = getOperands().drop_front(getAsyncDependencies().size());
  return KernelDim3{operands[0], operands[1], operands[2]};
}

KernelDim3 LaunchFuncOp::getBlockSizeOperandValues() {
  auto operands = getOperands().drop_front(getAsyncDependencies().size());
  return KernelDim3{operands[3], operands[4], operands[5]};
}

KernelDim3 LaunchFuncOp::getClusterSizeOperandValues() {
  assert(hasClusterSize() &&
         "cluster size is not set, check hasClusterSize() first");
  auto operands = getOperands().drop_front(getAsyncDependencies().size());
  return KernelDim3{operands[6], operands[7], operands[8]};
}

LogicalResult LaunchFuncOp::verify() {
  auto module = (*this)->getParentOfType<ModuleOp>();
  if (!module)
    return emitOpError("expected to belong to a module");

  if (!module->getAttrOfType<UnitAttr>(
          GPUDialect::getContainerModuleAttrName()))
    return emitOpError("expected the closest surrounding module to have the '" +
                       GPUDialect::getContainerModuleAttrName() +
                       "' attribute");

  if (hasClusterSize()) {
    if (getClusterSizeY().getType() != getClusterSizeX().getType() ||
        getClusterSizeZ().getType() != getClusterSizeX().getType())
      return emitOpError()
             << "expects types of the cluster dimensions must be the same";
  }

  return success();
}

static ParseResult
parseLaunchDimType(OpAsmParser &parser, Type &dimTy,
                   std::optional<OpAsmParser::UnresolvedOperand> clusterValue,
                   Type &clusterXTy, Type &clusterYTy, Type &clusterZTy) {
  if (succeeded(parser.parseOptionalColon())) {
    if (parser.parseType(dimTy))
      return failure();
  } else {
    dimTy = IndexType::get(parser.getContext());
  }
  if (clusterValue.has_value()) {
    clusterXTy = clusterYTy = clusterZTy = dimTy;
  }
  return success();
}

static void printLaunchDimType(OpAsmPrinter &printer, Operation *op, Type dimTy,
                               Value clusterValue, Type clusterXTy,
                               Type clusterYTy, Type clusterZTy) {
  if (!dimTy.isIndex())
    printer << ": " << dimTy;
}

static ParseResult parseLaunchFuncOperands(
    OpAsmParser &parser,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &argNames,
    SmallVectorImpl<Type> &argTypes) {
  if (parser.parseOptionalKeyword("args"))
    return success();

  auto parseElement = [&]() -> ParseResult {
    return failure(parser.parseOperand(argNames.emplace_back()) ||
                   parser.parseColonType(argTypes.emplace_back()));
  };

  return parser.parseCommaSeparatedList(OpAsmParser::Delimiter::Paren,
                                        parseElement, " in argument list");
}

static void printLaunchFuncOperands(OpAsmPrinter &printer, Operation *,
                                    OperandRange operands, TypeRange types) {
  if (operands.empty())
    return;
  printer << "args(";
  llvm::interleaveComma(llvm::zip(operands, types), printer,
                        [&](const auto &pair) {
                          printer.printOperand(std::get<0>(pair));
                          printer << " : ";
                          printer.printType(std::get<1>(pair));
                        });
  printer << ")";
}

//===----------------------------------------------------------------------===//
// ShuffleOp
//===----------------------------------------------------------------------===//

void ShuffleOp::build(OpBuilder &builder, OperationState &result, Value value,
                      int32_t offset, int32_t width, ShuffleMode mode) {
  build(builder, result, value,
        builder.create<arith::ConstantOp>(result.location,
                                          builder.getI32IntegerAttr(offset)),
        builder.create<arith::ConstantOp>(result.location,
                                          builder.getI32IntegerAttr(width)),
        mode);
}

//===----------------------------------------------------------------------===//
// BarrierOp
//===----------------------------------------------------------------------===//

namespace {

/// Remove gpu.barrier after gpu.barrier, the threads are already synchronized!
LogicalResult eraseRedundantGpuBarrierOps(BarrierOp op,
                                          PatternRewriter &rewriter) {
  if (isa_and_nonnull<BarrierOp>(op->getNextNode())) {
    rewriter.eraseOp(op);
    return success();
  }
  return failure();
}

} // end anonymous namespace

void BarrierOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                            MLIRContext *context) {
  results.add(eraseRedundantGpuBarrierOps);
}

//===----------------------------------------------------------------------===//
// GPUFuncOp
//===----------------------------------------------------------------------===//

/// Adds a new block argument that corresponds to buffers located in
/// workgroup memory.
BlockArgument GPUFuncOp::addWorkgroupAttribution(Type type, Location loc) {
  auto attrName = getNumWorkgroupAttributionsAttrName();
  auto attr = (*this)->getAttrOfType<IntegerAttr>(attrName);
  (*this)->setAttr(attrName,
                   IntegerAttr::get(attr.getType(), attr.getValue() + 1));
  return getBody().insertArgument(
      getFunctionType().getNumInputs() + attr.getInt(), type, loc);
}

/// Adds a new block argument that corresponds to buffers located in
/// private memory.
BlockArgument GPUFuncOp::addPrivateAttribution(Type type, Location loc) {
  // Buffers on the private memory always come after buffers on the workgroup
  // memory.
  return getBody().addArgument(type, loc);
}

void GPUFuncOp::build(OpBuilder &builder, OperationState &result,
                      StringRef name, FunctionType type,
                      TypeRange workgroupAttributions,
                      TypeRange privateAttributions,
                      ArrayRef<NamedAttribute> attrs) {
  result.addAttribute(SymbolTable::getSymbolAttrName(),
                      builder.getStringAttr(name));
  result.addAttribute(getFunctionTypeAttrName(result.name),
                      TypeAttr::get(type));
  result.addAttribute(getNumWorkgroupAttributionsAttrName(),
                      builder.getI64IntegerAttr(workgroupAttributions.size()));
  result.addAttributes(attrs);
  Region *body = result.addRegion();
  Block *entryBlock = new Block;

  // TODO: Allow passing in proper locations here.
  for (Type argTy : type.getInputs())
    entryBlock->addArgument(argTy, result.location);
  for (Type argTy : workgroupAttributions)
    entryBlock->addArgument(argTy, result.location);
  for (Type argTy : privateAttributions)
    entryBlock->addArgument(argTy, result.location);

  body->getBlocks().push_back(entryBlock);
}

/// Parses a GPU function memory attribution.
///
/// memory-attribution ::= (`workgroup` `(` ssa-id-and-type-list `)`)?
///                        (`private` `(` ssa-id-and-type-list `)`)?
///
/// Note that this function parses only one of the two similar parts, with the
/// keyword provided as argument.
static ParseResult
parseAttributions(OpAsmParser &parser, StringRef keyword,
                  SmallVectorImpl<OpAsmParser::Argument> &args,
                  Attribute &attributionAttrs) {
  // If we could not parse the keyword, just assume empty list and succeed.
  if (failed(parser.parseOptionalKeyword(keyword)))
    return success();

  size_t existingArgs = args.size();
  ParseResult result =
      parser.parseArgumentList(args, OpAsmParser::Delimiter::Paren,
                               /*allowType=*/true, /*allowAttrs=*/true);
  if (failed(result))
    return result;

  bool hadAttrs = llvm::any_of(ArrayRef(args).drop_front(existingArgs),
                               [](const OpAsmParser::Argument &arg) -> bool {
                                 return arg.attrs && !arg.attrs.empty();
                               });
  if (!hadAttrs) {
    attributionAttrs = nullptr;
    return result;
  }

  Builder &builder = parser.getBuilder();
  SmallVector<Attribute> attributionAttrsVec;
  for (const auto &argument : ArrayRef(args).drop_front(existingArgs)) {
    if (!argument.attrs)
      attributionAttrsVec.push_back(builder.getDictionaryAttr({}));
    else
      attributionAttrsVec.push_back(argument.attrs);
  }
  attributionAttrs = builder.getArrayAttr(attributionAttrsVec);
  return result;
}

/// Parses a GPU function.
///
/// <operation> ::= `gpu.func` symbol-ref-id `(` argument-list `)`
///                 (`->` function-result-list)? memory-attribution `kernel`?
///                 function-attributes? region
ParseResult GPUFuncOp::parse(OpAsmParser &parser, OperationState &result) {
  SmallVector<OpAsmParser::Argument> entryArgs;
  SmallVector<DictionaryAttr> resultAttrs;
  SmallVector<Type> resultTypes;
  bool isVariadic;

  // Parse the function name.
  StringAttr nameAttr;
  if (parser.parseSymbolName(nameAttr, ::mlir::SymbolTable::getSymbolAttrName(),
                             result.attributes))
    return failure();

  auto signatureLocation = parser.getCurrentLocation();
  if (failed(function_interface_impl::parseFunctionSignature(
          parser, /*allowVariadic=*/false, entryArgs, isVariadic, resultTypes,
          resultAttrs)))
    return failure();

  if (!entryArgs.empty() && entryArgs[0].ssaName.name.empty())
    return parser.emitError(signatureLocation)
           << "gpu.func requires named arguments";

  // Construct the function type. More types will be added to the region, but
  // not to the function type.
  Builder &builder = parser.getBuilder();

  SmallVector<Type> argTypes;
  for (auto &arg : entryArgs)
    argTypes.push_back(arg.type);
  auto type = builder.getFunctionType(argTypes, resultTypes);
  result.addAttribute(getFunctionTypeAttrName(result.name),
                      TypeAttr::get(type));

  function_interface_impl::addArgAndResultAttrs(
      builder, result, entryArgs, resultAttrs, getArgAttrsAttrName(result.name),
      getResAttrsAttrName(result.name));

  Attribute workgroupAttributionAttrs;
  // Parse workgroup memory attributions.
  if (failed(parseAttributions(parser, GPUFuncOp::getWorkgroupKeyword(),
                               entryArgs, workgroupAttributionAttrs)))
    return failure();

  // Store the number of operands we just parsed as the number of workgroup
  // memory attributions.
  unsigned numWorkgroupAttrs = entryArgs.size() - type.getNumInputs();
  result.addAttribute(GPUFuncOp::getNumWorkgroupAttributionsAttrName(),
                      builder.getI64IntegerAttr(numWorkgroupAttrs));
  if (workgroupAttributionAttrs)
    result.addAttribute(GPUFuncOp::getWorkgroupAttribAttrsAttrName(result.name),
                        workgroupAttributionAttrs);

  Attribute privateAttributionAttrs;
  // Parse private memory attributions.
  if (failed(parseAttributions(parser, GPUFuncOp::getPrivateKeyword(),
                               entryArgs, privateAttributionAttrs)))
    return failure();
  if (privateAttributionAttrs)
    result.addAttribute(GPUFuncOp::getPrivateAttribAttrsAttrName(result.name),
                        privateAttributionAttrs);

  // Parse the kernel attribute if present.
  if (succeeded(parser.parseOptionalKeyword(GPUFuncOp::getKernelKeyword())))
    result.addAttribute(GPUDialect::getKernelFuncAttrName(),
                        builder.getUnitAttr());

  // Parse attributes.
  if (failed(parser.parseOptionalAttrDictWithKeyword(result.attributes)))
    return failure();

  // Parse the region. If no argument names were provided, take all names
  // (including those of attributions) from the entry block.
  auto *body = result.addRegion();
  return parser.parseRegion(*body, entryArgs);
}

static void printAttributions(OpAsmPrinter &p, StringRef keyword,
                              ArrayRef<BlockArgument> values,
                              ArrayAttr attributes) {
  if (values.empty())
    return;

  p << ' ' << keyword << '(';
  llvm::interleaveComma(
      llvm::enumerate(values), p, [&p, attributes](auto pair) {
        BlockArgument v = pair.value();
        p << v << " : " << v.getType();

        size_t attributionIndex = pair.index();
        DictionaryAttr attrs;
        if (attributes && attributionIndex < attributes.size())
          attrs = llvm::cast<DictionaryAttr>(attributes[attributionIndex]);
        if (attrs)
          p.printOptionalAttrDict(attrs.getValue());
      });
  p << ')';
}

void GPUFuncOp::print(OpAsmPrinter &p) {
  p << ' ';
  p.printSymbolName(getName());

  FunctionType type = getFunctionType();
  function_interface_impl::printFunctionSignature(p, *this, type.getInputs(),
                                                  /*isVariadic=*/false,
                                                  type.getResults());

  printAttributions(p, getWorkgroupKeyword(), getWorkgroupAttributions(),
                    getWorkgroupAttribAttrs().value_or(nullptr));
  printAttributions(p, getPrivateKeyword(), getPrivateAttributions(),
                    getPrivateAttribAttrs().value_or(nullptr));
  if (isKernel())
    p << ' ' << getKernelKeyword();

  function_interface_impl::printFunctionAttributes(
      p, *this,
      {getNumWorkgroupAttributionsAttrName(),
       GPUDialect::getKernelFuncAttrName(), getFunctionTypeAttrName(),
       getArgAttrsAttrName(), getResAttrsAttrName(),
       getWorkgroupAttribAttrsAttrName(), getPrivateAttribAttrsAttrName()});
  p << ' ';
  p.printRegion(getBody(), /*printEntryBlockArgs=*/false);
}

static DictionaryAttr getAttributionAttrs(GPUFuncOp op, unsigned index,
                                          StringAttr attrName) {
  auto allAttrs = llvm::dyn_cast_or_null<ArrayAttr>(op->getAttr(attrName));
  if (!allAttrs || index >= allAttrs.size())
    return DictionaryAttr();
  return llvm::cast<DictionaryAttr>(allAttrs[index]);
}

DictionaryAttr GPUFuncOp::getworkgroupAttributionAttrs(unsigned index) {
  return getAttributionAttrs(*this, index, getWorkgroupAttribAttrsAttrName());
}

DictionaryAttr GPUFuncOp::getPrivateAttributionAttrs(unsigned index) {
  return getAttributionAttrs(*this, index, getPrivateAttribAttrsAttrName());
}

static void setAttributionAttrs(GPUFuncOp op, unsigned index,
                                DictionaryAttr value, StringAttr attrName) {
  MLIRContext *ctx = op.getContext();
  auto allAttrs = llvm::dyn_cast_or_null<ArrayAttr>(op->getAttr(attrName));
  SmallVector<Attribute> elements;
  if (allAttrs)
    elements.append(allAttrs.begin(), allAttrs.end());
  while (elements.size() <= index)
    elements.push_back(DictionaryAttr::get(ctx));
  if (!value)
    elements[index] = DictionaryAttr::get(ctx);
  else
    elements[index] = value;
  ArrayAttr newValue = ArrayAttr::get(ctx, elements);
  op->setAttr(attrName, newValue);
}

void GPUFuncOp::setworkgroupAttributionAttrs(unsigned index,
                                             DictionaryAttr value) {
  setAttributionAttrs(*this, index, value, getWorkgroupAttribAttrsAttrName());
}

void GPUFuncOp::setPrivateAttributionAttrs(unsigned int index,
                                           DictionaryAttr value) {
  setAttributionAttrs(*this, index, value, getPrivateAttribAttrsAttrName());
}

static Attribute getAttributionAttr(GPUFuncOp op, unsigned index,
                                    StringAttr name, StringAttr attrsName) {
  DictionaryAttr dict = getAttributionAttrs(op, index, attrsName);
  if (!dict)
    return Attribute();
  return dict.get(name);
}

Attribute GPUFuncOp::getWorkgroupAttributionAttr(unsigned index,
                                                 StringAttr name) {
  assert(index < getNumWorkgroupAttributions() &&
         "index must map to a workgroup attribution");
  return getAttributionAttr(*this, index, name,
                            getWorkgroupAttribAttrsAttrName());
}

Attribute GPUFuncOp::getPrivateAttributionAttr(unsigned index,
                                               StringAttr name) {
  assert(index < getNumPrivateAttributions() &&
         "index must map to a private attribution");
  return getAttributionAttr(*this, index, name,
                            getPrivateAttribAttrsAttrName());
}

static void setAttributionAttr(GPUFuncOp op, unsigned index, StringAttr name,
                               Attribute value, StringAttr attrsName) {
  MLIRContext *ctx = op.getContext();
  SmallVector<NamedAttribute> elems;
  DictionaryAttr oldDict = getAttributionAttrs(op, index, attrsName);
  if (oldDict)
    elems.append(oldDict.getValue().begin(), oldDict.getValue().end());

  bool found = false;
  bool mustSort = true;
  for (unsigned i = 0, e = elems.size(); i < e; ++i) {
    if (elems[i].getName() == name) {
      found = true;
      if (!value) {
        std::swap(elems[i], elems[elems.size() - 1]);
        elems.pop_back();
      } else {
        mustSort = false;
        elems[i] = NamedAttribute(elems[i].getName(), value);
      }
      break;
    }
  }
  if (!found) {
    if (!value)
      return;
    elems.emplace_back(name, value);
  }
  if (mustSort) {
    DictionaryAttr::sortInPlace(elems);
  }
  auto newDict = DictionaryAttr::getWithSorted(ctx, elems);
  setAttributionAttrs(op, index, newDict, attrsName);
}

void GPUFuncOp::setWorkgroupAttributionAttr(unsigned index, StringAttr name,
                                            Attribute value) {
  assert(index < getNumWorkgroupAttributions() &&
         "index must map to a workgroup attribution");
  setAttributionAttr(*this, index, name, value,
                     getWorkgroupAttribAttrsAttrName());
}

void GPUFuncOp::setPrivateAttributionAttr(unsigned index, StringAttr name,
                                          Attribute value) {
  assert(index < getNumPrivateAttributions() &&
         "index must map to a private attribution");
  setAttributionAttr(*this, index, name, value,
                     getPrivateAttribAttrsAttrName());
}

LogicalResult GPUFuncOp::verifyType() {
  if (isKernel() && getFunctionType().getNumResults() != 0)
    return emitOpError() << "expected void return type for kernel function";

  return success();
}

/// Verifies the body of the function.
LogicalResult GPUFuncOp::verifyBody() {
  if (empty())
    return emitOpError() << "expected body with at least one block";
  unsigned numFuncArguments = getNumArguments();
  unsigned numWorkgroupAttributions = getNumWorkgroupAttributions();
  unsigned numBlockArguments = front().getNumArguments();
  if (numBlockArguments < numFuncArguments + numWorkgroupAttributions)
    return emitOpError() << "expected at least "
                         << numFuncArguments + numWorkgroupAttributions
                         << " arguments to body region";

  ArrayRef<Type> funcArgTypes = getFunctionType().getInputs();
  for (unsigned i = 0; i < numFuncArguments; ++i) {
    Type blockArgType = front().getArgument(i).getType();
    if (funcArgTypes[i] != blockArgType)
      return emitOpError() << "expected body region argument #" << i
                           << " to be of type " << funcArgTypes[i] << ", got "
                           << blockArgType;
  }

  if (failed(verifyAttributions(getOperation(), getWorkgroupAttributions(),
                                GPUDialect::getWorkgroupAddressSpace())) ||
      failed(verifyAttributions(getOperation(), getPrivateAttributions(),
                                GPUDialect::getPrivateAddressSpace())))
    return failure();

  return success();
}

static LogicalResult verifyKnownLaunchSizeAttr(gpu::GPUFuncOp op,
                                               StringRef attrName) {
  auto maybeAttr = op->getAttr(attrName);
  if (!maybeAttr)
    return success();
  auto array = llvm::dyn_cast<DenseI32ArrayAttr>(maybeAttr);
  if (!array)
    return op.emitOpError(attrName + " must be a dense i32 array");
  if (array.size() != 3)
    return op.emitOpError(attrName + " must contain exactly 3 elements");
  return success();
}

LogicalResult GPUFuncOp::verify() {
  if (failed(verifyKnownLaunchSizeAttr(*this, getKnownBlockSizeAttrName())))
    return failure();
  if (failed(verifyKnownLaunchSizeAttr(*this, getKnownGridSizeAttrName())))
    return failure();
  return success();
}

//===----------------------------------------------------------------------===//
// ReturnOp
//===----------------------------------------------------------------------===//

LogicalResult gpu::ReturnOp::verify() {
  GPUFuncOp function = (*this)->getParentOfType<GPUFuncOp>();

  FunctionType funType = function.getFunctionType();

  if (funType.getNumResults() != getOperands().size())
    return emitOpError()
        .append("expected ", funType.getNumResults(), " result operands")
        .attachNote(function.getLoc())
        .append("return type declared here");

  for (const auto &pair : llvm::enumerate(
           llvm::zip(function.getFunctionType().getResults(), getOperands()))) {
    auto [type, operand] = pair.value();
    if (type != operand.getType())
      return emitOpError() << "unexpected type `" << operand.getType()
                           << "' for operand #" << pair.index();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// GPUModuleOp
//===----------------------------------------------------------------------===//

void GPUModuleOp::build(OpBuilder &builder, OperationState &result,
                        StringRef name, ArrayAttr targets,
                        Attribute offloadingHandler) {
  ensureTerminator(*result.addRegion(), builder, result.location);
  result.attributes.push_back(builder.getNamedAttr(
      ::mlir::SymbolTable::getSymbolAttrName(), builder.getStringAttr(name)));

  Properties &props = result.getOrAddProperties<Properties>();
  if (targets)
    props.targets = targets;
  props.offloadingHandler = offloadingHandler;
}

void GPUModuleOp::build(OpBuilder &builder, OperationState &result,
                        StringRef name, ArrayRef<Attribute> targets,
                        Attribute offloadingHandler) {
  build(builder, result, name,
        targets.empty() ? ArrayAttr() : builder.getArrayAttr(targets),
        offloadingHandler);
}

ParseResult GPUModuleOp::parse(OpAsmParser &parser, OperationState &result) {
  StringAttr nameAttr;
  ArrayAttr targetsAttr;

  if (parser.parseSymbolName(nameAttr, mlir::SymbolTable::getSymbolAttrName(),
                             result.attributes))
    return failure();

  Properties &props = result.getOrAddProperties<Properties>();

  // Parse the optional offloadingHandler
  if (succeeded(parser.parseOptionalLess())) {
    if (parser.parseAttribute(props.offloadingHandler))
      return failure();
    if (parser.parseGreater())
      return failure();
  }

  // Parse the optional array of target attributes.
  OptionalParseResult targetsAttrResult =
      parser.parseOptionalAttribute(targetsAttr, Type{});
  if (targetsAttrResult.has_value()) {
    if (failed(*targetsAttrResult)) {
      return failure();
    }
    props.targets = targetsAttr;
  }

  // If module attributes are present, parse them.
  if (parser.parseOptionalAttrDictWithKeyword(result.attributes))
    return failure();

  // Parse the module body.
  auto *body = result.addRegion();
  if (parser.parseRegion(*body, {}))
    return failure();

  // Ensure that this module has a valid terminator.
  GPUModuleOp::ensureTerminator(*body, parser.getBuilder(), result.location);
  return success();
}

void GPUModuleOp::print(OpAsmPrinter &p) {
  p << ' ';
  p.printSymbolName(getName());

  if (Attribute attr = getOffloadingHandlerAttr()) {
    p << " <";
    p.printAttribute(attr);
    p << ">";
  }

  if (Attribute attr = getTargetsAttr()) {
    p << ' ';
    p.printAttribute(attr);
    p << ' ';
  }

  p.printOptionalAttrDictWithKeyword((*this)->getAttrs(),
                                     {mlir::SymbolTable::getSymbolAttrName(),
                                      getTargetsAttrName(),
                                      getOffloadingHandlerAttrName()});
  p << ' ';
  p.printRegion(getRegion(), /*printEntryBlockArgs=*/false,
                /*printBlockTerminators=*/false);
}

bool GPUModuleOp::hasTarget(Attribute target) {
  if (ArrayAttr targets = getTargetsAttr())
    return llvm::count(targets.getValue(), target);
  return false;
}

void GPUModuleOp::setTargets(ArrayRef<TargetAttrInterface> targets) {
  ArrayAttr &targetsAttr = getProperties().targets;
  SmallVector<Attribute> targetsVector(targets);
  targetsAttr = ArrayAttr::get(getContext(), targetsVector);
}

//===----------------------------------------------------------------------===//
// GPUBinaryOp
//===----------------------------------------------------------------------===//
void BinaryOp::build(OpBuilder &builder, OperationState &result, StringRef name,
                     Attribute offloadingHandler, ArrayAttr objects) {
  auto &properties = result.getOrAddProperties<Properties>();
  result.attributes.push_back(builder.getNamedAttr(
      SymbolTable::getSymbolAttrName(), builder.getStringAttr(name)));
  properties.objects = objects;
  if (offloadingHandler)
    properties.offloadingHandler = offloadingHandler;
  else
    properties.offloadingHandler = builder.getAttr<SelectObjectAttr>(nullptr);
}

void BinaryOp::build(OpBuilder &builder, OperationState &result, StringRef name,
                     Attribute offloadingHandler, ArrayRef<Attribute> objects) {
  build(builder, result, name, offloadingHandler,
        objects.empty() ? ArrayAttr() : builder.getArrayAttr(objects));
}

static ParseResult parseOffloadingHandler(OpAsmParser &parser,
                                          Attribute &offloadingHandler) {
  if (succeeded(parser.parseOptionalLess())) {
    if (parser.parseAttribute(offloadingHandler))
      return failure();
    if (parser.parseGreater())
      return failure();
  }
  if (!offloadingHandler)
    offloadingHandler = parser.getBuilder().getAttr<SelectObjectAttr>(nullptr);
  return success();
}

static void printOffloadingHandler(OpAsmPrinter &printer, Operation *op,
                                   Attribute offloadingHandler) {
  if (offloadingHandler != SelectObjectAttr::get(op->getContext(), nullptr))
    printer << '<' << offloadingHandler << '>';
}

//===----------------------------------------------------------------------===//
// GPUMemcpyOp
//===----------------------------------------------------------------------===//

LogicalResult MemcpyOp::verify() {
  auto srcType = getSrc().getType();
  auto dstType = getDst().getType();

  if (getElementTypeOrSelf(srcType) != getElementTypeOrSelf(dstType))
    return emitOpError("arguments have incompatible element type");

  if (failed(verifyCompatibleShape(srcType, dstType)))
    return emitOpError("arguments have incompatible shape");

  return success();
}

namespace {

/// Erases a common case of copy ops where a destination value is used only by
/// the copy op, alloc and dealloc ops.
struct EraseTrivialCopyOp : public OpRewritePattern<MemcpyOp> {
  using OpRewritePattern<MemcpyOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(MemcpyOp op,
                                PatternRewriter &rewriter) const override {
    Value dest = op.getDst();
    Operation *destDefOp = dest.getDefiningOp();
    // `dest` must be defined by an op having Allocate memory effect in order to
    // perform the folding.
    if (!destDefOp ||
        !hasSingleEffect<MemoryEffects::Allocate>(destDefOp, dest))
      return failure();
    // We can erase `op` iff `dest` has no other use apart from its
    // use by `op` and dealloc ops.
    if (llvm::any_of(dest.getUsers(), [op, dest](Operation *user) {
          return user != op &&
                 !hasSingleEffect<MemoryEffects::Free>(user, dest);
        }))
      return failure();
    // We can perform the folding if and only if op has a single async
    // dependency and produces an async token as result, or if it does not have
    // any async dependency and does not produce any async token result.
    if (op.getAsyncDependencies().size() > 1 ||
        ((op.getAsyncDependencies().empty() && op.getAsyncToken()) ||
         (!op.getAsyncDependencies().empty() && !op.getAsyncToken())))
      return failure();
    rewriter.replaceOp(op, op.getAsyncDependencies());
    return success();
  }
};

} // end anonymous namespace

void MemcpyOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                           MLIRContext *context) {
  results.add<EraseTrivialCopyOp>(context);
}

//===----------------------------------------------------------------------===//
// GPU_SubgroupMmaLoadMatrixOp
//===----------------------------------------------------------------------===//

LogicalResult SubgroupMmaLoadMatrixOp::verify() {
  auto srcType = getSrcMemref().getType();
  auto resType = getRes().getType();
  auto resMatrixType = llvm::cast<gpu::MMAMatrixType>(resType);
  auto operand = resMatrixType.getOperand();
  auto srcMemrefType = llvm::cast<MemRefType>(srcType);

  if (!isLastMemrefDimUnitStride(srcMemrefType))
    return emitError(
        "expected source memref most minor dim must have unit stride");

  if (!operand.equals("AOp") && !operand.equals("BOp") &&
      !operand.equals("COp"))
    return emitError("only AOp, BOp and COp can be loaded");

  return success();
}

//===----------------------------------------------------------------------===//
// GPU_SubgroupMmaStoreMatrixOp
//===----------------------------------------------------------------------===//

LogicalResult SubgroupMmaStoreMatrixOp::verify() {
  auto srcType = getSrc().getType();
  auto dstType = getDstMemref().getType();
  auto srcMatrixType = llvm::cast<gpu::MMAMatrixType>(srcType);
  auto dstMemrefType = llvm::cast<MemRefType>(dstType);

  if (!isLastMemrefDimUnitStride(dstMemrefType))
    return emitError(
        "expected destination memref most minor dim must have unit stride");

  if (!srcMatrixType.getOperand().equals("COp"))
    return emitError(
        "expected the operand matrix being stored to have 'COp' operand type");

  return success();
}

//===----------------------------------------------------------------------===//
// GPU_SubgroupMmaComputeOp
//===----------------------------------------------------------------------===//

LogicalResult SubgroupMmaComputeOp::verify() {
  enum OperandMap { A, B, C };
  SmallVector<MMAMatrixType, 3> opTypes;
  opTypes.push_back(llvm::cast<MMAMatrixType>(getOpA().getType()));
  opTypes.push_back(llvm::cast<MMAMatrixType>(getOpB().getType()));
  opTypes.push_back(llvm::cast<MMAMatrixType>(getOpC().getType()));

  if (!opTypes[A].getOperand().equals("AOp") ||
      !opTypes[B].getOperand().equals("BOp") ||
      !opTypes[C].getOperand().equals("COp"))
    return emitError("operands must be in the order AOp, BOp, COp");

  ArrayRef<int64_t> aShape, bShape, cShape;
  aShape = opTypes[A].getShape();
  bShape = opTypes[B].getShape();
  cShape = opTypes[C].getShape();

  if (aShape[1] != bShape[0] || aShape[0] != cShape[0] ||
      bShape[1] != cShape[1])
    return emitError("operand shapes do not satisfy matmul constraints");

  return success();
}

LogicalResult MemcpyOp::fold(FoldAdaptor adaptor,
                             SmallVectorImpl<::mlir::OpFoldResult> &results) {
  return memref::foldMemRefCast(*this);
}

LogicalResult MemsetOp::fold(FoldAdaptor adaptor,
                             SmallVectorImpl<::mlir::OpFoldResult> &results) {
  return memref::foldMemRefCast(*this);
}

//===----------------------------------------------------------------------===//
// GPU_WaitOp
//===----------------------------------------------------------------------===//

namespace {

/// Remove gpu.wait op use of gpu.wait op def without async dependencies.
/// %t = gpu.wait async []       // No async dependencies.
/// ...  gpu.wait ... [%t, ...]  // %t can be removed.
struct EraseRedundantGpuWaitOpPairs : public OpRewritePattern<WaitOp> {
public:
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(WaitOp op,
                                PatternRewriter &rewriter) const final {
    auto predicate = [](Value value) {
      auto waitOp = value.getDefiningOp<WaitOp>();
      return waitOp && waitOp->getNumOperands() == 0;
    };
    if (llvm::none_of(op.getAsyncDependencies(), predicate))
      return failure();
    SmallVector<Value> validOperands;
    for (Value operand : op->getOperands()) {
      if (predicate(operand))
        continue;
      validOperands.push_back(operand);
    }
    rewriter.modifyOpInPlace(op, [&]() { op->setOperands(validOperands); });
    return success();
  }
};

/// Simplify trivial gpu.wait ops for the following patterns.
/// 1. %t = gpu.wait async ... ops, where %t has no uses (regardless of async
/// dependencies).
/// 2. %t1 = gpu.wait async [%t0], in this case, we can replace uses of %t1 with
/// %t0.
/// 3. gpu.wait [] ops, i.e gpu.wait ops that neither have any async
/// dependencies nor return any token.
struct SimplifyGpuWaitOp : public OpRewritePattern<WaitOp> {
public:
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(WaitOp op,
                                PatternRewriter &rewriter) const final {
    // Erase gpu.wait ops that neither have any async dependencies nor return
    // any async token.
    if (op.getAsyncDependencies().empty() && !op.getAsyncToken()) {
      rewriter.eraseOp(op);
      return success();
    }
    // Replace uses of %t1 = gpu.wait async [%t0] ops with %t0 and erase the op.
    if (llvm::hasSingleElement(op.getAsyncDependencies()) &&
        op.getAsyncToken()) {
      rewriter.replaceOp(op, op.getAsyncDependencies());
      return success();
    }
    // Erase %t = gpu.wait async ... ops, where %t has no uses.
    if (op.getAsyncToken() && op.getAsyncToken().use_empty()) {
      rewriter.eraseOp(op);
      return success();
    }
    return failure();
  }
};

} // end anonymous namespace

void WaitOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                         MLIRContext *context) {
  results.add<EraseRedundantGpuWaitOpPairs, SimplifyGpuWaitOp>(context);
}

//===----------------------------------------------------------------------===//
// GPU_AllocOp
//===----------------------------------------------------------------------===//

LogicalResult AllocOp::verify() {
  auto memRefType = llvm::cast<MemRefType>(getMemref().getType());

  if (static_cast<int64_t>(getDynamicSizes().size()) !=
      memRefType.getNumDynamicDims())
    return emitOpError("dimension operand count does not equal memref "
                       "dynamic dimension count");

  unsigned numSymbols = 0;
  if (!memRefType.getLayout().isIdentity())
    numSymbols = memRefType.getLayout().getAffineMap().getNumSymbols();
  if (getSymbolOperands().size() != numSymbols) {
    return emitOpError(
        "symbol operand count does not equal memref symbol count");
  }

  return success();
}

namespace {

/// Folding of memref.dim(gpu.alloc(%size), %idx) -> %size similar to
/// `memref::AllocOp`.
struct SimplifyDimOfAllocOp : public OpRewritePattern<memref::DimOp> {
  using OpRewritePattern<memref::DimOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(memref::DimOp dimOp,
                                PatternRewriter &rewriter) const override {
    std::optional<int64_t> index = dimOp.getConstantIndex();
    if (!index)
      return failure();

    auto memrefType = llvm::dyn_cast<MemRefType>(dimOp.getSource().getType());
    if (!memrefType || !memrefType.isDynamicDim(index.value()))
      return failure();

    auto alloc = dimOp.getSource().getDefiningOp<AllocOp>();
    if (!alloc)
      return failure();

    Value substituteOp = *(alloc.getDynamicSizes().begin() +
                           memrefType.getDynamicDimIndex(index.value()));
    rewriter.replaceOp(dimOp, substituteOp);
    return success();
  }
};

} // namespace

void AllocOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                          MLIRContext *context) {
  results.add<SimplifyDimOfAllocOp>(context);
}

//===----------------------------------------------------------------------===//
// GPU object attribute
//===----------------------------------------------------------------------===//

LogicalResult ObjectAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                                 Attribute target, CompilationTarget format,
                                 StringAttr object, DictionaryAttr properties) {
  if (!target)
    return emitError() << "the target attribute cannot be null";
  if (target.hasPromiseOrImplementsInterface<TargetAttrInterface>())
    return success();
  return emitError() << "the target attribute must implement or promise the "
                        "`gpu::TargetAttrInterface`";
}

namespace {
LogicalResult parseObject(AsmParser &odsParser, CompilationTarget &format,
                          StringAttr &object) {
  std::optional<CompilationTarget> formatResult;
  StringRef enumKeyword;
  auto loc = odsParser.getCurrentLocation();
  if (failed(odsParser.parseOptionalKeyword(&enumKeyword)))
    formatResult = CompilationTarget::Fatbin;
  if (!formatResult &&
      (formatResult =
           gpu::symbolizeEnum<gpu::CompilationTarget>(enumKeyword)) &&
      odsParser.parseEqual())
    return odsParser.emitError(loc, "expected an equal sign");
  if (!formatResult)
    return odsParser.emitError(loc, "expected keyword for GPU object format");
  FailureOr<StringAttr> objectResult =
      FieldParser<StringAttr>::parse(odsParser);
  if (failed(objectResult))
    return odsParser.emitError(odsParser.getCurrentLocation(),
                               "failed to parse GPU_ObjectAttr parameter "
                               "'object' which is to be a `StringAttr`");
  format = *formatResult;
  object = *objectResult;
  return success();
}

void printObject(AsmPrinter &odsParser, CompilationTarget format,
                 StringAttr object) {
  if (format != CompilationTarget::Fatbin)
    odsParser << stringifyEnum(format) << " = ";
  odsParser << object;
}
} // namespace

//===----------------------------------------------------------------------===//
// GPU select object attribute
//===----------------------------------------------------------------------===//

LogicalResult
gpu::SelectObjectAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                              Attribute target) {
  // Check `target`, it can be null, an integer attr or a GPU Target attribute.
  if (target) {
    if (auto intAttr = mlir::dyn_cast<IntegerAttr>(target)) {
      if (intAttr.getInt() < 0) {
        return emitError() << "the object index must be positive";
      }
    } else if (!target.hasPromiseOrImplementsInterface<TargetAttrInterface>()) {
      return emitError()
             << "the target attribute must be a GPU Target attribute";
    }
  }
  return success();
}

//===----------------------------------------------------------------------===//
// DynamicSharedMemoryOp
//===----------------------------------------------------------------------===//

LogicalResult gpu::DynamicSharedMemoryOp::verify() {
  if (!getOperation()->getParentWithTrait<OpTrait::SymbolTable>())
    return emitOpError() << "must be inside an op with symbol table";

  MemRefType memrefType = getResultMemref().getType();
  // Check address space
  if (!GPUDialect::hasWorkgroupMemoryAddressSpace(memrefType)) {
    return emitOpError() << "address space must be "
                         << gpu::AddressSpaceAttr::getMnemonic() << "<"
                         << stringifyEnum(gpu::AddressSpace::Workgroup) << ">";
  }
  if (memrefType.hasStaticShape()) {
    return emitOpError() << "result memref type must be memref<?xi8, "
                            "#gpu.address_space<workgroup>>";
  }
  return success();
}

//===----------------------------------------------------------------------===//
// GPU target options
//===----------------------------------------------------------------------===//

TargetOptions::TargetOptions(
    StringRef toolkitPath, ArrayRef<std::string> linkFiles,
    StringRef cmdOptions, CompilationTarget compilationTarget,
    function_ref<SymbolTable *()> getSymbolTableCallback)
    : TargetOptions(TypeID::get<TargetOptions>(), toolkitPath, linkFiles,
                    cmdOptions, compilationTarget, getSymbolTableCallback) {}

TargetOptions::TargetOptions(
    TypeID typeID, StringRef toolkitPath, ArrayRef<std::string> linkFiles,
    StringRef cmdOptions, CompilationTarget compilationTarget,
    function_ref<SymbolTable *()> getSymbolTableCallback)
    : toolkitPath(toolkitPath.str()), linkFiles(linkFiles),
      cmdOptions(cmdOptions.str()), compilationTarget(compilationTarget),
      getSymbolTableCallback(getSymbolTableCallback), typeID(typeID) {}

TypeID TargetOptions::getTypeID() const { return typeID; }

StringRef TargetOptions::getToolkitPath() const { return toolkitPath; }

ArrayRef<std::string> TargetOptions::getLinkFiles() const { return linkFiles; }

StringRef TargetOptions::getCmdOptions() const { return cmdOptions; }

SymbolTable *TargetOptions::getSymbolTable() const {
  return getSymbolTableCallback ? getSymbolTableCallback() : nullptr;
}

CompilationTarget TargetOptions::getCompilationTarget() const {
  return compilationTarget;
}

CompilationTarget TargetOptions::getDefaultCompilationTarget() {
  return CompilationTarget::Fatbin;
}

std::pair<llvm::BumpPtrAllocator, SmallVector<const char *>>
TargetOptions::tokenizeCmdOptions() const {
  std::pair<llvm::BumpPtrAllocator, SmallVector<const char *>> options;
  llvm::StringSaver stringSaver(options.first);
  StringRef opts = cmdOptions;
  // For a correct tokenization of the command line options `opts` must be
  // unquoted, otherwise the tokenization function returns a single string: the
  // unquoted `cmdOptions` -which is not the desired behavior.
  // Remove any quotes if they are at the beginning and end of the string:
  if (!opts.empty() && opts.front() == '"' && opts.back() == '"')
    opts.consume_front("\""), opts.consume_back("\"");
  if (!opts.empty() && opts.front() == '\'' && opts.back() == '\'')
    opts.consume_front("'"), opts.consume_back("'");
#ifdef _WIN32
  llvm::cl::TokenizeWindowsCommandLine(opts, stringSaver, options.second,
                                       /*MarkEOLs=*/false);
#else
  llvm::cl::TokenizeGNUCommandLine(opts, stringSaver, options.second,
                                   /*MarkEOLs=*/false);
#endif // _WIN32
  return options;
}

MLIR_DEFINE_EXPLICIT_TYPE_ID(::mlir::gpu::TargetOptions)

#include "mlir/Dialect/GPU/IR/GPUOpInterfaces.cpp.inc"
#include "mlir/Dialect/GPU/IR/GPUOpsEnums.cpp.inc"

#define GET_ATTRDEF_CLASSES
#include "mlir/Dialect/GPU/IR/GPUOpsAttributes.cpp.inc"

#define GET_OP_CLASSES
#include "mlir/Dialect/GPU/IR/GPUOps.cpp.inc"

#include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.cpp.inc"