ACCComputeLowering.cpp

Go to the documentation of this file.

//===- ACCComputeLowering.cpp - Lower ACC compute to compute_region -------===//
//
// 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 pass decomposes OpenACC compute constructs into a representation that
// separates the data environment from the compute portion and prepares for
// parallelism assignment and privatization at the appropriate level.
//
// Overview:
// ---------
// Each compute construct (`acc.parallel`, `acc.serial`, `acc.kernels`) is
// lowered to (1) `acc.kernel_environment`, which captures the data environment
// and (2) `acc.compute_region`, which holds the compute body. Inside the
// compute region, acc.loop is converted to SCF loops (`scf.parallel` or
// `scf.for`) with any predetermined parallelism expressed as `par_dims`. This
// decomposition allows later phases to assign parallelism and handle
// privatization at the right granularity.
//
// Transformations:
// ----------------
// 1. Compute constructs: acc.parallel, acc.serial, and acc.kernels are
//    replaced by acc.kernel_environment containing a single acc.compute_region.
//    For acc.parallel / acc.kernels, launch arguments (num_gangs, num_workers,
//    vector_length) become acc.par_width ops (each result is `index`) and are
//    passed as compute_region launch operands. Compute regions with
//    num_gangs(1), num_workers(1), and vector_length(1) and acc serial use a
//    single sequential acc.par_width launch operand.
//
// 2. acc.loop: Converted according to context and attributes:
//    - Unstructured: body wrapped in scf.execute_region.
//    - Sequential (serial region, seq clause, or compute region with
//    num_gangs(1), num_workers(1), and vector_length(1)):
//      scf.parallel with par_dims = sequential.
//    - Auto (in parallel/kernels): scf.for with collapse when
//    multi-dimensional.
//    - Orphan (not inside a compute construct): scf.for, no collapse.
//    - Independent (in parallel/kernels): scf.parallel with par_dims from
//      gang/worker/vector mapping (e.g. block_x).
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Dialect/OpenACC/Transforms/Passes.h"
 
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/OpenACC/Analysis/OpenACCSupport.h"
#include "mlir/Dialect/OpenACC/OpenACC.h"
#include "mlir/Dialect/OpenACC/OpenACCParMapping.h"
#include "mlir/Dialect/OpenACC/OpenACCUtils.h"
#include "mlir/Dialect/OpenACC/OpenACCUtilsCG.h"
#include "mlir/Dialect/OpenACC/OpenACCUtilsLoop.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/Matchers.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/RegionUtils.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
 
namespace mlir {
namespace acc {
#define GEN_PASS_DEF_ACCCOMPUTELOWERING
#include "mlir/Dialect/OpenACC/Transforms/Passes.h.inc"
} // namespace acc
} // namespace mlir
 
#define DEBUG_TYPE "acc-compute-lowering"
 
using namespace mlir;
using namespace mlir::acc;
 
