如何使用选项CU_JIT_LTO与CUDA JIT链接？

Question

我想知道我是否能够在与选项CU_JIT_LTO链接期间改进链接时间优化(LTO)。如果是，我如何指定此选项？

我在NVIDIA开发者博客中找到了下面的代码，但我不明白为什么壁时间会给CU_JIT_LTO。在博客中没有定义壁时变量。当我尝试类似的东西时，它对我的内核性能没有任何影响。

options[0] = CU_JIT_LTO;
values[0] = (void*)&walltime;
...
cuLinkCreate(..., options, values, &linkState);

来源：https://developer.nvidia.com/blog/discovering-new-features-in-cuda-11-4/

我的示例使用输入选项CU_JIT_INPUT_NVVM链接使用LTO标志(-dlto或-code=lto_80)创建的对象。看来链接器已经做了一些LTO，因为内核比没有LTO的对象文件“性能更好”，但不如使用NVCC链接LTO那么好。(详细结果和讨论见示例)

例例

为了检查链接时间优化(LTO)的有效性，我使用4种不同的方法创建了一个简单的程序，并使用每个线程的寄存器数作为指示。这在我的系统上给出了以下结果(OS: Ubuntu20.04，CUDA工具包: 11.5.1，NVIDIA驱动程序: 495.44，GPU: NVIDIA RTX 3080)。

                       method                                registers/thread
Create program using a single translation file           :         30
Link files using NVCC without link time optimization     :         44
Link files using NVCC with link time optimization        :         30
Link files using NVRTC/JIT with link time optimization   :         38

结果解释：

从一个翻译文件创建程序应该会给出最好的结果。编译器可以看到所有的函数实现，并使用它来优化内核。这将导致30个寄存器/线程。

使用NVCC与LTO连接绝对有效。它使用与从单个.cu文件编译的程序相同数量的寄存器(30)，如果没有LTO (使用44个寄存器)，情况就不是这样。

使用NVRTC/JIT链接文件比使用没有LTO的NVCC链接“更好”(当我们只关注注册使用)，但不如与NVCC和LTO链接。内核使用38个寄存器/线程。

注意:我的目标不是减少寄存器的使用，我只使用它作为一个指标。因为来自单个翻译文件的程序使用30个寄存器/线程，我假设一个完全优化的链接程序将具有相同的“最终可执行代码”，从而使用相同数量的寄存器。由于情况并非如此，我开始研究JIT选项。

CU_JIT_LTO选项：

我试图进一步优化NVRTC/JIT案例中与JIT_option CU_JIT_LTO的链接。但是，我不知道如何使用这个选项。我尝试了以下两种方法(有关更多上下文，请参见下面的文件cuda代码)。连接代码从第41行开始：

方法1:将选项CU_JIT_LTO添加到cuLinkCreate(...)。这似乎没有任何效果。当int lto = 0和int lto = 1时，代码使用相同数量的寄存器。

方法2:将选项CU_JIT_LTO添加到cuLinkAddFile(...)和cuLinkAddData(...)中。这立即给出了错误CUDA_ERROR_INVALID_VALUE。

那么现在我的问题是:如何使用CU_JIT_LTO选项？

文件：

下面是两个文件。按照以下步骤运行示例(在linux操作系统上)：

将code代码保存在具有code.cu)

