Use nvptxcompile library.

[jdoerfert]

is this still supposed to land?

[tra]

In D145527#4303252, @jdoerfert wrote:

is this still supposed to land?

It's up to MLIR folks.

I did this patch as a prototype when I was trying to decouple MLIR compilation passes from hard dependency on NVIDIA GPU driver. The hard dependency has been sufficiently mitigated w/o this patch.

[mehdi_amini]

@tra : do you have thoughts on nvptxcompiler library vs "writing out a temp file and shelling out invoke ptxas" ?

[tra]

In D145527#4547804, @mehdi_amini wrote:

@tra : do you have thoughts on nvptxcompiler library vs "writing out a temp file and shelling out invoke ptxas" ?

The main benefit is that it makes things more hermetic. Less of a chance that things would go wrong due to the end-user setup (e.g. no space in /tmp, picking up wrong ptxas). Easier to reproduce issues when we need to -- official builds would be guaranteed to use the same PTX compiler as the end user.
If there's a bug in the ptxas compiler, we only need to rebuild with the fixed library. It gives clear idea which version is broken and which is fixed. With ptxas the answer would depend on the specific set up for the individual user, which makes triaging/support much more of a hassle.

mlir/lib/Dialect/GPU/CMakeLists.txt

 
 add_subdirectory(TransformOps)
 
 if(MLIR_ENABLE_CUDA_RUNNER)
   if(NOT MLIR_ENABLE_CUDA_CONVERSIONS)
     message(SEND_ERROR
       "Building mlir with cuda support requires the NVPTX backend")
   endif()
   
+  find_package(CUDAToolkit)
   # Configure CUDA language support. Using check_language first allows us to
   # give a custom error message.
   include(CheckLanguage)
   check_language(CUDA)
   if (CMAKE_CUDA_COMPILER)
     enable_language(CUDA)
   else()
     message(SEND_ERROR
       "Building mlir with cuda support requires a working CUDA install")
   endif()
 
   # Enable gpu-to-cubin pass.
   target_compile_definitions(obj.MLIRGPUTransforms
     PRIVATE
     MLIR_GPU_TO_CUBIN_PASS_ENABLE=1
   )
 
   # Add CUDA headers includes and the libcuda.so library.
   target_include_directories(obj.MLIRGPUTransforms
     PRIVATE
     ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}
   )
 
-  find_library(CUDA_DRIVER_LIBRARY cuda)
+  #find_library(CUDA_NVPTX_LIBRARY nvptxcompiler_static)
 
   target_link_libraries(MLIRGPUTransforms
     PRIVATE
     MLIRNVVMToLLVMIRTranslation
-    ${CUDA_DRIVER_LIBRARY}
+    CUDA::nvptxcompiler_static
   )
 
 endif()
 
 if(MLIR_ENABLE_ROCM_CONVERSIONS)
   if (NOT ("AMDGPU" IN_LIST LLVM_TARGETS_TO_BUILD))
     message(SEND_ERROR
       "Building mlir with ROCm support requires the AMDGPU backend")

mlir/lib/Dialect/GPU/Transforms/SerializeToCubin.cpp

 
 #if MLIR_GPU_TO_CUBIN_PASS_ENABLE
 #include "mlir/Pass/Pass.h"
 #include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Export.h"
 #include "llvm/Support/TargetSelect.h"
 
 #include <cuda.h>
+#include <nvPTXCompiler.h>
 
 using namespace mlir;
 