namespace {
 
//===----------------------------------------------------------------------===//
// Helper functions
//===----------------------------------------------------------------------===//
 
static bool isOpInComputeRegion(Operation *op) {
  Region *region = op->getBlock()->getParent();
  return getEnclosingComputeOp(*region) != nullptr;
}
 
static bool isOpInSerialRegion(Operation *op) {
  if (auto parallelOp = op->getParentOfType<ParallelOp>())
    return parallelOp.isEffectivelySerial();
  if (auto kernelsOp = op->getParentOfType<KernelsOp>())
    return kernelsOp.isEffectivelySerial();
  if (op->getParentOfType<SerialOp>())
    return true;
  if (auto computeRegion = op->getParentOfType<ComputeRegionOp>())
    return computeRegion.isEffectivelySerial();
  if (auto funcOp = op->getParentOfType<FunctionOpInterface>()) {
    if (isSpecializedAccRoutine(funcOp)) {
      auto attr = funcOp->getDiscardableAttrOfType<SpecializedRoutineAttr>(
          getSpecializedRoutineAttrName());
      if (attr && attr.getLevel().getValue() == ParLevel::seq)
        return true;
    }
  }
  return false;
}
 
/// Clone defining ops of constant live-in values into `region`, rewrite uses
/// inside the region to the clones, and remove those values from
/// `liveInValues` so they are not threaded through `acc.compute_region` ins.
static void materializeConstantLiveInsIntoRegion(Region &region,
                                                 SetVector<Value> &liveInValues,
                                                 RewriterBase &rewriter) {
  SmallVector<Value> constantLiveIns;
  for (Value v : liveInValues) {
    Operation *defOp = v.getDefiningOp();
    if (defOp && matchPattern(defOp, m_Constant())) {
      // As per the definition of ConstantLike trait, constants must have a
      // single result.
      assert(defOp->getNumResults() == 1 &&
             "constants must have a single result");
      constantLiveIns.push_back(v);
    }
  }
  if (constantLiveIns.empty())
    return;
 
  OpBuilder::InsertionGuard guard(rewriter);
  rewriter.setInsertionPointToStart(&region.front());
 
  for (Value v : constantLiveIns) {
    Value newV = rewriter.clone(*v.getDefiningOp())->getResult(0);
    replaceAllUsesInRegionWith(v, newV, region);
    liveInValues.remove(v);
  }
}
 
/// Return the device type from which gang/worker/vector clauses should be read.
/// If the requested device type has any such clauses, use that exclusively;
/// otherwise fall back to the default (DeviceType::None).
static DeviceType getGangWorkerVectorDeviceType(LoopOp loopOp,
                                                DeviceType deviceType) {
  if (deviceType != DeviceType::None &&
      loopOp.hasAnyGangWorkerVector(deviceType))
    return deviceType;
  return DeviceType::None;
}
 
template <typename ComputeConstructT>
static DeviceType getParDimsDeviceType(ComputeConstructT computeOp,
                                       DeviceType deviceType) {
  if (deviceType != DeviceType::None &&
      computeOp.hasAnyGangWorkerVector(deviceType))
    return deviceType;
  return DeviceType::None;
}
 
/// Constant sized gang/worker/vector clauses collected per compute construct.
struct SizedLevel {
  ParLevel level;
  int64_t size;
};
using SizedLevelMap = DenseMap<Operation *, SmallVector<SizedLevel>>;
 
/// Record a sized clause if `size` is a constant; NYI otherwise.
static LogicalResult tryAddSizedLevel(SizedLevelMap &sizedLevelMap,
                                      Operation *computeOp, LoopOp loopOp,
                                      ParLevel level, Value size,
                                      OpenACCSupport &accSupport) {
  if (!size)
    return success();
  std::optional<int64_t> constSize = getConstantIntValue(size);
  if (!constSize) {
    accSupport.emitNYI(loopOp.getLoc(),
                       "non-constant sized parallelism clause");
    return failure();
  }
  sizedLevelMap[computeOp].push_back({level, *constSize});
  return success();
}
 
/// Collect constant sized levels from loops in `acc.kernels` regions.
static LogicalResult fillSizedLevelMap(Operation *op, DeviceType deviceType,
                                       SizedLevelMap &sizedLevelMap,
                                       OpenACCSupport &accSupport) {
  WalkResult result = op->walk([&](LoopOp loopOp) {
    Operation *computeOp =
        getEnclosingComputeOp(*loopOp->getBlock()->getParent());
    if (!computeOp || !isa<KernelsOp>(computeOp))
      return WalkResult::advance();
    DeviceType loopDeviceType =
        getGangWorkerVectorDeviceType(loopOp, deviceType);
    if (failed(tryAddSizedLevel(
            sizedLevelMap, computeOp, loopOp, ParLevel::vector,
            loopOp.getVectorValue(loopDeviceType), accSupport)) ||
        failed(tryAddSizedLevel(
            sizedLevelMap, computeOp, loopOp, ParLevel::worker,
            loopOp.getWorkerValue(loopDeviceType), accSupport)) ||
        failed(tryAddSizedLevel(
            sizedLevelMap, computeOp, loopOp, ParLevel::gang_dim1,
            loopOp.getGangValue(GangArgType::Num, loopDeviceType), accSupport)))
      return WalkResult::interrupt();
    return WalkResult::advance();
  });
  return failure(result.wasInterrupted());
}
 
/// Map loop parallelism clauses (gang/worker/vector) to GPU parallel
/// dimensions using the given mapping policy. Sized clauses (e.g. vector(n))
/// count as the corresponding level.
static SmallVector<GPUParallelDimAttr>
getParallelDimensions(LoopOp loopOp, const ACCToGPUMappingPolicy &policy,
                      DeviceType deviceType) {
  deviceType = getGangWorkerVectorDeviceType(loopOp, deviceType);
  SmallVector<GPUParallelDimAttr> parDims;
  auto *ctx = loopOp->getContext();
 
  if (loopOp.hasVector(deviceType) || loopOp.getVectorValue(deviceType))
    insertParDim(parDims, policy.vectorDim(ctx));
  if (loopOp.hasWorker(deviceType) || loopOp.getWorkerValue(deviceType))
    insertParDim(parDims, policy.workerDim(ctx));
  if (auto gangDimValue = loopOp.getGangValue(GangArgType::Dim, deviceType)) {
    if (auto gangDimDefOp =
            gangDimValue.getDefiningOp<arith::ConstantIntOp>()) {
      auto gangLevel = getGangParLevel(gangDimDefOp.value());
      insertParDim(parDims, policy.gangDim(ctx, gangLevel));
    }
  } else if (loopOp.hasGang(deviceType) ||
             loopOp.getGangValue(GangArgType::Num, deviceType) ||
             loopOp.getGangValue(GangArgType::Static, deviceType)) {
    insertParDim(parDims, policy.gangDim(ctx, ParLevel::gang_dim1));
  }
  return parDims;
}
 
/// If \p loopOp has a `gang(static:)` value, return the corresponding
/// `acc.chunk_size` attribute. A constant size is stored directly; a
/// non-constant size is recorded as -1, the same encoding as `static:*`.
/// Returns a null attribute when the clause is absent.
static ChunkSizeAttr getChunkSizeFromStatic(LoopOp loopOp,
                                            DeviceType deviceType) {
  DeviceType loopDeviceType = getGangWorkerVectorDeviceType(loopOp, deviceType);
  Value gangStatic = loopOp.getGangValue(GangArgType::Static, loopDeviceType);
  if (!gangStatic)
    return {};
  std::optional<int64_t> chunkSize = getConstantIntValue(gangStatic);
  // Non-constant static is treated like `static:*` (implementation-chosen).
  return ChunkSizeAttr::get(loopOp->getContext(), chunkSize.value_or(-1));
}
 
/// Build `acc.compute_region` launch operands: one sequential `acc.par_width`
/// for `acc.serial`, for `acc.parallel` / `acc.kernels` when every num_gangs
/// operand and num_workers / vector_length are the constant 1, and otherwise
/// `acc.par_width` from gang/worker/vector (device-type operands first, then
/// default DeviceType::None).
template <typename ComputeConstructT>
static SmallVector<Value> assignKnownLaunchArgs(
    ComputeConstructT computeOp, DeviceType deviceType, RewriterBase &rewriter,
    const ACCToGPUMappingPolicy &policy, const SizedLevelMap &sizedLevelMap) {
  auto *ctx = rewriter.getContext();
  auto loc = computeOp->getLoc();
 
  if constexpr (std::is_same_v<ComputeConstructT, SerialOp>) {
    return {ParWidthOp::create(rewriter, loc, Value(), policy.seqDim(ctx))};
  } else if constexpr (llvm::is_one_of<ComputeConstructT, ParallelOp,
                                       KernelsOp>::value) {
    if (computeOp.isEffectivelySerial())
      return {ParWidthOp::create(rewriter, loc, Value(), policy.seqDim(ctx))};
 
    deviceType = getParDimsDeviceType(computeOp, deviceType);
 
    SmallVector<Value> values;
    auto indexTy = rewriter.getIndexType();
 
    auto numGangs = computeOp.getNumGangsValues(deviceType);
    for (auto [gangDimIdx, gangSize] : llvm::enumerate(numGangs)) {
      auto gangLevel = getGangParLevel(gangDimIdx + 1);
      values.push_back(ParWidthOp::create(
          rewriter, loc,
          getValueOrCreateCastToIndexLike(rewriter, gangSize.getLoc(), indexTy,
                                          gangSize),
          policy.gangDim(ctx, gangLevel)));
    }
 
    Value numWorkers = computeOp.getNumWorkersValue(deviceType);
    if (numWorkers) {
      values.push_back(ParWidthOp::create(
          rewriter, loc,
          getValueOrCreateCastToIndexLike(rewriter, numWorkers.getLoc(),
                                          indexTy, numWorkers),
          policy.workerDim(ctx)));
    }
 
    Value vectorLength = computeOp.getVectorLengthValue(deviceType);
    if (vectorLength) {
      values.push_back(ParWidthOp::create(
          rewriter, loc,
          getValueOrCreateCastToIndexLike(rewriter, vectorLength.getLoc(),
                                          indexTy, vectorLength),
          policy.vectorDim(ctx)));
    }
 
    // Loop-level sized clauses. Skip a dim already set on the construct.
    // Rematerialize the constant here so it dominates the compute region.
    auto sizedLevels = sizedLevelMap.find(computeOp.getOperation());
    if (sizedLevels != sizedLevelMap.end()) {
      for (const SizedLevel &sizedLevel : sizedLevels->second) {
        GPUParallelDimAttr dim = policy.map(ctx, sizedLevel.level);
        bool exists = llvm::any_of(values, [&](Value v) {
          auto parWidth = v.getDefiningOp<ParWidthOp>();
          return parWidth && parWidth.getParDim() == dim;
        });
        if (exists)
          continue;
        Value sizeVal =
            arith::ConstantIndexOp::create(rewriter, loc, sizedLevel.size);
        values.push_back(ParWidthOp::create(rewriter, loc, sizeVal, dim));
      }
    }
    return values;
  } else {
    llvm_unreachable("assignKnownLaunchArgs: expected parallel, kernels, or "
                     "serial");
  }
}
 
//===----------------------------------------------------------------------===//
// Loop conversion pattern
//===----------------------------------------------------------------------===//
 
class ACCLoopConversion : public OpRewritePattern<LoopOp> {
public:
  ACCLoopConversion(MLIRContext *ctx, const ACCToGPUMappingPolicy &policy,
                    DeviceType deviceType)
      : OpRewritePattern<LoopOp>(ctx), policy(policy), deviceType(deviceType) {}
 
