OpenACCCG.cpp

Go to the documentation of this file.

//===- 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/Arith/IR/Arith.h"
#include "mlir/Dialect/GPU/IR/GPUBase.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/DenseMap.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 continue after the 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(op));
}
 
static ValueRange getSingleRegionSuccessorInputs(Operation *op,
                                                 RegionSuccessor successor) {
  return successor.isOperation() ? 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();
  }
};
 
struct KernelEnvironmentDedupDataOperands
    : public OpRewritePattern<acc::KernelEnvironmentOp> {
  using OpRewritePattern<acc::KernelEnvironmentOp>::OpRewritePattern;
 
  static bool isReductionDataOperand(Value operand) {
    Operation *def = operand.getDefiningOp();
    if (!def)
      return false;
    std::optional<DataClause> clause = getDataClause(def);
    if (!clause || *clause != DataClause::acc_reduction)
      return false;
    assert(isa<CopyinOp>(def) && "reduction clause must be an acc.copyin");
    return true;
  }
 
  LogicalResult matchAndRewrite(acc::KernelEnvironmentOp op,
                                PatternRewriter &rewriter) const override {
    SmallVector<Value> operands(op.getDataClauseOperands());
    const unsigned numOperands = operands.size();
    if (numOperands < 2)
      return failure();
 
    SmallVector<Value> kept;
    SmallVector<std::pair<Value, Value>> useRewrites;
    SmallVector<SmallVector<Value>> boundsStorage;
    boundsStorage.reserve(numOperands);
    llvm::DenseMap<std::pair<Value, ArrayRef<Value>>, unsigned>
        keptIndexByVarAndBounds;
    for (Value operand : operands) {
      Operation *def = operand.getDefiningOp();
      // acc.map_info is a finalized mapping. Its order and semantics stand;
      // matching only its variable and bounds would ignore flags and
      // descriptors.
      if (!def || isa<MapInfoOp>(def)) {
        kept.push_back(operand);
        continue;
      }
 
      Value var = getVar(def);
      if (!var) {
        kept.push_back(operand);
        continue;
      }
 
      boundsStorage.push_back(getBounds(def));
      auto [it, inserted] = keptIndexByVarAndBounds.try_emplace(
          std::make_pair(var, ArrayRef<Value>(boundsStorage.back())),
          kept.size());
      if (inserted) {
        kept.push_back(operand);
        continue;
      }
 
      Value &survivor = kept[it->second];
      // A reduction acc.copyin is that copyin plus the reduction. Keeping
      // another clause would drop the reduction.
      if (isReductionDataOperand(operand) && !isReductionDataOperand(survivor))
        std::swap(survivor, operand);
      useRewrites.emplace_back(operand, survivor);
    }
    if (kept.size() == numOperands)
      return failure();
 
    for (auto [from, to] : useRewrites) {
      rewriter.replaceUsesWithIf(from, to, [&](OpOperand &use) {
        return op->isProperAncestor(use.getOwner());
      });
    }
    rewriter.modifyOpInPlace(
        op, [&] { op.getDataClauseOperandsMutable().assign(kept); });
    return success();
  }
};
 
// Capture `hasStream` before erasing inputs because segment metadata is only
// repaired at the end of the modification transaction.
static void setComputeRegionInputOperandSegments(ComputeRegionOp op,
                                                 PatternRewriter &rewriter,
                                                 size_t numInput,
                                                 bool hasStream) {
  const size_t numLaunch = op.getLaunchArgs().size();
  op->setInherentAttr(
      rewriter.getStringAttr(ComputeRegionOp::getOperandSegmentSizeAttr()),
      rewriter.getDenseI32ArrayAttr({static_cast<int32_t>(numLaunch),
                                     static_cast<int32_t>(numInput),
                                     hasStream ? 1 : 0}));
}
 
struct ComputeRegionRemoveDuplicateArgs
    : public OpRewritePattern<ComputeRegionOp> {
  using Base::Base;
 
  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 hasDuplicate = false;
    for (size_t j = 1; j < numInput && !hasDuplicate; ++j)
      for (size_t i = 0; i < j; ++i)
        if (op->getOperand(static_cast<unsigned>(numLaunch + i)) ==
            op->getOperand(static_cast<unsigned>(numLaunch + j))) {
          hasDuplicate = true;
          break;
        }
    if (!hasDuplicate)
      return failure();
 
    const bool hasStream = static_cast<bool>(op.getStream());
    rewriter.modifyOpInPlace(op, [&] {
      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;
            break;
          }
        }
        if (!merged)
          break;
      }
      setComputeRegionInputOperandSegments(op, rewriter, numInput, hasStream);
    });
    return success();
  }
};
 
