[mlir][gpu] Introduce mlir-rocm-runner.

[whchung]

Summary

mlir-rocm-runner is introduced in this commit to execute GPU modules on ROCm
platform. A small wrapper to encapsulate ROCm's HIP runtime API is also inside
the commit.

Due to behavior of ROCm, raw pointers inside memrefs passed to gpu.launch
must be modified on the host side to properly capture the pointer values
addressable on the GPU.

LLVM MC is used to assemble AMD GCN ISA coming out from
ConvertGPUKernelToBlobPass to binary form, and LLD is used to produce a shared
ELF object which could be loaded by ROCm HIP runtime.

gfx900 is the default target be used right now, although it could be altered via
an option in mlir-rocm-runner. Future revisions may consider using ROCm Agent
Enumerator to detect the right target on the system.

Notice AMDGPU Code Object V2 is used in this revision. Future enhancements may
upgrade to AMDGPU Code Object V3.

Bitcode libraries in ROCm-Device-Libs, which implements math routines exposed in
rocdl dialect are not yet linked, and is left as a TODO in the logic.

[whchung]

This patch depends on D80142.

[whchung]

Tame clang-tidy.

[whchung]

Improve ROCm HIP library search logic.

[whchung]

Add CMake checks to ensure lld is enabled building mlir-rocm-runner.

[whchung]

Fix build errors in a shared lib build.

[rriddle]

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

49

StringRef?

[whchung]

std::string -> StringRef.

[jerryyin]

mlir/tools/mlir-rocm-runner/CMakeLists.txt

73

When building with the following option: -DBUILD_SHARED_LIBS=ON -DLLVM_BUILD_LLVM_DYLIB=ON

Build failed with linker issues like below:

mlir-rocm-runner.cpp:(.text.startup.main+0x5e): undefined reference to `LLVMInitializeNVPTXTargetInfo' 
mlir-rocm-runner.cpp:(.text.startup.main+0x63): undefined reference to `LLVMInitializeAMDGPUTargetInfo'
mlir-rocm-runner.cpp:(.text.startup.main+0x6d): undefined reference to `LLVMInitializeNVPTXTargetMC'
mlir-rocm-runner.cpp:(.text.startup.main+0x72): undefined reference to `LLVMInitializeAMDGPUTargetMC'
mlir-rocm-runner.cpp:(.text.startup.main+0x86): undefined reference to `LLVMInitializeAMDGPUTargetInfo'
mlir-rocm-runner.cpp:(.text.startup.main+0x8b): undefined reference to `LLVMInitializeAMDGPUTargetMC'

Needs the following targets: LLVMAMDGPUInfo, LLVMAMDGPUDesc, LLVMNVPTXInfo, LLVMNVPTXDesc.

[whchung]

mlir/tools/mlir-rocm-runner/CMakeLists.txt

73

@jerryyin The issue should be no more should you adopt D80739.

[jerryyin]

Good job, mostly minor issues, except the one related with the difference between mlir-rocm-runner and mlir-cuda-runner. Giving an accept to unblock future development, but we need to figure out the best way to resolve the differences between runner wrapper interfaces.

mlir/tools/mlir-rocm-runner/mlir-rocm-runner.cpp

72

Adding a TODO regarding on supporting other target chips?

81

I only see you use the raw os stream, not the bos. Should you swap os with bos in the lines follows?

139

Nit: Could you add a const to &isaBlob just so that it is clear it is input

191

Nit: Use the stack allocated llvm::FileRemover instead of having to manually track/remove it?

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

123

I'm a bit concerned that we get diverged between different user APIs. It is probably fine now, but will be a maintenance burden later. I will have to, in the future, duplicate and sync every cuda-rocm-runner test with us.

Desirably, we can reach to a shape where cuda-runner and rocm-runner tests can be merged into one.

[whchung]

mlir/tools/mlir-rocm-runner/mlir-rocm-runner.cpp

72

Other targets such as gfx906 gfx908 can already work via this option. I'll add a TODO to detect native AMD GPU ISA version on the system.

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

123

This only has to deal with host pinned memrefs. In real-world applications memrefs could be explicitly allocated on GPU, and transferred between CPU and GPU.

[whchung]

Address code review comments.

[whchung]

@jerryyin I've revised the commit based on your comments. Could you give this patch another around of review? Thanks.

[jerryyin]

LGTM

mlir/tools/mlir-rocm-runner/mlir-rocm-runner.cpp

72

Thanks for clarifying

[whchung]

Apply NOLINT.

[whchung]

Remove NOLINT.

