CIRGenCUDANV.cpp

Go to the documentation of this file.

//========- CIRGenCUDANV.cpp - Interface to NVIDIA CUDA Runtime -----=========//
//
// 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 provides a class for CUDA code generation targeting the NVIDIA CUDA
// runtime library.
//
//===----------------------------------------------------------------------===//
 
#include "CIRGenCUDARuntime.h"
#include "CIRGenCXXABI.h"
#include "CIRGenFunction.h"
#include "CIRGenModule.h"
#include "mlir/IR/Operation.h"
#include "clang/AST/ASTContext.h"
#include "clang/AST/Attrs.inc"
#include "clang/AST/Decl.h"
#include "clang/AST/GlobalDecl.h"
#include "clang/Basic/AddressSpaces.h"
#include "clang/Basic/Cuda.h"
#include "clang/Basic/DiagnosticFrontend.h"
#include "clang/Basic/FileManager.h"
#include "clang/Basic/SourceManager.h"
#include "clang/CIR/Dialect/IR/CIRDialect.h"
#include "clang/CIR/Dialect/IR/CIRTypes.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/VirtualFileSystem.h"
 
using namespace clang;
using namespace clang::CIRGen;
 
namespace {
 
class CIRGenNVCUDARuntime : public CIRGenCUDARuntime {
protected:
  StringRef prefix;
 
  // Map a device stub function to a symbol for identifying kernel in host
  // code. For CUDA, the symbol for identifying the kernel is the same as the
  // device stub function. For HIP, they are different.
  llvm::StringMap<mlir::Operation *> kernelHandles;
 
  // Map a kernel handle to the kernel stub.
  llvm::DenseMap<mlir::Operation *, mlir::Operation *> kernelStubs;
 
  struct VarInfo {
    cir::GlobalOp var;
    const VarDecl *d;
    cir::CUDADeviceVarKind flags;
  };
  llvm::SmallVector<VarInfo, 16> deviceVars;
 
  // Mangle context for device.
  std::unique_ptr<MangleContext> deviceMC;
 
private:
  void emitDeviceStubBodyNew(CIRGenFunction &cgf, cir::FuncOp fn,
                             FunctionArgList &args);
  void recordDeviceBinary();
  void transformManagedVars();
  mlir::Value prepareKernelArgs(CIRGenFunction &cgf, mlir::Location loc,
                                FunctionArgList &args);
  mlir::Operation *getKernelHandle(cir::FuncOp fn, GlobalDecl gd) override;
 
  mlir::Operation *getKernelStub(mlir::Operation *handle) override {
    auto it = kernelStubs.find(handle);
    assert(it != kernelStubs.end());
    return it->second;
  }
  std::string addPrefixToName(StringRef funcName) const;
  std::string addUnderscoredPrefixToName(StringRef funcName) const;
 
public:
  CIRGenNVCUDARuntime(CIRGenModule &cgm);
  ~CIRGenNVCUDARuntime();
 
  void emitDeviceStub(CIRGenFunction &cgf, cir::FuncOp fn,
                      FunctionArgList &args) override;
 
  void handleVarRegistration(const VarDecl *vd, cir::GlobalOp var) override;
  void finalizeModule() override;
  void handleGlobalReplace(cir::GlobalOp oldGV, cir::GlobalOp newGV) override;
 
  void internalizeDeviceSideVar(const VarDecl *d,
                                cir::GlobalLinkageKind &linkage) override;
 
  std::string getDeviceSideName(const NamedDecl *nd) override;
 
  void registerDeviceVar(const VarDecl *vd, cir::GlobalOp &var, bool isExtern,
                         bool isConstant) {
    // Attach the device var attribute to the GlobalOp
    auto &builder = cgm.getBuilder();
    var->setAttr(cir::CUDAVarRegistrationInfoAttr::getMnemonic(),
                 cir::CUDAVarRegistrationInfoAttr::get(
                     builder.getContext(),
                     getDeviceSideName(cast<NamedDecl>(vd)),
                     cir::CUDADeviceVarKind::Variable, isExtern, isConstant,
                     vd->hasAttr<HIPManagedAttr>(),
                     /*surfTexType=*/0));
    deviceVars.push_back({
        var,
        vd,
        cir::CUDADeviceVarKind::Variable,
    });
  }
 