struct ComputeRegionRemoveUnusedArgs
    : public OpRewritePattern<ComputeRegionOp> {
  using Base::Base;
 
  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 hasUnused = false;
    for (size_t k = numLaunch; k < numLaunch + numInput; ++k)
      if (body->getArgument(static_cast<unsigned>(k)).use_empty()) {
        hasUnused = true;
        break;
      }
    if (!hasUnused)
      return failure();
 
    const bool hasStream = static_cast<bool>(op.getStream());
    rewriter.modifyOpInPlace(op, [&] {
      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;
      }
      setComputeRegionInputOperandSegments(op, rewriter, numInput, hasStream);
    });
    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);
}
 
static FailureOr<SmallVector<GPUParallelDimAttr>>
parseGPUParallelDimList(AsmParser &parser) {
  SmallVector<GPUParallelDimAttr> parDims;
  auto parseParDim = [&]() -> ParseResult {
    GPUParallelDimAttr dim;
    if (parseProcessorValue(parser, dim))
      return failure();
    parDims.push_back(dim);
    return success();
  };
  if (parser.parseCommaSeparatedList(AsmParser::Delimiter::Square, parseParDim,
                                     "list of OpenACC GPU parallel dimensions"))
    return failure();
  return parDims;
}
 
static void printGPUParallelDimList(AsmPrinter &printer,
                                    ArrayRef<GPUParallelDimAttr> dims) {
  printer << "[";
  llvm::interleaveComma(dims, printer, [&printer](const GPUParallelDimAttr &p) {
    printProcessorValue(printer, p);
  });
  printer << "]";
}
} // 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, KernelEnvironmentDedupDataOperands>(
      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");
  if ((*this)->getParentOfType<KernelEnvironmentOp>())
    return emitError(
        "must not be nested within another acc.kernel_environment");
  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();
}
 
//===----------------------------------------------------------------------===//
// ReductionAccumulateOp
//===----------------------------------------------------------------------===//
 
LogicalResult ReductionAccumulateOp::verify() {
  Type valueType = getValue().getType();
  auto ptrLikeTy = cast<PointerLikeType>(getMemref().getType());
  Type elementType = ptrLikeTy.getElementType();
  if (!elementType)
    return emitOpError("pointer-like destination must have an element type");
  if (elementType != valueType)
    return emitOpError("pointer-like element type must match value type");
  if (getParDims().getArray().empty())
    return emitOpError("par_dims must specify at least one parallel dimension");
  return success();
}
 
//===----------------------------------------------------------------------===//
// ReductionAccumulateArrayOp
//===----------------------------------------------------------------------===//
 
LogicalResult ReductionAccumulateArrayOp::verify() {
  if (getParDims().getArray().empty())
    return emitOpError("par_dims must specify at least one parallel dimension");
  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);
  ComputeRegionOp::printProperties(getContext(), p, getProperties(),
                                   /*elidedProps=*/getOperandSegmentSizeAttr());
  p.printOptionalAttrDict((*this)->getDiscardableAttrDictionary().getValue());
}
 
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");
 
  DenseI32ArrayAttr operandSegmentSizes = 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();
  }
 
  Attribute parsedProperties;
  if (ComputeRegionOp::genericParseProperties(parser, parsedProperties))
    return failure();
  auto propertyDictionary = dyn_cast_or_null<DictionaryAttr>(parsedProperties);
  if (parsedProperties && !propertyDictionary)
    return parser.emitError(parser.getNameLoc(),
                            "expected properties dictionary");
 
  NamedAttrList properties(propertyDictionary ? propertyDictionary
                                              : builder.getDictionaryAttr({}));
  properties.set(ComputeRegionOp::getOperandSegmentSizeAttr(),
                 operandSegmentSizes);
  propertyDictionary = properties.getDictionary(builder.getContext());
  auto emitError = [&]() {
    return mlir::emitError(result.location, "invalid properties ")
           << propertyDictionary << " for op " << result.name.getStringRef()
           << ": ";
  };
  if (failed(ComputeRegionOp::setPropertiesFromParsedAttr(
          result.getOrAddProperties<Properties>(), propertyDictionary,
          emitError)))
    return failure();
 
  auto attrsLoc = parser.getCurrentLocation();
  if (parser.parseOptionalAttrDict(result.attributes))
    return failure();
  for (StringRef attrName : ComputeRegionOp::getAttributeNames()) {
    if (result.attributes.get(attrName))
      return parser.emitError(attrsLoc)
             << "inherent attribute '" << attrName
             << "' cannot be parsed from attr-dict when strict properties in "
                "assembly format is enabled";
  }
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// GPUSharedMemoryOp
//===----------------------------------------------------------------------===//
 