mlir/CMakeLists.txt

 if ("AMDGPU" IN_LIST LLVM_TARGETS_TO_BUILD)
   set(MLIR_ROCM_CONVERSIONS_ENABLED 1)
 else()
   set(MLIR_ROCM_CONVERSIONS_ENABLED 0)
 endif()
 add_definitions(-DMLIR_ROCM_CONVERSIONS_ENABLED=${MLIR_ROCM_CONVERSIONS_ENABLED})
 
 set(MLIR_CUDA_RUNNER_ENABLED 0 CACHE BOOL "Enable building the mlir CUDA runner")
+set(MLIR_ROCM_RUNNER_ENABLED 0 CACHE BOOL "Enable building the mlir ROCm runner")
 set(MLIR_VULKAN_RUNNER_ENABLED 0 CACHE BOOL "Enable building the mlir Vulkan runner")
 
 option(MLIR_INCLUDE_TESTS
        "Generate build targets for the MLIR unit tests."
        ${LLVM_INCLUDE_TESTS})
 
 include_directories( "include")
 include_directories( ${MLIR_INCLUDE_DIR})

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

 /// 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();
+std::unique_ptr<OperationPass<ModuleOp>>
+createConvertGpuLaunchFuncToGpuRuntimeCallsPass(
+    std::string gpuBinaryAnnotation);
 
 /// Creates a pass to convert kernel functions into GPU target object blobs.
 ///
 /// This transformation takes the body of each function that is annotated with
 /// the 'gpu.kernel' attribute, copies it to a new LLVM module, compiles the
 /// module with help of the GPU backend to target object and then invokes
 /// the provided blobGenerator to produce a binary blob. Such blob is then
 /// attached as a string attribute to the kernel function.

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToRuntimeCalls.cpp

private:
   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:
+  GpuLaunchFuncToGpuRuntimeCallsPass() = default;
+  GpuLaunchFuncToGpuRuntimeCallsPass(std::string gpuBinaryAnnotation) {
+    this->gpuBinaryAnnotation = gpuBinaryAnnotation;
+  }
+
   // 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(
builder.create<LLVM::CallOp>(loc, ArrayRef<Type>{getGpuRuntimeResultType()},
                                ArrayRef<Value>(gpuStream.getResult(0)));
   launchOp.erase();
 }
 
 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
 mlir::createConvertGpuLaunchFuncToGpuRuntimeCallsPass() {
   return std::make_unique<GpuLaunchFuncToGpuRuntimeCallsPass>();
 }
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+mlir::createConvertGpuLaunchFuncToGpuRuntimeCallsPass(
+    std::string gpuBinaryAnnotation) {
+  return std::make_unique<GpuLaunchFuncToGpuRuntimeCallsPass>(
+      gpuBinaryAnnotation);
+}

mlir/test/CMakeLists.txt

 add_subdirectory(EDSC)
 add_subdirectory(mlir-cpu-runner)
 add_subdirectory(SDBM)
 add_subdirectory(lib)
 
 llvm_canonicalize_cmake_booleans(
   LLVM_BUILD_EXAMPLES
   )
 
 # Passed to lit.site.cfg.py.in to set up the path where to find the libraries
 # for linalg integration tests.
 set(MLIR_DIALECT_LINALG_INTEGRATION_TEST_LIB_DIR ${CMAKE_LIBRARY_OUTPUT_DIRECTORY})
 set(MLIR_RUNNER_UTILS_DIR ${CMAKE_LIBRARY_OUTPUT_DIRECTORY})
 
 # Passed to lit.site.cfg.py.in to set up the path where to find the libraries
-# for the mlir cuda runner tests.
+# for the mlir cuda / rocm / vulkan runner tests.
 set(MLIR_CUDA_WRAPPER_LIBRARY_DIR ${CMAKE_LIBRARY_OUTPUT_DIRECTORY})
+set(MLIR_ROCM_WRAPPER_LIBRARY_DIR ${CMAKE_LIBRARY_OUTPUT_DIRECTORY})
 set(MLIR_VULKAN_WRAPPER_LIBRARY_DIR ${CMAKE_LIBRARY_OUTPUT_DIRECTORY})
 
 configure_lit_site_cfg(
   ${CMAKE_CURRENT_SOURCE_DIR}/lit.site.cfg.py.in
   ${CMAKE_CURRENT_BINARY_DIR}/lit.site.cfg.py
   MAIN_CONFIG
   ${CMAKE_CURRENT_SOURCE_DIR}/lit.cfg.py
   )

 endif()
 
 if(MLIR_CUDA_RUNNER_ENABLED)
   list(APPEND MLIR_TEST_DEPENDS
     mlir-cuda-runner
   )
 endif()
 