  void registerDeviceSurf(const VarDecl *vd, cir::GlobalOp &var, bool isExtern,
                          int32_t surfaceType) {
    auto &builder = cgm.getBuilder();
 
    var->setAttr(cir::CUDAVarRegistrationInfoAttr::getMnemonic(),
                 cir::CUDAVarRegistrationInfoAttr::get(
                     builder.getContext(),
                     getDeviceSideName(cast<NamedDecl>(vd)),
                     cir::CUDADeviceVarKind::Surface, isExtern,
                     /*isConstant=*/false,
                     /*isManaged=*/false, surfaceType));
 
    deviceVars.push_back({
        var,
        vd,
        cir::CUDADeviceVarKind::Surface,
    });
  }
 
  void registerDeviceTex(const VarDecl *vd, cir::GlobalOp &var, bool isExtern) {
    auto &builder = cgm.getBuilder();
 
    var->setAttr(cir::CUDAVarRegistrationInfoAttr::getMnemonic(),
                 cir::CUDAVarRegistrationInfoAttr::get(
                     builder.getContext(),
                     getDeviceSideName(cast<NamedDecl>(vd)),
                     cir::CUDADeviceVarKind::Texture, isExtern,
                     /*isConstant=*/false,
                     /*isManaged=*/false,
                     /*surfTexType=*/0));
 
    deviceVars.push_back({
        var,
        vd,
        cir::CUDADeviceVarKind::Texture,
    });
  }
};
 
} // namespace
 
std::string CIRGenNVCUDARuntime::addPrefixToName(StringRef funcName) const {
  return (prefix + funcName).str();
}
 
std::string
CIRGenNVCUDARuntime::addUnderscoredPrefixToName(StringRef funcName) const {
  return ("__" + prefix + funcName).str();
}
 
CIRGenNVCUDARuntime::CIRGenNVCUDARuntime(CIRGenModule &cgm)
    : CIRGenCUDARuntime(cgm),
      deviceMC(cgm.getASTContext().cudaNVInitDeviceMC()) {
  if (cgm.getLangOpts().OffloadViaLLVM)
    cgm.errorNYI("CIRGenNVCUDARuntime: Offload via LLVM");
  else if (cgm.getLangOpts().HIP)
    prefix = "hip";
  else
    prefix = "cuda";
}
 
mlir::Value CIRGenNVCUDARuntime::prepareKernelArgs(CIRGenFunction &cgf,
                                                   mlir::Location loc,
                                                   FunctionArgList &args) {
  CIRGenBuilderTy &builder = cgm.getBuilder();
 
  // Build void *args[] and populate with the addresses of kernel arguments.
  auto voidPtrArrayTy = cir::ArrayType::get(cgm.voidPtrTy, args.size());
  mlir::Value kernelArgs =
      builder.createAlloca(loc, cir::PointerType::get(voidPtrArrayTy),
                           "kernel_args", CharUnits::fromQuantity(16));
 
  mlir::Value kernelArgsDecayed =
      builder.createCast(cir::CastKind::array_to_ptrdecay, kernelArgs,
                         cir::PointerType::get(cgm.voidPtrTy));
 
  for (const auto &[i, arg] : llvm::enumerate(args)) {
    mlir::Value index =
        builder.getConstInt(loc, llvm::APInt(/*numBits=*/32, i));
    mlir::Value storePos =
        builder.createPtrStride(loc, kernelArgsDecayed, index);
    mlir::Value argAddr = cgf.getAddrOfLocalVar(arg).getPointer();
    mlir::Value argAsVoid = builder.createBitcast(argAddr, cgm.voidPtrTy);
 
    builder.CIRBaseBuilderTy::createStore(loc, argAsVoid, storePos);
  }
 
  return kernelArgsDecayed;
}
 
