TheRock/compiler/amd-llvm/mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp

//===- OpenACCCG.cpp - OpenACC codegen ops, attributes, and types ---------===//
//
// Part of the MLIR 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
//
//===----------------------------------------------------------------------===//
//
// Implementation for OpenACC codegen operations, attributes, and types.
// These correspond to the definitions in OpenACCCG*.td tablegen files
// and are kept in a separate file because they do not represent direct mappings
// of OpenACC language constructs; they are intermediate representations used
// when decomposing and lowering primary `acc` dialect operations.
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/OpenACC/OpenACC.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/Region.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Transforms/RegionUtils.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"

using namespace mlir;
using namespace acc;

namespace {

/// Generic helper for single-region OpenACC ops that execute their body once
/// and then return to the parent operation with their results (if any).
static void
getSingleRegionOpSuccessorRegions(Operation *op, Region &region,
                                  RegionBranchPoint point,
                                  SmallVectorImpl<RegionSuccessor> &regions) {
  if (point.isParent()) {
    regions.push_back(RegionSuccessor(&region));
    return;
  }
  regions.push_back(RegionSuccessor::parent());
}

static ValueRange getSingleRegionSuccessorInputs(Operation *op,
                                                 RegionSuccessor successor) {
  return successor.isParent() ? ValueRange(op->getResults()) : ValueRange();
}

/// Remove empty acc.kernel_environment operations. If the operation has wait
/// operands, create a acc.wait operation to preserve synchronization.
struct RemoveEmptyKernelEnvironment
    : public OpRewritePattern<acc::KernelEnvironmentOp> {
  using OpRewritePattern<acc::KernelEnvironmentOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(acc::KernelEnvironmentOp op,
                                PatternRewriter &rewriter) const override {
    assert(op->getNumRegions() == 1 && "expected op to have one region");

    Block &block = op.getRegion().front();
    if (!block.empty())
      return failure();

    // Remove empty kernel environment.
    // Preserve synchronization by creating acc.wait operation if needed.
    if (!op.getWaitOperands().empty() || op.getWaitOnlyAttr())
      rewriter.replaceOpWithNewOp<acc::WaitOp>(
          op, op.getWaitOperands(), /*asyncOperand=*/Value(),
          op.getWaitDevnum(), /*async=*/nullptr, /*ifCond=*/Value());
    else
      rewriter.eraseOp(op);

    return success();
  }
};

static void updateComputeRegionInputOperandSegments(ComputeRegionOp op,
                                                    PatternRewriter &rewriter,
                                                    size_t numInput) {
  const size_t numLaunch = op.getLaunchArgs().size();
  op->setAttr(ComputeRegionOp::getOperandSegmentSizeAttr(),
              rewriter.getDenseI32ArrayAttr({static_cast<int32_t>(numLaunch),
                                             static_cast<int32_t>(numInput),
                                             op.getStream() ? 1 : 0}));
}

struct ComputeRegionRemoveDuplicateArgs
    : public OpRewritePattern<ComputeRegionOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(ComputeRegionOp op,
                                PatternRewriter &rewriter) const override {
    Block *body = op.getBody();
    const size_t numLaunch = op.getLaunchArgs().size();
    size_t numInput = op.getInputArgs().size();
    assert(body->getNumArguments() == numLaunch + numInput &&
           "region args mismatch");

    bool mergedAny = false;
    while (true) {
      bool merged = false;
      for (size_t j = 1; j < numInput && !merged; ++j) {
        for (size_t i = 0; i < j; ++i) {
          if (op->getOperand(static_cast<unsigned>(numLaunch + i)) !=
              op->getOperand(static_cast<unsigned>(numLaunch + j)))
            continue;
          unsigned keepIdx = static_cast<unsigned>(numLaunch + i);
          unsigned dropIdx = static_cast<unsigned>(numLaunch + j);
          rewriter.replaceAllUsesWith(body->getArgument(dropIdx),
                                      body->getArgument(keepIdx));
          body->eraseArgument(dropIdx);
          op->eraseOperand(dropIdx);
          --numInput;
          merged = true;
          mergedAny = true;
          break;
        }
      }
      if (!merged)
        break;
    }

    if (!mergedAny)
      return failure();
    updateComputeRegionInputOperandSegments(op, rewriter, numInput);
    return success();
  }
};

