[mlir][gpu] Refactor ConvertGpuLaunchFuncToCudaCalls pass.

[whchung]

Summary

Due to similar APIs between CUDA and ROCm (HIP),
ConvertGpuLaunchFuncToCudaCalls pass could be used on both platforms with some
refactoring.

In this commit:

• Migrate ConvertLaunchFuncToCudaCalls from GPUToCUDA to GPUCommon, and rename.
• Change CUDA-specific runtime wrapper APIs be specifiable as PassOptions.
• Naming changes within the implementation and tests.

Subsequent patches would introduce ROCm-specific tests and runtime wrapper
APIs.

[herhut]

As an alternative, we could stick with fixed names, replace the mcu prefix by something neutral and then for AMD use case we link a different helper library in the tests. WDYT?

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToRuntimeCalls.cpp

103–105

There is a CUDA left-over here.

[whchung]

Revise the patch addressing part of the code review comments.

[whchung]

In D80167#2043575, @herhut wrote:

As an alternative, we could stick with fixed names, replace the mcu prefix by something neutral and then for AMD use case we link a different helper library in the tests. WDYT?

@herhut I revised the patch so:

• runtime wrapper APIs have more platform-neutral names.
• GPU binary blobs still carry platform-specific attributes, and could be specified as a PassOption. If none specified, default would still be nvvm.cubin.
• Also it seems nvvm.cubingetter is no longer used so I removed it from this patch.

[whchung]

Amend unit test for ROCm-specific checks.

[herhut]

Thank you!

mlir/include/mlir/Conversion/GPUCommon/GPUCommonPass.h

1

Please fix comment.

mlir/lib/Conversion/GPUCommon/CMakeLists.txt

19

Please remove commented-out lines.

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToRuntimeCalls.cpp

189

Maybe remove the CUDA reference here, too.

mlir/test/Conversion/GPUCommon/lower-launch-func-to-gpu-runtime-calls.mlir

1

Move this file to GPUCommon directory.

[whchung]

Review code comments.

[whchung]

@herhut Revised the patch per your review comments. I left gpuBinaryAnnotation as a PassOption because in my use case I would invoke LLVM multiple times to trigger code generation for different targets / features, and I use this attribute to help me tell which binary blob comes from which path.

mlir/test/Conversion/GPUCommon/lower-launch-func-to-gpu-runtime-calls.mlir

1

this file is in GPUCommon now.

[herhut]

Thanks!

[nicolasvasilache]

mlir/lib/Conversion/GPUToCUDA/CMakeLists.txt

5

This seems to break the build, SOURCES can be empty when I try to build it.

[whchung]

mlir/lib/Conversion/GPUToCUDA/CMakeLists.txt

5

@nicolasvasilache in your setup, MLIR_CUDA_CONVERSIONS_ENABLED is OFF, correct? Let me see how to address this.

[whchung]

@nicolasvasilache could you check if D80343 solves your issue? Sorry for the confusion.

[rsmith]

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

40

I'm seeing test failures caused by undefined references to mgpuModuleGetFunction. Did you intend to rename this function too?

[mehdi_amini]

I revert this change because the bot was broken, cmake fails when -DLLVM_TARGETS_TO_BUILD=host is passed in.

[whchung]

Fix build errors. Adopting changes in D80343 and D80353.

Address review comments from @mehdi_amini in D80343.

[rriddle]

mlir/include/mlir/Conversion/GPUCommon/GPUCommonPass.h

13

Where is functional used?

15

Why are these necessary?

[whchung]

mlir/include/mlir/Conversion/GPUCommon/GPUCommonPass.h

15

These headers were introduced back in commit c72c6c390710 , but they don't really seem necessary now. I can modify D80142 and remove them.

mlir/include/mlir/Conversion/GPUCommon/GPUCommonPass.h