// CUDA 9.0+ uses new way to launch kernels. Parameters are packed in a local
// array and kernels are launched using cudaLaunchKernel().
void CIRGenNVCUDARuntime::emitDeviceStubBodyNew(CIRGenFunction &cgf,
                                                cir::FuncOp fn,
                                                FunctionArgList &args) {
 
  // This requires arguments to be sent to kernels in a different way.
  if (cgm.getLangOpts().OffloadViaLLVM)
    cgm.errorNYI("CIRGenNVCUDARuntime: Offload via LLVM");
 
  CIRGenBuilderTy &builder = cgm.getBuilder();
  mlir::Location loc = fn.getLoc();
 
  // For [cuda|hip]LaunchKernel, we must add another layer of indirection
  // to arguments. For example, for function `add(int a, float b)`,
  // we need to pass it as `void *args[2] = { &a, &b }`.
  mlir::Value kernelArgs = prepareKernelArgs(cgf, loc, args);
 
  // Lookup cudaLaunchKernel/hipLaunchKernel function.
  // HIP kernel launching API name depends on -fgpu-default-stream option. For
  // the default value 'legacy', it is hipLaunchKernel. For 'per-thread',
  // it is hipLaunchKernel_spt.
  // cudaError_t cudaLaunchKernel(const void *func, dim3 gridDim, dim3 blockDim,
  //                              void **args, size_t sharedMem,
  //                              cudaStream_t stream);
  // hipError_t hipLaunchKernel[_spt](const void *func, dim3 gridDim,
  //                                  dim3 blockDim, void **args,
  //                                  size_t sharedMem, hipStream_t stream);
  TranslationUnitDecl *tuDecl = cgm.getASTContext().getTranslationUnitDecl();
  DeclContext *dc = TranslationUnitDecl::castToDeclContext(tuDecl);
 
  // The default stream is usually stream 0 (the legacy default stream).
  // For per-thread default stream, we need a different LaunchKernel function.
  std::string kernelLaunchAPI = "LaunchKernel";
  if (cgm.getLangOpts().GPUDefaultStream ==
      LangOptions::GPUDefaultStreamKind::PerThread) {
    if (cgm.getLangOpts().HIP)
      kernelLaunchAPI += "_spt";
    else if (cgm.getLangOpts().CUDA)
      kernelLaunchAPI += "_ptsz";
  }
 
  std::string launchKernelName = addPrefixToName(kernelLaunchAPI);
  const IdentifierInfo &launchII =
      cgm.getASTContext().Idents.get(launchKernelName);
  FunctionDecl *cudaLaunchKernelFD = nullptr;
  for (NamedDecl *result : dc->lookup(&launchII)) {
    if (FunctionDecl *fd = dyn_cast<FunctionDecl>(result))
      cudaLaunchKernelFD = fd;
  }
 
  if (cudaLaunchKernelFD == nullptr) {
    cgm.error(cgf.curFuncDecl->getLocation(),
              "Can't find declaration for " + launchKernelName);
    return;
  }
 
  // Use this function to retrieve arguments for cudaLaunchKernel:
  // int __[cuda|hip]PopCallConfiguration(dim3 *gridDim, dim3 *blockDim, size_t
  //                                *sharedMem, cudaStream_t *stream)
  //
  // Here [cuda|hip]Stream_t, while also being the 6th argument of
  // [cuda|hip]LaunchKernel, is a pointer to some opaque struct.
 
  mlir::Type dim3Ty = cgf.getTypes().convertType(
      cudaLaunchKernelFD->getParamDecl(1)->getType());
  mlir::Type streamTy = cgf.getTypes().convertType(
      cudaLaunchKernelFD->getParamDecl(5)->getType());
 
  mlir::Value gridDim =
      builder.createAlloca(loc, cir::PointerType::get(dim3Ty), "grid_dim",
                           CharUnits::fromQuantity(8));
  mlir::Value blockDim =
      builder.createAlloca(loc, cir::PointerType::get(dim3Ty), "block_dim",
                           CharUnits::fromQuantity(8));
  mlir::Value sharedMem = builder.createAlloca(
      loc, cir::PointerType::get(cgm.sizeTy), "shared_mem", cgm.getSizeAlign());
  mlir::Value stream = builder.createAlloca(
      loc, cir::PointerType::get(streamTy), "stream", cgm.getPointerAlign());
 
  cir::FuncOp popConfig = cgm.createRuntimeFunction(
      cir::FuncType::get({gridDim.getType(), blockDim.getType(),
                          sharedMem.getType(), stream.getType()},
                         cgm.sInt32Ty),
      addUnderscoredPrefixToName("PopCallConfiguration"));
  cgf.emitRuntimeCall(loc, popConfig, {gridDim, blockDim, sharedMem, stream});
 
  // Now emit the call to cudaLaunchKernel
  // [cuda|hip]Error_t [cuda|hip]LaunchKernel(const void *func, dim3 gridDim,
  // dim3 blockDim,
  //                              void **args, size_t sharedMem,
  //                              [cuda|hip]Stream_t stream);
 
  // We now either pick the function or the stub global for cuda, hip
  // respectively.
  mlir::Value kernel = [&]() -> mlir::Value {
    if (cir::GlobalOp globalOp = llvm::dyn_cast_or_null<cir::GlobalOp>(
            kernelHandles[fn.getSymName()])) {
      cir::PointerType kernelTy = cir::PointerType::get(globalOp.getSymType());
      mlir::Value kernelVal = cir::GetGlobalOp::create(builder, loc, kernelTy,
                                                       globalOp.getSymName());
      mlir::Value func = builder.createBitcast(kernelVal, cgm.voidPtrTy);
      return func;
    }
    if (cir::FuncOp funcOp = llvm::dyn_cast_or_null<cir::FuncOp>(
            kernelHandles[fn.getSymName()])) {
      cir::PointerType kernelTy =
          cir::PointerType::get(funcOp.getFunctionType());
      mlir::Value kernelVal =
          cir::GetGlobalOp::create(builder, loc, kernelTy, funcOp.getSymName());
      mlir::Value func = builder.createBitcast(kernelVal, cgm.voidPtrTy);
      return func;
    }
    llvm_unreachable("Expected stub handle to be cir::GlobalOp or FuncOp");
  }();
 
  CallArgList launchArgs;
  launchArgs.add(RValue::get(kernel),
                 cudaLaunchKernelFD->getParamDecl(0)->getType());
  launchArgs.add(
      RValue::getAggregate(Address(gridDim, CharUnits::fromQuantity(8))),
      cudaLaunchKernelFD->getParamDecl(1)->getType());
  launchArgs.add(
      RValue::getAggregate(Address(blockDim, CharUnits::fromQuantity(8))),
      cudaLaunchKernelFD->getParamDecl(2)->getType());
  launchArgs.add(RValue::get(kernelArgs),
                 cudaLaunchKernelFD->getParamDecl(3)->getType());
  launchArgs.add(
      RValue::get(builder.CIRBaseBuilderTy::createLoad(loc, sharedMem)),
      cudaLaunchKernelFD->getParamDecl(4)->getType());
  launchArgs.add(RValue::get(builder.CIRBaseBuilderTy::createLoad(loc, stream)),
                 cudaLaunchKernelFD->getParamDecl(5)->getType());
 
  mlir::Type launchTy =
      cgm.getTypes().convertType(cudaLaunchKernelFD->getType());
  mlir::Operation *cudaKernelLauncherFn = cgm.createRuntimeFunction(
      cast<cir::FuncType>(launchTy), launchKernelName);
  const CIRGenFunctionInfo &callInfo =
      cgm.getTypes().arrangeFunctionDeclaration(cudaLaunchKernelFD);
  cgf.emitCall(callInfo, CIRGenCallee::forDirect(cudaKernelLauncherFn),
               ReturnValueSlot(), launchArgs, /*isMustTail=*/false);
 
  if (cgm.getASTContext().getTargetInfo().getCXXABI().isMicrosoft() &&
      !cgf.getLangOpts().HIP)
    cgm.errorNYI("MSVC CUDA stub handling");
}
 