struct ComputeRegionRemoveUnusedArgs
    : public OpRewritePattern<ComputeRegionOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(ComputeRegionOp op,
                                PatternRewriter &rewriter) const override {
    Block *body = op.getBody();
    const size_t numLaunch = op.getLaunchArgs().size();
    size_t numInput = op.getInputArgs().size();
    assert(body->getNumArguments() == numLaunch + numInput &&
           "region args mismatch");

    bool changed = false;
    for (size_t k = numLaunch; k < numLaunch + numInput;) {
      if (!body->getArgument(static_cast<unsigned>(k)).use_empty()) {
        ++k;
        continue;
      }
      body->eraseArgument(static_cast<unsigned>(k));
      op->eraseOperand(static_cast<unsigned>(k));
      --numInput;
      changed = true;
    }

    if (!changed)
      return failure();
    updateComputeRegionInputOperandSegments(op, rewriter, numInput);
    return success();
  }
};

template <typename EffectTy>
static void addOperandEffect(
    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
        &effects,
    const MutableOperandRange &operand) {
  for (unsigned i = 0, e = operand.size(); i < e; ++i)
    effects.emplace_back(EffectTy::get(), &operand[i]);
}

template <typename EffectTy>
static void addResultEffect(
    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
        &effects,
    Value result) {
  effects.emplace_back(EffectTy::get(), mlir::cast<mlir::OpResult>(result));
}

static int64_t gpuProcessorIndex(gpu::Processor p) {
  switch (p) {
  case gpu::Processor::Sequential:
    return 0;
  case gpu::Processor::ThreadX:
    return 1;
  case gpu::Processor::ThreadY:
    return 2;
  case gpu::Processor::ThreadZ:
    return 3;
  case gpu::Processor::BlockX:
    return 4;
  case gpu::Processor::BlockY:
    return 5;
  case gpu::Processor::BlockZ:
    return 6;
  }
  llvm_unreachable("unhandled gpu::Processor");
}

static gpu::Processor indexToGpuProcessor(int64_t idx) {
  switch (idx) {
  case 0:
    return gpu::Processor::Sequential;
  case 1:
    return gpu::Processor::ThreadX;
  case 2:
    return gpu::Processor::ThreadY;
  case 3:
    return gpu::Processor::ThreadZ;
  case 4:
    return gpu::Processor::BlockX;
  case 5:
    return gpu::Processor::BlockY;
  case 6:
    return gpu::Processor::BlockZ;
  default:
    return gpu::Processor::Sequential;
  }
}

static GPUParallelDimAttr intToParDim(MLIRContext *context, int64_t dimInt) {
  return GPUParallelDimAttr::get(
      context, IntegerAttr::get(IndexType::get(context), dimInt));
}

static GPUParallelDimAttr processorParDim(MLIRContext *context,
                                          gpu::Processor proc) {
  return GPUParallelDimAttr::get(
      context,
      IntegerAttr::get(IndexType::get(context), gpuProcessorIndex(proc)));
}

static ParseResult parseProcessorValue(AsmParser &parser,
                                       GPUParallelDimAttr &dim) {
  std::string keyword;
  llvm::SMLoc loc = parser.getCurrentLocation();
  if (failed(parser.parseKeywordOrString(&keyword)))
    return failure();
  auto maybeProcessor = gpu::symbolizeProcessor(keyword);
  if (!maybeProcessor)
    return parser.emitError(loc)
           << "expected one of ::mlir::gpu::Processor enum names";
  dim = intToParDim(parser.getContext(), gpuProcessorIndex(*maybeProcessor));
  return success();
}

static void printProcessorValue(AsmPrinter &printer,
                                const GPUParallelDimAttr &attr) {
  gpu::Processor processor = indexToGpuProcessor(attr.getValue().getInt());
  printer << gpu::stringifyProcessor(processor);
}

} // namespace

//===----------------------------------------------------------------------===//
// KernelEnvironmentOp
//===----------------------------------------------------------------------===//

void KernelEnvironmentOp::getSuccessorRegions(
    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
  getSingleRegionOpSuccessorRegions(getOperation(), getRegion(), point,
                                    regions);
}

ValueRange KernelEnvironmentOp::getSuccessorInputs(RegionSuccessor successor) {
  return getSingleRegionSuccessorInputs(getOperation(), successor);
}

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