• Save扩展名的文件中(例如：bash run.sh code.cu

，文件中的bash脚本(例如:run.sh)，

• 从终端：bash run.sh code.cu

运行此命令

code代码：

#include <iostream>
#include <stdio.h>

#ifdef RTC
#include <cuda.h>
#include <nvrtc.h>
#define NVRTC_CHECK(x)                                                                            \
  do {                                                                                            \
    nvrtcResult result = x;                                                                       \
    if (result != NVRTC_SUCCESS) {                                                                \
      std::cerr << "\nerror: " #x " failed with error " << nvrtcGetErrorString(result) << '\n';   \
      exit(1);                                                                                    \
    }                                                                                             \
  } while (0)
#define CUDA_CHECK(x)                                                                             \
  do {                                                                                            \
    CUresult result = x;                                                                          \
    if (result != CUDA_SUCCESS) {                                                                 \
      const char* msg;                                                                            \
      cuGetErrorName(result, &msg);                                                               \
      std::cerr << "\nerror: " #x " failed with error " << msg << '\n';                           \
      exit(1);                                                                                    \
    }                                                                                             \
  } while (0)

CUmodule compileModule(std::string program)
{
  // Compile nvvm from program string ===============
  nvrtcProgram prog;
  NVRTC_CHECK(nvrtcCreateProgram(&prog, program.c_str(), "programRTC.cu", 0, NULL, NULL));

  const char* opts[] = {"-arch=compute_80", "-dlto", "-dc"};
  nvrtcResult compileResult = nvrtcCompileProgram(prog, 3, opts);

  // Obtain NVVM from the program.
  size_t nvvmSize;
  NVRTC_CHECK(nvrtcGetNVVMSize(prog, &nvvmSize));
  char* nvvm = new char[nvvmSize];
  NVRTC_CHECK(nvrtcGetNVVM(prog, nvvm));

  // Link files ===============
  CUlinkState linker;

  // ARE THE OPTIONS SPECIFIED CORRECTLY?
  int lto = 1;
  CUjit_option options[] = {CU_JIT_LTO};
  void* values[] = {(void*)&lto};

  // METHOD 1: GIVE THE OPTIONS TO 'cuLinkCreate(...)'
  //           -> HAS NO EFFECT ON THE AMOUNT OF REGISTERS USED
  // -------------------------------------------------------------------------------------------
  // CUDA_CHECK(cuLinkCreate(0, NULL, NULL, &linker));
  CUDA_CHECK(cuLinkCreate(1, options, values, &linker));
  // -------------------------------------------------------------------------------------------


  // METHOD 2: GIVE THE OPTIONS TO 'cuLinkAddFile(...)' and 'cuLinkAddData(...)'
  //           -> FUNCTION FAILS WITH ERROR 'CUDA_ERROR_INVALID_VALUE'
  // -------------------------------------------------------------------------------------------
  CUDA_CHECK(cuLinkAddFile(linker, CU_JIT_INPUT_NVVM, "func_link_nvrtc_lto.o", 0, NULL, NULL));
  CUDA_CHECK(cuLinkAddData(linker, CU_JIT_INPUT_NVVM, (void*)nvvm, nvvmSize, "programRTC.o", 0,
                           NULL, NULL));

  // CUDA_CHECK(cuLinkAddFile(linker, CU_JIT_INPUT_NVVM, "func_link_nvrtc_lto.o", 1, options, values));
  // CUDA_CHECK(cuLinkAddData(linker, CU_JIT_INPUT_NVVM, (void*)nvvm, nvvmSize, "programRTC.o", 1,
  //                          options, values));
  // -------------------------------------------------------------------------------------------

  // Create module ===============
  void* cubin;
  CUmodule module;
  CUDA_CHECK(cuLinkComplete(linker, &cubin, NULL));
  CUDA_CHECK(cuModuleLoadDataEx(&module, cubin, 0, NULL, NULL));

  // Cleanup
  NVRTC_CHECK(nvrtcDestroyProgram(&prog));
  CUDA_CHECK(cuLinkDestroy(linker));

  return module;
}
#endif // RTC

__device__ double func(double a, double b);
#ifdef FUNC
__device__ double func(double a, double b)
{
  return pow(a, b);
}
#endif

#ifdef MAIN
#ifdef RTC
std::string the_program = R"===(
__device__ double func(double a, double b);

extern "C" __global__ void kernel(double* out, double* a, double* b)
{
  size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
  if(tid >= 1){
    return;
  }
  a[tid] = 2;
  b[tid] = 3;
  out[tid] = func(a[tid], b[tid]);
  printf("out[%lu] = %f\n", tid, out[tid]);
})===";
#else  // RTC
__global__ void kernel(double* out, double* a, double* b)
{
  size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
  if (tid >= 1) {
    return;
  }
  a[tid] = 2;
  b[tid] = 3;
  out[tid] = func(a[tid], b[tid]);
  printf("out[%lu] = %f\n", tid, out[tid]);
}
#endif // RTC

int main()
{
  double* a;
  double* b;
  double* out;
  cudaMalloc((void**)&a, sizeof(double));
  cudaMalloc((void**)&b, sizeof(double));
  cudaMalloc((void**)&out, sizeof(double));

#ifdef RTC
  // Create context
  CUdevice cuDevice;
  CUcontext context;
  CUDA_CHECK(cuInit(0));
  CUDA_CHECK(cuDeviceGet(&cuDevice, 0));
  CUDA_CHECK(cuCtxCreate(&context, 0, cuDevice));

  CUmodule module = compileModule(the_program);

  CUfunction kernel;
  CUDA_CHECK(cuModuleGetFunction(&kernel, module, "kernel"));

  size_t n_blocks = 1;
  size_t n_threads = 1;
  void* args[] = {&out, &a, &b};
  CUDA_CHECK(cuLaunchKernel(kernel, n_blocks, 1, 1, // grid dim
                            n_threads, 1, 1,        // block dim
                            0, NULL,                // shared mem and stream
                            args, 0));              // arguments
  CUDA_CHECK(cuCtxSynchronize());

  // Cleanup
  CUDA_CHECK(cuModuleUnload(module));
  CUDA_CHECK(cuCtxDestroy(context));
#else  // RTC
  kernel<<<1, 1>>>(out, a, b);
  cudaDeviceSynchronize();
#endif // RTC

  return 0;
}
#endif // MAIN