void CIRGenNVCUDARuntime::emitDeviceStub(CIRGenFunction &cgf, cir::FuncOp fn,
                                         FunctionArgList &args) {
  CIRGenBuilderTy &builder = cgm.getBuilder();
  StringRef kernelName = cgm.getMangledName(
      cgf.curGD.getWithKernelReferenceKind(KernelReferenceKind::Kernel));
  auto kernelNameAttr = cir::CUDAKernelNameAttr::get(
      builder.getContext(), builder.getStringAttr(kernelName));
  fn->setDiscardableAttr(cir::CUDAKernelNameAttr::getMnemonic(),
                         kernelNameAttr);
 
  if (auto globalOp =
          llvm::dyn_cast<cir::GlobalOp>(kernelHandles[fn.getSymName()])) {
    mlir::Type fnPtrTy = globalOp.getSymType();
    auto sym = mlir::FlatSymbolRefAttr::get(fn.getSymNameAttr());
    auto gv = cir::GlobalViewAttr::get(fnPtrTy, sym);
 
    globalOp->setAttr("initial_value", gv);
    globalOp->setAttr("alignment", builder.getI64IntegerAttr(
                                       cgm.getPointerAlign().getQuantity()));
 
    // The stub's linkage is only final once it is defined, so the handle
    // takes it over here.
    globalOp.setLinkage(fn.getLinkage());
    mlir::SymbolTable::setSymbolVisibility(
        globalOp, cgm.getMLIRVisibilityFromCIRLinkage(fn.getLinkage()));
    cgm.setDSOLocal(static_cast<mlir::Operation *>(globalOp));
  }
 
  // CUDA 9.0 changed the way to launch kernels.
  if (CudaFeatureEnabled(cgm.getTarget().getSDKVersion(),
                         CudaFeature::CUDA_USES_NEW_LAUNCH) ||
      (cgm.getLangOpts().HIP && cgm.getLangOpts().HIPUseNewLaunchAPI) ||
      cgm.getLangOpts().OffloadViaLLVM)
    emitDeviceStubBodyNew(cgf, fn, args);
  else
    cgm.errorNYI("Emit Stub Body Legacy");
}
 
