OpenACCUtils.cpp

Go to the documentation of this file.

//===- OpenACCUtils.cpp ---------------------------------------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Dialect/OpenACC/OpenACCUtils.h"
 
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/OpenACC/OpenACC.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Dominance.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Interfaces/ViewLikeInterface.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/IR/Intrinsics.h"
#include "llvm/Support/Casting.h"
 
mlir::Operation *mlir::acc::getEnclosingComputeOp(mlir::Region &region) {
  return region
      .getParentOfType<ACC_COMPUTE_CONSTRUCT_OPS, mlir::acc::ComputeRegionOp>();
}
 
mlir::Value mlir::acc::getACCOperandForBlockArg(mlir::Value v) {
  auto barg = mlir::dyn_cast<mlir::BlockArgument>(v);
  if (!barg)
    return nullptr;
 
  mlir::Block *block = barg.getOwner();
  auto computeReg =
      mlir::dyn_cast<mlir::acc::ComputeRegionOp>(block->getParentOp());
  if (!computeReg)
    return nullptr;
  assert(block == computeReg.getBody() &&
         "block must be the body of acc.compute_region");
  return computeReg.getOperand(barg);
}
 
mlir::Operation *mlir::acc::getACCDataClauseOpForBlockArg(mlir::Value v) {
  mlir::Value orig = getACCOperandForBlockArg(v);
  if (!orig)
    return nullptr;
  mlir::Operation *def = orig.getDefiningOp();
  return mlir::isa_and_nonnull<ACC_DATA_ENTRY_OPS>(def) ? def : nullptr;
}
 
template <typename OpTy>
static bool isOnlyUsedByOpClauses(mlir::Value val, mlir::Region &region) {
  auto checkIfUsedOnlyByOpInside = [&](mlir::Operation *user) {
    // For any users which are not in the current acc region, we can ignore.
    // Return true so that it can be used in a `all_of` check.
    if (!region.isAncestor(user->getParentRegion()))
      return true;
    return mlir::isa<OpTy>(user);
  };
 
  return llvm::all_of(val.getUsers(), checkIfUsedOnlyByOpInside);
}
 
bool mlir::acc::isOnlyUsedByPrivateClauses(mlir::Value val,
                                           mlir::Region &region) {
  return isOnlyUsedByOpClauses<mlir::acc::PrivateOp>(val, region);
}
 
bool mlir::acc::isOnlyUsedByReductionClauses(mlir::Value val,
                                             mlir::Region &region) {
  return isOnlyUsedByOpClauses<mlir::acc::ReductionOp>(val, region);
}
 
std::optional<mlir::acc::ClauseDefaultValue>
mlir::acc::getDefaultAttr(Operation *op) {
  std::optional<mlir::acc::ClauseDefaultValue> defaultAttr;
  Operation *currOp = op;
 
  // Iterate outwards until a default clause is found (since OpenACC
  // specification notes that a visible default clause is the nearest default
  // clause appearing on the compute construct or a lexically containing data
  // construct.
  while (!defaultAttr.has_value() && currOp) {
    defaultAttr =
        llvm::TypeSwitch<mlir::Operation *,
                         std::optional<mlir::acc::ClauseDefaultValue>>(currOp)
            .Case<ACC_COMPUTE_CONSTRUCT_OPS, mlir::acc::DataOp>(
                [&](auto op) { return op.getDefaultAttr(); })
            .Default([&](Operation *) { return std::nullopt; });
    currOp = currOp->getParentOp();
  }
 
  return defaultAttr;
}
 
mlir::acc::VariableTypeCategory mlir::acc::getTypeCategory(mlir::Value var) {
  mlir::acc::VariableTypeCategory typeCategory =
      mlir::acc::VariableTypeCategory::uncategorized;
  if (auto mappableTy = dyn_cast<mlir::acc::MappableType>(var.getType()))
    typeCategory = mappableTy.getTypeCategory(var);
  else if (auto pointerLikeTy =
               dyn_cast<mlir::acc::PointerLikeType>(var.getType()))
    typeCategory = pointerLikeTy.getPointeeTypeCategory(
        cast<TypedValue<mlir::acc::PointerLikeType>>(var),
        pointerLikeTy.getElementType());
  return typeCategory;
}
 
llvm::StringLiteral mlir::acc::getVarNamePlaceholder() {
  return llvm::StringLiteral("<acc.varname.placeholder>");
}
 
