[mlir] Specify cuda-runner pass pipeline as command line options.

[csigg]

Summary

The cuda-runner registers two pass pipelines for nested passes,
so that we don't have to use verbose textual pass pipeline specification.

[mehdi_amini]

mlir/tools/mlir-cuda-runner/mlir-cuda-runner.cpp

175

I'm not fond of exposing arbitrary pipeline to the runner, can we leave this to mlir-opt and pipe it to the runner instead? There is value in keeping the tools focused in the Unix-style, without loss of generality I think.

[csigg]

mlir/tools/mlir-cuda-runner/mlir-cuda-runner.cpp

175

I see your point, and this was intended as a step in that direction, at least allowing textual pass pipeline specification on the command line instead of hardcoding it in the cuda-runner.

• I added gpu-to-async because I didn't know about PassManager::Nesting::Implicit. I can removed that again.
• I added gpu-to-cuda-runtime because -gpu-to-llvm=nvvm.cubin looks pretty arbitrary (to be fair, nvvm.cubin is arbitraty too, but it needs to match the name used in gpu-to-cubin -- that pass pipeline should probably have the name as an option).

We will presumably need a cuda-opt though which runs or provides the gpu-to-cubin pipeline added here, because mlir-opt would not want to depend on CUDA.

I will make the above two changes here and leave the cuda-opt > mlir-opt > cpu-runner piping as a follow-up.

As a side note, the downside of tests piping code through tools is that it's harder to setup IDEs for debugging. When I debug an mlir-opt > cpu-runner test, I think twice before setting up proper debugging, add some printfs, waste time, and then do it anyway ;-)

[csigg]

Only register the one required pass pipeline.

[herhut]

lgtm for a first step in making the pass pipeline more configurable.

mlir/lib/Conversion/GPUCommon/ConvertKernelFuncToBlob.cpp

private:
   /// Translates llvmModule to a blob with target instructions and returns the
   /// result as attribute.
   StringAttr translateGPUModuleToBinaryAnnotation(llvm::Module &llvmModule,
                                                   Location loc, StringRef name);
 
   LoweringCallback loweringCallback;
   BlobGenerator blobGenerator;
   llvm::Triple triple;
