ACCRecipeMaterialization.cpp

Go to the documentation of this file.

//===- ACCRecipeMaterialization.cpp - Materialize ACC recipes -------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Overview:
// ---------
// OpenACC compute constructs (acc.parallel, acc.serial, acc.kernels) and
// acc.loop can carry data clauses (acc.private, acc.firstprivate,
// acc.reduction) that refer to recipes (acc.private.recipe,
// acc.firstprivate.recipe, acc.reduction.recipe). Recipes define how to
// initialize, copy, combine, or destroy a particular variable. This pass clones
// those regions into the construct and ensures the materialized SSA values are
// used instead.
//
// Transforms:
// -----------
// 1. Firstprivate: Inserts acc.firstprivate_map so the initial value is
//    available on the device, then clones the recipe init and copy regions
//    into the construct and replaces uses with the materialized alloca.
//    Optional destroy region is cloned before the region terminator.
//
// 2. Private: Clones the recipe init region into the construct (at the
//    region entry or at the loop op for acc.loop private). Replaces uses
//    of the recipe result with the materialized alloca. Optional destroy
//    region is cloned before the region terminator.
//
// 3. Reduction: Creates acc.reduction_init (init region inlined) and
//    acc.reduction_combine_region (combiner region inlined). Uses within
//    the region are updated to the reduction init result.
//    In addition, creates appropriate acc.copyin/copyout around
//    the compute region; the reduction's initial value is taken
//    from the copied in variable, and the final reduction value
//    is copied out to the variable. If there are existing
//    data operations for the reduction variable, no new data
//    operations are added.
//
// Requirements:
// -------------
// 1. OpenACCSupport: The pass uses the `acc::OpenACCSupport` analysis
//    including emitNYI for unsupported cases.
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/OpenACC/Analysis/OpenACCSupport.h"
#include "mlir/Dialect/OpenACC/Extensions/AllExtensions.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/OpenACC/Transforms/Passes.h"
#include "mlir/IR/Block.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/IR/Value.h"
#include "mlir/IR/ValueRange.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/RegionUtils.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/ErrorHandling.h"
 
namespace mlir {
namespace acc {
#define GEN_PASS_DEF_ACCRECIPEMATERIALIZATION
#include "mlir/Dialect/OpenACC/Transforms/Passes.h.inc"
} // namespace acc
} // namespace mlir
 
#define DEBUG_TYPE "acc-recipe-materialization"
 