+//===- GPUToCUDAPass.h - MLIR CUDA runtime support --------------*- 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_GPUCOMMON_GPUCOMMONPASS_H_
+#define MLIR_CONVERSION_GPUCOMMON_GPUCOMMONPASS_H_
+
+#include "mlir/Support/LLVM.h"
+#include <functional>
+#include <memory>
+#include <string>
+#include <vector>
+
+namespace mlir {
+
+class Location;
+class ModuleOp;
+
+template <typename T>
+class OperationPass;
+
+/// Creates a pass to convert a gpu.launch_func operation into a sequence of
+/// GPU runtime calls.
+///
+/// This pass does not generate code to call GPU runtime APIs directly but
+/// instead uses a small wrapper library that exports a stable and conveniently
+/// typed ABI on top of GPU runtimes such as CUDA or ROCm (HIP).
+std::unique_ptr<OperationPass<ModuleOp>>
+createConvertGpuLaunchFuncToGpuRuntimeCallsPass();
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_GPUCOMMON_GPUCOMMONPASS_H_

mlir/include/mlir/Conversion/GPUToCUDA/GPUToCUDAPass.h

 /// module with help of the nvptx backend to PTX and then invokes the provided
 /// cubinGenerator to produce a binary blob (the cubin). Such blob is then
 /// attached as a string attribute named 'nvvm.cubin' to the kernel function.
 /// After the transformation, the body of the kernel function is removed (i.e.,
 /// it is turned into a declaration).
 std::unique_ptr<OperationPass<gpu::GPUModuleOp>>
 createConvertGPUKernelToCubinPass(CubinGenerator cubinGenerator);
 
-/// Creates a pass to convert a gpu.launch_func operation into a sequence of
-/// CUDA calls.
-///
-/// This pass does not generate code to call CUDA directly but instead uses a
-/// small wrapper library that exports a stable and conveniently typed ABI
-/// on top of CUDA.
-std::unique_ptr<OperationPass<ModuleOp>>
-createConvertGpuLaunchFuncToCudaCallsPass();
-
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_GPUTOCUDA_GPUTOCUDAPASS_H_

mlir/include/mlir/Conversion/Passes.td

 
 def ConvertAVX512ToLLVM : Pass<"convert-avx512-to-llvm", "ModuleOp"> {
   let summary = "Convert the operations from the avx512 dialect into the LLVM "
                 "dialect";
   let constructor = "mlir::createConvertAVX512ToLLVMPass()";
 }
 
 //===----------------------------------------------------------------------===//
-// GPUToCUDA
+// GPUCommon
 //===----------------------------------------------------------------------===//
 
-def ConvertGpuLaunchFuncToCudaCalls : Pass<"launch-func-to-cuda", "ModuleOp"> {
-  let summary = "Convert all launch_func ops to CUDA runtime calls";
-  let constructor = "mlir::createConvertGpuLaunchFuncToCudaCallsPass()";
+def ConvertGpuLaunchFuncToGpuRuntimeCalls : Pass<"launch-func-to-gpu-runtime",
+                                                 "ModuleOp"> {
+  let summary = "Convert all launch_func ops to GPU runtime calls";
+  let constructor = "mlir::createConvertGpuLaunchFuncToGpuRuntimeCallsPass()";
+  let options = [
+    Option<"gpuBinaryAnnotation", "gpu-binary-annotation", "std::string",
+           "\"nvvm.cubin\"",
+           "Annotation attribute string for GPU binary">,
+  ];
 }
 
 //===----------------------------------------------------------------------===//
 // GPUToNVVM
 //===----------------------------------------------------------------------===//
 
 def ConvertGpuOpsToNVVMOps : Pass<"convert-gpu-to-nvvm", "gpu::GPUModuleOp"> {
   let summary = "Generate NVVM operations for gpu operations";

mlir/include/mlir/InitAllPasses.h

 // passes to the system.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_INITALLPASSES_H_
 #define MLIR_INITALLPASSES_H_
 
 #include "mlir/Conversion/AVX512ToLLVM/ConvertAVX512ToLLVM.h"
+#include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #include "mlir/Conversion/GPUToCUDA/GPUToCUDAPass.h"
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
 #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h"
 #include "mlir/Conversion/GPUToSPIRV/ConvertGPUToSPIRVPass.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"

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(AffineToStandard)
 add_subdirectory(AVX512ToLLVM)
+add_subdirectory(GPUCommon)
 add_subdirectory(GPUToCUDA)
 add_subdirectory(GPUToNVVM)
 add_subdirectory(GPUToROCDL)
 add_subdirectory(GPUToSPIRV)
 add_subdirectory(GPUToVulkan)
 add_subdirectory(LinalgToLLVM)
 add_subdirectory(LinalgToSPIRV)
 add_subdirectory(LinalgToStandard)
 add_subdirectory(SCFToGPU)
 add_subdirectory(SCFToStandard)
 add_subdirectory(StandardToLLVM)
 add_subdirectory(StandardToSPIRV)
 add_subdirectory(VectorToLLVM)
 add_subdirectory(VectorToSCF)

mlir/lib/Conversion/GPUCommon/CMakeLists.txt

+set(SOURCES
+  ConvertLaunchFuncToRuntimeCalls.cpp
+)
+
+add_mlir_conversion_library(MLIRGPUtoGPURuntimeTransforms
+  ${SOURCES}
+
+  DEPENDS
+  MLIRConversionPassIncGen
+  intrinsics_gen
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIRGPU
+  MLIRIR
+  MLIRLLVMIR
+  #MLIRNVVMIR
+  MLIRPass
+  MLIRSupport
+  #MLIRTargetNVVMIR
+)

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToRuntimeCalls.cpp

-//===- ConvertLaunchFuncToCudaCalls.cpp - MLIR CUDA lowering passes -------===//
+//===- ConvertLaunchFuncToGpuRuntimeCalls.cpp - MLIR GPU lowering passes --===//
 //
 // 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 convert gpu.launch_func op into a sequence of
-// CUDA runtime calls. As the CUDA runtime does not have a stable published ABI,
-// this pass uses a slim runtime layer that builds on top of the public API from
-// the CUDA headers.
+// GPU runtime calls. As most of GPU runtimes does not have a stable published
+// ABI, this pass uses a slim runtime layer that builds on top of the public
+// API from GPU runtime headers.
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Conversion/GPUToCUDA/GPUToCUDAPass.h"
+#include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 
 #include "../PassDetail.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/StandardTypes.h"
 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FormatVariadic.h"
 
 using namespace mlir;
 
 // To avoid name mangling, these are defined in the mini-runtime file.
-static constexpr const char *cuModuleLoadName = "mcuModuleLoad";
-static constexpr const char *cuModuleGetFunctionName = "mcuModuleGetFunction";
-static constexpr const char *cuLaunchKernelName = "mcuLaunchKernel";
-static constexpr const char *cuGetStreamHelperName = "mcuGetStreamHelper";
-static constexpr const char *cuStreamSynchronizeName = "mcuStreamSynchronize";
-static constexpr const char *kMcuMemHostRegister = "mcuMemHostRegister";
-
-static constexpr const char *kCubinAnnotation = "nvvm.cubin";
-static constexpr const char *kCubinStorageSuffix = "_cubin_cst";
+static constexpr const char *kGpuModuleLoadName = "mgpuModuleLoad";
+static constexpr const char *kGpuModuleGetFunctionName =
+    "mgpuModuleGetFunction";
+static constexpr const char *kGpuLaunchKernelName = "mgpuLaunchKernel";
+static constexpr const char *kGpuGetStreamHelperName = "mgpuGetStreamHelper";
+static constexpr const char *kGpuStreamSynchronizeName =
+    "mgpuStreamSynchronize";
+static constexpr const char *kGpuMemHostRegisterName = "mgpuMemHostRegister";
+static constexpr const char *kGpuBinaryStorageSuffix = "_gpubin_cst";
 
 namespace {
 
-/// A pass to convert gpu.launch_func operations into a sequence of CUDA
-/// runtime calls.
+/// A pass to convert gpu.launch_func operations into a sequence of GPU
+/// runtime calls. Currently it supports CUDA and ROCm (HIP).
 ///
 /// In essence, a gpu.launch_func operations gets compiled into the following
 /// sequence of runtime calls:
 ///
-/// * mcuModuleLoad        -- loads the module given the cubin data
-/// * mcuModuleGetFunction -- gets a handle to the actual kernel function
-/// * mcuGetStreamHelper   -- initializes a new CUDA stream
-/// * mcuLaunchKernelName  -- launches the kernel on a stream
-/// * mcuStreamSynchronize -- waits for operations on the stream to finish
+/// * moduleLoad        -- loads the module given the cubin / hsaco data
+/// * moduleGetFunction -- gets a handle to the actual kernel function
+/// * getStreamHelper   -- initializes a new compute stream on GPU
+/// * launchKernel      -- launches the kernel on a stream
+/// * streamSynchronize -- waits for operations on the stream to finish
 ///
 /// Intermediate data structures are allocated on the stack.
-class GpuLaunchFuncToCudaCallsPass
-    : public ConvertGpuLaunchFuncToCudaCallsBase<GpuLaunchFuncToCudaCallsPass> {
+class GpuLaunchFuncToGpuRuntimeCallsPass
+    : public ConvertGpuLaunchFuncToGpuRuntimeCallsBase<
+          GpuLaunchFuncToGpuRuntimeCallsPass> {
 private:
   LLVM::LLVMDialect *getLLVMDialect() { return llvmDialect; }
 
   llvm::LLVMContext &getLLVMContext() {
     return getLLVMDialect()->getLLVMContext();
   }
 
   void initializeCachedTypes() {
private:
   LLVM::LLVMType getInt64Type() { return llvmInt64Type; }
 
   LLVM::LLVMType getIntPtrType() {
     const llvm::Module &module = getLLVMDialect()->getLLVMModule();
     return LLVM::LLVMType::getIntNTy(
         getLLVMDialect(), module.getDataLayout().getPointerSizeInBits());
   }
 
-  LLVM::LLVMType getCUResultType() {
-    // This is declared as an enum in CUDA but helpers use i32.
+  LLVM::LLVMType getGpuRuntimeResultType() {
+    // This is declared as an enum in both CUDA and ROCm (HIP), but helpers
+    // use i32.
     return getInt32Type();
   }
 
   // Allocate a void pointer on the stack.
   Value allocatePointer(OpBuilder &builder, Location loc) {
     auto one = builder.create<LLVM::ConstantOp>(loc, getInt32Type(),
                                                 builder.getI32IntegerAttr(1));
     return builder.create<LLVM::AllocaOp>(loc, getPointerPointerType(), one,
                                           /*alignment=*/0);
   }
 
-  void declareCudaFunctions(Location loc);
+  void declareGpuRuntimeFunctions(Location loc);
   void addParamToList(OpBuilder &builder, Location loc, Value param, Value list,
                       unsigned pos, Value one);
   Value setupParamsArray(gpu::LaunchFuncOp launchOp, OpBuilder &builder);
   Value generateKernelNameConstant(StringRef moduleName, StringRef name,
                                    Location loc, OpBuilder &builder);
   void translateGpuLaunchCalls(mlir::gpu::LaunchFuncOp launchOp);
 
 public:
   // Run the dialect converter on the module.
   void runOnOperation() override {
     // Cache the LLVMDialect for the current module.
     llvmDialect = getContext().getRegisteredDialect<LLVM::LLVMDialect>();
     // Cache the used LLVM types.
     initializeCachedTypes();
 
     getOperation().walk(
         [this](mlir::gpu::LaunchFuncOp op) { translateGpuLaunchCalls(op); });
 
     // GPU kernel modules are no longer necessary since we have a global
-    // constant with the CUBIN data.
+    // constant with the CUBIN, or HSACO data.
     for (auto m :
          llvm::make_early_inc_range(getOperation().getOps<gpu::GPUModuleOp>()))
       m.erase();
   }
 
 private:
   LLVM::LLVMDialect *llvmDialect;
   LLVM::LLVMType llvmVoidType;
   LLVM::LLVMType llvmPointerType;
   LLVM::LLVMType llvmPointerPointerType;
   LLVM::LLVMType llvmInt8Type;
   LLVM::LLVMType llvmInt32Type;
   LLVM::LLVMType llvmInt64Type;
   LLVM::LLVMType llvmIntPtrType;
 };
 
 } // anonymous namespace
 
-// Adds declarations for the needed helper functions from the CUDA wrapper.
+// Adds declarations for the needed helper functions from the runtime wrappers.
 // The types in comments give the actual types expected/returned but the API
 // uses void pointers. This is fine as they have the same linkage in C.
-void GpuLaunchFuncToCudaCallsPass::declareCudaFunctions(Location loc) {
+void GpuLaunchFuncToGpuRuntimeCallsPass::declareGpuRuntimeFunctions(
+    Location loc) {
   ModuleOp module = getOperation();
   OpBuilder builder(module.getBody()->getTerminator());
-  if (!module.lookupSymbol(cuModuleLoadName)) {
+  if (!module.lookupSymbol(kGpuModuleLoadName)) {
     builder.create<LLVM::LLVMFuncOp>(
-        loc, cuModuleLoadName,
+        loc, kGpuModuleLoadName,
         LLVM::LLVMType::getFunctionTy(
-            getCUResultType(),
+            getGpuRuntimeResultType(),
             {
                 getPointerPointerType(), /* CUmodule *module */
                 getPointerType()         /* void *cubin */
             },
             /*isVarArg=*/false));
   }
-  if (!module.lookupSymbol(cuModuleGetFunctionName)) {
+  if (!module.lookupSymbol(kGpuModuleGetFunctionName)) {
     // The helper uses void* instead of CUDA's opaque CUmodule and
     // CUfunction.
     builder.create<LLVM::LLVMFuncOp>(
-        loc, cuModuleGetFunctionName,
+        loc, kGpuModuleGetFunctionName,
         LLVM::LLVMType::getFunctionTy(
-            getCUResultType(),
+            getGpuRuntimeResultType(),
             {
                 getPointerPointerType(), /* void **function */
                 getPointerType(),        /* void *module */
                 getPointerType()         /* char *name */
             },
             /*isVarArg=*/false));
   }
-  if (!module.lookupSymbol(cuLaunchKernelName)) {
+  if (!module.lookupSymbol(kGpuLaunchKernelName)) {
     // Other than the CUDA api, the wrappers use uintptr_t to match the
     // LLVM type if MLIR's index type, which the GPU dialect uses.
     // Furthermore, they use void* instead of CUDA's opaque CUfunction and
     // CUstream.
     builder.create<LLVM::LLVMFuncOp>(
-        loc, cuLaunchKernelName,
+        loc, kGpuLaunchKernelName,
         LLVM::LLVMType::getFunctionTy(
-            getCUResultType(),
+            getGpuRuntimeResultType(),
             {
                 getPointerType(),        /* void* f */
                 getIntPtrType(),         /* intptr_t gridXDim */
                 getIntPtrType(),         /* intptr_t gridyDim */
                 getIntPtrType(),         /* intptr_t gridZDim */
                 getIntPtrType(),         /* intptr_t blockXDim */
                 getIntPtrType(),         /* intptr_t blockYDim */
                 getIntPtrType(),         /* intptr_t blockZDim */
                 getInt32Type(),          /* unsigned int sharedMemBytes */
                 getPointerType(),        /* void *hstream */
                 getPointerPointerType(), /* void **kernelParams */
                 getPointerPointerType()  /* void **extra */
             },
             /*isVarArg=*/false));
   }
-  if (!module.lookupSymbol(cuGetStreamHelperName)) {
+  if (!module.lookupSymbol(kGpuGetStreamHelperName)) {
     // Helper function to get the current CUDA stream. Uses void* instead of
     // CUDAs opaque CUstream.
     builder.create<LLVM::LLVMFuncOp>(
-        loc, cuGetStreamHelperName,
+        loc, kGpuGetStreamHelperName,
         LLVM::LLVMType::getFunctionTy(getPointerType(), /*isVarArg=*/false));
   }
-  if (!module.lookupSymbol(cuStreamSynchronizeName)) {
+  if (!module.lookupSymbol(kGpuStreamSynchronizeName)) {
     builder.create<LLVM::LLVMFuncOp>(
-        loc, cuStreamSynchronizeName,
-        LLVM::LLVMType::getFunctionTy(getCUResultType(),
+        loc, kGpuStreamSynchronizeName,
+        LLVM::LLVMType::getFunctionTy(getGpuRuntimeResultType(),
                                       getPointerType() /* CUstream stream */,
                                       /*isVarArg=*/false));
   }
-  if (!module.lookupSymbol(kMcuMemHostRegister)) {
+  if (!module.lookupSymbol(kGpuMemHostRegisterName)) {
     builder.create<LLVM::LLVMFuncOp>(
-        loc, kMcuMemHostRegister,
+        loc, kGpuMemHostRegisterName,
         LLVM::LLVMType::getFunctionTy(getVoidType(),
                                       {
                                           getPointerType(), /* void *ptr */
                                           getInt64Type()    /* int64 sizeBytes*/
                                       },
                                       /*isVarArg=*/false));
   }
 }
 
 /// Emits the IR with the following structure:
 ///
 ///   %data = llvm.alloca 1 x type-of(<param>)
 ///   llvm.store <param>, %data
 ///   %typeErased = llvm.bitcast %data to !llvm<"i8*">
 ///   %addr = llvm.getelementptr <list>[<pos>]
 ///   llvm.store %typeErased, %addr
 ///
 /// This is necessary to construct the list of arguments passed to the kernel
 /// function as accepted by cuLaunchKernel, i.e. as a void** that points to list
 /// of stack-allocated type-erased pointers to the actual arguments.
-void GpuLaunchFuncToCudaCallsPass::addParamToList(OpBuilder &builder,
-                                                  Location loc, Value param,
-                                                  Value list, unsigned pos,
+void GpuLaunchFuncToGpuRuntimeCallsPass::addParamToList(OpBuilder &builder,
+                                                        Location loc,
+                                                        Value param, Value list,
+                                                        unsigned pos,
                                                         Value one) {
   auto memLocation = builder.create<LLVM::AllocaOp>(
       loc, param.getType().cast<LLVM::LLVMType>().getPointerTo(), one,
       /*alignment=*/1);
   builder.create<LLVM::StoreOp>(loc, param, memLocation);
   auto casted =
       builder.create<LLVM::BitcastOp>(loc, getPointerType(), memLocation);
 
   auto index = builder.create<LLVM::ConstantOp>(loc, getInt32Type(),
                                                 builder.getI32IntegerAttr(pos));
   auto gep = builder.create<LLVM::GEPOp>(loc, getPointerPointerType(), list,
                                          ArrayRef<Value>{index});
   builder.create<LLVM::StoreOp>(loc, casted, gep);
 }
 
-// Generates a parameters array to be used with a CUDA kernel launch call. The
-// arguments are extracted from the launchOp.
+// Generates a parameters array to be used with a CUDA / ROCm (HIP) kernel
+// launch call. The arguments are extracted from the launchOp.
 // The generated code is essentially as follows:
 //
 // %array = alloca(numparams * sizeof(void *))
 // for (i : [0, NumKernelOperands))
 //   %array[i] = cast<void*>(KernelOperand[i])
 // return %array
-Value GpuLaunchFuncToCudaCallsPass::setupParamsArray(gpu::LaunchFuncOp launchOp,
-                                                     OpBuilder &builder) {
+Value GpuLaunchFuncToGpuRuntimeCallsPass::setupParamsArray(
+    gpu::LaunchFuncOp launchOp, OpBuilder &builder) {
 
   // Get the launch target.
   auto gpuFunc = SymbolTable::lookupNearestSymbolFrom<LLVM::LLVMFuncOp>(
       launchOp, launchOp.kernel());
   if (!gpuFunc)
     return {};
 
   unsigned numArgs = gpuFunc.getNumArguments();

 // The code is essentially:
 //
 // llvm.global constant @kernel_name("function_name\00")
 // func(...) {
 //   %0 = llvm.addressof @kernel_name
 //   %1 = llvm.constant (0 : index)
 //   %2 = llvm.getelementptr %0[%1, %1] : !llvm<"i8*">
 // }
-Value GpuLaunchFuncToCudaCallsPass::generateKernelNameConstant(
+Value GpuLaunchFuncToGpuRuntimeCallsPass::generateKernelNameConstant(
     StringRef moduleName, StringRef name, Location loc, OpBuilder &builder) {
   // Make sure the trailing zero is included in the constant.
   std::vector<char> kernelName(name.begin(), name.end());
   kernelName.push_back('\0');
 
   std::string globalName =
       std::string(llvm::formatv("{0}_{1}_kernel_name", moduleName, name));
   return LLVM::createGlobalString(
       loc, builder, globalName, StringRef(kernelName.data(), kernelName.size()),
       LLVM::Linkage::Internal, llvmDialect);
 }
 
 // Emits LLVM IR to launch a kernel function. Expects the module that contains
-// the compiled kernel function as a cubin in the 'nvvm.cubin' attribute of the
-// kernel function in the IR.
-// While MLIR has no global constants, also expects a cubin getter function in
-// an 'nvvm.cubingetter' attribute. Such function is expected to return a
-// pointer to the cubin blob when invoked.
-// With these given, the generated code in essence is
+// the compiled kernel function as a cubin in the 'nvvm.cubin' attribute, or a
+// hsaco in the 'rocdl.hsaco' attribute of the kernel function in the IR.
 //
-// %0 = call %cubingetter
+// %0 = call %binarygetter
 // %1 = alloca sizeof(void*)
-// call %mcuModuleLoad(%2, %1)
+// call %moduleLoad(%2, %1)
 // %2 = alloca sizeof(void*)
 // %3 = load %1
 // %4 = <see generateKernelNameConstant>
-// call %mcuModuleGetFunction(%2, %3, %4)
-// %5 = call %mcuGetStreamHelper()
+// call %moduleGetFunction(%2, %3, %4)
+// %5 = call %getStreamHelper()
 // %6 = load %2
 // %7 = <see setupParamsArray>
-// call %mcuLaunchKernel(%6, <launchOp operands 0..5>, 0, %5, %7, nullptr)
-// call %mcuStreamSynchronize(%5)
-void GpuLaunchFuncToCudaCallsPass::translateGpuLaunchCalls(
+// call %launchKernel(%6, <launchOp operands 0..5>, 0, %5, %7, nullptr)
+// call %streamSynchronize(%5)
+void GpuLaunchFuncToGpuRuntimeCallsPass::translateGpuLaunchCalls(
     mlir::gpu::LaunchFuncOp launchOp) {
   OpBuilder builder(launchOp);
   Location loc = launchOp.getLoc();
-  declareCudaFunctions(loc);
+  declareGpuRuntimeFunctions(loc);
 
   auto zero = builder.create<LLVM::ConstantOp>(loc, getInt32Type(),
                                                builder.getI32IntegerAttr(0));
   // Create an LLVM global with CUBIN extracted from the kernel annotation and
   // obtain a pointer to the first byte in it.
   auto kernelModule = getOperation().lookupSymbol<gpu::GPUModuleOp>(
       launchOp.getKernelModuleName());
   assert(kernelModule && "expected a kernel module");
 
-  auto cubinAttr = kernelModule.getAttrOfType<StringAttr>(kCubinAnnotation);
-  if (!cubinAttr) {
+  auto binaryAttr = kernelModule.getAttrOfType<StringAttr>(gpuBinaryAnnotation);
+  if (!binaryAttr) {
     kernelModule.emitOpError()
-        << "missing " << kCubinAnnotation << " attribute";
+        << "missing " << gpuBinaryAnnotation << " attribute";
     return signalPassFailure();
   }
 
   SmallString<128> nameBuffer(kernelModule.getName());
-  nameBuffer.append(kCubinStorageSuffix);
+  nameBuffer.append(kGpuBinaryStorageSuffix);
   Value data = LLVM::createGlobalString(
-      loc, builder, nameBuffer.str(), cubinAttr.getValue(),
+      loc, builder, nameBuffer.str(), binaryAttr.getValue(),
       LLVM::Linkage::Internal, getLLVMDialect());
 
   // Emit the load module call to load the module data. Error checking is done
   // in the called helper function.
-  auto cuModule = allocatePointer(builder, loc);
-  auto cuModuleLoad =
-      getOperation().lookupSymbol<LLVM::LLVMFuncOp>(cuModuleLoadName);
-  builder.create<LLVM::CallOp>(loc, ArrayRef<Type>{getCUResultType()},
-                               builder.getSymbolRefAttr(cuModuleLoad),
-                               ArrayRef<Value>{cuModule, data});
+  auto gpuModule = allocatePointer(builder, loc);
+  auto gpuModuleLoad =
+      getOperation().lookupSymbol<LLVM::LLVMFuncOp>(kGpuModuleLoadName);
+  builder.create<LLVM::CallOp>(loc, ArrayRef<Type>{getGpuRuntimeResultType()},
+                               builder.getSymbolRefAttr(gpuModuleLoad),
+                               ArrayRef<Value>{gpuModule, data});
   // Get the function from the module. The name corresponds to the name of
   // the kernel function.
-  auto cuOwningModuleRef =
-      builder.create<LLVM::LoadOp>(loc, getPointerType(), cuModule);
+  auto gpuOwningModuleRef =
+      builder.create<LLVM::LoadOp>(loc, getPointerType(), gpuModule);
   auto kernelName = generateKernelNameConstant(
       launchOp.getKernelModuleName(), launchOp.getKernelName(), loc, builder);
-  auto cuFunction = allocatePointer(builder, loc);
-  auto cuModuleGetFunction =
-      getOperation().lookupSymbol<LLVM::LLVMFuncOp>(cuModuleGetFunctionName);
+  auto gpuFunction = allocatePointer(builder, loc);
+  auto gpuModuleGetFunction =
+      getOperation().lookupSymbol<LLVM::LLVMFuncOp>(kGpuModuleGetFunctionName);
   builder.create<LLVM::CallOp>(
-      loc, ArrayRef<Type>{getCUResultType()},
-      builder.getSymbolRefAttr(cuModuleGetFunction),
-      ArrayRef<Value>{cuFunction, cuOwningModuleRef, kernelName});
+      loc, ArrayRef<Type>{getGpuRuntimeResultType()},
+      builder.getSymbolRefAttr(gpuModuleGetFunction),
+      ArrayRef<Value>{gpuFunction, gpuOwningModuleRef, kernelName});
   // Grab the global stream needed for execution.
-  auto cuGetStreamHelper =
-      getOperation().lookupSymbol<LLVM::LLVMFuncOp>(cuGetStreamHelperName);
-  auto cuStream = builder.create<LLVM::CallOp>(
+  auto gpuGetStreamHelper =
+      getOperation().lookupSymbol<LLVM::LLVMFuncOp>(kGpuGetStreamHelperName);
+  auto gpuStream = builder.create<LLVM::CallOp>(
       loc, ArrayRef<Type>{getPointerType()},
-      builder.getSymbolRefAttr(cuGetStreamHelper), ArrayRef<Value>{});
+      builder.getSymbolRefAttr(gpuGetStreamHelper), ArrayRef<Value>{});
   // Invoke the function with required arguments.
-  auto cuLaunchKernel =
-      getOperation().lookupSymbol<LLVM::LLVMFuncOp>(cuLaunchKernelName);
-  auto cuFunctionRef =
-      builder.create<LLVM::LoadOp>(loc, getPointerType(), cuFunction);
+  auto gpuLaunchKernel =
+      getOperation().lookupSymbol<LLVM::LLVMFuncOp>(kGpuLaunchKernelName);
+  auto gpuFunctionRef =
+      builder.create<LLVM::LoadOp>(loc, getPointerType(), gpuFunction);
   auto paramsArray = setupParamsArray(launchOp, builder);
   if (!paramsArray) {
     launchOp.emitOpError() << "cannot pass given parameters to the kernel";
     return signalPassFailure();
   }
   auto nullpointer =
       builder.create<LLVM::IntToPtrOp>(loc, getPointerPointerType(), zero);
   builder.create<LLVM::CallOp>(
-      loc, ArrayRef<Type>{getCUResultType()},
-      builder.getSymbolRefAttr(cuLaunchKernel),
-      ArrayRef<Value>{cuFunctionRef, launchOp.getOperand(0),
+      loc, ArrayRef<Type>{getGpuRuntimeResultType()},
+      builder.getSymbolRefAttr(gpuLaunchKernel),
+      ArrayRef<Value>{gpuFunctionRef, launchOp.getOperand(0),
                       launchOp.getOperand(1), launchOp.getOperand(2),
                       launchOp.getOperand(3), launchOp.getOperand(4),
                       launchOp.getOperand(5), zero, /* sharedMemBytes */
-                      cuStream.getResult(0),        /* stream */
+                      gpuStream.getResult(0),       /* stream */
                       paramsArray,                  /* kernel params */
                       nullpointer /* extra */});
   // Sync on the stream to make it synchronous.
-  auto cuStreamSync =
-      getOperation().lookupSymbol<LLVM::LLVMFuncOp>(cuStreamSynchronizeName);
-  builder.create<LLVM::CallOp>(loc, ArrayRef<Type>{getCUResultType()},
-                               builder.getSymbolRefAttr(cuStreamSync),
-                               ArrayRef<Value>(cuStream.getResult(0)));
+  auto gpuStreamSync =
+      getOperation().lookupSymbol<LLVM::LLVMFuncOp>(kGpuStreamSynchronizeName);
+  builder.create<LLVM::CallOp>(loc, ArrayRef<Type>{getGpuRuntimeResultType()},
+                               builder.getSymbolRefAttr(gpuStreamSync),
+                               ArrayRef<Value>(gpuStream.getResult(0)));
   launchOp.erase();
 }
 
 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
-mlir::createConvertGpuLaunchFuncToCudaCallsPass() {
-  return std::make_unique<GpuLaunchFuncToCudaCallsPass>();
+mlir::createConvertGpuLaunchFuncToGpuRuntimeCallsPass() {
+  return std::make_unique<GpuLaunchFuncToGpuRuntimeCallsPass>();
 }

mlir/lib/Conversion/GPUToCUDA/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   ConvertKernelFuncToCubin.cpp
 )
 
-set(SOURCES
-  ConvertLaunchFuncToCudaCalls.cpp
-)
-
 if (MLIR_CUDA_CONVERSIONS_ENABLED)
  list(APPEND SOURCES "ConvertKernelFuncToCubin.cpp")
   set(NVPTX_LIBS
     MC
     NVPTXCodeGen
     NVPTXDesc
     NVPTXInfo
   )

mlir/test/Conversion/GPUCommon/lower-launch-func-to-gpu-runtime-calls.mlir

-// RUN: mlir-opt -allow-unregistered-dialect %s --launch-func-to-cuda | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect %s --launch-func-to-gpu-runtime="gpu-binary-annotation=nvvm.cubin" | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect %s --launch-func-to-gpu-runtime="gpu-binary-annotation=rocdl.hsaco" | FileCheck %s --check-prefix=ROCDL
 
 module attributes {gpu.container_module} {
 
   // CHECK: llvm.mlir.global internal constant @[[kernel_name:.*]]("kernel\00")
   // CHECK: llvm.mlir.global internal constant @[[global:.*]]("CUBIN")
+  // ROCDL: llvm.mlir.global internal constant @[[global:.*]]("HSACO")
 
-  gpu.module @kernel_module attributes {nvvm.cubin = "CUBIN"} {
+  gpu.module @kernel_module attributes {nvvm.cubin = "CUBIN", rocdl.hsaco = "HSACO"} {
     llvm.func @kernel(%arg0: !llvm.float, %arg1: !llvm<"float*">) attributes {gpu.kernel} {
       llvm.return
     }
   }
 
   llvm.func @foo() {
     %0 = "op"() : () -> (!llvm.float)
     %1 = "op"() : () -> (!llvm<"float*">)
     %cst = llvm.mlir.constant(8 : index) : !llvm.i64
 
     // CHECK: %[[addressof:.*]] = llvm.mlir.addressof @[[global]]
     // CHECK: %[[c0:.*]] = llvm.mlir.constant(0 : index)
-    // CHECK: %[[cubin_ptr:.*]] = llvm.getelementptr %[[addressof]][%[[c0]], %[[c0]]]
+    // CHECK: %[[binary_ptr:.*]] = llvm.getelementptr %[[addressof]][%[[c0]], %[[c0]]]
     // CHECK-SAME: -> !llvm<"i8*">
     // CHECK: %[[module_ptr:.*]] = llvm.alloca {{.*}} x !llvm<"i8*"> : (!llvm.i32) -> !llvm<"i8**">
-    // CHECK: llvm.call @mcuModuleLoad(%[[module_ptr]], %[[cubin_ptr]]) : (!llvm<"i8**">, !llvm<"i8*">) -> !llvm.i32
+    // CHECK: llvm.call @mgpuModuleLoad(%[[module_ptr]], %[[binary_ptr]]) : (!llvm<"i8**">, !llvm<"i8*">) -> !llvm.i32
     // CHECK: %[[func_ptr:.*]] = llvm.alloca {{.*}} x !llvm<"i8*"> : (!llvm.i32) -> !llvm<"i8**">
-    // CHECK: llvm.call @mcuModuleGetFunction(%[[func_ptr]], {{.*}}, {{.*}}) : (!llvm<"i8**">, !llvm<"i8*">, !llvm<"i8*">) -> !llvm.i32
-    // CHECK: llvm.call @mcuGetStreamHelper
-    // CHECK: llvm.call @mcuLaunchKernel
-    // CHECK: llvm.call @mcuStreamSynchronize
+    // CHECK: llvm.call @mgpuModuleGetFunction(%[[func_ptr]], {{.*}}, {{.*}}) : (!llvm<"i8**">, !llvm<"i8*">, !llvm<"i8*">) -> !llvm.i32
+    // CHECK: llvm.call @mgpuGetStreamHelper
+    // CHECK: llvm.call @mgpuLaunchKernel
+    // CHECK: llvm.call @mgpuStreamSynchronize
     "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %0, %1) { kernel = @kernel_module::@kernel }
         : (!llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.float, !llvm<"float*">) -> ()
 
     llvm.return
   }
 
 }

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

 int32_t reportErrorIfAny(CUresult result, const char *where) {
   if (result != CUDA_SUCCESS) {
     llvm::errs() << "CUDA failed with " << result << " in " << where << "\n";
   }
   return result;
 }
 } // anonymous namespace
 
-extern "C" int32_t mcuModuleLoad(void **module, void *data) {
+extern "C" int32_t mgpuModuleLoad(void **module, void *data) {
   int32_t err = reportErrorIfAny(
       cuModuleLoadData(reinterpret_cast<CUmodule *>(module), data),
       "ModuleLoad");
   return err;
 }
 
 extern "C" int32_t mcuModuleGetFunction(void **function, void *module,
                                         const char *name) {
   return reportErrorIfAny(
       cuModuleGetFunction(reinterpret_cast<CUfunction *>(function),
                           reinterpret_cast<CUmodule>(module), name),
       "GetFunction");
 }
 
 // The wrapper uses intptr_t instead of CUDA's unsigned int to match
 // the type of MLIR's index type. This avoids the need for casts in the
 // generated MLIR code.
-extern "C" int32_t mcuLaunchKernel(void *function, intptr_t gridX,
+extern "C" int32_t mgpuLaunchKernel(void *function, intptr_t gridX,
                                     intptr_t gridY, intptr_t gridZ,
                                     intptr_t blockX, intptr_t blockY,
                                     intptr_t blockZ, int32_t smem, void *stream,
                                     void **params, void **extra) {
   return reportErrorIfAny(
       cuLaunchKernel(reinterpret_cast<CUfunction>(function), gridX, gridY,
                      gridZ, blockX, blockY, blockZ, smem,
                      reinterpret_cast<CUstream>(stream), params, extra),
       "LaunchKernel");
 }
 
-extern "C" void *mcuGetStreamHelper() {
+extern "C" void *mgpuGetStreamHelper() {
   CUstream stream;
   reportErrorIfAny(cuStreamCreate(&stream, CU_STREAM_DEFAULT), "StreamCreate");
   return stream;
 }
 
-extern "C" int32_t mcuStreamSynchronize(void *stream) {
+extern "C" int32_t mgpuStreamSynchronize(void *stream) {
   return reportErrorIfAny(
       cuStreamSynchronize(reinterpret_cast<CUstream>(stream)), "StreamSync");
 }
 
 /// Helper functions for writing mlir example code
 
 // Allows to register byte array with the CUDA runtime. Helpful until we have
 // transfer functions implemented.
-extern "C" void mcuMemHostRegister(void *ptr, uint64_t sizeBytes) {
+extern "C" void mgpuMemHostRegister(void *ptr, uint64_t sizeBytes) {
   reportErrorIfAny(cuMemHostRegister(ptr, sizeBytes, /*flags=*/0),
                    "MemHostRegister");
 }
 
 // Allows to register a MemRef with the CUDA runtime. Initializes array with
 // value. Helpful until we have transfer functions implemented.
 template <typename T>
 void mcuMemHostRegisterMemRef(T *pointer, llvm::ArrayRef<int64_t> sizes,
                               llvm::ArrayRef<int64_t> strides, T value) {
   assert(sizes.size() == strides.size());
   llvm::SmallVector<int64_t, 4> denseStrides(strides.size());
 
   std::partial_sum(sizes.rbegin(), sizes.rend(), denseStrides.rbegin(),
                    std::multiplies<int64_t>());
   auto count = denseStrides.front();
 
   // Only densely packed tensors are currently supported.
   std::rotate(denseStrides.begin(), denseStrides.begin() + 1,
               denseStrides.end());
   denseStrides.back() = 1;
   assert(strides == llvm::makeArrayRef(denseStrides));
 
   std::fill_n(pointer, count, value);
-  mcuMemHostRegister(pointer, count * sizeof(T));
+  mgpuMemHostRegister(pointer, count * sizeof(T));
 }
 
 extern "C" void mcuMemHostRegisterFloat(int64_t rank, void *ptr) {
   auto *desc = static_cast<StridedMemRefType<float, 1> *>(ptr);
   auto sizes = llvm::ArrayRef<int64_t>(desc->sizes, rank);
   auto strides = llvm::ArrayRef<int64_t>(desc->sizes + rank, rank);
   mcuMemHostRegisterMemRef(desc->data + desc->offset, sizes, strides, 1.23f);
 }
 
 extern "C" void mcuMemHostRegisterInt32(int64_t rank, void *ptr) {
   auto *desc = static_cast<StridedMemRefType<int32_t, 1> *>(ptr);
   auto sizes = llvm::ArrayRef<int64_t>(desc->sizes, rank);
   auto strides = llvm::ArrayRef<int64_t>(desc->sizes + rank, rank);
   mcuMemHostRegisterMemRef(desc->data + desc->offset, sizes, strides, 123);
 }

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/GPUCommon/GPUCommonPass.h"
 #include "mlir/Conversion/GPUToCUDA/GPUToCUDAPass.h"
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.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"
static LogicalResult runMLIRPasses(ModuleOp m) {
   applyPassManagerCLOptions(pm);
 
   pm.addPass(createGpuKernelOutliningPass());
   auto &kernelPm = pm.nest<gpu::GPUModuleOp>();
   kernelPm.addPass(createStripDebugInfoPass());
   kernelPm.addPass(createLowerGpuOpsToNVVMOpsPass());
   kernelPm.addPass(createConvertGPUKernelToCubinPass(&compilePtxToCubin));
   pm.addPass(createLowerToLLVMPass());
-  pm.addPass(createConvertGpuLaunchFuncToCudaCallsPass());
+  pm.addPass(createConvertGpuLaunchFuncToGpuRuntimeCallsPass());
 
   return pm.run(m);
 }
 
 int main(int argc, char **argv) {
   registerPassManagerCLOptions();
   mlir::registerAllDialects();
   llvm::InitLLVM y(argc, argv);
   llvm::InitializeNativeTarget();
   llvm::InitializeNativeTargetAsmPrinter();
   mlir::initializeLLVMPasses();
   return mlir::JitRunnerMain(argc, argv, &runMLIRPasses);
 }