-  StringRef targetChip;
-  StringRef features;
-  StringRef blobAnnotation;
+  std::string targetChip;
+  std::string features;
+  std::string blobAnnotation;
 };
 
 } // anonymous namespace
 
 std::string
 GpuKernelToBlobPass::translateModuleToISA(llvm::Module &module,
                                           llvm::TargetMachine &targetMachine) {
   std::string targetISA;

mlir/test/mlir-cuda-runner/all-reduce-and.mlir

 // RUN: mlir-cuda-runner %s \
+// RUN:   -gpu-to-cubin="gpu-binary-annotation=nvvm.cubin" \
+// RUN:   -gpu-to-llvm="gpu-binary-annotation=nvvm.cubin" \
 // RUN:   --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext \
 // RUN:   --entry-point-result=void \
 // RUN: | FileCheck %s
 
 func @main() {
   %data = alloc() : memref<2x6xi32>
   %sum = alloc() : memref<2xi32>

mlir/test/mlir-cuda-runner/all-reduce-max.mlir

 // RUN: mlir-cuda-runner %s \
+// RUN:   -gpu-to-cubin="gpu-binary-annotation=nvvm.cubin" \
+// RUN:   -gpu-to-llvm="gpu-binary-annotation=nvvm.cubin" \
 // RUN:   --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext \
 // RUN:   --entry-point-result=void \
 // RUN: | FileCheck %s
 
 func @main() {
   %data = alloc() : memref<2x6xi32>
   %sum = alloc() : memref<2xi32>

mlir/test/mlir-cuda-runner/all-reduce-min.mlir

 // RUN: mlir-cuda-runner %s \
+// RUN:   -gpu-to-cubin="gpu-binary-annotation=nvvm.cubin" \
+// RUN:   -gpu-to-llvm="gpu-binary-annotation=nvvm.cubin" \
 // RUN:   --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext \
 // RUN:   --entry-point-result=void \
 // RUN: | FileCheck %s
 
 func @main() {
   %data = alloc() : memref<2x6xi32>
   %sum = alloc() : memref<2xi32>

mlir/test/mlir-cuda-runner/all-reduce-op.mlir

 // RUN: mlir-cuda-runner %s \
+// RUN:   -gpu-to-cubin="gpu-binary-annotation=nvvm.cubin" \
+// RUN:   -gpu-to-llvm="gpu-binary-annotation=nvvm.cubin" \
 // RUN:   --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext \
 // RUN:   --entry-point-result=void \
 // RUN: | FileCheck %s
 
 // CHECK-COUNT-8: [{{(5356, ){12}5356}}]
 func @main() {
   %arg = alloc() : memref<2x4x13xf32>

mlir/test/mlir-cuda-runner/all-reduce-or.mlir

 // RUN: mlir-cuda-runner %s \
+// RUN:   -gpu-to-cubin="gpu-binary-annotation=nvvm.cubin" \
+// RUN:   -gpu-to-llvm="gpu-binary-annotation=nvvm.cubin" \
 // RUN:   --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext \
 // RUN:   --entry-point-result=void \
 // RUN: | FileCheck %s
 
 func @main() {
   %data = alloc() : memref<2x6xi32>
   %sum = alloc() : memref<2xi32>

mlir/test/mlir-cuda-runner/all-reduce-region.mlir

 // RUN: mlir-cuda-runner %s \
+// RUN:   -gpu-to-cubin="gpu-binary-annotation=nvvm.cubin" \
+// RUN:   -gpu-to-llvm="gpu-binary-annotation=nvvm.cubin" \
 // RUN:   --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext \
 // RUN:   --entry-point-result=void \
 // RUN: | FileCheck %s
 
 // CHECK: [{{(35, ){34}35}}]
 func @main() {
   %arg = alloc() : memref<35xf32>

mlir/test/mlir-cuda-runner/all-reduce-xor.mlir

 // RUN: mlir-cuda-runner %s \
+// RUN:   -gpu-to-cubin="gpu-binary-annotation=nvvm.cubin" \
+// RUN:   -gpu-to-llvm="gpu-binary-annotation=nvvm.cubin" \
 // RUN:   --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext \
 // RUN:   --entry-point-result=void \
 // RUN: | FileCheck %s
 
 func @main() {
   %data = alloc() : memref<2x6xi32>
   %sum = alloc() : memref<2xi32>

mlir/test/mlir-cuda-runner/async.mlir

-// RUN: mlir-cuda-runner %s --entry-point-result=void -O0 \
+// RUN: mlir-cuda-runner %s \
+// RUN:   -gpu-to-cubin="gpu-binary-annotation=nvvm.cubin" \
+// RUN:   -gpu-async-region -async-ref-counting \
+// RUN:   -gpu-to-llvm="gpu-binary-annotation=nvvm.cubin" \
+// RUN:   -async-to-async-runtime -convert-async-to-llvm -convert-std-to-llvm \
 // RUN:   --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_async_runtime%shlibext \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext \
+// RUN:   --entry-point-result=void -O0 \
 // RUN: | FileCheck %s
 
 func @main() {
   %c0    = constant 0 : index
   %c1    = constant 1 : index
   %count = constant 2 : index
 
   // initialize h0 on host

mlir/test/mlir-cuda-runner/gpu-to-cubin.mlir

 // RUN: mlir-cuda-runner %s \
+// RUN:   -gpu-to-cubin="gpu-binary-annotation=nvvm.cubin" \
+// RUN:   -gpu-to-llvm="gpu-binary-annotation=nvvm.cubin" \
 // RUN:   --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext \
 // RUN:   --entry-point-result=void \
 // RUN: | FileCheck %s
 
 func @other_func(%arg0 : f32, %arg1 : memref<?xf32>) {
   %cst = constant 1 : index
   %c0 = constant 0 : index

mlir/test/mlir-cuda-runner/multiple-all-reduce.mlir

 // RUN: mlir-cuda-runner %s \
+// RUN:   -gpu-to-cubin="gpu-binary-annotation=nvvm.cubin" \
+// RUN:   -gpu-to-llvm="gpu-binary-annotation=nvvm.cubin" \
 // RUN:   --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext \
 // RUN:   --entry-point-result=void \
 // RUN: | FileCheck %s
 
 func @main() {
   %data = alloc() : memref<2x6xf32>
   %sum = alloc() : memref<2xf32>

mlir/test/mlir-cuda-runner/shuffle.mlir

 // RUN: mlir-cuda-runner %s \
+// RUN:   -gpu-to-cubin="gpu-binary-annotation=nvvm.cubin" \
+// RUN:   -gpu-to-llvm="gpu-binary-annotation=nvvm.cubin" \
 // RUN:   --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%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>

mlir/test/mlir-cuda-runner/two-modules.mlir

 // RUN: mlir-cuda-runner %s \
+// RUN:   -gpu-to-cubin="gpu-binary-annotation=nvvm.cubin" \
+// RUN:   -gpu-to-llvm="gpu-binary-annotation=nvvm.cubin" \
 // RUN:   --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext \
 // RUN:   --shared-libs=%linalg_test_lib_dir/libmlir_runner_utils%shlibext \
 // RUN:   --entry-point-result=void \
 // RUN: | FileCheck %s
 
 // CHECK: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
 func @main() {
   %arg = alloc() : memref<13xi32>

mlir/tools/mlir-cuda-runner/mlir-cuda-runner.cpp

 //===- mlir-cuda-runner.cpp - MLIR CUDA Execution Driver-------------------===//
 //
 // 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 is a command line utility that executes an MLIR file on the GPU by
 // translating MLIR to NVVM/LVVM IR before JIT-compiling and executing the
 // latter.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/STLExtras.h"
 
-#include "mlir/Conversion/AsyncToLLVM/AsyncToLLVM.h"
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
-#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
-#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
+#include "mlir/Conversion/Passes.h"
 #include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/Async/Passes.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/ExecutionEngine/JitRunner.h"

 
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/TargetSelect.h"
 
 #include "cuda.h"
 
 using namespace mlir;
 
-inline void emit_cuda_error(const llvm::Twine &message, const char *buffer,
-                            CUresult error, Location loc) {
-  emitError(loc, message.concat(" failed with error code ")
+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, msg)                                        \
-  {                                                                            \
-    auto _cuda_error = (expr);                                                 \
-    if (_cuda_error != CUDA_SUCCESS) {                                         \
-      emit_cuda_error(msg, jitErrorBuffer, _cuda_error, loc);                  \
+#define RETURN_ON_CUDA_ERROR(expr)                                             \
+  do {                                                                         \
+    if (auto status = (expr)) {                                                \
+      emitCudaError(#expr, jitErrorBuffer, status, loc);                       \
       return {};                                                               \
     }                                                                          \
-  }
+  } while (false)
 
 OwnedBlob compilePtxToCubin(const std::string ptx, Location loc,
                             StringRef name) {
   char jitErrorBuffer[4096] = {0};
 
-  RETURN_ON_CUDA_ERROR(cuInit(0), "cuInit");
+  RETURN_ON_CUDA_ERROR(cuInit(0));
 
   // Linking requires a device context.
   CUdevice device;
-  RETURN_ON_CUDA_ERROR(cuDeviceGet(&device, 0), "cuDeviceGet");
+  RETURN_ON_CUDA_ERROR(cuDeviceGet(&device, 0));
   CUcontext context;
-  RETURN_ON_CUDA_ERROR(cuCtxCreate(&context, 0, device), "cuCtxCreate");
+  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),
-                       "cuLinkCreate");
+                                    &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 */
-                    ),
-      "cuLinkAddData");
+                    ));
 
   void *cubinData;
   size_t cubinSize;
-  RETURN_ON_CUDA_ERROR(cuLinkComplete(linkState, &cubinData, &cubinSize),
-                       "cuLinkComplete");
+  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), "cuLinkDestroy");
-  RETURN_ON_CUDA_ERROR(cuCtxDestroy(context), "cuCtxDestroy");
+  RETURN_ON_CUDA_ERROR(cuLinkDestroy(linkState));
+  RETURN_ON_CUDA_ERROR(cuCtxDestroy(context));
 
   return result;
 }
 
-static LogicalResult runMLIRPasses(ModuleOp m) {
-  PassManager pm(m.getContext());
-  applyPassManagerCLOptions(pm);
-
-  const char gpuBinaryAnnotation[] = "nvvm.cubin";
+struct GpuToCubinPipelineOptions
+    : public mlir::PassPipelineOptions<GpuToCubinPipelineOptions> {
+  Option<std::string> gpuBinaryAnnotation{
+      *this, "gpu-binary-annotation",
+      llvm::cl::desc("Annotation attribute string for GPU binary")};
+};
+
+// Register cuda-runner specific passes.
+static void registerCudaRunnerPasses() {
+  PassPipelineRegistration<GpuToCubinPipelineOptions> registerGpuToCubin(
+      "gpu-to-cubin", "Generate CUBIN from gpu.launch regions",
+      [&](OpPassManager &pm, const GpuToCubinPipelineOptions &options) {
         pm.addPass(createGpuKernelOutliningPass());
         auto &kernelPm = pm.nest<gpu::GPUModuleOp>();
         kernelPm.addPass(createStripDebugInfoPass());
         kernelPm.addPass(createLowerGpuOpsToNVVMOpsPass());
         kernelPm.addPass(createConvertGPUKernelToBlobPass(
             translateModuleToLLVMIR, compilePtxToCubin, "nvptx64-nvidia-cuda",
-      "sm_35", "+ptx60", gpuBinaryAnnotation));
-  auto &funcPm = pm.nest<FuncOp>();
-  funcPm.addPass(createGpuAsyncRegionPass());
-  funcPm.addPass(createAsyncRefCountingPass());
-  pm.addPass(createGpuToLLVMConversionPass(gpuBinaryAnnotation));
-  pm.addPass(createAsyncToAsyncRuntimePass());
-  pm.addPass(createConvertAsyncToLLVMPass());
-  mlir::LowerToLLVMOptions lower_to_llvm_opts;
-  pm.addPass(mlir::createLowerToLLVMPass(lower_to_llvm_opts));
+            "sm_35", "+ptx60", options.gpuBinaryAnnotation));
+      });
+  registerGPUPasses();
+  registerGpuToLLVMConversionPassPass();
+  registerAsyncPasses();
+  registerConvertAsyncToLLVMPass();
+  registerConvertStandardToLLVMPass();
+}
+
+static LogicalResult runMLIRPasses(ModuleOp module,
+                                   PassPipelineCLParser &passPipeline) {
+  PassManager pm(module.getContext(), PassManager::Nesting::Implicit);
+  applyPassManagerCLOptions(pm);
 
-  return pm.run(m);
+  auto errorHandler = [&](const Twine &msg) {
+    emitError(UnknownLoc::get(module.getContext())) << msg;
+    return failure();
+  };
+
+  // Build the provided pipeline.
+  if (failed(passPipeline.addToPipeline(pm, errorHandler)))
+    return failure();
+
+  // Run the pipeline.
+  return pm.run(module);
 }
 
 int main(int argc, char **argv) {
-  registerPassManagerCLOptions();
   llvm::InitLLVM y(argc, argv);
   llvm::InitializeNativeTarget();
   llvm::InitializeNativeTargetAsmPrinter();
 
   // Initialize LLVM NVPTX backend.
   LLVMInitializeNVPTXTarget();
   LLVMInitializeNVPTXTargetInfo();
   LLVMInitializeNVPTXTargetMC();
   LLVMInitializeNVPTXAsmPrinter();
 
   mlir::initializeLLVMPasses();
 
+  registerCudaRunnerPasses();
+  PassPipelineCLParser passPipeline("", "Compiler passes to run");
+  registerPassManagerCLOptions();
+
+  auto mlirTransformer = [&](ModuleOp module) {
+    return runMLIRPasses(module, passPipeline);
+  };
+
   mlir::JitRunnerConfig jitRunnerConfig;
-  jitRunnerConfig.mlirTransformer = runMLIRPasses;
+  jitRunnerConfig.mlirTransformer = mlirTransformer;
 
   mlir::DialectRegistry registry;
   registry.insert<mlir::LLVM::LLVMDialect, mlir::NVVM::NVVMDialect,
                   mlir::async::AsyncDialect, mlir::gpu::GPUDialect,
                   mlir::StandardOpsDialect>();
   registry.addDialectInterface<NVVM::NVVMDialect,
                                NVVMDialectLLVMIRTranslationInterface>();
   mlir::registerLLVMDialectTranslation(registry);
 
   return mlir::JitRunnerMain(argc, argv, registry, jitRunnerConfig);
 }