/// Extract async for `clauseDeviceType`. Returns true if a clause was found.
template <typename ComputeConstructT>
static bool
extractAsyncClause(ComputeConstructT computeConstruct,
                   DeviceType clauseDeviceType, MLIRContext *context,
                   std::optional<Value> &asyncOperand, UnitAttr &asyncOnly) {
  if (computeConstruct.hasAsyncOnly(clauseDeviceType)) {
    asyncOnly = UnitAttr::get(context);
    return true;
  }
  if (Value asyncValue = computeConstruct.getAsyncValue(clauseDeviceType)) {
    asyncOperand = asyncValue;
    return true;
  }
  return false;
}

/// Extract wait for `clauseDeviceType`. Returns true if a clause was found.
template <typename ComputeConstructT>
static bool extractWaitClause(ComputeConstructT computeConstruct,
                              DeviceType clauseDeviceType, MLIRContext *context,
                              std::optional<Value> &waitDevnum,
                              SmallVectorImpl<Value> &waitOperands,
                              UnitAttr &waitOnly) {
  if (computeConstruct.hasWaitOnly(clauseDeviceType)) {
    waitOnly = UnitAttr::get(context);
    return true;
  }
  Value devnum = computeConstruct.getWaitDevnum(clauseDeviceType);
  auto waitValues = computeConstruct.getWaitValues(clauseDeviceType);
  if (!devnum && waitValues.empty())
    return false;
  if (devnum)
    waitDevnum = devnum;
  waitOperands.append(waitValues.begin(), waitValues.end());
  return true;
}

template <typename ComputeConstructT>
static void populateKernelEnvironmentAsyncWait(
    ComputeConstructT computeConstruct, DeviceType deviceType,
    std::optional<Value> &asyncOperand, UnitAttr &asyncOnly,
    std::optional<Value> &waitDevnum, SmallVectorImpl<Value> &waitOperands,
    UnitAttr &waitOnly) {
  MLIRContext *context = computeConstruct->getContext();

  // Prefer device_type-specific clauses, then default ones.
  if (!extractAsyncClause(computeConstruct, deviceType, context, asyncOperand,
                          asyncOnly)) {
    if (deviceType != DeviceType::None)
      extractAsyncClause(computeConstruct, DeviceType::None, context,
                         asyncOperand, asyncOnly);
  }

  if (!extractWaitClause(computeConstruct, deviceType, context, waitDevnum,
                         waitOperands, waitOnly)) {
    if (deviceType != DeviceType::None)
      extractWaitClause(computeConstruct, DeviceType::None, context, waitDevnum,
                        waitOperands, waitOnly);
  }
}

template <typename ComputeConstructT>
KernelEnvironmentOp
KernelEnvironmentOp::createAndPopulate(ComputeConstructT computeConstruct,
                                       DeviceType deviceType,
                                       OpBuilder &builder) {
  std::optional<Value> asyncOperand;
  UnitAttr asyncOnly = nullptr;
  std::optional<Value> waitDevnum;
  SmallVector<Value> waitOperands;
  UnitAttr waitOnly = nullptr;
  populateKernelEnvironmentAsyncWait(computeConstruct, deviceType, asyncOperand,
                                     asyncOnly, waitDevnum, waitOperands,
                                     waitOnly);

  auto kernelEnvironment = KernelEnvironmentOp::create(
      builder, computeConstruct->getLoc(),
      computeConstruct.getDataClauseOperands(), asyncOperand.value_or(Value()),
      asyncOnly, waitDevnum.value_or(Value()), waitOperands, waitOnly);
  Block &block = kernelEnvironment.getRegion().emplaceBlock();
  builder.setInsertionPointToStart(&block);
  return kernelEnvironment;
}

template KernelEnvironmentOp
KernelEnvironmentOp::createAndPopulate<ParallelOp>(ParallelOp, DeviceType,
                                                   OpBuilder &);
template KernelEnvironmentOp
KernelEnvironmentOp::createAndPopulate<KernelsOp>(KernelsOp, DeviceType,
                                                  OpBuilder &);
template KernelEnvironmentOp
KernelEnvironmentOp::createAndPopulate<SerialOp>(SerialOp, DeviceType,
                                                 OpBuilder &);

LogicalResult KernelEnvironmentOp::verify() {
  if (getAsyncOnly() && getAsyncOperand())
    return emitError("async-only cannot appear with async operand");
  if (getWaitOnly() && (!getWaitOperands().empty() || getWaitDevnum()))
    return emitError("wait-only cannot appear with wait operands or devnum");
  return success();
}

//===----------------------------------------------------------------------===//
// FirstprivateMapInitialOp
//===----------------------------------------------------------------------===//

