[ROCm] Adding pass to generate the HSACO binary blob from the GPU kernel function

CMakeLists.txt

@@ -44,6 +44,13 @@ endif()
 
 set(MLIR_CUDA_RUNNER_ENABLED 0 CACHE BOOL "Enable building the mlir CUDA runner")
 
+# Build the ROCM conversions if the AMDGPU backend is available
+if ("AMDGPU" IN_LIST LLVM_TARGETS_TO_BUILD)
+  set(MLIR_ROCM_CONVERSIONS_ENABLED 1)
+else()
+  set(MLIR_ROCM_CONVERSIONS_ENABLED 0)
+endif()
+
 include_directories( "include")
 include_directories( ${MLIR_INCLUDE_DIR})
 

include/mlir/CMakeLists.txt

@@ -2,3 +2,4 @@ add_subdirectory(Analysis)
 add_subdirectory(Dialect)
 add_subdirectory(EDSC)
 add_subdirectory(Transforms)
+add_subdirectory(Conversion/GPUToROCM)

include/mlir/Conversion/GPUToROCM/CMakeLists.txt

@@ -0,0 +1,33 @@
+if(MLIR_ROCM_CONVERSIONS_ENABLED)
+
+  # Check whether the ROCm installation dir exists
+  set(ROCM_INSTALL_DIR "/opt/rocm" CACHE STRING "ROCm installation directory")
+  if (EXISTS ${ROCM_INSTALL_DIR})
+    message("-- ROCm Install Dir - ${ROCM_INSTALL_DIR}")
+  else()
+    message(SEND_ERROR "-- NOT FOUND : ROCm Install Dir - ${ROCM_INSTALL_DIR}")
+  endif()
+
+  # Check whether the ROCm device library dir exists
+  set(ROCM_DEVICE_LIB_DIR ${ROCM_INSTALL_DIR}/lib)
+  if (EXISTS ${ROCM_DEVICE_LIB_DIR})
+    message("-- ROCm Device Library Dir - ${ROCM_DEVICE_LIB_DIR}")
+  else ()
+    message(SEND_ERROR "-- NOT FOUND : ROCm Device Library Dir - ${ROCM_DEVICE_LIB_DIR}")
+  endif()
+
+  # Check whether the ROCm HCC linker exists
+  set(ROCM_HCC_LINKER ${ROCM_INSTALL_DIR}/hcc/bin/ld.lld)
+  if (EXISTS ${ROCM_HCC_LINKER})
+    message("-- ROCm HCC Linker - ${ROCM_HCC_LINKER}")
+  else ()
+    message(SEND_ERROR "-- NOT FOUND : ROCm HCC Linker - ${ROCM_HCC_LINKER}")
+  endif()
+
+  # Generate the ROCm Configuration header file
+  configure_file(
+    "${CMAKE_CURRENT_SOURCE_DIR}/ROCMConfig.h.in"
+    "${CMAKE_CURRENT_BINARY_DIR}/ROCMConfig.h"
+    )
+
+endif()

include/mlir/Conversion/GPUToROCM/GPUToROCMPass.h