-static void emitCudaError(const llvm::Twine &expr, const char *buffer,
-                          CUresult result, Location loc) {
+static void emitNvptxError(const llvm::Twine &expr,
+                           nvPTXCompilerHandle compiler,
+                           nvPTXCompileResult result, Location loc) {
   const char *error;
-  cuGetErrorString(result, &error);
+  auto GetErrMsg = [](nvPTXCompileResult result) -> const char * {
+    switch (result) {
+    case NVPTXCOMPILE_SUCCESS:
+      return "Success";
+    case NVPTXCOMPILE_ERROR_INVALID_COMPILER_HANDLE:
+      return "Invalid compiler handle";
+    case NVPTXCOMPILE_ERROR_INVALID_INPUT:
+      return "Invalid input";
+    case NVPTXCOMPILE_ERROR_COMPILATION_FAILURE:
+      return "Compilation failure";
+    case NVPTXCOMPILE_ERROR_INTERNAL:
+      return "Internal error";
+    case NVPTXCOMPILE_ERROR_OUT_OF_MEMORY:
+      return "Out of memory";
+    case NVPTXCOMPILE_ERROR_COMPILER_INVOCATION_INCOMPLETE:
+      return "Invocation incomplete";
+    case NVPTXCOMPILE_ERROR_UNSUPPORTED_PTX_VERSION:
+      return "Unsupported PTX version";
+    }
+  };
+  size_t errorSize;
+  auto status = nvPTXCompilerGetErrorLogSize(compiler, &errorSize);
+  std::string error_log;
+  if (status == NVPTXCOMPILE_SUCCESS) {
+    error_log.resize(errorSize);
+    status = nvPTXCompilerGetErrorLog(compiler, error_log.data());
+    if (status != NVPTXCOMPILE_SUCCESS)
+      error_log = "<failed to retrieve compiler error log>";
+  }
   emitError(loc, expr.concat(" failed with error code ")
-                     .concat(llvm::Twine{error})
+                     .concat(llvm::Twine{GetErrMsg(result)})
                      .concat("[")
-                     .concat(buffer)
+                     .concat(error_log)
                      .concat("]"));
 }
 
 #define RETURN_ON_CUDA_ERROR(expr)                                             \
   do {                                                                         \
     if (auto status = (expr)) {                                                \
-      emitCudaError(#expr, jitErrorBuffer, status, loc);                       \
+      emitNvptxError(#expr, compiler, status, loc);                            \
+      return {};                                                               \
+    }                                                                          \
+  } while (false)
+
+#define RETURN_ON_NVPTX_ERROR(expr)                                            \
+  do {                                                                         \
+    nvPTXCompileResult result = (expr);                                        \
+    if (result != NVPTXCOMPILE_SUCCESS) {                                      \
+      emitNvptxError(#expr, compiler, result, loc);                            \
       return {};                                                               \
     }                                                                          \
   } while (false)
 
 namespace {
 class SerializeToCubinPass
     : public PassWrapper<SerializeToCubinPass, gpu::SerializeToBlobPass> {
 public:
void SerializeToCubinPass::getDependentDialects(
   registerNVVMDialectTranslation(registry);
   gpu::SerializeToBlobPass::getDependentDialects(registry);
 }
 
 std::unique_ptr<std::vector<char>>
 SerializeToCubinPass::serializeISA(const std::string &isa) {
   Location loc = getOperation().getLoc();
   char jitErrorBuffer[4096] = {0};
+  nvPTXCompilerHandle compiler;
+  nvPTXCompilerCreate(&compiler, isa.length(), isa.c_str());
 
-  RETURN_ON_CUDA_ERROR(cuInit(0));
-
-  // Linking requires a device context.
-  CUdevice device;
-  RETURN_ON_CUDA_ERROR(cuDeviceGet(&device, 0));
-  CUcontext context;
-  RETURN_ON_CUDA_ERROR(cuCtxCreate(&context, 0, device));
-  CUlinkState linkState;
-
-  CUjit_option jitOptions[] = {CU_JIT_ERROR_LOG_BUFFER,
-                               CU_JIT_ERROR_LOG_BUFFER_SIZE_BYTES};
-  void *jitOptionsVals[] = {jitErrorBuffer,
-                            reinterpret_cast<void *>(sizeof(jitErrorBuffer))};
-
-  RETURN_ON_CUDA_ERROR(cuLinkCreate(2,              /* number of jit options */
-                                    jitOptions,     /* jit options */
-                                    jitOptionsVals, /* jit option values */
-                                    &linkState));
-
-  auto kernelName = getOperation().getName().str();
-  RETURN_ON_CUDA_ERROR(cuLinkAddData(
-      linkState, CUjitInputType::CU_JIT_INPUT_PTX,
-      const_cast<void *>(static_cast<const void *>(isa.c_str())), isa.length(),
-      kernelName.c_str(), 0, /* number of jit options */
-      nullptr,               /* jit options */
-      nullptr                /* jit option values */
-      ));
+  nvPTXCompilerCompile(compiler, 0, nullptr);
 
-  void *cubinData;
   size_t cubinSize;
-  RETURN_ON_CUDA_ERROR(cuLinkComplete(linkState, &cubinData, &cubinSize));
+  nvPTXCompilerGetCompiledProgramSize(compiler, &cubinSize);
 
-  char *cubinAsChar = static_cast<char *>(cubinData);
-  auto result =
-      std::make_unique<std::vector<char>>(cubinAsChar, cubinAsChar + cubinSize);
-
-  // This will also destroy the cubin data.
-  RETURN_ON_CUDA_ERROR(cuLinkDestroy(linkState));
-  RETURN_ON_CUDA_ERROR(cuCtxDestroy(context));
+  auto result = std::make_unique<std::vector<char>>(cubinSize);
+  nvPTXCompilerGetCompiledProgram(compiler, result->data());
+  nvPTXCompilerDestroy(&compiler);
 
   return result;
 }
 
 // Register pass to serialize GPU kernel functions to a CUBIN binary annotation.
 void mlir::registerGpuSerializeToCubinPass() {
   PassRegistration<SerializeToCubinPass> registerSerializeToCubin([] {
     // Initialize LLVM NVPTX backend.

mlir/lib/ExecutionEngine/CMakeLists.txt

add_mlir_library(mlir_async_runtime
 
     LINK_LIBS PUBLIC
     ${LLVM_PTHREAD_LIB}
   )
   set_property(TARGET mlir_async_runtime PROPERTY CXX_VISIBILITY_PRESET hidden)
   target_compile_definitions(mlir_async_runtime PRIVATE mlir_async_runtime_EXPORTS)
 
   if(MLIR_ENABLE_CUDA_RUNNER)
+    find_package(CUDAToolkit)
     # Configure CUDA support. Using check_language first allows us to give a
     # custom error message.
     include(CheckLanguage)
     check_language(CUDA)
     if (CMAKE_CUDA_COMPILER)
       enable_language(CUDA)
     else()
       message(SEND_ERROR
         "Building the mlir cuda runner requires a working CUDA install")
     endif()
 
-    # We need the libcuda.so library.
-    find_library(CUDA_RUNTIME_LIBRARY cuda)
-
     add_mlir_library(mlir_cuda_runtime
       SHARED
       CudaRuntimeWrappers.cpp
 
       EXCLUDE_FROM_LIBMLIR
     )
     set_property(TARGET mlir_cuda_runtime PROPERTY CXX_STANDARD 14)
     target_include_directories(mlir_cuda_runtime
       PRIVATE
       ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}
     )
     target_link_libraries(mlir_cuda_runtime
       PRIVATE
-      ${CUDA_RUNTIME_LIBRARY}
+      CUDA::cuda_driver
     )
   endif()
 
   if(MLIR_ENABLE_ROCM_RUNNER)
     # Configure ROCm support.
     if (NOT DEFINED ROCM_PATH)
       if (NOT DEFINED ENV{ROCM_PATH})
         set(ROCM_PATH "/opt/rocm" CACHE PATH "Path to which ROCm has been installed")