LogicalResult FirstprivateMapInitialOp::verify() {
  if (getDataClause() != acc::DataClause::acc_firstprivate)
    return emitError("data clause associated with firstprivate operation must "
                     "match its intent");
  if (!getVar())
    return emitError("must have var operand");
  if (!mlir::isa<mlir::acc::PointerLikeType>(getVar().getType()) &&
      !mlir::isa<mlir::acc::MappableType>(getVar().getType()))
    return emitError("var must be mappable or pointer-like");
  if (mlir::isa<mlir::acc::PointerLikeType>(getVar().getType()) &&
      getVarType() == getVar().getType())
    return emitError("varType must capture the element type of var");
  if (getModifiers() != acc::DataClauseModifier::none)
    return emitError("no data clause modifiers are allowed");
  return success();
}

void FirstprivateMapInitialOp::getEffects(
    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
        &effects) {
  effects.emplace_back(MemoryEffects::Read::get(),
                       acc::CurrentDeviceIdResource::get());
  addOperandEffect<MemoryEffects::Read>(effects, getVarMutable());
  addResultEffect<MemoryEffects::Write>(effects, getAccVar());
}

//===----------------------------------------------------------------------===//
// ReductionInitOp
//===----------------------------------------------------------------------===//

void ReductionInitOp::getSuccessorRegions(
    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
  getSingleRegionOpSuccessorRegions(getOperation(), getRegion(), point,
                                    regions);
}

void ReductionInitOp::getRegionInvocationBounds(
    ArrayRef<Attribute> operands,
    SmallVectorImpl<InvocationBounds> &invocationBounds) {
  invocationBounds.emplace_back(1, 1);
}

ValueRange ReductionInitOp::getSuccessorInputs(RegionSuccessor successor) {
  return getSingleRegionSuccessorInputs(getOperation(), successor);
}

LogicalResult ReductionInitOp::verify() {
  Block &block = getRegion().front();
  if (auto yieldOp = dyn_cast<acc::YieldOp>(block.getTerminator())) {
    if (yieldOp.getNumOperands() != 1)
      return emitOpError(
          "region must yield exactly one value (private storage)");
    if (yieldOp.getOperand(0).getType() != getVar().getType())
      return emitOpError("yielded value type must match var type");
  }
  return success();
}

//===----------------------------------------------------------------------===//
// ReductionCombineRegionOp
//===----------------------------------------------------------------------===//

void ReductionCombineRegionOp::getSuccessorRegions(
    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
  getSingleRegionOpSuccessorRegions(getOperation(), getRegion(), point,
                                    regions);
}

void ReductionCombineRegionOp::getRegionInvocationBounds(
    ArrayRef<Attribute> operands,
    SmallVectorImpl<InvocationBounds> &invocationBounds) {
  invocationBounds.emplace_back(1, 1);
}

ValueRange
ReductionCombineRegionOp::getSuccessorInputs(RegionSuccessor successor) {
  return getSingleRegionSuccessorInputs(getOperation(), successor);
}

LogicalResult ReductionCombineRegionOp::verify() {
  Block &block = getRegion().front();
  if (auto yieldOp = dyn_cast<acc::YieldOp>(block.getTerminator())) {
    if (yieldOp.getNumOperands() != 0)
      return emitOpError("region must be terminated by acc.yield with no "
                         "operands");
  }
  return success();
}

//===----------------------------------------------------------------------===//
// ReductionCombineOp
//===----------------------------------------------------------------------===//

void ReductionCombineOp::getEffects(
    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
        &effects) {
  effects.emplace_back(MemoryEffects::Read::get(), &getSrcMemrefMutable(),
                       SideEffects::DefaultResource::get());
  effects.emplace_back(MemoryEffects::Read::get(), &getDestMemrefMutable(),
                       SideEffects::DefaultResource::get());
  effects.emplace_back(MemoryEffects::Write::get(), &getDestMemrefMutable(),
                       SideEffects::DefaultResource::get());
}

//===----------------------------------------------------------------------===//
// ComputeRegionOp
//===----------------------------------------------------------------------===//

static ParWidthOp getParWidthOpForLaunchArg(ComputeRegionOp op,
                                            GPUParallelDimAttr parDim) {
  for (auto launchArg : op.getLaunchArgs()) {
    auto parOp = launchArg.getDefiningOp<ParWidthOp>();
    if (!parOp)
      continue;
    auto launchArgDim = cast<GPUParallelDimAttr>(parOp.getParDim());
    if (launchArgDim == parDim)
      return parOp;
  }
  return nullptr;
}