std::string mlir::acc::getVariableName(mlir::Value v) {
  Value current = v;
 
  // Walk through view operations until a name is found or can't go further
  while (Operation *definingOp = current.getDefiningOp()) {
    // For integer constants, return their value as a string.
    if (std::optional<int64_t> constVal = getConstantIntValue(current))
      return std::to_string(*constVal);
 
    // Check for `acc.var_name` attribute
    if (auto varNameAttr = definingOp->getDiscardableAttrOfType<VarNameAttr>(
            getVarNameAttrName()))
      return varNameAttr.getName().str();
 
    // If it is a data entry operation, get name via getVarName
    if (isa<ACC_DATA_ENTRY_OPS, MapInfoOp>(definingOp))
      if (auto name = acc::getVarName(definingOp))
        return name->str();
 
    // A global goes by the symbol it is addressed through.
    if (auto addressOf = dyn_cast<AddressOfGlobalOpInterface>(definingOp))
      return addressOf.getSymbol().getLeafReference().str();
 
    // If it's a view operation, continue to the source
    if (auto viewOp = dyn_cast<ViewLikeOpInterface>(definingOp)) {
      current = viewOp.getViewSource();
      continue;
    }
 
    break;
  }
 
  return "";
}
 
std::string mlir::acc::getRecipeName(mlir::acc::RecipeKind kind,
                                     mlir::Type type) {
  assert(kind == mlir::acc::RecipeKind::private_recipe ||
         kind == mlir::acc::RecipeKind::firstprivate_recipe ||
         kind == mlir::acc::RecipeKind::reduction_recipe);
  if (!llvm::isa<mlir::acc::PointerLikeType, mlir::acc::MappableType>(type))
    return "";
 
  std::string recipeName;
  llvm::raw_string_ostream ss(recipeName);
  ss << (kind == mlir::acc::RecipeKind::private_recipe ? "privatization_"
         : kind == mlir::acc::RecipeKind::firstprivate_recipe
             ? "firstprivatization_"
             : "reduction_");
 
  // Print the type using its dialect-defined textual format.
  type.print(ss);
  ss.flush();
 
  // Replace invalid characters (anything that's not a letter, number, or
  // period) since this needs to be a valid MLIR identifier.
  for (char &c : recipeName) {
    if (!std::isalnum(static_cast<unsigned char>(c)) && c != '.' && c != '_') {
      if (c == '?')
        c = 'U';
      else if (c == '*')
        c = 'Z';
      else if (c == '(' || c == ')' || c == '[' || c == ']' || c == '{' ||
               c == '}' || c == '<' || c == '>')
        c = '_';
      else
        c = 'X';
    }
  }
 
  return recipeName;
}
 
mlir::Value mlir::acc::getBaseEntity(mlir::Value val) {
  if (auto partialEntityAccessOp =
          val.getDefiningOp<PartialEntityAccessOpInterface>()) {
    if (!partialEntityAccessOp.isCompleteView())
      return partialEntityAccessOp.getBaseEntity();
  }
 
  return val;
}
 
/// Look up `symbol` in the `gpu.module`s of the enclosing module. A
/// `gpu.module` is its own symbol table, so `lookupNearestSymbolFrom` from a
/// user outside of it does not find definitions placed inside.
static mlir::Operation *lookupSymbolInGPUModules(mlir::Operation *user,
                                                 mlir::SymbolRefAttr symbol) {
  auto moduleOp = user->getParentOfType<mlir::ModuleOp>();
  if (!moduleOp)
    return nullptr;
  for (auto gpuModule : moduleOp.getOps<mlir::gpu::GPUModuleOp>()) {
    if (mlir::Operation *op = mlir::SymbolTable::lookupSymbolIn(
            gpuModule, symbol.getRootReference()))
      return op;
  }
  return nullptr;
}
 