+if(MLIR_ROCM_RUNNER_ENABLED)
+  list(APPEND MLIR_TEST_DEPENDS
+    mlir-rocm-runner
+  )
+endif()
+
 if(MLIR_VULKAN_RUNNER_ENABLED)
   list(APPEND MLIR_TEST_DEPENDS
     mlir-vulkan-runner
   )
 endif()
 
 add_lit_testsuite(check-mlir "Running the MLIR regression tests"
   ${CMAKE_CURRENT_BINARY_DIR}
   DEPENDS ${MLIR_TEST_DEPENDS}
   )
 set_target_properties(check-mlir PROPERTIES FOLDER "Tests")
 
 add_lit_testsuites(MLIR ${CMAKE_CURRENT_SOURCE_DIR}
   DEPENDS ${MLIR_TEST_DEPS}
 )

mlir/test/lit.cfg.py

tools.extend([
     ToolSubst('toy-ch1', unresolved='ignore'),
     ToolSubst('toy-ch2', unresolved='ignore'),
     ToolSubst('toy-ch3', unresolved='ignore'),
     ToolSubst('toy-ch4', unresolved='ignore'),
     ToolSubst('toy-ch5', unresolved='ignore'),
     ToolSubst('%cuda_wrapper_library_dir', config.cuda_wrapper_library_dir, unresolved='ignore'),
     ToolSubst('%linalg_test_lib_dir', config.linalg_test_lib_dir, unresolved='ignore'),
     ToolSubst('%mlir_runner_utils_dir', config.mlir_runner_utils_dir, unresolved='ignore'),
+    ToolSubst('%rocm_wrapper_library_dir', config.rocm_wrapper_library_dir, unresolved='ignore'),
     ToolSubst('%vulkan_wrapper_library_dir', config.vulkan_wrapper_library_dir, unresolved='ignore')
 ])
 
 llvm_config.add_tool_substitutions(tools, tool_dirs)

mlir/test/lit.site.cfg.py.in

 config.mlir_runner_utils_dir = "@MLIR_RUNNER_UTILS_DIR@"
 config.mlir_tools_dir = "@MLIR_TOOLS_DIR@"
 config.linalg_test_lib_dir = "@MLIR_DIALECT_LINALG_INTEGRATION_TEST_LIB_DIR@"
 config.build_examples = @LLVM_BUILD_EXAMPLES@
 config.run_cuda_tests = @MLIR_CUDA_CONVERSIONS_ENABLED@
 config.cuda_wrapper_library_dir = "@MLIR_CUDA_WRAPPER_LIBRARY_DIR@"
 config.enable_cuda_runner = @MLIR_CUDA_RUNNER_ENABLED@
 config.run_rocm_tests = @MLIR_ROCM_CONVERSIONS_ENABLED@
+config.rocm_wrapper_library_dir = "@MLIR_ROCM_WRAPPER_LIBRARY_DIR@"
+config.enable_rocm_runner = @MLIR_ROCM_RUNNER_ENABLED@
 config.vulkan_wrapper_library_dir = "@MLIR_VULKAN_WRAPPER_LIBRARY_DIR@"
 config.enable_vulkan_runner = @MLIR_VULKAN_RUNNER_ENABLED@
 
 # Support substitution of the tools_dir with user parameters. This is
 # used when we can't determine the tool dir at configuration time.
 try:
     config.llvm_tools_dir = config.llvm_tools_dir % lit_config.params
     config.llvm_shlib_dir = config.llvm_shlib_dir % lit_config.params

mlir/test/mlir-rocm-runner/gpu-to-hsaco.mlir

+// RUN: mlir-rocm-runner %s --shared-libs=%rocm_wrapper_library_dir/librocm-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s
+
+func @other_func(%arg0 : f32, %arg1 : memref<?xf32>) {
+  %cst = constant 1 : index
+  %cst2 = dim %arg1, 0 : memref<?xf32>
+  gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst, %grid_z = %cst)
+             threads(%tx, %ty, %tz) in (%block_x = %cst2, %block_y = %cst, %block_z = %cst) {
+    store %arg0, %arg1[%tx] : memref<?xf32>
+    gpu.terminator
+  }
+  return
+}
+
+// CHECK: [1, 1, 1, 1, 1]
+func @main() {
+  %arg0 = alloc() : memref<5xf32>
+  %21 = constant 5 : i32
+  %22 = memref_cast %arg0 : memref<5xf32> to memref<?xf32>
+  %cast = memref_cast %22 : memref<?xf32> to memref<*xf32>
+  call @mgpuMemHostRegisterFloat(%cast) : (memref<*xf32>) -> ()
+  %23 = memref_cast %22 : memref<?xf32> to memref<*xf32>
+  call @print_memref_f32(%23) : (memref<*xf32>) -> ()
+  %24 = constant 1.0 : f32
+  %25 = call @mgpuMemGetDeviceMemRef1dFloat(%22) : (memref<?xf32>) -> (memref<?xf32>)
+  call @other_func(%24, %25) : (f32, memref<?xf32>) -> ()
+  call @print_memref_f32(%23) : (memref<*xf32>) -> ()
+  return
+}
+
+func @mgpuMemHostRegisterFloat(%ptr : memref<*xf32>)
+func @mgpuMemGetDeviceMemRef1dFloat(%ptr : memref<?xf32>) -> (memref<?xf32>)
+func @print_memref_f32(%ptr : memref<*xf32>)