@@ -0,0 +1,92 @@
+//===- GPUToROCmPass.h - MLIR ROCm runtime support --------------*- C++ -*-===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+#ifndef MLIR_CONVERSION_GPUTOROCM_GPUTOROCMPASS_H_
+#define MLIR_CONVERSION_GPUTOROCM_GPUTOROCMPASS_H_
+
+#include <functional>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "mlir/Conversion/GPUToROCM/ROCMConfig.h"
+
+namespace mlir {
+
+namespace rocm {
+
+/// string constants used by the ROCM backend
+static constexpr const char *kHSACOAnnotation = "amdgpu.hsaco";
+static constexpr const char *kHSACOGetterAnnotation = "amdgpu.hsacogetter";
+static constexpr const char *kHSACOGetterSuffix = "_hsaco";
+static constexpr const char *kHSACOStorageSuffix = "_hsaco_cst";
+
+/// enum to represent the AMD GPU versions supported by the ROCM backend
+enum class AMDGPUVersion { GFX900 };
+
+/// enum to represent the HSA Code Object versions supported by the ROCM backend
+enum class HSACOVersion { V3 };
+
+/// Configurable parameters for generating the HSACO blobs from GPU Kernels
+struct HSACOGeneratorConfig {
+
+  /// Constructor - sets the default values for the configurable parameters
+  HSACOGeneratorConfig(bool isTestMode)
+      : testMode(isTestMode), amdgpuVersion(AMDGPUVersion::GFX900),
+        hsacoVersion(HSACOVersion::V3), rocdlDir(ROCM_DEVICE_LIB_DIR),
+        linkerPath(ROCM_HCC_LINKER) {}
+
+  /// testMode == true will result in skipping the HASCO generation process, and
+  /// simply return the string "HSACO" as the HSACO blob
+  bool testMode;
+
+  /// the AMDGPU version for which to generate the HSACO
+  AMDGPUVersion amdgpuVersion;
+
+  /// the code object version for the generated HSACO
+  HSACOVersion hsacoVersion;
+
+  /// the directory containing the ROCDL bitcode libraries
+  std::string rocdlDir;
+
+  /// the path the ld.lld linker to use when generating the HSACO
+  std::string linkerPath;
+};
+
+} // namespace rocm
+
+// unique pointer to the HSA Code Object (which is stored as char vector)
+using OwnedHSACO = std::unique_ptr<std::vector<char>>;
+
+class ModuleOp;
+template <typename T>
+class OpPassBase;
+
+/// Creates a pass to convert kernel functions into HSA Code Object blobs.
+///
+/// This transformation takes the body of each function that is annotated with
+/// the amdgpu_kernel calling convention, copies it to a new LLVM module,
+/// compiles the module with help of the AMDGPU backend to GCN ISA, and then
+/// invokes lld to produce a binary blob in HSA Code Object format. Such blob
+/// is then attached as a string attribute named 'amdgpu.hsaco' 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<OpPassBase<ModuleOp>> createConvertGPUKernelToHSACOPass(
+    rocm::HSACOGeneratorConfig hsacoGeneratorConfig);
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_GPUTOROCM_GPUTOROCMPASS_H_

include/mlir/Conversion/GPUToROCM/ROCMConfig.h.in

@@ -0,0 +1,30 @@
+//===- ROCMConfig.h - ROCm Configuration Header -----------------*- C++ -*-===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+#ifndef MLIR_CONVERSION_GPUTOROCM_ROCMCONFIG_H_
+#define MLIR_CONVERSION_GPUTOROCM_ROCMCONFIG_H_
+
+/// The code to generate the HSACO binary blobs (corresponding the GPU kernels)
+/// assumes the presense of ROCm libraries/utilities. The location of these
+/// tools is configured via cmake
+
+/// Path to the ROCm Device Library dir in the ROCM install
+#cmakedefine ROCM_DEVICE_LIB_DIR "@ROCM_DEVICE_LIB_DIR@"
+
+/// Path to the HCC Linker in the ROCM install
+#cmakedefine ROCM_HCC_LINKER "@ROCM_HCC_LINKER@"
+
+#endif // MLIR_CONVERSION_GPUTOROCM_ROCMCONFIG_H_

lib/Conversion/CMakeLists.txt

@@ -1,5 +1,6 @@
 add_subdirectory(GPUToCUDA)
 add_subdirectory(GPUToNVVM)
+add_subdirectory(GPUToROCM)
 add_subdirectory(GPUToROCDL)
 add_subdirectory(GPUToSPIRV)
 add_subdirectory(LoopsToGPU)

lib/Conversion/GPUToROCM/CMakeLists.txt

@@ -0,0 +1,15 @@
+if(MLIR_ROCM_CONVERSIONS_ENABLED)
+  llvm_map_components_to_libnames(amdgpu "AMDGPU")
+
+  add_llvm_library(MLIRGPUtoROCMTransforms
+    ConvertKernelFuncToHSACO.cpp
+  )
+  target_link_libraries(MLIRGPUtoROCMTransforms
+    MLIRGPU
+    MLIRLLVMIR
+    MLIRROCDLIR
+    MLIRPass
+    MLIRTargetROCDLIR
+    ${amdgpu}
+  )
+endif()