std::optional<Value> ComputeRegionOp::getLaunchArg(GPUParallelDimAttr parDim) {
  if (auto parWidthOp = getParWidthOpForLaunchArg(*this, parDim))
    return parWidthOp.getResult();
  return {};
}

std::optional<Value>
ComputeRegionOp::getKnownLaunchArg(GPUParallelDimAttr parDim) {
  if (auto parWidthOp = getParWidthOpForLaunchArg(*this, parDim))
    if (parWidthOp.getLaunchArg())
      return parWidthOp.getLaunchArg();
  return {};
}

std::optional<uint64_t>
ComputeRegionOp::getKnownConstantLaunchArg(GPUParallelDimAttr parDim) {
  auto knownParWidth = getKnownLaunchArg(parDim);
  if (knownParWidth.has_value())
    return getConstantIntValue(knownParWidth.value());
  return {};
}

BlockArgument ComputeRegionOp::appendInputArg(Value value) {
  getInputArgsMutable().append(value);
  return getBody()->addArgument(value.getType(), getLoc());
}

std::optional<BlockArgument>
ComputeRegionOp::wireHoistedValueThroughIns(Value value) {
  Region &region = getRegion();

  auto useIsInRegion = [&](OpOperand &use) -> bool {
    return region.isAncestor(use.getOwner()->getParentRegion());
  };

  if (!areValuesDefinedAbove(ValueRange(value), region) ||
      !llvm::any_of(value.getUses(), useIsInRegion))
    return std::nullopt;

  BlockArgument arg = appendInputArg(value);
  replaceAllUsesInRegionWith(value, arg, region);
  return arg;
}

bool ComputeRegionOp::isEffectivelySerial() {
  auto *ctx = getContext();

  if (getLaunchArg(GPUParallelDimAttr::seqDim(ctx)))
    return true;

  auto checkDim = [&](GPUParallelDimAttr dim) -> bool {
    auto val = getKnownConstantLaunchArg(dim);
    return val && *val == 1;
  };

  return checkDim(GPUParallelDimAttr::threadXDim(ctx)) &&
         checkDim(GPUParallelDimAttr::threadYDim(ctx)) &&
         checkDim(GPUParallelDimAttr::threadZDim(ctx)) &&
         checkDim(GPUParallelDimAttr::blockXDim(ctx)) &&
         checkDim(GPUParallelDimAttr::blockYDim(ctx)) &&
         checkDim(GPUParallelDimAttr::blockZDim(ctx));
}

BlockArgument ComputeRegionOp::parDimToWidth(GPUParallelDimAttr parDim) {
  for (auto [pos, launchArg] : llvm::enumerate(getLaunchArgs())) {
    auto parOp = launchArg.getDefiningOp<ParWidthOp>();
    assert(parOp);
    auto launchArgDim = cast<GPUParallelDimAttr>(parOp.getParDim());
    if (launchArgDim == parDim) {
      assert(pos < getRegion().front().getNumArguments() &&
             "launch arg position out of range");
      return getRegion().front().getArgument(pos);
    }
  }
  llvm_unreachable("attempting to get unspecified parDim");
}

SmallVector<GPUParallelDimAttr> ComputeRegionOp::getLaunchParDims() {
  SmallVector<GPUParallelDimAttr> parDims;
  for (auto launchArg : getLaunchArgs()) {
    auto parOp = launchArg.getDefiningOp<ParWidthOp>();
    auto launchArgDim = cast<GPUParallelDimAttr>(parOp.getParDim());
    int64_t dimInt = launchArgDim.getValue().getInt();
    parDims.push_back(intToParDim(getContext(), dimInt));
  }
  return parDims;
}

Value ComputeRegionOp::getOperand(BlockArgument blockArg) {
  Block *body = getBody();
  if (blockArg.getOwner() != body)
    return Value();
  unsigned argNumber = blockArg.getArgNumber();
  unsigned numLaunchArgs = getLaunchArgs().size();
  unsigned numInputArgs = getInputArgs().size();
  if (argNumber >= numLaunchArgs + numInputArgs)
    return Value();
  if (argNumber < numLaunchArgs)
    return getLaunchArgs()[argNumber];
  return getInputArgs()[argNumber - numLaunchArgs];
}