bash脚本：

#!/bin/bash

set -e # stop script when an error occurs

SCRIPT=$1
xCCx=80 # CUDA compute compatibility

# Create program using a single translation file
echo -e "\n---------- main_single ----------\n"
nvcc -DFUNC -DMAIN $SCRIPT -o main_single \
    -gencode arch=compute_$xCCx,code=sm_$xCCx
./main_single # should print 'out[0] = 8.0'
cuobjdump main_single -res-usage | grep kernel -A1

# Link files using NVCC without link time optimization (code=compute_...)
echo -e "\n---------- main_link_nvcc ----------\n"
nvcc -DFUNC $SCRIPT -o func_link_nvcc.o -dc \
    -gencode arch=compute_$xCCx,code=compute_$xCCx
nvcc -DMAIN $SCRIPT -o main_link_nvcc.o -dc \
    -gencode arch=compute_$xCCx,code=compute_$xCCx
nvcc func_link_nvcc.o main_link_nvcc.o -o main_link_nvcc \
    -gencode arch=compute_$xCCx,code=sm_$xCCx
./main_link_nvcc # should print 'out[0] = 8.0'
cuobjdump main_link_nvcc -res-usage | grep kernel -A1

# Link files using NVCC with link time optimization (code=lto_...)
echo -e "\n---------- main_link_nvcc_lto ----------\n"
nvcc -DFUNC $SCRIPT -o func_link_nvcc_lto.o -dc \
    -gencode arch=compute_$xCCx,code=lto_$xCCx
nvcc -DMAIN $SCRIPT -o main_link_nvcc_lto.o -dc \
    -gencode arch=compute_$xCCx,code=lto_$xCCx
nvcc func_link_nvcc_lto.o main_link_nvcc_lto.o -o main_link_nvcc_lto -dlto \
    -gencode arch=compute_$xCCx,code=sm_$xCCx
./main_link_nvcc_lto # should print 'out[0] = 8.0'
cuobjdump main_link_nvcc_lto -res-usage | grep kernel -A1

# Link files using NVRTC with link time optimization
echo -e "\n---------- main_link_nvrtc_lto ----------\n"
nvcc -DFUNC $SCRIPT -o func_link_nvrtc_lto.o -dc -ptx \
    -gencode arch=compute_$xCCx,code=lto_$xCCx
nvcc -DMAIN -DRTC $SCRIPT -o main_link_nvrtc_lto \
    -lnvrtc_static -lnvrtc-builtins_static -lnvptxcompiler_static -lcuda -lpthread \
    -gencode arch=compute_$xCCx,code=sm_$xCCx
./main_link_nvrtc_lto # should print 'out[0] = 8.0'
ncu main_link_nvrtc_lto | grep register/thread


# Registers/thread used on my system with an NVIDIA RTX 3080:
# main_single          : 30 registers/thread
# main_link_nvcc       : 44 registers/thread
# main_link_nvcc_lto   : 30 registers/thread
# main_link_nvrtc_lto  : 38 registers/thread

子问题:生成一个NVVM IR文件

要处理与命令func_link_nvrtc_lto.o一起工作的文件cuLinkAddFile(linker, CU_JIT_INPUT_NVVM, "func_link_nvrtc_lto.o", ...)，我必须添加-ptx标志，如下所示。我没有在任何文档中找到这一点，而是通过尝试和错误。我想知道是否有更好的方法来制作这样的文件。

nvcc -DFUNC $SCRIPT -o func_link_nvrtc_lto.o -dc -ptx \
-gencode arch=compute_$xCCx,code=lto_$xCCx

Answer 1

首先，不幸的是，在带有CU_JIT_LTO值的博客文章中出现了错误。相反，应该是：

values[0] = (void*)1;

但是，这并不重要，因为值被忽略了--使用的仅仅是CU_JIT_LTO的存在。正如您所发现的那样，CU_JIT_LTO确实应该传递给cuLinkCreate。

对于您的子问题，您的-ptx所做的是在生成nvvm之后停止编译，但这是一个没有文档记录的副作用。更简单、更安全的方法是只使用：

nvcc -dc -arch=compute_XX,code=lto_XX

它创建一个包含nvvm的主机对象。然后将其传递为：

CU_JIT_INPUT_OBJECT to cuLinkAddFile().