CIRGenCUDARuntime *clang::CIRGen::createNVCUDARuntime(CIRGenModule &cgm) {
  return new CIRGenNVCUDARuntime(cgm);
}
 
CIRGenNVCUDARuntime::~CIRGenNVCUDARuntime() {}
 
mlir::Operation *CIRGenNVCUDARuntime::getKernelHandle(cir::FuncOp fn,
                                                      GlobalDecl gd) {
 
  // Check if we already have a kernel handle for this function
  auto it = kernelHandles.find(fn.getSymName());
  if (it != kernelHandles.end()) {
    mlir::Operation *oldHandle = it->second;
    // Here we know that the fn did not change. Return it
    if (kernelStubs[oldHandle] == fn)
      return oldHandle;
 
    // We've found the function name, but F itself has changed, so we need to
    // update the references.
    if (cgm.getLangOpts().HIP) {
      // For HIP compilation the handle itself does not change, so we only need
      // to update the Stub value.
      kernelStubs[oldHandle] = fn;
      return oldHandle;
    }
    // For non-HIP compilation, erase the old Stub and fall-through to creating
    // new entries.
    kernelStubs.erase(oldHandle);
  }
 
  // If not targeting HIP, store the function itself
  if (!cgm.getLangOpts().HIP) {
    kernelHandles[fn.getSymName()] = fn;
    kernelStubs[fn] = fn;
    return fn;
  }
 
  // Create a new CIR global variable to represent the kernel handle
  CIRGenBuilderTy &builder = cgm.getBuilder();
  StringRef globalName = cgm.getMangledName(
      gd.getWithKernelReferenceKind(KernelReferenceKind::Kernel));
  cir::PointerType fnPtrTy = builder.getPointerTo(fn.getFunctionType());
  cir::GlobalOp globalOp =
      cgm.createGlobalOp(fn.getLoc(), globalName, fnPtrTy, /*isConstant=*/true);
 
  globalOp->setAttr("alignment", builder.getI64IntegerAttr(
                                     cgm.getPointerAlign().getQuantity()));
 
  // Inherit visibility and comdat from the kernel's declaration.
  auto *fd = cast<FunctionDecl>(gd.getDecl());
  cgm.setGVProperties(globalOp, fd);
  FunctionTemplateDecl *ft = fd->getPrimaryTemplate();
  if (!ft || ft->isThisDeclarationADefinition())
    cgm.maybeSetTrivialComdat(*fd, globalOp);
 
  // Store references
  kernelHandles[fn.getSymName()] = globalOp;
  kernelStubs[globalOp] = fn;
 
  return globalOp;
}
 