std::optional<BlockArgument> ComputeRegionOp::getBlockArg(Value value) {
  Block *body = getBody();
  for (auto [idx, launchVal] : llvm::enumerate(getLaunchArgs())) {
    if (launchVal == value)
      return body->getArgument(idx);
  }
  unsigned numLaunch = getLaunchArgs().size();
  for (auto [idx, inputVal] : llvm::enumerate(getInputArgs())) {
    if (inputVal == value)
      return body->getArgument(numLaunch + idx);
  }
  return std::nullopt;
}

void ComputeRegionOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                  MLIRContext *context) {
  results.add<ComputeRegionRemoveDuplicateArgs, ComputeRegionRemoveUnusedArgs>(
      context);
}

BlockArgument ComputeRegionOp::gpuParWidth(gpu::Processor processor) {
  return parDimToWidth(GPUParallelDimAttr::get(getContext(), processor));
}

LogicalResult ComputeRegionOp::verify() {
  for (auto op : getLaunchArgs())
    if (!op.getDefiningOp<acc::ParWidthOp>())
      return emitOpError(
          "launch arguments must be results of acc.par_width operations");

  unsigned expectedBlockArgs = getLaunchArgs().size() + getInputArgs().size();
  unsigned actualBlockArgs = getRegion().front().getNumArguments();
  if (expectedBlockArgs != actualBlockArgs)
    return emitOpError("expected ")
           << expectedBlockArgs << " block arguments (launch + input), got "
           << actualBlockArgs;

  return success();
}

void ComputeRegionOp::print(OpAsmPrinter &p) {
  ValueRange regionArgs = getBody()->getArguments();
  ValueRange launchArgs = getLaunchArgs();
  ValueRange inputArgs = getInputArgs();

  assert(regionArgs.size() == (launchArgs.size() + inputArgs.size()) &&
         "region args mismatch");

  if (getStream())
    p << " stream(" << getStream() << " : " << getStream().getType() << ")";

  size_t i = 0;
  if (!launchArgs.empty()) {
    p << " launch(";
    for (size_t j = 0; j < launchArgs.size(); ++j, ++i) {
      p << regionArgs[i] << " = " << launchArgs[j];
      if (j < launchArgs.size() - 1)
        p << ", ";
    }
    p << ")";
  }
  if (!inputArgs.empty()) {
    p << " ins(";
    for (size_t j = 0; j < inputArgs.size(); ++j, ++i) {
      p << regionArgs[i] << " = " << inputArgs[j];
      if (j < inputArgs.size() - 1)
        p << ", ";
    }
    p << ") : (";
    for (size_t j = 0; j < inputArgs.size(); ++j) {
      p << inputArgs[j].getType();
      if (j < inputArgs.size() - 1)
        p << ", ";
    }
    p << ")";
  }
  p.printOptionalArrowTypeList(getResultTypes());
  p << " ";
  p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);
  p.printOptionalAttrDict((*this)->getAttrs(),
                          /*elidedAttrs=*/getOperandSegmentSizeAttr());
}

