mlir/lib/Dialect/OpenMP/Utils/Utils.cpp

//===- StackToShared.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
//
//===----------------------------------------------------------------------===//
//
// This file implements various OpenMP dialect utilities.
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/OpenMP/Utils/Utils.h"

#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/OpenMP/OpenMPDialect.h"

using namespace mlir;

void mlir::omp::setOffloadModuleInterfaceAttributes(
    ModuleOp module, const OffloadModuleOpts &opts) {
  if (auto offloadMod = llvm::dyn_cast<OffloadModuleInterface>(*module)) {
    offloadMod.setIsTargetDevice(opts.isTargetDevice);
    offloadMod.setIsGPU(opts.isGPU);
    if (opts.forceUSM)
      offloadMod.setRequires(offloadMod.getRequires() |
                             ClauseRequires::unified_shared_memory);
    offloadMod.setFlags(opts.targetDebugKind, opts.assumeTeamsOversubscription,
                        opts.assumeThreadsOversubscription,
                        opts.assumeNoThreadState,
                        opts.assumeNoNestedParallelism,
                        opts.openMPDeviceVersion, opts.noGPULib);
    if (opts.isTargetDevice && !opts.hostIRFile.empty())
      offloadMod.setHostIRFilePath(opts.hostIRFile);

    auto strTriples = llvm::to_vector(
        llvm::map_range(opts.targetTriples, [](const llvm::Triple &triple) {
          return triple.normalize();
        }));
    offloadMod.setTargetTriples(strTriples);
  }
}

void mlir::omp::setOpenMPVersionAttribute(ModuleOp module, int64_t version) {
  module->setDiscardableAttr(
      StringAttr::get(module.getContext(), llvm::Twine{"omp.version"}),
      VersionAttr::get(module.getContext(), version));
}

void mlir::omp::setOpenMPIntegerWrapAround(ModuleOp module, bool value) {
  module->setAttr(StringAttr::get(module.getContext(),
                                  llvm::Twine{"omp.integer_wrap_around"}),
                  IntegerWrapAroundAttr::get(module.getContext(), value));
}

int64_t mlir::omp::getOpenMPVersionAttribute(ModuleOp module,
                                             int64_t fallback) {
  if (Attribute verAttr = module->getDiscardableAttr("omp.version"))
    return llvm::cast<VersionAttr>(verAttr).getVersion();
  return fallback;
}

bool mlir::omp::isOpenMPModule(ModuleOp module) {
  return module->hasDiscardableAttr("omp.version");
}

static bool allocaUseRequiresSharedMem(const OpOperand &use) {
  Operation *owner = use.getOwner();
  if (auto parallelOp = dyn_cast<omp::ParallelOp>(owner)) {
    if (llvm::is_contained(parallelOp.getReductionVars(), use.get()))
      return true;
  } else if (auto callOp = dyn_cast<CallOpInterface>(owner)) {
    if (llvm::is_contained(callOp.getArgOperands(), use.get()))
      return true;
  }

  // If it is used directly inside of a parallel region, it has to be replaced
  // unless the use is a private clause.
  if (owner->getParentOfType<omp::ParallelOp>()) {
    if (auto argIface = dyn_cast<omp::BlockArgOpenMPOpInterface>(owner)) {
      OperandRange privateVars = argIface.getPrivateVars();
      auto it = llvm::find(privateVars, use.get());
      if (it != privateVars.end()) {
        ArrayAttr privateSyms = *argIface.getPrivateSyms();
        size_t idx = std::distance(privateVars.begin(), it);
        auto privateOp =
            SymbolTable::lookupNearestSymbolFrom<omp::PrivateClauseOp>(
                owner, cast<SymbolRefAttr>(privateSyms[idx]));
        return privateOp.getDataSharingType() !=
               omp::DataSharingClauseType::Private;
      }
    }
    return true;
  }
  return false;
}

bool mlir::omp::allocaUsesRequireSharedMem(Value alloc) {
  for (const OpOperand &use : alloc.getUses()) {
    Operation *owner = use.getOwner();
    if (isa<LLVM::AddrSpaceCastOp, LLVM::GEPOp>(owner)) {
      if (llvm::any_of(owner->getResults(), [&](Value result) {
            return allocaUsesRequireSharedMem(result);
          }))
        return true;
    } else if (allocaUseRequiresSharedMem(use)) {
      return true;
    }
  }
  return false;
}

bool mlir::omp::opInSharedDeviceContext(Operation &op) {
  if (isa<omp::ParallelOp>(op))
    return false;

  auto offloadIface = op.getParentOfType<omp::OffloadModuleInterface>();
  if (!offloadIface || !offloadIface.getIsTargetDevice())
    return false;

  auto targetOp = op.getParentOfType<omp::TargetOp>();

  // It must be inside of a generic omp.target or in a target device function,
  // and not inside of omp.parallel.
  if (auto parallelOp = op.getParentOfType<omp::ParallelOp>()) {
    if (!targetOp || targetOp->isProperAncestor(parallelOp))
      return false;
  }

  // The omp.target operation itself is considered in a shared device context in
  // order to properly process its own allocation-defining entry block
  // arguments.
  if (!targetOp)
    targetOp = dyn_cast<omp::TargetOp>(op);

  if (targetOp) {
    if (targetOp.getKernelType() != omp::TargetExecMode::generic)
      return false;
  } else {
    auto declTargetIface = op.getParentOfType<omp::DeclareTargetInterface>();
    omp::DeclareTargetAttr declTargetAttr =
        declTargetIface ? declTargetIface.getDeclareTarget() : nullptr;
    if (!declTargetAttr ||
        declTargetAttr.getDeviceType() == omp::DeclareTargetDeviceType::host)
      return false;
  }
  return true;
}