bool mlir::acc::isValidSymbolUse(mlir::Operation *user,
                                 mlir::SymbolRefAttr symbol,
                                 mlir::Operation **definingOpPtr) {
  // A pass may prepare the device-side definition of a symbol inside a
  // `gpu.module` while the use is still a reference from host IR. Such a
  // symbol is meant to be used on device, so the use is valid.
  if (mlir::Operation *gpuOp = lookupSymbolInGPUModules(user, symbol)) {
    if (definingOpPtr)
      *definingOpPtr = gpuOp;
    return true;
  }
 
  mlir::Operation *definingOp =
      mlir::SymbolTable::lookupNearestSymbolFrom(user, symbol);
 
  // If there are no defining ops, we have no way to ensure validity because
  // we cannot check for any attributes.
  if (!definingOp)
    return false;
 
  if (definingOpPtr)
    *definingOpPtr = definingOp;
 
  // Check if the defining op is a recipe.
  // Recipes are valid as they get materialized before being offloaded to
  // device. They are only instructions for how to materialize.
  if (mlir::isa<mlir::accomp::RecipeInterface>(definingOp))
    return true;
 
  // Check if the defining op is a global variable that is device data.
  // Device data is already resident on the device and does not need mapping.
  if (auto globalVar =
          mlir::dyn_cast<mlir::acc::GlobalVariableOpInterface>(definingOp))
    if (globalVar.isDeviceAccessible())
      return true;
 
  // Check if the defining op is a function
  if (auto func =
          mlir::dyn_cast_if_present<mlir::FunctionOpInterface>(definingOp)) {
    // If this symbol is actually an acc routine or a specialized acc routine -
    // then it is expected for it to be offloaded - therefore it is valid.
    if (func->hasDiscardableAttr(mlir::acc::getRoutineInfoAttrName()) ||
        func->hasDiscardableAttr(mlir::acc::getSpecializedRoutineAttrName()))
      return true;
 
    // If this symbol is a call to an LLVM intrinsic, then it is likely valid.
    // Check the following:
    // 1. The function is private
    // 2. The function has no body
    // 3. Name starts with "llvm."
    // 4. The function's name is a valid LLVM intrinsic name
    if (func.getVisibility() == mlir::SymbolTable::Visibility::Private &&
        func.getFunctionBody().empty() && func.getName().starts_with("llvm.") &&
        llvm::Intrinsic::lookupIntrinsicID(func.getName()) !=
            llvm::Intrinsic::not_intrinsic)
      return true;
  }
 
  // A declare attribute is needed for symbol references.
  bool hasDeclare =
      definingOp->hasDiscardableAttr(mlir::acc::getDeclareAttrName());
  return hasDeclare;
}
 
/// Finds the global variable that `addrOf` refers to, through `symbolTables`
/// when provided (cached lookup) and by scanning the nearest symbol table
/// otherwise.
static mlir::acc::GlobalVariableOpInterface
lookupGlobalVariable(mlir::Operation *addrOf, mlir::SymbolRefAttr symbol,
                     mlir::SymbolTableCollection *symbolTables) {
  if (symbolTables)
    return symbolTables
        ->lookupNearestSymbolFrom<mlir::acc::GlobalVariableOpInterface>(addrOf,
                                                                        symbol);
  return mlir::SymbolTable::lookupNearestSymbolFrom<
      mlir::acc::GlobalVariableOpInterface>(addrOf, symbol);
}
 
bool mlir::acc::isDeviceAccessibleValue(
    mlir::Value val, mlir::SymbolTableCollection *symbolTables) {
  // Check if the value is device data via type interfaces.
  // Device data is already resident on the device and does not need mapping.
  if (auto mappableTy = dyn_cast<mlir::acc::MappableType>(val.getType()))
    if (mappableTy.isDeviceAccessible(val, symbolTables))
      return true;
 
  if (auto pointerLikeTy = dyn_cast<mlir::acc::PointerLikeType>(val.getType()))
    if (pointerLikeTy.isDeviceAccessible(val, symbolTables))
      return true;
 
  mlir::Operation *defOp = val.getDefiningOp();
  if (!defOp)
    return false;
 
  // `acc.declare` with deviceptr marks data that is already associated with
  // the device.
  if (auto declareAttr =
          defOp->getDiscardableAttrOfType<mlir::acc::DeclareAttr>(
              mlir::acc::getDeclareAttrName()))
    if (declareAttr.getDataClause().getValue() ==
        mlir::acc::DataClause::acc_deviceptr)
      return true;
 
  // Handle operations that access a partial entity - check if the base entity
  // is device data.
  if (auto partialAccess =
          dyn_cast<mlir::acc::PartialEntityAccessOpInterface>(defOp)) {
    if (mlir::Value base = partialAccess.getBaseEntity())
      return isDeviceAccessibleValue(base, symbolTables);
  }
 
  // Handle address_of - check if the referenced global is device data.
  if (auto addrOfIface =
          dyn_cast<mlir::acc::AddressOfGlobalOpInterface>(defOp)) {
    auto symbol = addrOfIface.getSymbol();
    if (auto global = lookupGlobalVariable(defOp, symbol, symbolTables))
      return global.isDeviceAccessible();
  }
 
  return false;
}
 
