[mlir] Add NVVM to CUBIN conversion to mlir-opt

[csigg]

Summary

If MLIR_CUDA_RUNNER_ENABLED, register a 'gpu-to-cubin' conversion pass to mlir-opt.

The next step is to switch CUDA integration tests from mlir-cuda-runner to mlir-opt + mlir-cpu-runner and remove mlir-cuda-runner.

Depends On D98279

[rriddle]

mlir/include/mlir/Conversion/GPUToCUBIN/GPUToCUBINPass.h

1

The Conversion/ directory is intended for dialect->dialect conversions, that doesn't seem to be the case here.

[csigg]

mlir/include/mlir/Conversion/GPUToCUBIN/GPUToCUBINPass.h

1

This change registers the pre-existing GpuKernelToBlobPass pass in mlir/lib/Conversion/GPUCommon/ConvertKernelFuncToBlob.cpp. Where do you think this registration should live, and should the pass move as well?

[mehdi_amini]

mlir/lib/Conversion/GPUToCUBIN/LowerGPUToCUBIN.cpp

96

I don't think this safe to do in a multi-threaded environment.

I suspect we're missing an assert for this in the context @ftynse?

mlir/test/Integration/GPU/CUDA/shuffle.mlir

1

I think this is a leftover debugging option

[rriddle]

mlir/include/mlir/Conversion/GPUToCUBIN/GPUToCUBINPass.h

1

The pass looks like a much more natural fit for the Target/ directory, given that it is translating an MLIR module to an external format.

[mehdi_amini]

mlir/include/mlir/Conversion/GPUToCUBIN/GPUToCUBINPass.h

1

Not exactly: it takes MLIR and generates MLIR, it only takes a subset of the IR and turn it into a CUBIN serialized into the IR :)

I suspect it may better be a GPU/Transforms pass?

[bondhugula]

switch CUDA integration tests from mlir-cuda-runner to mlir-opt + mlir-cpu-runner

If mlir-cpu-runner can also run things on the GPUs, it's name should be changed (to mlir-runner?). Or there could be an mlir-gpu-runner in the interest of reducing the link time for those who may just want to run on the CPU and not link in many other passes and dependences?

[csigg]

mlir/include/mlir/Conversion/GPUToCUBIN/GPUToCUBINPass.h

1

Thanks River and Mehdi. I will move it to GPU/Transforms. For that, I would like to first land a prep revision.

mlir/lib/Conversion/GPUToCUBIN/LowerGPUToCUBIN.cpp

96

Good point. I will change GpuKernelToBlobPass to a base class (and move it to GPU/Transforms), so that we can add those as dependent dialects.

mlir/test/Integration/GPU/CUDA/shuffle.mlir

1

Indeed. Thanks for spotting it.

[mehdi_amini]

mlir/lib/Conversion/GPUToCUBIN/LowerGPUToCUBIN.cpp

96

We can't rely on "dependent dialects": these registration are separate from the dialect registration.

I would suggest using the Pass initialization that is invoked at the beginning of the pipeline before the pass manager runs for this purpose.

[csigg]

Use base class added in D98279.

[csigg]

Rebase, cmake fixes.

[csigg]

This change should be ready for review again.

mlir/lib/Conversion/GPUToCUBIN/LowerGPUToCUBIN.cpp

96

Thanks Mehdi. Does the reworked implementation look correct?

[mehdi_amini]

LGTM, but please give a chance to @herhut to have a general look since he reviewed the previous revisions in this series!

mlir/lib/Conversion/GPUToCUBIN/LowerGPUToCUBIN.cpp

96