LogicalResult GPUSharedMemoryOp::verify() {
  if (getNumCopies() <= 0)
    return emitOpError("num_copies must be positive");
  if (getStaticUpperBoundBytes() <= 0)
    return emitOpError("static_upper_bound_bytes must be positive");
 
  bool hasScaling = static_cast<bool>(getDynamicSharedMemoryScalingBytes());
  bool hasFixed = static_cast<bool>(getDynamicSharedMemoryFixedBytes());
  if (hasScaling != hasFixed)
    return emitOpError(
        "dynamic_shared_memory_scaling_bytes and "
        "dynamic_shared_memory_fixed_bytes must both be present or both be "
        "absent");
  if (auto scalingAttr = getDynamicSharedMemoryScalingBytesAttr())
    if (scalingAttr.getValue().isNegative())
      return emitOpError("dynamic_shared_memory_scaling_bytes must be "
                         "non-negative");
  if (auto fixedAttr = getDynamicSharedMemoryFixedBytesAttr())
    if (fixedAttr.getValue().isNegative())
      return emitOpError("dynamic_shared_memory_fixed_bytes must be "
                         "non-negative");
 
  auto resultTy = cast<MemRefType>(getResult().getType());
  auto addrSpace =
      dyn_cast_if_present<gpu::AddressSpaceAttr>(resultTy.getMemorySpace());
  if (!addrSpace || addrSpace.getValue() != gpu::AddressSpace::Workgroup)
    return emitOpError("result memref must use #gpu.address_space<workgroup>");
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// PredicateRegionOp
//===----------------------------------------------------------------------===//
 
LogicalResult PredicateRegionOp::verify() {
  if (getRegion().empty())
    return emitOpError("region needs to have at least one block");
  if (getRegion().front().getNumArguments() > 0)
    return emitOpError("region cannot have any arguments");
  if (!getOperation()->getParentOfType<ComputeRegionOp>())
    return emitOpError("must be nested within an acc.compute_region operation");
  return success();
}
 
//===----------------------------------------------------------------------===//
// MapInfoOp
//===----------------------------------------------------------------------===//
 
LogicalResult MapInfoOp::verify() {
  // A pointer-like var addresses the mapped object, so varType has to name that
  // object rather than the address of it.
  if (mlir::isa<acc::PointerLikeType>(getVar().getType()) &&
      getVarType() == getVar().getType())
    return emitOpError("varType must capture the element type of var");
 
  // A descriptor operand is only meaningful together with the layout it
  // follows. Without a kind, consumers have no way to interpret it.
  if (getDesc() && getDescKind() == DataDescKind::none)
    return emitOpError("desc requires a descKind other than none");
 
  // Bounds count elements of the mapped object, which is the OpenACC
  // descriptor layout, so it must be among the kinds named here.
  if (!getBounds().empty() &&
      !acc::bitEnumContainsAny(getDescKind(), DataDescKind::openacc))
    return emitOpError("bounds require descKind openacc");
 
  // Anything below -1 has no meaning: -1 states that the size is unknown at
  // compile time and 0 defers it to bounds or to a descriptor.
  if (getSize()) {
    std::optional<int64_t> constantSize = getConstantIntValue(getSize());
    if (constantSize && *constantSize < -1)
      return emitOpError("size must be -1, 0, or a positive byte count");
  }
 
  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) {
  FailureOr<SmallVector<GPUParallelDimAttr>> parDims =
      parseGPUParallelDimList(parser);
  if (failed(parDims))
    return {};
  return GPUParallelDimsAttr::get(parser.getContext(), *parDims);
}
 
void GPUParallelDimsAttr::print(AsmPrinter &printer) const {
  printGPUParallelDimList(printer, getArray());
}
 
//===----------------------------------------------------------------------===//
// ActiveParDimsAttr
//===----------------------------------------------------------------------===//
 
Attribute ActiveParDimsAttr::parse(AsmParser &parser, Type type) {
  FailureOr<SmallVector<GPUParallelDimAttr>> parDims =
      parseGPUParallelDimList(parser);
  if (failed(parDims))
    return {};
  return ActiveParDimsAttr::get(parser.getContext(), *parDims);
}
 
void ActiveParDimsAttr::print(AsmPrinter &printer) const {
  printGPUParallelDimList(printer, getArray());
}