mlir/test/mlir-rocm-runner/lit.local.cfg

+if not config.enable_rocm_runner:
+  config.unsupported = True

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

+// RUN: mlir-rocm-runner %s --print-ir-after-all --shared-libs=%rocm_wrapper_library_dir/librocm-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s --dump-input=always
+
+// CHECK: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
+func @main() {
+  %arg = alloc() : memref<13xi32>
+  %dst = memref_cast %arg : memref<13xi32> to memref<?xi32>
+  %one = constant 1 : index
+  %sx = dim %dst, 0 : memref<?xi32>
+  %cast_dst = memref_cast %dst : memref<?xi32> to memref<*xi32>
+  call @mgpuMemHostRegisterInt32(%cast_dst) : (memref<*xi32>) -> ()
+  %dst_device = call @mgpuMemGetDeviceMemRef1dInt32(%dst) : (memref<?xi32>) -> (memref<?xi32>)
+  gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %one, %grid_y = %one, %grid_z = %one)
+             threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %one, %block_z = %one) {
+    %t0 = index_cast %tx : index to i32
+    store %t0, %dst_device[%tx] : memref<?xi32>
+    gpu.terminator
+  }
+  gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %one, %grid_y = %one, %grid_z = %one)
+             threads(%tx, %ty, %tz) in (%block_x = %sx, %block_y = %one, %block_z = %one) {
+    %t0 = index_cast %tx : index to i32
+    store %t0, %dst_device[%tx] : memref<?xi32>
+    gpu.terminator
+  }
+  call @print_memref_i32(%cast_dst) : (memref<*xi32>) -> ()
+  return
+}
+
+func @mgpuMemHostRegisterInt32(%ptr : memref<*xi32>)
+func @mgpuMemGetDeviceMemRef1dInt32(%ptr : memref<?xi32>) -> (memref<?xi32>)
+func @print_memref_i32(%ptr : memref<*xi32>)

mlir/test/mlir-rocm-runner/vecadd.mlir

+// RUN: mlir-rocm-runner %s --shared-libs=%rocm_wrapper_library_dir/librocm-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s
+
+func @vecadd(%arg0 : memref<?xf32>, %arg1 : memref<?xf32>, %arg2 : memref<?xf32>) {
+  %cst = constant 1 : index
+  %cst2 = dim %arg0, 0 : memref<?xf32>
+  gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst, %grid_z = %cst)
+             threads(%tx, %ty, %tz) in (%block_x = %cst2, %block_y = %cst, %block_z = %cst) {
+    %a = load %arg0[%tx] : memref<?xf32>
+    %b = load %arg1[%tx] : memref<?xf32>
+    %c = addf %a, %b : f32
+    store %c, %arg2[%tx] : memref<?xf32>
+    gpu.terminator
+  }
+  return
+}
+
+// CHECK: [2.46, 2.46, 2.46, 2.46, 2.46]
+func @main() {
+  %0 = alloc() : memref<5xf32>
+  %1 = alloc() : memref<5xf32>
+  %2 = alloc() : memref<5xf32>
+  %3 = memref_cast %0 : memref<5xf32> to memref<?xf32>
+  %4 = memref_cast %1 : memref<5xf32> to memref<?xf32>
+  %5 = memref_cast %2 : memref<5xf32> to memref<?xf32>
+  %6 = memref_cast %3 : memref<?xf32> to memref<*xf32>
+  %7 = memref_cast %4 : memref<?xf32> to memref<*xf32>
+  %8 = memref_cast %5 : memref<?xf32> to memref<*xf32>
+  call @mgpuMemHostRegisterFloat(%6) : (memref<*xf32>) -> ()
+  call @mgpuMemHostRegisterFloat(%7) : (memref<*xf32>) -> ()
+  call @mgpuMemHostRegisterFloat(%8) : (memref<*xf32>) -> ()
+  %9 = call @mgpuMemGetDeviceMemRef1dFloat(%3) : (memref<?xf32>) -> (memref<?xf32>)
+  %10 = call @mgpuMemGetDeviceMemRef1dFloat(%4) : (memref<?xf32>) -> (memref<?xf32>)
+  %11 = call @mgpuMemGetDeviceMemRef1dFloat(%5) : (memref<?xf32>) -> (memref<?xf32>)
+
+  call @vecadd(%9, %10, %11) : (memref<?xf32>, memref<?xf32>, memref<?xf32>) -> ()
+  call @print_memref_f32(%8) : (memref<*xf32>) -> ()
+  return
+}
+
+func @mgpuMemHostRegisterFloat(%ptr : memref<*xf32>)
+func @mgpuMemGetDeviceMemRef1dFloat(%ptr : memref<?xf32>) -> (memref<?xf32>)
+func @print_memref_f32(%ptr : memref<*xf32>)