ParseResult ComputeRegionOp::parse(OpAsmParser &parser,
                                   OperationState &result) {
  auto &builder = parser.getBuilder();

  SmallVector<OpAsmParser::Argument> regionArgs;
  OpAsmParser::UnresolvedOperand streamOperand;
  Type streamType;
  SmallVector<OpAsmParser::UnresolvedOperand> launchOperands;
  SmallVector<OpAsmParser::UnresolvedOperand> inputOperands;
  SmallVector<Type> types;

  bool hasStream = false;
  if (succeeded(parser.parseOptionalKeyword("stream"))) {
    hasStream = true;
    if (parser.parseLParen() || parser.parseOperand(streamOperand) ||
        parser.parseColon() || parser.parseType(streamType) ||
        parser.parseRParen())
      return failure();
  }

  if (succeeded(parser.parseOptionalKeyword("launch"))) {
    if (parser.parseAssignmentList(regionArgs, launchOperands))
      return failure();
    Type indexType = builder.getIndexType();
    for (size_t i = 0; i < regionArgs.size(); ++i)
      types.push_back(indexType);
  }

  if (succeeded(parser.parseOptionalKeyword("ins"))) {
    if (parser.parseAssignmentList(regionArgs, inputOperands) ||
        parser.parseColon() || parser.parseLParen() ||
        parser.parseTypeList(types) || parser.parseRParen())
      return failure();
  }

  if (parser.parseOptionalArrowTypeList(result.types))
    return failure();

  for (auto [iterArg, type] : llvm::zip_equal(regionArgs, types))
    iterArg.type = type;

  Region *body = result.addRegion();
  if (parser.parseRegion(*body, regionArgs))
    return failure();
  ComputeRegionOp::ensureTerminator(*body, parser.getBuilder(),
                                    result.location);

  const size_t numLaunchOperands = launchOperands.size();
  const size_t numInputOperands = inputOperands.size();
  assert(numLaunchOperands + numInputOperands == regionArgs.size() &&
         "compute region args mismatch");

  result.addAttribute(
      ComputeRegionOp::getOperandSegmentSizeAttr(),
      builder.getDenseI32ArrayAttr({static_cast<int32_t>(numLaunchOperands),
                                    static_cast<int32_t>(numInputOperands),
                                    hasStream ? 1 : 0}));

  for (size_t i = 0; i < numLaunchOperands; ++i) {
    if (parser.resolveOperand(launchOperands[i], types[i], result.operands))
      return failure();
  }

  for (size_t i = numLaunchOperands; i < regionArgs.size(); ++i) {
    if (parser.resolveOperand(inputOperands[i - numLaunchOperands], types[i],
                              result.operands))
      return failure();
  }

  if (hasStream) {
    if (parser.resolveOperand(streamOperand, streamType, result.operands))
      return failure();
  }

  if (parser.parseOptionalAttrDict(result.attributes))
    return failure();

  return success();
}

//===----------------------------------------------------------------------===//
// GPUParallelDimAttr
//===----------------------------------------------------------------------===//

GPUParallelDimAttr GPUParallelDimAttr::get(MLIRContext *context,
                                           gpu::Processor proc) {
  return processorParDim(context, proc);
}

GPUParallelDimAttr GPUParallelDimAttr::seqDim(MLIRContext *context) {
  return processorParDim(context, gpu::Processor::Sequential);
}

GPUParallelDimAttr GPUParallelDimAttr::threadXDim(MLIRContext *context) {
  return processorParDim(context, gpu::Processor::ThreadX);
}

GPUParallelDimAttr GPUParallelDimAttr::threadYDim(MLIRContext *context) {
  return processorParDim(context, gpu::Processor::ThreadY);
}

GPUParallelDimAttr GPUParallelDimAttr::threadZDim(MLIRContext *context) {
  return processorParDim(context, gpu::Processor::ThreadZ);
}

GPUParallelDimAttr GPUParallelDimAttr::blockXDim(MLIRContext *context) {
  return processorParDim(context, gpu::Processor::BlockX);
}

GPUParallelDimAttr GPUParallelDimAttr::blockYDim(MLIRContext *context) {
  return processorParDim(context, gpu::Processor::BlockY);
}

GPUParallelDimAttr GPUParallelDimAttr::blockZDim(MLIRContext *context) {
  return processorParDim(context, gpu::Processor::BlockZ);
}

Attribute GPUParallelDimAttr::parse(AsmParser &parser, Type type) {
  GPUParallelDimAttr dim;
  if (parser.parseLess() || parseProcessorValue(parser, dim) ||
      parser.parseGreater()) {
    parser.emitError(parser.getCurrentLocation(),
                     "expected format `<` processor_name `>`");
    return {};
  }
  return dim;
}

void GPUParallelDimAttr::print(AsmPrinter &printer) const {
  printer << "<";
  printProcessorValue(printer, *this);
  printer << ">";
}

GPUParallelDimAttr GPUParallelDimAttr::threadDim(MLIRContext *context,
                                                 unsigned index) {
  assert(index <= 2 && "thread dimension index must be 0, 1, or 2");
  switch (index) {
  case 0:
    return threadXDim(context);
  case 1:
    return threadYDim(context);
  case 2:
    return threadZDim(context);
  }
  llvm_unreachable("validated thread dimension index");
}

GPUParallelDimAttr GPUParallelDimAttr::blockDim(MLIRContext *context,
                                                unsigned index) {
  assert(index <= 2 && "block dimension index must be 0, 1, or 2");
  switch (index) {
  case 0:
    return blockXDim(context);
  case 1:
    return blockYDim(context);
  case 2:
    return blockZDim(context);
  }
  llvm_unreachable("validated block dimension index");
}

gpu::Processor GPUParallelDimAttr::getProcessor() const {
  return indexToGpuProcessor(getValue().getInt());
}