I had another hook in mind ( https://mlir.llvm.org/docs/PassManagement/#initialization ) but what you wrote is just equally fine in this case I think!

[rriddle]

mlir/include/mlir/Dialect/GPU/Passes.h

89 ↗

(On Diff #329650)

nit: ///

[csigg]

Fix nit.

[herhut]

Thanks!

mlir/include/mlir/Conversion/GPUToCUBIN/GPUToCUBINPass.h

+//===- GPUToCUBINPass.h - Convert GPU kernel to CUBIN blob ------*- C++ -*-===//
+//
+// 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 MLIR_CONVERSION_GPUTOCUBIN_GPUTOCUBINPASS_H_
+#define MLIR_CONVERSION_GPUTOCUBIN_GPUTOCUBINPASS_H_
+
+namespace mlir {
+namespace gpu {
+
+// Register pass to lower a GPU kernel function to a CUBIN binary annotation.
+void registerGpuKernelToCubinPass();
+
+} // namespace gpu
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_GPUTOCUBIN_GPUTOCUBINPASS_H_

mlir/include/mlir/Conversion/Passes.h

 //===- Passes.h - Conversion Pass Construction and Registration -----------===//
 //
 // 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 MLIR_CONVERSION_PASSES_H
 #define MLIR_CONVERSION_PASSES_H
 
 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Conversion/AsyncToLLVM/AsyncToLLVM.h"
 #include "mlir/Conversion/ComplexToLLVM/ComplexToLLVM.h"
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
+#include "mlir/Conversion/GPUToCUBIN/GPUToCUBINPass.h"
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
 #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h"
 #include "mlir/Conversion/GPUToSPIRV/GPUToSPIRVPass.h"
 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
 #include "mlir/Conversion/LinalgToLLVM/LinalgToLLVM.h"
 #include "mlir/Conversion/LinalgToSPIRV/LinalgToSPIRVPass.h"
 #include "mlir/Conversion/LinalgToStandard/LinalgToStandard.h"
 #include "mlir/Conversion/OpenMPToLLVM/ConvertOpenMPToLLVM.h"

mlir/include/mlir/InitAllPasses.h

 // individual passes.
 // The global registry is interesting to interact with the command-line tools.
 inline void registerAllPasses() {
   // General passes
   registerTransformsPasses();
 
   // Conversion passes
   registerConversionPasses();
+  gpu::registerGpuKernelToCubinPass();
 
   // Dialect passes
   registerAffinePasses();
   registerAsyncPasses();
   registerGPUPasses();
   registerLinalgPasses();
   LLVM::registerLLVMPasses();
   quant::registerQuantPasses();

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(AffineToStandard)
 add_subdirectory(AsyncToLLVM)
 add_subdirectory(AVX512ToLLVM)
 add_subdirectory(ComplexToLLVM)
 add_subdirectory(GPUCommon)
+add_subdirectory(GPUToCUBIN)
 add_subdirectory(GPUToNVVM)
 add_subdirectory(GPUToROCDL)
 add_subdirectory(GPUToSPIRV)
 add_subdirectory(GPUToVulkan)
 add_subdirectory(LinalgToLLVM)
 add_subdirectory(LinalgToSPIRV)
 add_subdirectory(LinalgToStandard)
 add_subdirectory(OpenMPToLLVM)

mlir/lib/Conversion/GPUToCUBIN/CMakeLists.txt

+if (MLIR_CUDA_CONVERSIONS_ENABLED)
+  add_mlir_conversion_library(MLIRGPUToCUBIN
+    LowerGPUToCUBIN.cpp
+
+    LINK_COMPONENTS
+    NVPTXCodeGen
+    NVPTXDesc
+    NVPTXInfo
+
+    LINK_LIBS PUBLIC
+    MLIRGPUToGPURuntimeTransforms
+    MLIRGPUToNVVMTransforms
+    MLIRLLVMToLLVMIRTranslation
+    MLIRNVVMToLLVMIRTranslation
+  )
+endif()

mlir/lib/Conversion/GPUToCUBIN/LowerGPUToCUBIN.cpp

+//===- LowerGPUToCUBIN.cpp - Convert GPU kernel to CUBIN blob -------------===//
+//
+// 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 a pass to generate CUBIN blob from a GPU kernel
+// function.
+//
+//===----------------------------------------------------------------------===//
+#if MLIR_CUDA_CONVERSIONS_ENABLED
+#include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
+#include "mlir/Conversion/GPUToCUBIN/GPUToCUBINPass.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Export.h"
+#include "llvm/Support/TargetSelect.h"
+
+#include <cuda.h>
+
+using namespace mlir;
+
+static void emitCudaError(const llvm::Twine &expr, const char *buffer,
+                          CUresult result, Location loc) {
+  const char *error;
+  cuGetErrorString(result, &error);
+  emitError(loc, expr.concat(" failed with error code ")
+                     .concat(llvm::Twine{error})
+                     .concat("[")
+                     .concat(buffer)
+                     .concat("]"));
+}
+
+#define RETURN_ON_CUDA_ERROR(expr)                                             \
+  do {                                                                         \
+    if (auto status = (expr)) {                                                \
+      emitCudaError(#expr, jitErrorBuffer, status, loc);                       \
+      return {};                                                               \
+    }                                                                          \
+  } while (false)
+
+static OwnedBlob compilePtxToCubin(const std::string ptx, Location loc,
+                                   StringRef name) {
+  char jitErrorBuffer[4096] = {0};
+
+  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));
+
+  RETURN_ON_CUDA_ERROR(
+      cuLinkAddData(linkState, CUjitInputType::CU_JIT_INPUT_PTX,
+                    const_cast<void *>(static_cast<const void *>(ptx.c_str())),
+                    ptx.length(), name.str().data(), /* kernel name */
+                    0,                               /* number of jit options */
+                    nullptr,                         /* jit options */
+                    nullptr                          /* jit option values */
+                    ));
+
+  void *cubinData;
+  size_t cubinSize;
+  RETURN_ON_CUDA_ERROR(cuLinkComplete(linkState, &cubinData, &cubinSize));
+
+  char *cubinAsChar = static_cast<char *>(cubinData);
+  OwnedBlob 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));
+
+  return result;
+}
+
+static std::unique_ptr<llvm::Module>
+translateModuleToNVVMIR(Operation *m, llvm::LLVMContext &llvmContext,
+                        StringRef moduleName) {
+  registerLLVMDialectTranslation(*m->getContext());
+  registerNVVMDialectTranslation(*m->getContext());
+  return translateModuleToLLVMIR(m, llvmContext, moduleName);
+}
+
+// Register pass to lower a GPU kernel function to a CUBIN binary annotation.
+void mlir::gpu::registerGpuKernelToCubinPass() {
+  PassRegistration<OperationPass<gpu::GPUModuleOp>> registerGpuToCubin(
+      "gpu-to-cubin", "Lower GPU kernel function to CUBIN binary annotations",
+      [] {
+        // Initialize LLVM NVPTX backend.
+        LLVMInitializeNVPTXTarget();
+        LLVMInitializeNVPTXTargetInfo();
+        LLVMInitializeNVPTXTargetMC();
+        LLVMInitializeNVPTXAsmPrinter();
+
+        return createConvertGPUKernelToBlobPass(
+            translateModuleToNVVMIR, compilePtxToCubin, "nvptx64-nvidia-cuda",
+            "sm_35", "+ptx60", "nvvm.cubin");
+      });
+}
+#else  // MLIR_CUDA_CONVERSIONS_ENABLED
+void mlir::gpu::registerGpuKernelToCubinPass() {}
+#endif // MLIR_CUDA_CONVERSIONS_ENABLED

mlir/test/Integration/GPU/CUDA/shuffle.mlir

-// RUN: mlir-cuda-runner %s \
-// RUN:   -gpu-to-cubin="gpu-binary-annotation=nvvm.cubin" \
+// RUN: mlir-opt %s --print-ir-after-all \
+// RUN:   -gpu-kernel-outlining \
+// RUN:   -pass-pipeline='gpu.module(strip-debuginfo,convert-gpu-to-nvvm,gpu-to-cubin)' \
 // RUN:   -gpu-to-llvm="gpu-binary-annotation=nvvm.cubin" \
+// RUN: | mlir-cpu-runner \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_cuda_runtime%shlibext \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext \
 // RUN:   --entry-point-result=void \
 // RUN: | FileCheck %s
 
 // CHECK: [4, 5, 6, 7, 0, 1, 2, 3, 12, -1, -1, -1, 8]
 func @main() {
   %arg = alloc() : memref<13xf32>