mlir/tools/CMakeLists.txt

 add_subdirectory(mlir-cuda-runner)
 add_subdirectory(mlir-cpu-runner)
 add_subdirectory(mlir-linalg-ods-gen)
 add_subdirectory(mlir-opt)
+add_subdirectory(mlir-rocm-runner)
 add_subdirectory(mlir-translate)
 add_subdirectory(mlir-vulkan-runner)
 add_subdirectory(mlir-shlib)

mlir/tools/mlir-rocm-runner/CMakeLists.txt

+set(LLVM_OPTIONAL_SOURCES
+  rocm-runtime-wrappers.cpp
+  mlir-rocm-runner.cpp
+  )
+
+if(MLIR_ROCM_RUNNER_ENABLED)
+  if (NOT ("AMDGPU" IN_LIST LLVM_TARGETS_TO_BUILD))
+    message(SEND_ERROR
+      "Building the mlir rocm runner requires the AMDGPU backend")
+  endif()
+
+  # lld header files.
+  include_directories(${MLIR_SOURCE_DIR}/../lld/include)
+
+  # 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")
+    else()
+      set(ROCM_PATH $ENV{ROCM_PATH} CACHE PATH "Path to which ROCm has been installed")
+    endif()
+    set(HIP_PATH "${ROCM_PATH}/hip" CACHE PATH " Path to which HIP has been installed")
+  endif()
+  set(CMAKE_MODULE_PATH "${HIP_PATH}/cmake" ${CMAKE_MODULE_PATH})
+  find_package(HIP)
+  if (NOT HIP_FOUND)
+    message(SEND_ERROR "Build the mlir rocm runner requires a working ROCm and HIP install")
+  else()
+    message(STATUS "ROCm HIP version: ${HIP_VERSION}")
+  endif()
+
+  # Locate HIP runtime library.
+  find_library(ROCM_RUNTIME_LIBRARY hip_hcc)
+
+  # Set HIP compile-time flags.
+  add_definitions(-D__HIP_PLATFORM_HCC__)
+
+  add_llvm_library(rocm-runtime-wrappers SHARED
+    rocm-runtime-wrappers.cpp
+  )
+  target_include_directories(rocm-runtime-wrappers
+    PRIVATE
+    "${HIP_PATH}/../include"
+    "${HIP_PATH}/include"
+    LLVMSupport
+  )
+  target_link_libraries(rocm-runtime-wrappers
+    PUBLIC
+    LLVMSupport
+    ${ROCM_RUNTIME_LIBRARY}
+  )
+
+  get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+  get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
+  set(LIBS
+    ${dialect_libs}
+    ${conversion_libs}
+    lldCommon
+    lldDriver
+    lldELF
+    LLVMCore
+    LLVMLTO
+    LLVMOption
+    LLVMSupport
+    MLIRJitRunner
+    MLIRAnalysis
+    MLIREDSC
+    MLIRExecutionEngine
+    MLIRIR
+    MLIRParser
+    MLIRROCDLIR
+    MLIRSupport
+    MLIRTargetLLVMIR
+    MLIRTargetROCDLIR
+    MLIRTransforms
+    MLIRTranslation
+    ${ROCM_RUNTIME_LIBRARY}
+  )
+
+  # Manually expand the target library, since our MLIR libraries
+  # aren't plugged into the LLVM dependency tracking. If we don't
+  # do this then we can't insert the CodeGen library after ourselves
+  llvm_expand_pseudo_components(TARGET_LIBS AllTargetsCodeGens AllTargetsAsmParsers)
+  # Prepend LLVM in front of every target, this is how the library
+  # are named with CMake
+  SET(targets_to_link)
+  FOREACH(t ${TARGET_LIBS})
+    LIST(APPEND targets_to_link "LLVM${t}")
+  ENDFOREACH(t)
+
+  add_llvm_tool(mlir-rocm-runner
+    mlir-rocm-runner.cpp
+
+    DEPENDS
+    rocm-runtime-wrappers
+    )
+  llvm_update_compile_flags(mlir-rocm-runner)
+  target_link_libraries(mlir-rocm-runner PRIVATE ${LIBS} ${targets_to_link})
+
+endif()

mlir/tools/mlir-rocm-runner/mlir-rocm-runner.cpp