int GPUParallelDimAttr::getOrder() const {
  return gpuProcessorIndex(getProcessor());
}

GPUParallelDimAttr GPUParallelDimAttr::getOneHigher() const {
  int order = getOrder();
  if (order >= 6) // BlockZ is the highest
    return *this;
  return get(getContext(), indexToGpuProcessor(order + 1));
}

GPUParallelDimAttr GPUParallelDimAttr::getOneLower() const {
  int order = getOrder();
  if (order <= 0) // Sequential is the lowest
    return *this;
  return get(getContext(), indexToGpuProcessor(order - 1));
}

bool GPUParallelDimAttr::isSeq() const {
  return getProcessor() == gpu::Processor::Sequential;
}
bool GPUParallelDimAttr::isThreadX() const {
  return getProcessor() == gpu::Processor::ThreadX;
}
bool GPUParallelDimAttr::isThreadY() const {
  return getProcessor() == gpu::Processor::ThreadY;
}
bool GPUParallelDimAttr::isThreadZ() const {
  return getProcessor() == gpu::Processor::ThreadZ;
}
bool GPUParallelDimAttr::isBlockX() const {
  return getProcessor() == gpu::Processor::BlockX;
}
bool GPUParallelDimAttr::isBlockY() const {
  return getProcessor() == gpu::Processor::BlockY;
}
bool GPUParallelDimAttr::isBlockZ() const {
  return getProcessor() == gpu::Processor::BlockZ;
}
bool GPUParallelDimAttr::isAnyThread() const {
  return isThreadX() || isThreadY() || isThreadZ();
}
bool GPUParallelDimAttr::isAnyBlock() const {
  return isBlockX() || isBlockY() || isBlockZ();
}

//===----------------------------------------------------------------------===//
// GPUParallelDimsAttr
//===----------------------------------------------------------------------===//

GPUParallelDimsAttr GPUParallelDimsAttr::seq(MLIRContext *ctx) {
  return GPUParallelDimsAttr::get(ctx, {GPUParallelDimAttr::seqDim(ctx)});
}

bool GPUParallelDimsAttr::isSeq() const {
  assert(!getArray().empty() && "no par_dims found");
  if (getArray().size() == 1) {
    auto parDim = dyn_cast<GPUParallelDimAttr>(getArray()[0]);
    assert(parDim && "expected GPUParallelDimAttr");
    return parDim.isSeq();
  }
  return false;
}

bool GPUParallelDimsAttr::isParallel() const { return !isSeq(); }

bool GPUParallelDimsAttr::isMultiDim() const { return getArray().size() > 1; }

bool GPUParallelDimsAttr::hasAnyBlockLevel() const {
  return llvm::any_of(
      getArray(), [](const GPUParallelDimAttr &p) { return p.isAnyBlock(); });
}

bool GPUParallelDimsAttr::hasOnlyBlockLevel() const {
  return !getArray().empty() &&
         llvm::all_of(getArray(), [](const GPUParallelDimAttr &p) {
           return p.isAnyBlock();
         });
}

bool GPUParallelDimsAttr::hasOnlyThreadYLevel() const {
  return !getArray().empty() &&
         llvm::all_of(getArray(), [](const GPUParallelDimAttr &p) {
           return p.isThreadY();
         });
}

bool GPUParallelDimsAttr::hasOnlyThreadXLevel() const {
  return !getArray().empty() &&
         llvm::all_of(getArray(), [](const GPUParallelDimAttr &p) {
           return p.isThreadX();
         });
}

Attribute GPUParallelDimsAttr::parse(AsmParser &parser, Type type) {
  auto delimiter = AsmParser::Delimiter::Square;
  SmallVector<GPUParallelDimAttr> parDims;
  auto parseParDim = [&]() -> ParseResult {
    GPUParallelDimAttr dim;
    if (parseProcessorValue(parser, dim))
      return failure();
    parDims.push_back(dim);
    return success();
  };
  if (parser.parseCommaSeparatedList(delimiter, parseParDim,
                                     "list of OpenACC GPU parallel dimensions"))
    return {};
  return GPUParallelDimsAttr::get(parser.getContext(), parDims);
}

void GPUParallelDimsAttr::print(AsmPrinter &printer) const {
  printer << "[";
  llvm::interleaveComma(getArray(), printer,
                        [&printer](const GPUParallelDimAttr &p) {
                          printProcessorValue(printer, p);
                        });
  printer << "]";
}