void CIRGenNVCUDARuntime::internalizeDeviceSideVar(
    const VarDecl *d, cir::GlobalLinkageKind &linkage) {
  if (cgm.getLangOpts().GPURelocatableDeviceCode)
    cgm.errorNYI(d->getSourceRange(),
                 "internalizeDeviceSideVar: GPU Relocatable Device Code (RDC)");
 
  // __shared__ variables are odd. Shadows do get created, but
  // they are not registered with the CUDA runtime, so they
  // can't really be used to access their device-side
  // counterparts. It's not clear yet whether it's nvcc's bug or
  // a feature, but we've got to do the same for compatibility.
  if (d->hasAttr<CUDADeviceAttr>() || d->hasAttr<CUDAConstantAttr>() ||
      d->hasAttr<CUDASharedAttr>() ||
      d->getType()->isCUDADeviceBuiltinSurfaceType() ||
      d->getType()->isCUDADeviceBuiltinTextureType()) {
    linkage = cir::GlobalLinkageKind::InternalLinkage;
  }
 
  if (d->getType()->isCUDADeviceBuiltinTextureType())
    cgm.errorNYI(d->getSourceRange(),
                 "internalizeDeviceSideVar: CUDA Texture support");
}
 
std::string CIRGenNVCUDARuntime::getDeviceSideName(const NamedDecl *nd) {
  GlobalDecl gd;
  // nd could be either a kernel or a variable.
  if (auto *fd = dyn_cast<FunctionDecl>(nd))
    gd = GlobalDecl(fd, KernelReferenceKind::Kernel);
  else
    gd = GlobalDecl(nd);
  std::string deviceSideName;
  MangleContext *mc;
  if (cgm.getLangOpts().CUDAIsDevice)
    mc = &cgm.getCXXABI().getMangleContext();
  else
    mc = deviceMC.get();
  if (mc->shouldMangleDeclName(nd)) {
    SmallString<256> buffer;
    llvm::raw_svector_ostream out(buffer);
    mc->mangleName(gd, out);
    deviceSideName = std::string(out.str());
  } else
    deviceSideName = std::string(nd->getIdentifier()->getName());
 
  // Make unique name for device side static file-scope variable for HIP.
  if (cgm.getASTContext().shouldExternalize(nd) &&
      cgm.getLangOpts().GPURelocatableDeviceCode) {
    SmallString<256> buffer;
    llvm::raw_svector_ostream out(buffer);
    out << deviceSideName;
    cgm.printPostfixForExternalizedDecl(out, nd);
    deviceSideName = std::string(out.str());
  }
  return deviceSideName;
}
 
void CIRGenNVCUDARuntime::handleVarRegistration(const VarDecl *vd,
                                                cir::GlobalOp var) {
  if (vd->hasAttr<CUDADeviceAttr>() || vd->hasAttr<CUDAConstantAttr>()) {
    // Shadow variables and their properties must be registered with CUDA
    // runtime. Skip Extern global variables, which will be registered in
    // the TU where they are defined.
    //
    // Don't register a C++17 inline variable. The local symbol can be
    // discarded and referencing a discarded local symbol from outside the
    // comdat (__cuda_register_globals) is disallowed by the ELF spec.
    //
    // HIP managed variables need to be always recorded in device and host
    // compilations for transformation.
    //
    // HIP managed variables and variables in CUDADeviceVarODRUsedByHost are
    // added to llvm.compiler-used, therefore they are safe to be registered.
    if ((!vd->hasExternalStorage() && !vd->isInline()) ||
        cgm.getASTContext().CUDADeviceVarODRUsedByHost.contains(vd) ||
        vd->hasAttr<HIPManagedAttr>()) {
      registerDeviceVar(vd, var, !vd->hasDefinition(),
                        vd->hasAttr<CUDAConstantAttr>());
    }
  } else if (vd->getType()->isCUDADeviceBuiltinSurfaceType()) {
    // Builtin surfaces and their template arguments are also registered
    // with CUDA runtime.
    const auto *td = cast<ClassTemplateSpecializationDecl>(
        vd->getType()->castAsCXXRecordDecl());
    const TemplateArgumentList &args = td->getTemplateArgs();
 
    assert(args.size() == 2 &&
           "Unexpected number of template arguments of CUDA device "
           "builtin surface type.");
 
    llvm::APSInt surfaceType = args[1].getAsIntegral();
 
    if (!vd->hasExternalStorage())
      registerDeviceSurf(vd, var, !vd->hasDefinition(),
                         surfaceType.getSExtValue());
 
  } else if (vd->getType()->isCUDADeviceBuiltinTextureType()) {
    // Builtin textures and their template arguments are also registered
    // with CUDA runtime.
    if (!vd->hasExternalStorage())
      registerDeviceTex(vd, var, !vd->hasDefinition());
  }
}
 