+//===- mlir-rocm-runner.cpp - MLIR ROCM 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 ROCDL/LLVM IR before JIT-compiling and executing the
+// latter.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/STLExtras.h"
+
+#include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
+#include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.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/ROCDLDialect.h"
+#include "mlir/ExecutionEngine/JitRunner.h"
+#include "mlir/ExecutionEngine/OptUtils.h"
+#include "mlir/IR/Function.h"
+#include "mlir/IR/Module.h"
+#include "mlir/InitAllDialects.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Support/FileUtilities.h"
+#include "mlir/Target/ROCDLIR.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/Passes.h"
+#include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/TargetRegistry.h"
+#include "llvm/Support/TargetSelect.h"
+
+// MC headers.
+#include "llvm/MC/MCAsmBackend.h"
+#include "llvm/MC/MCAsmInfo.h"
+#include "llvm/MC/MCCodeEmitter.h"
+#include "llvm/MC/MCContext.h"
+#include "llvm/MC/MCInstPrinter.h"
+#include "llvm/MC/MCInstrInfo.h"
+#include "llvm/MC/MCObjectFileInfo.h"
+#include "llvm/MC/MCObjectWriter.h"
+#include "llvm/MC/MCParser/AsmLexer.h"
+#include "llvm/MC/MCParser/MCTargetAsmParser.h"
+#include "llvm/MC/MCRegisterInfo.h"
+#include "llvm/MC/MCStreamer.h"
+#include "llvm/MC/MCSubtargetInfo.h"
+#include "llvm/MC/MCTargetOptionsCommandFlags.h"
+
+// lld headers.
+#include "lld/Common/Driver.h"
+
+using namespace mlir;
+using namespace llvm;
+
+using Blob = SmallVector<char, 0>;
+
+static cl::opt<std::string> tripleName("triple", cl::desc("target triple"),
+                                       cl::value_desc("triple string"),
+                                       cl::init("amdgcn-amd-amdhsa"));
+
+static cl::opt<std::string> targetChip("target", cl::desc("target chip"),
+                                       cl::value_desc("AMDGPU ISA version"),
+                                       cl::init("gfx900"));
+
+static cl::opt<std::string> features("feature", cl::desc("target features"),
+                                     cl::value_desc("AMDGPU target features"),
+                                     cl::init("-code-object-v3"));
+
+static LogicalResult assembleIsa(const std::string isa, StringRef name,
+                                 Blob &result) {
+  raw_svector_ostream os(result);
+  std::unique_ptr<buffer_ostream> bos = std::make_unique<buffer_ostream>(os);
+
+  std::string error;
+  Triple theTriple(Triple::normalize(tripleName));
+  const Target *theTarget =
+      TargetRegistry::lookupTarget(theTriple.normalize(), error);
+  if (!theTarget) {
+    WithColor::error(errs(), name) << error;
+    return failure();
+  }
+
+  SourceMgr srcMgr;
+  srcMgr.AddNewSourceBuffer(MemoryBuffer::getMemBuffer(isa), SMLoc());
+
+  const MCTargetOptions mcOptions;
+  std::unique_ptr<MCRegisterInfo> mri(theTarget->createMCRegInfo(tripleName));
+  std::unique_ptr<MCAsmInfo> mai(
+      theTarget->createMCAsmInfo(*mri, tripleName, mcOptions));
+  mai->setRelaxELFRelocations(true);
+
+  MCObjectFileInfo mofi;
+  MCContext ctx(mai.get(), mri.get(), &mofi, &srcMgr, &mcOptions);
+  mofi.InitMCObjectFileInfo(theTriple, false, ctx, false);
+
+  SmallString<128> cwd;
+  if (!sys::fs::current_path(cwd))
+    ctx.setCompilationDir(cwd);
+
+  std::unique_ptr<MCStreamer> mcStreamer;
+  std::unique_ptr<MCInstrInfo> mcii(theTarget->createMCInstrInfo());
+  std::unique_ptr<MCSubtargetInfo> sti(
+      theTarget->createMCSubtargetInfo(tripleName, targetChip, features));
+
+  MCCodeEmitter *ce = theTarget->createMCCodeEmitter(*mcii, *mri, ctx);
+  MCAsmBackend *mab = theTarget->createMCAsmBackend(*sti, *mri, mcOptions);
+  mcStreamer.reset(theTarget->createMCObjectStreamer(
+      theTriple, ctx, std::unique_ptr<MCAsmBackend>(mab),
+      mab->createObjectWriter(os), std::unique_ptr<MCCodeEmitter>(ce), *sti,
+      mcOptions.MCRelaxAll, mcOptions.MCIncrementalLinkerCompatible,
+      /*DWARFMustBeAtTheEnd*/ false));
+  mcStreamer->setUseAssemblerInfoForParsing(true);
+
+  std::unique_ptr<MCAsmParser> parser(
+      createMCAsmParser(srcMgr, ctx, *mcStreamer, *mai));
+  std::unique_ptr<MCTargetAsmParser> tap(
+      theTarget->createMCAsmParser(*sti, *parser, *mcii, mcOptions));
+
+  if (!tap) {
+    WithColor::error(errs(), name) << "assembler initialization error.\n";
+    return failure();
+  }
+
+  parser->setTargetParser(*tap);
+  parser->Run(false);
+
+  return success();
+}
+
+static LogicalResult createHsaco(Blob &isaBlob, StringRef name,
+                                 Blob &hsacoBlob) {
+  // Save the ISA binary to a temp file.
+  int tempIsaBinaryFd = -1;
+  SmallString<128> tempIsaBinaryFilename;
+  std::error_code ec = sys::fs::createTemporaryFile(
+      "kernel", "o", tempIsaBinaryFd, tempIsaBinaryFilename);
+  if (ec) {
+    WithColor::error(errs(), name)
+        << "temporary file for ISA binary creation error.\n";
+    return failure();
+  }
+  raw_fd_ostream tempIsaBinaryOs(tempIsaBinaryFd, true);
+  tempIsaBinaryOs << isaBlob;
+  tempIsaBinaryOs.close();
+
+  // Create a temp file for HSA code object.
+  int tempHsacoFD = -1;
+  SmallString<128> temphsacoFilename;
+  ec = sys::fs::createTemporaryFile("kernel", "hsaco", tempHsacoFD,
+                                    temphsacoFilename);
+  if (ec) {
+    WithColor::error(errs(), name)
+        << "temporary file for HSA code object creation error.\n";
+    return failure();
+  }
+
+  // Invoke lld. Expect a true return value from lld.
+  bool ret = lld::elf::link({"ld.lld", "-shared", tempIsaBinaryFilename.c_str(),
+                             "-o", temphsacoFilename.c_str()},
+                            /*canEarlyExit=*/false, llvm::outs(), llvm::errs());
+  if (!ret) {
+    WithColor::error(errs(), name) << "lld invocation error.\n";
+    return failure();
+  }
+
+  // Load the HSA code object.
+  auto hsacoFile = mlir::openInputFile(temphsacoFilename);
+  if (!hsacoFile) {
+    WithColor::error(errs(), name)
+        << "read HSA code object from temp file error.\n";
+    return failure();
+  }
+  hsacoBlob.assign(hsacoFile->getBuffer().begin(),
+                   hsacoFile->getBuffer().end());
+
+  // Remove temp files.
+  ec = sys::fs::remove(tempIsaBinaryFilename);
+  if (ec) {
+    WithColor::error(errs(), name) << "remove ISA binary temp file error.\n";
+    return failure();
+  }
+  ec = sys::fs::remove(temphsacoFilename);
+  if (ec) {
+    WithColor::error(errs(), name)
+        << "remove HSA code object temp file error.\n";
+    return failure();
+  }
+  return success();
+}
+
+static std::unique_ptr<llvm::Module> compileModuleToROCDLIR(Operation *m) {
+  auto llvmModule = translateModuleToROCDLIR(m);
+  // TODO(whchung): Link with ROCm-Device-Libs in case needed (ex: the Module
+  // depends on math functions).
+  return llvmModule;
+}
+
+static OwnedBlob compileISAToHsaco(const std::string isa, Location loc,
+                                   StringRef name) {
+  // ISA -> ISA in binary form via MC.
+  // Use lld to create HSA code object.
+  Blob isaBlob;
+  Blob hsacoBlob;
+
+  if (succeeded(assembleIsa(isa, name, isaBlob)) &&
+      succeeded(createHsaco(isaBlob, name, hsacoBlob)))
+    return std::make_unique<std::vector<char>>(hsacoBlob.begin(),
+                                               hsacoBlob.end());
+
+  WithColor::error(errs(), name) << "producing HSA code object error.\n";
+  return {};
+}
+
+static LogicalResult runMLIRPasses(ModuleOp m) {
+  PassManager pm(m.getContext());
+  applyPassManagerCLOptions(pm);
+
+  pm.addPass(createGpuKernelOutliningPass());
+  auto &kernelPm = pm.nest<gpu::GPUModuleOp>();
+  kernelPm.addPass(createStripDebugInfoPass());
+  kernelPm.addPass(createLowerGpuOpsToROCDLOpsPass());
+  kernelPm.addPass(createConvertGPUKernelToBlobPass(
+      compileModuleToROCDLIR, compileISAToHsaco, tripleName, targetChip,
+      features, /*gpuBinaryAnnotation=*/"rocdl.hsaco"));
+  pm.addPass(createLowerToLLVMPass());
+  pm.addPass(createConvertGpuLaunchFuncToGpuRuntimeCallsPass(
+      /*gpuBinaryAnnotation=*/"rocdl.hsaco"));
+
+  return pm.run(m);
+}
+
+int main(int argc, char **argv) {
+  registerPassManagerCLOptions();
+  mlir::registerAllDialects();
+  llvm::InitLLVM y(argc, argv);
+  llvm::InitializeAllTargetInfos();
+  llvm::InitializeAllTargetMCs();
+  llvm::InitializeAllAsmParsers();
+
+  // Initialize LLVM AMDGPU backend.
+  LLVMInitializeAMDGPUTarget();
+  LLVMInitializeAMDGPUTargetInfo();
+  LLVMInitializeAMDGPUTargetMC();
+  LLVMInitializeAMDGPUAsmPrinter();
+
+  mlir::initializeLLVMPasses();
+  return mlir::JitRunnerMain(argc, argv, &runMLIRPasses);
+}

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