  LogicalResult matchAndRewrite(LoopOp loopOp,
                                PatternRewriter &rewriter) const override {
    if (loopOp.getUnstructured()) {
      auto executeRegion =
          convertUnstructuredACCLoopToSCFExecuteRegion(loopOp, rewriter);
      if (!executeRegion)
        return failure();
      rewriter.replaceOp(loopOp, executeRegion);
      return success();
    }
 
    LoopParMode parMode = loopOp.getDefaultOrDeviceTypeParallelism(deviceType);
 
    if (parMode == LoopParMode::loop_seq || isOpInSerialRegion(loopOp)) {
      // Use scf.for with sequential loops, because the loop's parallelism is
      // already determined.
      auto forOp =
          convertACCLoopToSCFFor(loopOp, rewriter, /*enableCollapse=*/true);
      if (!forOp)
        return failure();
      setParDimsAttr(forOp, GPUParallelDimsAttr::seq(loopOp->getContext()));
      rewriter.replaceOp(loopOp, forOp);
    } else if (parMode == LoopParMode::loop_auto) {
      // All loops in serial regions should have already been handled.
      assert(!isOpInSerialRegion(loopOp) &&
             "Expected loop to be in non-serial region");
      // Mark as scf.for to allow auto-parallelization analysis later.
      auto forOp =
          convertACCLoopToSCFFor(loopOp, rewriter, /*enableCollapse=*/true);
      if (!forOp)
        return failure();
      SmallVector<GPUParallelDimAttr> parDims =
          getParallelDimensions(loopOp, policy, deviceType);
      if (!parDims.empty()) {
        auto parDimsAttr =
            GPUParallelDimsAttr::get(loopOp->getContext(), parDims);
        setParDimsAttr(forOp, parDimsAttr);
        if (ChunkSizeAttr chunkSize =
                getChunkSizeFromStatic(loopOp, deviceType))
          setChunkSizeAttr(forOp, chunkSize);
      }
      rewriter.replaceOp(loopOp, forOp);
    } else if (!isOpInComputeRegion(loopOp) &&
               !isSpecializedAccRoutine(
                   loopOp->getParentOfType<FunctionOpInterface>())) {
      // This loop is an orphan `acc loop` but it is not in any sort
      // of compute region. Thus it is just a sequential non-accelerator loop.
      auto forOp =
          convertACCLoopToSCFFor(loopOp, rewriter, /*enableCollapse=*/false);
      if (!forOp)
        return failure();
      rewriter.replaceOp(loopOp, forOp);
    } else {
      assert(parMode == LoopParMode::loop_independent &&
             "Expected loop to be independent");
      auto parallelOp = convertACCLoopToSCFParallel(loopOp, rewriter);
      if (!parallelOp)
        return failure();
 
      SmallVector<GPUParallelDimAttr> parDims =
          getParallelDimensions(loopOp, policy, deviceType);
      if (!parDims.empty()) {
        auto parDimsAttr =
            GPUParallelDimsAttr::get(loopOp->getContext(), parDims);
        setParDimsAttr(parallelOp, parDimsAttr);
        if (ChunkSizeAttr chunkSize =
                getChunkSizeFromStatic(loopOp, deviceType))
          setChunkSizeAttr(parallelOp, chunkSize);
      }
      rewriter.replaceOp(loopOp, parallelOp);
    }
    return success();
  }
 
private:
  const ACCToGPUMappingPolicy &policy;
  DeviceType deviceType;
};
 
//===----------------------------------------------------------------------===//
// Compute construct conversion pattern
//===----------------------------------------------------------------------===//
 
template <typename ComputeConstructT>
class ComputeOpConversion : public OpRewritePattern<ComputeConstructT> {
public:
  ComputeOpConversion(MLIRContext *ctx, const ACCToGPUMappingPolicy &policy,
                      DeviceType deviceType, const SizedLevelMap &sizedLevelMap)
      : OpRewritePattern<ComputeConstructT>(ctx), policy(policy),
        deviceType(deviceType), sizedLevelMap(sizedLevelMap) {}
 