bool mlir::acc::isInDeviceMemoryValue(
    mlir::Value val, mlir::SymbolTableCollection *symbolTables) {
  // In-device-memory data is a subset of device-accessible data: it must be
  // accessible and its storage must physically reside in device memory.
  if (auto mappableTy = dyn_cast<mlir::acc::MappableType>(val.getType()))
    if (mappableTy.isInDeviceMemory(val, symbolTables))
      return true;
 
  if (auto pointerLikeTy = dyn_cast<mlir::acc::PointerLikeType>(val.getType()))
    if (pointerLikeTy.isInDeviceMemory(val, symbolTables))
      return true;
 
  mlir::Operation *defOp = val.getDefiningOp();
  if (!defOp)
    return false;
 
  // `acc.declare` with deviceptr marks data whose storage is already in device
  // memory.
  if (auto declareAttr =
          defOp->getDiscardableAttrOfType<mlir::acc::DeclareAttr>(
              mlir::acc::getDeclareAttrName()))
    if (declareAttr.getDataClause().getValue() ==
        mlir::acc::DataClause::acc_deviceptr)
      return true;
 
  // Handle operations that access a partial entity - check if the base entity
  // is in device memory.
  if (auto partialAccess =
          dyn_cast<mlir::acc::PartialEntityAccessOpInterface>(defOp)) {
    if (mlir::Value base = partialAccess.getBaseEntity())
      return isInDeviceMemoryValue(base, symbolTables);
  }
 
  // Handle address_of - check if the referenced global is in device memory.
  if (auto addrOfIface =
          dyn_cast<mlir::acc::AddressOfGlobalOpInterface>(defOp)) {
    auto symbol = addrOfIface.getSymbol();
    if (auto global = lookupGlobalVariable(defOp, symbol, symbolTables))
      return global.isInDeviceMemory();
  }
 
  return false;
}
 
bool mlir::acc::isValidValueUse(mlir::Value val, mlir::Region &region,
                                mlir::SymbolTableCollection *symbolTables) {
  // Types that can be passed by value are legal.
  Type type = val.getType();
  if (type.isIntOrIndexOrFloat() || isa<mlir::ComplexType>(type) ||
      llvm::isa<mlir::VectorType>(type))
    return true;
 
  // If this is produced by an ACC data entry operation, it is valid.
  if (isa_and_nonnull<ACC_DATA_ENTRY_OPS>(val.getDefiningOp()))
    return true;
 
  // If the value is only used by private clauses, it is not a live-in.
  if (isOnlyUsedByPrivateClauses(val, region))
    return true;
 
  // If this is device data, it is valid.
  if (isDeviceAccessibleValue(val, symbolTables))
    return true;
 
  // Arguments of an enclosing acc routine are already on the device.
  if (mlir::Operation *parent = region.getParentOp()) {
    if (auto func = parent->getParentOfType<mlir::FunctionOpInterface>()) {
      if ((mlir::acc::isAccRoutine(func) ||
           mlir::acc::isSpecializedAccRoutine(func)) &&
          llvm::is_contained(func.getArguments(), val))
        return true;
    }
  }
 
  return false;
}
 