+//===- rocm-runtime-wrappers.cpp - MLIR ROCM runner wrapper library -------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Implements C wrappers around the ROCM library for easy linking in ORC jit.
+// Also adds some debugging helpers that are helpful when writing MLIR code to
+// run on GPUs.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cassert>
+#include <numeric>
+
+#include "mlir/ExecutionEngine/CRunnerUtils.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include "hip/hip_runtime.h"
+
+namespace {
+int32_t reportErrorIfAny(hipError_t result, const char *where) {
+  if (result != hipSuccess) {
+    llvm::errs() << "HIP failed with " << result << " in " << where << "\n";
+  }
+  return result;
+}
+} // anonymous namespace
+
+extern "C" int32_t mgpuModuleLoad(void **module, void *data) {
+  int32_t err = reportErrorIfAny(
+      hipModuleLoadData(reinterpret_cast<hipModule_t *>(module), data),
+      "ModuleLoad");
+  return err;
+}
+
+extern "C" int32_t mgpuModuleGetFunction(void **function, void *module,
+                                         const char *name) {
+  return reportErrorIfAny(
+      hipModuleGetFunction(reinterpret_cast<hipFunction_t *>(function),
+                           reinterpret_cast<hipModule_t>(module), name),
+      "GetFunction");
+}
+
+// The wrapper uses intptr_t instead of ROCM'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 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(
+      hipModuleLaunchKernel(reinterpret_cast<hipFunction_t>(function), gridX,
+                            gridY, gridZ, blockX, blockY, blockZ, smem,
+                            reinterpret_cast<hipStream_t>(stream), params,
+                            extra),
+      "LaunchKernel");
+}
+
+extern "C" void *mgpuGetStreamHelper() {
+  hipStream_t stream;
+  reportErrorIfAny(hipStreamCreate(&stream), "StreamCreate");
+  return stream;
+}
+
+extern "C" int32_t mgpuStreamSynchronize(void *stream) {
+  return reportErrorIfAny(
+      hipStreamSynchronize(reinterpret_cast<hipStream_t>(stream)),
+      "StreamSync");
+}
+
+/// Helper functions for writing mlir example code
+
+// Allows to register byte array with the ROCM runtime. Helpful until we have
+// transfer functions implemented.
+extern "C" void mgpuMemHostRegister(void *ptr, uint64_t sizeBytes) {
+  reportErrorIfAny(hipHostRegister(ptr, sizeBytes, /*flags=*/0),
+                   "MemHostRegister");
+}
+
+// Allows to register a MemRef with the ROCM runtime. Initializes array with
+// value. Helpful until we have transfer functions implemented.
+template <typename T>
+void mgpuMemHostRegisterMemRef(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);
+  mgpuMemHostRegister(pointer, count * sizeof(T));
+}
+
+extern "C" void mgpuMemHostRegisterFloat(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);
+  mgpuMemHostRegisterMemRef(desc->data + desc->offset, sizes, strides, 1.23f);
+}
+
+extern "C" void mgpuMemHostRegisterInt32(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);
+  mgpuMemHostRegisterMemRef(desc->data + desc->offset, sizes, strides, 123);
+}
+
+template <typename T>
+void mgpuMemGetDevicePointer(T *hostPtr, T **devicePtr) {
+  reportErrorIfAny(hipSetDevice(0), "hipSetDevice");
+  reportErrorIfAny(
+      hipHostGetDevicePointer((void **)devicePtr, hostPtr, /*flags=*/0),
+      "hipHostGetDevicePointer");
+}
+
+extern "C" StridedMemRefType<float, 1>
+mgpuMemGetDeviceMemRef1dFloat(float *allocated, float *aligned, int64_t offset,
+                              int64_t size, int64_t stride) {
+  float *devicePtr = nullptr;
+  mgpuMemGetDevicePointer(aligned, &devicePtr);
+  return {devicePtr, devicePtr, offset, {size}, {stride}};
+}
+
+extern "C" StridedMemRefType<int32_t, 1>
+mgpuMemGetDeviceMemRef1dInt32(int32_t *allocated, int32_t *aligned,
+                              int64_t offset, int64_t size, int64_t stride) {
+  int32_t *devicePtr = nullptr;
+  mgpuMemGetDevicePointer(aligned, &devicePtr);
+  return {devicePtr, devicePtr, offset, {size}, {stride}};
+}