  LogicalResult matchAndRewrite(ComputeConstructT computeOp,
                                PatternRewriter &rewriter) const override {
    rewriter.setInsertionPoint(computeOp);
    auto kernelEnv =
        KernelEnvironmentOp::createAndPopulate(computeOp, deviceType, rewriter);
    auto launchArgs = assignKnownLaunchArgs(computeOp, deviceType, rewriter,
                                            policy, sizedLevelMap);
    Region &region = computeOp.getRegion();
    SetVector<Value> liveInValues;
    getUsedValuesDefinedAbove(region, region, liveInValues);
    materializeConstantLiveInsIntoRegion(region, liveInValues, rewriter);
    IRMapping mapping;
    auto computeRegion = buildComputeRegion(
        computeOp->getLoc(), launchArgs, liveInValues.getArrayRef(),
        ComputeConstructT::getOperationName(), region, rewriter, mapping);
    if (!computeRegion) {
      rewriter.eraseOp(kernelEnv);
      return failure();
    }
    rewriter.eraseOp(computeOp);
    return success();
  }
 
private:
  const ACCToGPUMappingPolicy &policy;
  DeviceType deviceType;
  const SizedLevelMap &sizedLevelMap;
};
 
//===----------------------------------------------------------------------===//
// Pass implementation
//===----------------------------------------------------------------------===//
 
class ACCComputeLowering
    : public acc::impl::ACCComputeLoweringBase<ACCComputeLowering> {
public:
  using ACCComputeLoweringBase::ACCComputeLoweringBase;
 
  void runOnOperation() override {
    auto op = getOperation();
    auto *context = op.getContext();
 
    DefaultACCToGPUMappingPolicy policy;
    // Collect loop sized levels before loops are rewritten away.
    SizedLevelMap sizedLevelMap;
    OpenACCSupport &accSupport = getAnalysis<OpenACCSupport>();
    if (failed(fillSizedLevelMap(op, deviceType, sizedLevelMap, accSupport)))
      return signalPassFailure();
 
    // Part 1: Convert acc.loop to scf.parallel/scf.for while the parent
    // compute construct is still present (needed to determine conversion
    // strategy).
    RewritePatternSet loopPatterns(context);
    loopPatterns.insert<ACCLoopConversion>(context, policy, deviceType);
    if (failed(applyPatternsGreedily(op, std::move(loopPatterns))))
      return signalPassFailure();
 
    // Part 2: Convert acc.parallel, acc.kernels, and acc.serial to
    // acc.kernel_environment { acc.compute_region { ... } }.
    RewritePatternSet computePatterns(context);
    computePatterns
        .insert<ComputeOpConversion<ParallelOp>, ComputeOpConversion<KernelsOp>,
                ComputeOpConversion<SerialOp>>(context, policy, deviceType,
                                               sizedLevelMap);
    if (failed(applyPatternsGreedily(op, std::move(computePatterns))))
      return signalPassFailure();
  }
};
 
} // namespace