namespace {
using namespace mlir;
 
static void setLocation(Region &region, Location loc) {
  // Since recipes are generated per type and not per variable, the location
  // of the recipe operations which get inlined will not necessarily be the
  // same as the location of the op that is being materialized. Force an update
  // of the location of the recipe operations to the location of the op that is
  // being materialized.
  region.walk([&](Operation *op) { op->setLoc(loc); });
}
 
static void saveVarName(StringRef name, Value dst) {
  if (name.empty())
    return;
  if (Operation *dstOp = dst.getDefiningOp()) {
    if (dstOp->getDiscardableAttrOfType<acc::VarNameAttr>(
            acc::getVarNameAttrName()))
      return;
    if (isa<ACC_DATA_ENTRY_OPS>(dstOp))
      return;
    dstOp->setDiscardableAttr(acc::getVarNameAttrName(),
                              acc::VarNameAttr::get(dstOp->getContext(), name));
    return;
  }
  auto blockArg = dyn_cast<BlockArgument>(dst);
  if (!blockArg)
    return;
  Block *block = blockArg.getOwner();
  Region *region = block ? block->getParent() : nullptr;
  if (!region || !block->isEntryBlock())
    return;
  Operation *parent = region->getParentOp();
  if (!parent)
    return;
  auto funcOp = dyn_cast<FunctionOpInterface>(parent);
  if (!funcOp)
    return;
  unsigned argIdx = blockArg.getArgNumber();
  if (argIdx >= funcOp.getNumArguments())
    return;
  if (funcOp.getArgAttr(argIdx, acc::getVarNameAttrName()))
    return;
  funcOp.setArgAttr(argIdx, acc::getVarNameAttrName(),
                    acc::VarNameAttr::get(parent->getContext(), name));
}
 
static void saveVarName(Value src, Value dst) {
  saveVarName(acc::getVariableName(src), dst);
}
 
static void resolveVarNamePlaceholders(Block *block, Block::iterator ip,
                                       StringRef name) {
  StringRef placeholder = acc::getVarNamePlaceholder();
  for (auto it = block->begin(); it != std::next(ip); ++it) {
    it->walk([&](Operation *op) {
      auto attr = op->getDiscardableAttrOfType<acc::VarNameAttr>(
          acc::getVarNameAttrName());
      if (!attr || attr.getName() != placeholder)
        return;
      if (name.empty())
        op->removeDiscardableAttr(acc::getVarNameAttrName());
      else
        op->setDiscardableAttr(acc::getVarNameAttrName(),
                               acc::VarNameAttr::get(op->getContext(), name));
    });
  }
}
 
// Clone the destroy region of the recipe before the terminator of the provided
// block. Values must be provided for the destroy region block arguments
// according to the recipe specifications.
template <typename RecipeOpTy>
static void cloneDestroy(Location loc, RecipeOpTy recipe, mlir::Block *block,
                         Block::iterator ip,
                         const llvm::SmallVector<mlir::Value> &arguments) {
  IRMapping mapping{};
  Region &destroyRegion = recipe.getDestroyRegion();
  assert(destroyRegion.getBlocks().front().getNumArguments() ==
             arguments.size() &&
         "unexpected acc recipe destroy block arguments");
 
  setLocation(destroyRegion, loc);
 
  mapping.map(destroyRegion.getBlocks().front().getArguments(), arguments);
  acc::cloneACCRegionInto(&destroyRegion, block, ip, mapping,
                          /*resultsToReplace=*/{});
}
 
class ACCRecipeMaterialization
    : public acc::impl::ACCRecipeMaterializationBase<ACCRecipeMaterialization> {
public:
  using acc::impl::ACCRecipeMaterializationBase<
      ACCRecipeMaterialization>::ACCRecipeMaterializationBase;
 
  void getDependentDialects(DialectRegistry &registry) const override {
    acc::impl::ACCRecipeMaterializationBase<
        ACCRecipeMaterialization>::getDependentDialects(registry);
    acc::registerAllExtensions(registry);
  }
  void runOnOperation() override;
 
private:
  // When the recipe reads the original variable, its initial value needs to be
  // available on the GPU. One way to get that value there is to map the
  // variable through global memory.
  // Thus, when we materialize a firstprivate, we materialize it into
  // a mapping action first. This function ends up with doing the following:
  // %dev = acc.firstprivate var(%var)
  // =>
  // %copy = acc.firstprivate_map var(%var)
  // %dev = acc.firstprivate var(%copy)
  // When the recipe materialization happens, the `acc.firstprivate` ends up
  // being removed. But because of the way we chain it to the
  // `acc.firstprivate_map`, then its result becomes live-in to the
  // compute region and used as the variable the initial value is loaded from.
  template <typename OpTy>
  void handleInitialValueMapping(OpTy op) const;
  template <typename OpTy>
  void removeRecipe(
      OpTy op, ModuleOp moduleOp,
      const std::optional<llvm::DenseSet<StringAttr>> &usedSymbols) const;
  template <typename OpTy, typename RecipeOpTy, typename AccOpTy>
  LogicalResult materialize(OpTy op, RecipeOpTy recipe, AccOpTy accOp,
                            acc::OpenACCSupport &accSupport,
                            acc::ACCToGPUMappingPolicy &policy,
                            Value materializationVar = {}) const;
  template <typename OpTy>
  LogicalResult materializeForACCOp(OpTy accOp, acc::OpenACCSupport &accSupport,
                                    acc::ACCToGPUMappingPolicy &policy) const;
};
 
template <typename OpTy>
void ACCRecipeMaterialization::handleInitialValueMapping(OpTy op) const {
  OpBuilder builder(op);
  auto mapInitialOp = acc::FirstprivateMapInitialOp::create(
      builder, op.getLoc(), op.getVar(), op.getStructured(), op.getImplicit(),
      op.getBounds());
  mapInitialOp.setName(op.getName());
  op.getVarMutable().assign(mapInitialOp.getAccVar());
}
 
// Whether a recipe region reads the variable it privatizes - a descriptor
// recipe loads it for the bounds, while a scalar one ignores it. Both init and
// destroy receive it as their first argument.
static bool readsVar(Region &region) {
  if (region.empty() || region.getNumArguments() == 0)
    return false;
  return !region.getArgument(0).use_empty();
}
 
template <typename OpTy>
void ACCRecipeMaterialization::removeRecipe(
    OpTy op, ModuleOp moduleOp,
    const std::optional<llvm::DenseSet<StringAttr>> &usedSymbols) const {
  auto recipeName = op.getNameAttr();
  // Fall back to scanning the module when the symbol uses could not be
  // gathered up front.
  bool useEmpty = usedSymbols
                      ? !usedSymbols->contains(recipeName)
                      : SymbolTable::symbolKnownUseEmpty(recipeName, moduleOp);
  if (useEmpty) {
    LLVM_DEBUG(llvm::dbgs() << "erasing recipe: " << recipeName << "\n");
    op.erase();
  } else {
    LLVM_DEBUG({
      std::optional<SymbolTable::UseRange> symbolUses =
          op.getSymbolUses(moduleOp);
      if (symbolUses.has_value()) {
        for (SymbolTable::SymbolUse symbolUse : *symbolUses) {
          llvm::dbgs() << "symbol use: ";
          symbolUse.getUser()->dump();
        }
      }
    });
    llvm_unreachable("expected no use of recipe symbol");
  }
}
 
template <typename OpTy, typename RecipeOpTy, typename AccOpTy>
LogicalResult ACCRecipeMaterialization::materialize(
    OpTy op, RecipeOpTy recipe, AccOpTy accOp, acc::OpenACCSupport &accSupport,
    acc::ACCToGPUMappingPolicy &policy, Value materializationVar) const {
  Region &region = accOp.getRegion();
  Value origPtr = materializationVar ? materializationVar : op.getVar();
  Value accPtr = op.getAccVar();
  assert(accPtr && "invalid op: null acc var");
 
  OpBuilder b(op);
  SmallVector<Value> triples;
 
  // Clone init block into the region at the insertion point specified.
  Region &initRegion = recipe.getInitRegion();
  unsigned initNumArguments =
      initRegion.getBlocks().front().getArguments().size();
  if (initNumArguments > 1) {
    // Code from C/C++ will most likely only provide extent arguments to the
    // recipe arguments.
    if ((initNumArguments - 1) % 3 != 0) {
      (void)accSupport.emitNYI(recipe.getLoc(),
                               "privatization of array section with extents");
      return failure();
    }
    // The remaining arguments must be the bounds triples
    // (lower-bound, upper-bound, step), ...
    unsigned argIdx = 1;
    // Cast the given value to the type of the combiner region's argument
    // at position argIdx, and increment argIdx.
    auto castValueToArgType = [&](Location loc, Value v) {
      return convertScalarToDtype(
          b, loc, v,
          initRegion.getBlocks().front().getArgument(argIdx++).getType(),
          /*isUnsignedCast=*/false);
    };
    for (Value bound : acc::getBounds(op)) {
      auto dataBound = bound.getDefiningOp<acc::DataBoundsOp>();
      assert(dataBound &&
             "acc.reduction's bound must be defined by acc.bounds");
      // NOTE: we should probably generate get_lowerbound, get_upperbound
      // and get_stride here, so that we can stop looking for the acc.bounds
      // operation above, and just use the `bound` value.
      Value lb =
          castValueToArgType(dataBound.getLoc(), dataBound.getLowerbound());
      Value ub =
          castValueToArgType(dataBound.getLoc(), dataBound.getUpperbound());
      Value step =
          castValueToArgType(dataBound.getLoc(), dataBound.getStride());
      triples.append({lb, ub, step});
    }
    assert(triples.size() + 1 == initNumArguments &&
           "mismatch between number bounds and number of recipe init block "
           "arguments");
  }
 
  IRMapping mapping;
  SmallVector<Value> initArgs{origPtr};
  initArgs.append(triples);
  mapping.map(initRegion.getBlocks().front().getArguments(), initArgs);
 
  Location loc = op.getLoc();
  setLocation(initRegion, loc);
 
  if constexpr (std::is_same_v<OpTy, acc::PrivateOp>) {
    // Clone the init region for a private.
    Block *block = &region.front();
    auto [results, ip] = acc::cloneACCRegionInto(
        &initRegion, block, block->begin(), mapping, {accPtr});
    assert(!results.empty() && "expected a result from init region");
    saveVarName(op.getAccVar(), results[0]);
    resolveVarNamePlaceholders(block, ip, acc::getVariableName(op.getAccVar()));
    // Clone the destroy region for a private, if it exists.
    if (!recipe.getDestroyRegion().empty()) {
      results.insert(results.begin(), origPtr);
      results.append(triples);
      cloneDestroy(loc, recipe, block, std::prev(block->end()), results);
    }
  } else if constexpr (std::is_same_v<OpTy, acc::FirstprivateOp>) {
    // Clone the init region for a firstprivate.
    Block *block = &region.front();
    auto [results, ip] = acc::cloneACCRegionInto(
        &initRegion, block, block->begin(), mapping, {accPtr});
    assert(!results.empty() && "expected a result from init region");
    saveVarName(op.getAccVar(), results[0]);
    resolveVarNamePlaceholders(block, ip, acc::getVariableName(op.getAccVar()));
    // The copy only consumes the original and user-visible private value.
    SmallVector<Value> copyArgs{origPtr, results.front()};
    copyArgs.append(triples);
    // Destruction also consumes any cleanup values yielded by init.
    SmallVector<Value> destroyArgs{origPtr};
    destroyArgs.append(results);
    destroyArgs.append(triples);
 
    // Clone the copy region for a firstprivate
    mapping.clear();
    mapping.map(recipe.getCopyRegion().front().getArguments(), copyArgs);
    // Clone the copy region for a firstprivate.
    Region &copyRegion = recipe.getCopyRegion();
    setLocation(copyRegion, loc);
    acc::cloneACCRegionInto(&copyRegion, block, std::next(ip), mapping, {});
    if (!recipe.getDestroyRegion().empty()) {
      cloneDestroy(loc, recipe, block, std::prev(block->end()), destroyArgs);
    }
  } else if constexpr (std::is_same_v<OpTy, acc::ReductionOp>) {
    auto cloneRegionIntoAccRegion = [&](Region *src, Region *dest,
                                        bool hasResult) {
      src->cloneInto(dest, mapping);
      Block *block = &dest->front();
      Operation *terminator = block->getTerminator();
      b.setInsertionPoint(terminator);
      if (hasResult)
        acc::YieldOp::create(b, op.getLoc(), terminator->getOperands());
      else
        acc::YieldOp::create(b, op.getLoc(), ValueRange{});
      terminator->erase();
    };
 
    // Clone the init region into acc.reduction_init.
    if constexpr (std::is_same_v<AccOpTy, acc::ParallelOp>)
      b.setInsertionPointToStart(&region.front());
    else if constexpr (std::is_same_v<AccOpTy, acc::LoopOp>)
      b.setInsertionPoint(op);
    else
      llvm_unreachable("unexpected acc op with reduction recipe");
 
    SmallVector<Value> reductionBounds(acc::getBounds(op));
    auto reductionOp =
        acc::ReductionInitOp::create(b, op.getLoc(), origPtr, reductionBounds,
                                     recipe.getReductionOperatorAttr());
    saveVarName(op.getAccVar(), reductionOp.getResult());
    cloneRegionIntoAccRegion(&initRegion, &reductionOp.getRegion(),
                             /*hasResult=*/true);
    Block *initBlock = &reductionOp.getRegion().front();
    resolveVarNamePlaceholders(initBlock, std::prev(initBlock->end()),
                               acc::getVariableName(op.getAccVar()));
 
    // Update the uses within the loop to use the reduction op result.
    replaceAllUsesInRegionWith(accPtr, reductionOp.getResult(), region);
 
    // Clone the combiner region into acc.reduction_combine_region.
    Region &combinerRegion = recipe.getCombinerRegion();
    setLocation(combinerRegion, loc);
 
    Block *entryBlock = &combinerRegion.front();
 
    if constexpr (std::is_same_v<AccOpTy, acc::ParallelOp>)
      b.setInsertionPoint(region.back().getTerminator());
    else if constexpr (std::is_same_v<AccOpTy, acc::LoopOp>)
      b.setInsertionPointAfter(accOp);
    else
      llvm_unreachable("unexpected acc op with reduction recipe");
 
    // Map the first two block arguments to the original and private
    // reduction variables. If the recipe's combiner region has the bounds
    // arguments, we have to map them to the corresponding operands of
    // acc.reduction operation.
    mapping.clear();
    SmallVector<Value, 2> argsRemapping{origPtr, reductionOp.getResult()};
    argsRemapping.append(triples);
    mapping.map(entryBlock->getArguments(), argsRemapping);
 
    auto combineRegionOp = acc::ReductionCombineRegionOp::create(
        b, op.getLoc(), origPtr, reductionOp.getResult());
    cloneRegionIntoAccRegion(&combinerRegion, &combineRegionOp.getRegion(),
                             /*hasResult=*/false);
 
    auto *ctx = b.getContext();
 
    // For reductions that come from parallel constructs, explicitly set the
    // GPU parallel dimensions attribute to blockXDim since they will always be
    // gang private. GPU parallel dimensions cannot be determined for acc.loop
    // at this point.
    if constexpr (std::is_same_v<AccOpTy, acc::ParallelOp>) {
      acc::GPUParallelDimsAttr parDimsAttr;
      if (accOp.isEffectivelySerial()) {
        // If acc.serial has been lowered to a parallel op that is effectively
        // sequential
        parDimsAttr = acc::getSeqParDimsAttr(ctx, policy);
      } else {
        parDimsAttr = acc::getGangDim1ParDimsAttr(ctx, policy);
      }
      acc::setParDimsAttr(reductionOp, parDimsAttr);
      acc::setParDimsAttr(combineRegionOp, parDimsAttr);
    }
 
    // Set sequential parallel dimensions attribute for loops in the recipe.
    auto setSeqParDimsForRecipeLoops = [&](Region *r) {
      r->walk([&](LoopLikeOpInterface loopLike) {
        acc::setParDimsAttr(loopLike, acc::getSeqParDimsAttr(ctx, policy));
      });
    };
    setSeqParDimsForRecipeLoops(&reductionOp.getRegion());
    setSeqParDimsForRecipeLoops(&combineRegionOp.getRegion());
 
    if (!recipe.getDestroyRegion().empty()) {
      SmallVector<Value> results{origPtr, reductionOp.getResult()};
      results.append(triples);
      Block::iterator ip = std::next(Block::iterator(combineRegionOp));
      cloneDestroy(loc, recipe, combineRegionOp->getBlock(), ip, results);
    }
  } else {
    llvm_unreachable("unexpected op type");
  }
 
  op.erase();
  return success();
}
 
template <typename OpTy>
LogicalResult ACCRecipeMaterialization::materializeForACCOp(
    OpTy accOp, acc::OpenACCSupport &accSupport,
    acc::ACCToGPUMappingPolicy &policy) const {
  assert(isa<ACC_COMPUTE_CONSTRUCT_AND_LOOP_OPS>(accOp));
 
  // Reduction recipes use the original variable both to initialize the
  // private reduction value and to combine the result. Preserve copy semantics
  // when materializing reductions on compute constructs, before the
  // acc.reduction operation carrying that intent is erased. Loop reductions
  // are excluded: their original variable can be an outer private reduction
  // value rather than a host variable. Keep acc.reduction referring to its
  // original host variable and pass the mapped value to materialization
  // separately.
  struct ReductionMapping {
    Value originalVar;
    Value mappedVar;
  };
  SmallVector<ReductionMapping> mappedReductionVars;
  if constexpr (!std::is_same_v<OpTy, acc::LoopOp>) {
    for (Value dataOperand : accOp.getDataClauseOperands()) {
      Operation *dataOp = dataOperand.getDefiningOp();
      if (dataOp && isa<ACC_DATA_ENTRY_OPS>(dataOp))
        mappedReductionVars.push_back({acc::getVar(dataOp), dataOperand});
    }
  }
 
  auto getMappedReductionVar = [&](acc::ReductionOp reductionOp) -> Value {
    if constexpr (std::is_same_v<OpTy, acc::LoopOp>) {
      return reductionOp.getVar();
    } else {
      Value originalVar = reductionOp.getVar();
      if (isa_and_nonnull<ACC_DATA_ENTRY_OPS>(originalVar.getDefiningOp()))
        return originalVar;
 
      // Note that we do not require matching bounds here.
      // Bounds may be represented by different SSA values while evaluating to
      // the same values at runtime. Data clauses for the same variable are
      // expected to specify matching bounds.
      auto existing = llvm::find_if(mappedReductionVars,
                                    [&](const ReductionMapping &mapping) {
                                      return mapping.originalVar == originalVar;
                                    });
      if (existing != mappedReductionVars.end())
        return existing->mappedVar;
 
      OpBuilder builder(reductionOp);
      acc::CopyinOp copyinOp;
      if (std::optional<StringRef> name = reductionOp.getName())
        copyinOp =
            acc::CopyinOp::create(builder, reductionOp.getLoc(), originalVar,
                                  /*structured=*/true, /*implicit=*/true, *name,
                                  reductionOp.getBounds());
      else
        copyinOp = acc::CopyinOp::create(
            builder, reductionOp.getLoc(), originalVar,
            /*structured=*/true, /*implicit=*/true, reductionOp.getBounds());
      copyinOp.setDataClause(acc::DataClause::acc_reduction);
      accOp.getDataClauseOperandsMutable().append(copyinOp.getAccVar());
 
      builder.setInsertionPointAfter(accOp);
      acc::CopyoutOp copyoutOp;
      if (std::optional<StringRef> name = reductionOp.getName())
        copyoutOp = acc::CopyoutOp::create(
            builder, reductionOp.getLoc(), copyinOp.getAccVar(), originalVar,
            /*structured=*/true, /*implicit=*/true, *name,
            reductionOp.getBounds());
      else
        copyoutOp = acc::CopyoutOp::create(
            builder, reductionOp.getLoc(), copyinOp.getAccVar(), originalVar,
            /*structured=*/true, /*implicit=*/true, reductionOp.getBounds());
      copyoutOp.setDataClause(acc::DataClause::acc_reduction);
 
      mappedReductionVars.push_back({originalVar, copyinOp.getAccVar()});
      return copyinOp.getAccVar();
    }
  };
 
  if (!accOp.getFirstprivateOperands().empty()) {
    // Clear the firstprivate operands list so there will be no uses after
    // the recipe is materialized.
    SmallVector<Value> operands(accOp.getFirstprivateOperands());
    accOp.getFirstprivateOperandsMutable().clear();
    for (Value operand : operands) {
      auto firstprivateOp = cast<acc::FirstprivateOp>(operand.getDefiningOp());
      auto symbolRef = cast<SymbolRefAttr>(firstprivateOp.getRecipeAttr());
      auto decl = SymbolTable::lookupNearestSymbolFrom(accOp, symbolRef);
      auto recipeOp = cast<acc::FirstprivateRecipeOp>(decl);
      LLVM_DEBUG(llvm::dbgs() << "materializing: " << firstprivateOp << "\n"
                              << symbolRef << "\n");
      handleInitialValueMapping(firstprivateOp);
      if (failed(
              materialize(firstprivateOp, recipeOp, accOp, accSupport, policy)))
        return failure();
    }
  }
 
  if (!accOp.getPrivateOperands().empty()) {
    // Clear the private operands list so there will be no uses after
    // the recipe is materialized.
    SmallVector<Value> operands(accOp.getPrivateOperands());
    accOp.getPrivateOperandsMutable().clear();
    for (Value operand : operands) {
      auto privateOp = cast<acc::PrivateOp>(operand.getDefiningOp());
      auto symbolRef = cast<SymbolRefAttr>(privateOp.getRecipeAttr());
      auto decl = SymbolTable::lookupNearestSymbolFrom(accOp, symbolRef);
      auto recipeOp = cast<acc::PrivateRecipeOp>(decl);
      LLVM_DEBUG(llvm::dbgs() << "materializing: " << privateOp << "\n"
                              << symbolRef << "\n");
      if (readsVar(recipeOp.getInitRegion()) ||
          readsVar(recipeOp.getDestroyRegion()))
        handleInitialValueMapping(privateOp);
      if (failed(materialize(privateOp, recipeOp, accOp, accSupport, policy)))
        return failure();
    }
  }
 
  if (!accOp.getReductionOperands().empty()) {
    // Clear the reduction operands list so there will be no uses after
    // the recipe is materialized.
    SmallVector<Value> operands(accOp.getReductionOperands());
    accOp.getReductionOperandsMutable().clear();
    for (Value operand : operands) {
      auto reductionOp = cast<acc::ReductionOp>(operand.getDefiningOp());
      auto symbolRef = cast<SymbolRefAttr>(reductionOp.getRecipeAttr());
      auto decl = SymbolTable::lookupNearestSymbolFrom(accOp, symbolRef);
      auto recipeOp = cast<acc::ReductionRecipeOp>(decl);
      LLVM_DEBUG(llvm::dbgs() << "materializing: " << reductionOp << "\n"
                              << symbolRef << "\n");
      Value mappedVar = getMappedReductionVar(reductionOp);
      if (failed(materialize(reductionOp, recipeOp, accOp, accSupport, policy,
                             mappedVar)))
        return failure();
    }
  }
  return success();
}
 
void ACCRecipeMaterialization::runOnOperation() {
  ModuleOp moduleOp = getOperation();
  acc::OpenACCSupport &accSupport = getAnalysis<acc::OpenACCSupport>();
 
  acc::DefaultACCToGPUMappingPolicy policy;
 
  // Materialize all recipes for all compute constructs and loop constructs.
  bool anyFailed = false;
  moduleOp.walk([&](Operation *op) {
    if (anyFailed)
      return;
    TypeSwitch<Operation *>(op).Case<ACC_COMPUTE_CONSTRUCT_AND_LOOP_OPS>(
        [&](auto constructOp) {
          if (failed(materializeForACCOp(constructOp, accSupport, policy)))
            anyFailed = true;
        });
  });
  if (anyFailed) {
    signalPassFailure();
    return;
  }
 
  // Remove all recipes. Gather the symbol uses that are left with a single
  // walk: asking whether each recipe is still referenced walks the whole module
  // again for every recipe, and recipes are generated per type, so there can be
  // many of them.
  std::optional<llvm::DenseSet<StringAttr>> usedSymbols;
  // The module region, not the module op, is the symbol table scope: asking
  // for the uses on the op itself would not walk into the body.
  if (std::optional<SymbolTable::UseRange> uses =
          SymbolTable::getSymbolUses(&moduleOp.getBodyRegion())) {
    usedSymbols.emplace();
    for (const SymbolTable::SymbolUse &use : *uses)
      usedSymbols->insert(use.getSymbolRef().getLeafReference());
  }
 
  moduleOp.walk([&](Operation *op) {
    if (auto recipe = dyn_cast<acc::ReductionRecipeOp>(op))
      removeRecipe(recipe, moduleOp, usedSymbols);
    else if (auto recipe = dyn_cast<acc::PrivateRecipeOp>(op))
      removeRecipe(recipe, moduleOp, usedSymbols);
    else if (auto recipe = dyn_cast<acc::FirstprivateRecipeOp>(op))
      removeRecipe(recipe, moduleOp, usedSymbols);
  });
}
} // namespace