CUFCommon.h

//===-- CUFCommon.h -------------------------------------------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
 
#ifndef FORTRAN_OPTIMIZER_TRANSFORMS_CUFCOMMON_H_
#define FORTRAN_OPTIMIZER_TRANSFORMS_CUFCOMMON_H_
 
#include "flang/Optimizer/Dialect/FIROps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/IR/BuiltinOps.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/TargetParser/Triple.h"
 
static constexpr llvm::StringRef cudaDeviceModuleName = "cuda_device_mod";
static constexpr llvm::StringRef cudaSharedMemSuffix = "__shared_mem__";
static constexpr llvm::StringRef cudaDeviceCopySuffix = ".device";
static constexpr llvm::StringRef cudaSharedTypeInfoSection = "__nv_type_info";
static constexpr llvm::StringRef cudaSharedTypeInfoEndPrefix =
    "__nv_type_info_end_";
inline int64_t getCudaSharedTypeInfoAlignment(const llvm::Triple &triple) {
  if (triple.isAArch64() || triple.isPPC64())
    return 65536;
  return 4096;
}
static constexpr llvm::StringRef cudaSharedTypeDescsAttrName =
    "cuf.shared_type_descs";
static constexpr llvm::StringRef cudaHostTypeDescAttrName =
    "cuf.host_type_desc";
 
namespace fir {
class FirOpBuilder;
class KindMapping;
} // namespace fir
 
namespace cuf {

mlir::gpu::GPUModuleOp getOrCreateGPUModule(mlir::ModuleOp mod,
                                            mlir::SymbolTable &symTab);
 
bool isCUDADeviceContext(mlir::Operation *op);
bool isCUDADeviceContext(mlir::Region &,
                         bool isDoConcurrentOffloadEnabled = false);

bool isExecutingOnDevice(mlir::Operation *op);
bool isRegisteredDeviceGlobal(fir::GlobalOp op);
bool isRegisteredDeviceAttr(std::optional<cuf::DataAttribute> attr);

bool isTypeInfoGlobal(fir::GlobalOp op);

bool isTypeDescriptorGlobal(fir::GlobalOp op);

bool isDeviceProcedure(mlir::func::FuncOp funcOp);

struct DeviceCodeSet {
  llvm::SetVector<mlir::func::FuncOp> deviceFuncs;
  llvm::SetVector<mlir::func::FuncOp> calledFromDevice;
  llvm::SetVector<mlir::func::FuncOp> keepInModule;
};

DeviceCodeSet collectDeviceCode(mlir::ModuleOp mod, mlir::SymbolTable &symTab,
                                bool rejectDynamicDispatch = false);

void remapProcedureSymbols(
    mlir::Operation *root,
    const llvm::DenseMap<mlir::StringAttr, mlir::FlatSymbolRefAttr> &map);

void setDeviceAllocationPolicy(mlir::Operation *func);

void genPointerSync(mlir::Location loc, const mlir::Value box,
                    fir::FirOpBuilder &builder);
 
int computeElementByteSize(mlir::Location loc, mlir::Type type,
                           fir::KindMapping &kindMap,
                           bool emitErrorOnFailure = true);
 
mlir::Value computeElementCount(mlir::PatternRewriter &rewriter,
                                mlir::Location loc, mlir::Value shapeOperand,
                                mlir::Type seqType, mlir::Type targetType);
 
} // namespace cuf
 
#endif // FORTRAN_OPTIMIZER_TRANSFORMS_CUFCOMMON_H_