llvm::SmallVector<mlir::Value>
mlir::acc::getDominatingDataClauses(mlir::Operation *computeConstructOp,
                                    mlir::DominanceInfo &domInfo,
                                    mlir::PostDominanceInfo &postDomInfo) {
  llvm::SmallSetVector<mlir::Value, 8> dominatingDataClauses;
 
  llvm::TypeSwitch<mlir::Operation *>(computeConstructOp)
      .Case<mlir::acc::ParallelOp, mlir::acc::KernelsOp, mlir::acc::SerialOp>(
          [&](auto op) {
            for (auto dataClause : op.getDataClauseOperands()) {
              dominatingDataClauses.insert(dataClause);
            }
          })
      .Default([](mlir::Operation *) {});
 
  // Collect the data clauses from enclosing data constructs.
  mlir::Operation *currParentOp = computeConstructOp->getParentOp();
  while (currParentOp) {
    if (mlir::isa<mlir::acc::DataOp>(currParentOp)) {
      for (auto dataClause : mlir::dyn_cast<mlir::acc::DataOp>(currParentOp)
                                 .getDataClauseOperands()) {
        dominatingDataClauses.insert(dataClause);
      }
    }
    currParentOp = currParentOp->getParentOp();
  }
 
  // Find the enclosing function/subroutine
  auto funcOp =
      computeConstructOp->getParentOfType<mlir::FunctionOpInterface>();
  if (!funcOp)
    return dominatingDataClauses.takeVector();
 
  // Walk the function to find `acc.declare_enter`/`acc.declare_exit` pairs that
  // dominate and post-dominate the compute construct and add their data
  // clauses to the list.
  funcOp->walk([&](mlir::acc::DeclareEnterOp declareEnterOp) {
    if (domInfo.dominates(declareEnterOp.getOperation(), computeConstructOp)) {
      // Collect all `acc.declare_exit` ops for this token.
      llvm::SmallVector<mlir::acc::DeclareExitOp> exits;
      for (auto *user : declareEnterOp.getToken().getUsers())
        if (auto declareExit = mlir::dyn_cast<mlir::acc::DeclareExitOp>(user))
          exits.push_back(declareExit);
 
      // Only add clauses if every `acc.declare_exit` op post-dominates the
      // compute construct.
      if (!exits.empty() &&
          llvm::all_of(exits, [&](mlir::acc::DeclareExitOp exitOp) {
            return postDomInfo.postDominates(exitOp, computeConstructOp);
          })) {
        for (auto dataClause : declareEnterOp.getDataClauseOperands())
          dominatingDataClauses.insert(dataClause);
      }
    }
  });
 
  return dominatingDataClauses.takeVector();
}
 
static mlir::Value getIfCondition(mlir::Operation *op) {
  if (auto ifOp = mlir::dyn_cast<mlir::scf::IfOp>(op))
    return ifOp.getCondition();
 
  if (!mlir::isa<mlir::RegionBranchOpInterface>(op) ||
      mlir::isa<mlir::LoopLikeOpInterface>(op))
    return {};
  if (op->getNumOperands() != 1 || op->getNumRegions() != 2 ||
      !op->getOperand(0).getType().isSignlessInteger(1))
    return {};
  return op->getOperand(0);
}
 
bool mlir::acc::isInOffTargetBranch(mlir::Operation *op,
                                    llvm::ArrayRef<int64_t> deviceTypes) {
  // No device type is known to be true, so no branch can be classified.
  if (deviceTypes.empty())
    return false;
 
  for (mlir::Operation *parent = op->getParentOp();
       parent &&
       !mlir::isa<ACC_COMPUTE_CONSTRUCT_OPS, mlir::acc::ComputeRegionOp,
                  mlir::FunctionOpInterface>(parent);
       parent = parent->getParentOp()) {
    mlir::Value condition = getIfCondition(parent);
    if (!condition)
      continue;
 
    // The condition must be the `acc.on_device` result itself.
    auto onDeviceOp = condition.getDefiningOp<mlir::acc::OnDeviceOp>();
    if (!onDeviceOp)
      continue;
 
    std::optional<int64_t> deviceTypeValue =
        mlir::getConstantIntValue(onDeviceOp.getDeviceType());
    if (!deviceTypeValue)
      continue;
 
    bool onTarget = llvm::is_contained(deviceTypes, *deviceTypeValue);
    bool inThen = parent->getRegion(0).isAncestor(op->getParentRegion());
    bool inElse = parent->getRegion(1).isAncestor(op->getParentRegion());
    // Off the target: the then of a device type outside `deviceTypes`, or
    // the else of a device type in `deviceTypes`.
    if ((!onTarget && inThen) || (onTarget && inElse))
      return true;
  }
  return false;
}
 
mlir::remark::detail::InFlightRemark
mlir::acc::emitRemark(mlir::Operation *op,
                      const std::function<std::string()> &messageFn,
                      llvm::StringRef category) {
  using namespace mlir::remark;
  mlir::Location loc = op->getLoc();
  auto *engine = loc->getContext()->getRemarkEngine();
  if (!engine)
    return remark::detail::InFlightRemark{};
 
  llvm::StringRef funcName;
  if (auto func = dyn_cast<mlir::FunctionOpInterface>(op))
    funcName = func.getName();
  else if (auto funcOp = op->getParentOfType<mlir::FunctionOpInterface>())
    funcName = funcOp.getName();
 
  auto opts = RemarkOpts::name("openacc").category(category);
  if (!funcName.empty())
    opts = opts.function(funcName);
 
  auto remark = engine->emitOptimizationRemark(loc, opts);
  if (remark)
    remark << messageFn();
  return remark;
}