void CIRGenNVCUDARuntime::handleGlobalReplace(cir::GlobalOp oldGV,
                                              cir::GlobalOp newGV) {
  for (auto &info : deviceVars) {
    if (info.var == oldGV)
      info.var = newGV;
  }
}
 
/// Whether this translation unit has anything for the CUDA runtime to register.
/// Like LoweringPrepare, count only kernel stub definitions and device
/// variables carrying registration metadata, not kernel declarations or
/// launch-site calls.
static bool hasEntitiesToRegister(mlir::ModuleOp module) {
  // A walk, not a scan of the module body, so this stays in agreement with the
  // recursive walk LoweringPrepare collects them with.
  return module
      ->walk([](mlir::Operation *op) {
        if (auto fn = mlir::dyn_cast<cir::FuncOp>(op);
            fn && !fn.isDeclaration() &&
            fn->hasAttr(cir::CUDAKernelNameAttr::getMnemonic()))
          return mlir::WalkResult::interrupt();
        if (op->hasAttr(cir::CUDAVarRegistrationInfoAttr::getMnemonic()))
          return mlir::WalkResult::interrupt();
        return mlir::WalkResult::advance();
      })
      .wasInterrupted();
}
 
/// Read the device-side fat binary and record its contents on the module as
/// `cir.cu.device_binary`, for LoweringPrepare to build the fatbin global from.
///
/// This mirrors the read in CGNVCUDARuntime::makeModuleCtorFunction, guards
/// included: nothing is read in a compilation that would build no module
/// constructor.
void CIRGenNVCUDARuntime::recordDeviceBinary() {
  StringRef binaryName = cgm.getCodeGenOpts().OffloadBinaryToEmbedFile;
  if (binaryName.empty())
    return;
 
  const LangOptions &langOpts = cgm.getLangOpts();
  if ((langOpts.HIP || !langOpts.GPURelocatableDeviceCode) &&
      !hasEntitiesToRegister(cgm.getModule()))
    return;
 
  llvm::vfs::FileSystem &fs = cgm.getASTContext()
                                  .getSourceManager()
                                  .getFileManager()
                                  .getVirtualFileSystem();
  llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> binaryOrErr =
      fs.getBufferForFile(binaryName, /*FileSize=*/-1,
                          /*RequiresNullTerminator=*/false);
  if (std::error_code ec = binaryOrErr.getError()) {
    cgm.getDiags().Report(diag::err_cannot_open_file)
        << binaryName << ec.message();
    return;
  }
 
  // Typed as the fatbin global's array type so LoweringPrepare can use this
  // attribute as the initializer as-is: attributes are uniqued on {value, type}
  // and never freed, so building a second, typed copy there would keep the fat
  // binary in memory twice.
  StringRef bytes = binaryOrErr.get()->getBuffer();
  mlir::MLIRContext &ctx = cgm.getMLIRContext();
  auto charTy = cir::IntType::get(&ctx, cgm.getTarget().getCharWidth(),
                                  /*isSigned=*/false);
  auto fatbinTy = cir::ArrayType::get(charTy, bytes.size());
  cgm.getModule()->setAttr(cir::CIRDialect::getCUDADeviceBinaryAttrName(),
                           mlir::StringAttr::get(bytes, fatbinTy));
}
 
// Uses of a managed variable load its address from a pointer the runtime fills.
void CIRGenNVCUDARuntime::transformManagedVars() {
  CIRGenBuilderTy &builder = cgm.getBuilder();
  mlir::ModuleOp module = cgm.getModule();
  for (auto &&info : deviceVars) {
    cir::GlobalOp var = info.var;
    if (info.flags != cir::CUDADeviceVarKind::Variable ||
        !info.d->hasAttr<HIPManagedAttr>())
      continue;
 
    cir::PointerType varPtrTy =
        builder.getPointerTo(var.getSymType(), var.getAddrSpaceAttr());
    mlir::ptr::MemorySpaceAttrInterface managedAS = cir::toCIRAddressSpaceAttr(
        cgm.getMLIRContext(),
        cgm.getLangOpts().CUDAIsDevice ? LangAS::cuda_device : LangAS::Default);
    std::optional<mlir::SymbolTable::UseRange> uses = var.getSymbolUses(module);
 
    std::string name = var.getSymName().str();
    // An empty comdat would follow the rename.
    std::optional<llvm::StringRef> comdat = var.getComdat();
    if (comdat && comdat->empty())
      var.setComdat(llvm::StringRef(name));
    cgm.eraseGlobalSymbol(var);
    var.setSymName(name + ".managed");
    cgm.insertGlobalSymbol(var);
 
    cir::GlobalOp managedVar =
        cgm.createGlobalOp(var.getLoc(), name, varPtrTy,
                           /*isConstant=*/false, managedAS);
    managedVar.setLinkage(var.getLinkage());
    if (!var.isDeclaration())
      managedVar.setInitialValueAttr(builder.getConstNullPtrAttr(varPtrTy));
    managedVar.setDsoLocal(var.getDsoLocal());
    managedVar.setGlobalVisibility(var.getGlobalVisibility());
    mlir::SymbolTable::setSymbolVisibility(
        managedVar, mlir::SymbolTable::getSymbolVisibility(var));
    managedVar->setAttr(
        cir::CUDAExternallyInitializedAttr::getMnemonic(),
        cir::CUDAExternallyInitializedAttr::get(&cgm.getMLIRContext()));
 
    // The uses now name the pointer global.
    cir::PointerType managedPtrTy =
        builder.getPointerTo(varPtrTy, managedVar.getAddrSpaceAttr());
    for (mlir::SymbolTable::SymbolUse use : *uses) {
      auto getGlobalOp = dyn_cast<cir::GetGlobalOp>(use.getUser());
      if (!getGlobalOp) {
        cgm.errorNYI(info.d->getSourceRange(),
                     "transformManagedVars: non-instruction use");
        continue;
      }
      mlir::Value addr = getGlobalOp.getAddr();
      addr.setType(managedPtrTy);
      mlir::OpBuilder::InsertionGuard guard(builder);
      builder.setInsertionPointAfter(getGlobalOp);
      mlir::Value ld = builder.createAlignedLoad(
          getGlobalOp.getLoc(), varPtrTy, addr,
          CharUnits::fromQuantity(var.getAlignment().value_or(1)));
      addr.replaceAllUsesExcept(ld, ld.getDefiningOp());
    }
 
    // The runtime allocates them even if device code never uses them.
    if (cgm.getLangOpts().CUDAIsDevice && !var.isDeclaration()) {
      cgm.addCompilerUsedGlobal(var);
      cgm.addCompilerUsedGlobal(managedVar);
    }
  }
}
 
void CIRGenNVCUDARuntime::finalizeModule() {
  transformManagedVars();
  if (!cgm.getLangOpts().CUDAIsDevice) {
    recordDeviceBinary();
    return;
  }
 
  // Mark ODR-used device variables as compiler used to prevent them from being
  // eliminated by optimization. This is necessary for device variables
  // ODR-used by host functions. Sema correctly marks them as ODR-used no
  // matter whether they are ODR-used by device or host functions.
  //
  // We do not need to do this if the variable has used attribute since it
  // has already been added.
  //
  // Static device variables have been externalized at this point, therefore
  // variables with private or internal linkage need not be added.
  for (auto &&info : deviceVars) {
    auto kind = info.flags;
    bool isDecl = info.var.isDeclaration();
    bool isLocalLinkage = cir::isLocalLinkage(info.var.getLinkage());
    bool isVarOrSurfaceOrTexture = (kind == cir::CUDADeviceVarKind::Variable ||
                                    kind == cir::CUDADeviceVarKind::Surface ||
                                    kind == cir::CUDADeviceVarKind::Texture);
    bool isUsed = info.d->isUsed();
    bool hasUsedAttr = info.d->hasAttr<UsedAttr>();
    if (!isDecl && !isLocalLinkage && isVarOrSurfaceOrTexture && isUsed &&
        !hasUsedAttr) {
      if (auto globalValue = mlir::dyn_cast<cir::CIRGlobalValueInterface>(
              info.var.getOperation())) {
        cgm.addCompilerUsedGlobal(globalValue);
      }
    }
  }
}