[ROCM] Build rocm device bc files from source and installation setup. (…

…iree-org#15096)

* Adds a scheme whereby compiler targets can tack dependencies onto
libIREECompiler.so by adding targets to the global property
`IREE_COMPILER_DYLIB_DEPENDS` (allowing arbitrary required artifacts to
be guaranteed to be generated when the main shared library is built).
* Added a global property `IREE_COMPILER_DYLIB_RELPATHS` which always
contains the relative paths of all artifacts that must exist along with
libIREECompiler.so (including itself). This is used for installing into
bindings (like Python), etc without all sorts of conditional code going
the wrong way in the project.
* Changed the main shared library to
`CMAKE_PLATFORM_NO_VERSIONED_SONAME` mode always (we were manually
setting this only for pip builds of Python bindings). This makes a lot
of things easier and more consistent. In reality, if ever major
versioning, we will bake a version number into the name, since that is
portable to Windows, versus relying on Unix symlink/versioning arcana.
* Set `CMAKE_RUNTIME_OUTPUT_DIRECTORY`,
`CMAKE_LIBRARY_OUTPUT_DIRECTORY`, and `CMAKE_ARCHIVE_OUTPUT_DIRECTORY`
in the top level compiler/ directory vs having them set inconsistently
throughout the tree. As a side-effect, all libraries in the compiler now
land in lib/. Since they all have fully qualified names, this is fine.
* Define a `IREE_COMPILER_DYLIB_DIR` and
`IREE_COMPILER_DYLIB_INSTALL_PREFIX` at the top level compiler/ dir vs
ad-hoc code throughout checking to see if Windows (in order to special
case DLLs installing next to binaries on that platform).
* Reworked the Python compiler build so that it always builds a symlink
tree in the build directory that matches how it gets installed. While
not done in this PR, this should let us remove all of the weird special
casing for how to find libraries if in a build tree vs deployed. Might
need a bit more massaging to delete the code.
* Adds a `findPlatformLibDirectory(StringRef platformName)` helper that
can always find the `iree_platform_libs/{platformName}` directory
located adjacent to the shared-library/DLL. Uses this to initialize the
flag for finding ROCM bc files. We may want to do this with other
prebuilts too.
* Builds the ROCM bitcode files with our clang and includes them
adjacent to the shared-lib/DLL under iree_platform_libs/rocm, which
should always be available in all build/deploy/install modality.

ci-extra: build_test_all_windows

---------

Co-authored-by: Stanley Winata <[email protected]>
Co-authored-by: Stella Laurenzo <[email protected]>
Co-authored-by: Scott Todd <[email protected]>

.gitmodules

@@ -44,3 +44,6 @@
 [submodule "third_party/torch-mlir"]
 	path = third_party/torch-mlir
 	url = https://github.com/shark-infra/torch-mlir.git
+[submodule "third_party/ROCm-Device-Libs"]
+	path = third_party/ROCm-Device-Libs
+	url = https://github.com/shark-infra/ROCm-Device-Libs

CMakeLists.txt

@@ -812,6 +812,9 @@ else()
   # Add default external projects.
   iree_llvm_add_external_project(mlir-iree-dialects ${CMAKE_CURRENT_SOURCE_DIR}/llvm-external-projects/iree-dialects)
   iree_llvm_add_external_project(stablehlo ${CMAKE_CURRENT_SOURCE_DIR}/third_party/stablehlo)
+  if(IREE_TARGET_BACKEND_ROCM)
+    iree_llvm_add_external_project(ROCm-Device-Libs ${CMAKE_CURRENT_SOURCE_DIR}/third_party/ROCm-Device-Libs)
+  endif()
 
   # Ensure that LLVM-based dependencies needed for testing are included.
   add_dependencies(iree-test-deps FileCheck)

compiler/CMakeLists.txt

@@ -4,6 +4,23 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+# On Windows, DLLs go to the runtime directory and import
+# libraries go to the lib directory.
+# TODO: We should really be dumping binaries into bin/ not
+# tools/. This must line up with binaries built this way because
+# DLLs must be in the same directory as the binary.
+# See: https://github.com/openxla/iree/issues/11297
+set(CMAKE_RUNTIME_OUTPUT_DIRECTORY "${PROJECT_BINARY_DIR}/tools")
+set(CMAKE_LIBRARY_OUTPUT_DIRECTORY "${PROJECT_BINARY_DIR}/lib")
+set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY "${PROJECT_BINARY_DIR}/lib")
+if(WIN32)
+  set(IREE_COMPILER_DYLIB_DIR "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}")
+  set(IREE_COMPILER_DYLIB_INSTALL_PREFIX "bin/")
+else()
+  set(IREE_COMPILER_DYLIB_DIR "${CMAKE_LIBRARY_OUTPUT_DIRECTORY}")
+  set(IREE_COMPILER_DYLIB_INSTALL_PREFIX "lib/")
+endif()
+
 # Always build the C bindings, since the API is available apart from
 # actually building the compiler.
 add_subdirectory(bindings/c)
@@ -22,4 +39,11 @@ if(IREE_BUILD_COMPILER)
     configure_file(setup.py setup.py @ONLY)
     add_subdirectory(bindings/python)
   endif()
+
+  # Post processing.
+  get_property(_iree_compiler_dylib_depends GLOBAL PROPERTY IREE_COMPILER_DYLIB_DEPENDS)
+  if(_iree_compiler_dylib_depends)
+    add_dependencies(iree_compiler_API_SharedImpl ${_iree_compiler_dylib_depends})
+  endif()
 endif()
+

compiler/bindings/python/CMakeLists.txt

@@ -214,42 +214,56 @@ if(TARGET lld)
   )
 endif()
 
-# Install shared libraries that the extension depends on. This uses
-# CMake's defer feature to evaluate the install directive once everything
-# has been evaluated (because there is no guarantee that this file evaluates
-# before the API libraries are defined). While deferred calls are generally
-# fragile, this install is purely a leaf feature (with no other calls
-# depending on its sequencing), so this use is okay.
-# We defer it to the compiler/ directory so that cmake_install.cmake targets
-# are all self contained to the compiler, which is the most common spanning
-# parent.
-cmake_language(EVAL CODE "
-cmake_language(DEFER DIRECTORY \"${IREE_SOURCE_DIR}/compiler\"
-  CALL install
-  TARGETS
-    iree_compiler_API_SharedImpl
-  DESTINATION \"${_PYTHON_INSTALL_PREFIX}/iree/compiler/_mlir_libs\"
-)
-")
-
-# On Windows, the IREECompiler.dll must be physically adjacent to the python
-# extensions. This is only an issue for the build tree (the install tree
-# places things appropriately). There is a lot of unavoidable fragility
-# here. Most notably, the actual name of the DLL must be static and computed
-# the same as where it is built.
-if(WIN32)
-  set(_local_dll_copy
-    "${_PYTHON_BUILD_PREFIX}/iree/compiler/_mlir_libs/IREECompiler.dll")
-  add_custom_command(
-    OUTPUT "${_local_dll_copy}"
-    COMMAND "${CMAKE_COMMAND}" -E copy_if_different
-      "$<TARGET_FILE:iree_compiler_API_SharedImpl>"
-      "${_local_dll_copy}"
-    DEPENDS iree_compiler_API_SharedImpl
+################################################################################
+# libIREECompiler.so dylib tree
+# We copy the compiler shared library and its related files to both the build
+# and install tree for consistency. Note that on Windows, this is load bearing
+# since the DLL must be colocated in the same directory. On Unix, CMake's
+# automatic RPATH handling in the build tree will cause binaries to resolve
+# against the project level library, so this is just cosmetic, but it does
+# have the side effect of making the Python build tree relocatable and making
+# logic that expects to find the compiler libraries in a consistent place
+# simple. Since a symlink is used on Unix and it is not the versioned dylib, 
+# the realpath will still be the build-tree wide shared library. This also means
+# that any tools or files that must be colocated with the dylib will come from
+# the project-wide location.
+################################################################################
+
+# Copy compiler dylib files into the python _mlir_libs tree so that the
+# binaries can find them with a relative rpath.
+get_property(_dylib_relpaths GLOBAL PROPERTY IREE_COMPILER_DYLIB_RELPATHS)
+set(_dylib_copy_commands)
+set(_dylib_target_files)
+set(_dylib_src_files)
+foreach (_dylib_relpath ${_dylib_relpaths})
+  set(_dylib_srcpath "${IREE_COMPILER_DYLIB_DIR}/${_dylib_relpath}")
+  set(_dylib_outputpath "${_PYTHON_BUILD_PREFIX}/iree/compiler/_mlir_libs/${_dylib_relpath}")
+  cmake_path(GET _dylib_outputpath PARENT_PATH _parent_dir)
+  file(MAKE_DIRECTORY "${_parent_dir}")
+  list(APPEND _dylib_copy_commands
+    COMMAND ${CMAKE_COMMAND} -E create_symlink
+      "${_dylib_srcpath}" "${_dylib_outputpath}"
   )
-  add_custom_target(IREEPythonCopyDLL ALL
-    DEPENDS "${_local_dll_copy}")
-endif()
+  list(APPEND _dylib_target_files "${_dylib_outputpath}")
+  list(APPEND _dylib_src_files "${_dylib_srcpath}")
+  cmake_path(GET _dylib_relpath PARENT_PATH _dylib_install_destination)
+  install(FILES "${_dylib_srcpath}"
+    DESTINATION "${_PYTHON_INSTALL_PREFIX}/iree/compiler/_mlir_libs/${_dylib_install_destination}"
+  )
+endforeach()
+
+add_custom_command(
+  OUTPUT ${_dylib_target_files}
+  DEPENDS ${_dylib_src_files}
+  POST_BUILD
+    ${_dylib_copy_commands}
+)
+
+add_custom_target(IREECompilerPythonDylibFiles
+  DEPENDS ${_dylib_target_files}
+)
+
+add_dependencies(IREECompilerPythonModules IREECompilerPythonDylibFiles)
 
 ################################################################################
 # Subdirectories

compiler/src/iree/compiler/API/CMakeLists.txt

@@ -139,12 +139,16 @@ foreach(_object_lib ${_EXPORT_OBJECT_LIBS})
   # another level of fix upstream if you like pain.
   list(APPEND _EXPORT_OBJECT_DEPS "$<GENEX_EVAL:$<GENEX_EVAL:$<GENEX_EVAL:$<GENEX_EVAL:$<TARGET_PROPERTY:${_object_lib},LINK_LIBRARIES>>>>>")
 endforeach()
-if(NOT GENERATOR_IS_MULTI_CONFIG)
-  # NOTE: multi-config generators (xcode, visual studio, ninja in multi-config
-  # mode, etc) don't support writing files with fixed names with contents that
-  # differ based on the available variants.
-  file(GENERATE OUTPUT export_objects_debug.txt CONTENT "OBJECTS:${_EXPORT_OBJECT_SRCS}\n\nDEPS:${_EXPORT_OBJECT_DEPS}")
-endif()
+
+# UNCOMMENT TO DEBUG WHAT IS GOING ON.
+# file(GENERATE OUTPUT export_objects_debug.txt CONTENT "OBJECTS:${_EXPORT_OBJECT_SRCS}\n\nDEPS:${_EXPORT_OBJECT_DEPS}")
+
+# Disable .so.0 style naming/linking. In order to be consistent across platforms
+# and bindings, we will embed a major version in the library name when it is time.
+# OS packagers can re-enable this or take other steps to get canonical shared
+# library structures.
+set(CMAKE_PLATFORM_NO_VERSIONED_SONAME ON)
+
 iree_cc_library(
   SHARED
   NAME
@@ -195,18 +199,18 @@ set_target_properties(
     OUTPUT_NAME "IREECompiler"
     VERSION "0"
     SOVERSION "0"
-
-    LIBRARY_OUTPUT_DIRECTORY "${PROJECT_BINARY_DIR}/lib"
-    # On Windows, DLLs go to the runtime directory and import
-    # libraries go to the lib directory.
-    # TODO: We should really be dumping binaries into bin/ not
-    # tools/. This must line up with binaries built this way because
-    # DLLs must be in the same directory as the binary.
-    # See: https://github.com/openxla/iree/issues/11297
-    RUNTIME_OUTPUT_DIRECTORY "${PROJECT_BINARY_DIR}/tools"
-    ARCHIVE_OUTPUT_DIRECTORY "${PROJECT_BINARY_DIR}/lib"
 )
 
+# Signal the main shared library/dll as part of the bundle of files that
+# must be included with compiler dylibs.
+# What we really want to be using here is the TARGET_LINKER_FILE_NAME
+# generator expression. But since this is used in places where target
+# dependent generator expressions are not legal, we guesstimate it with
+# platform specific logic.
+set(_target_linker_file_name "${CMAKE_SHARED_LIBRARY_PREFIX}IREECompiler${CMAKE_SHARED_LIBRARY_SUFFIX}")
+set_property(GLOBAL APPEND PROPERTY 
+  IREE_COMPILER_DYLIB_RELPATHS "${_target_linker_file_name}")
+
 install(
   TARGETS iree_compiler_API_SharedImpl
   COMPONENT Compiler

compiler/src/iree/compiler/Dialect/HAL/Target/ROCM/CMakeLists.txt

@@ -48,3 +48,54 @@ iree_cc_library(
 )
 
 ### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
+
+# Query for list of device libs to build.
+get_property(AMD_DEVICE_LIBS GLOBAL PROPERTY AMD_DEVICE_LIBS)
+
+set(_platform_lib_reldir "iree_platform_libs/rocm")
+file(MAKE_DIRECTORY "${IREE_COMPILER_DYLIB_DIR}/${_platform_lib_reldir}")
+
+# Transform device lib targets to the generated bc file of each.
+set(_all_device_bc_deps)
+set(_all_device_bc_copy_commands)
+set(_all_device_bc_files)
+foreach (_device_lib_target ${AMD_DEVICE_LIBS})
+  get_target_property(_device_basename ${_device_lib_target} ARCHIVE_OUTPUT_NAME)
+  get_target_property(_device_output_path ${_device_lib_target} OUTPUT_NAME)
+  set(_device_bc_relpath "${_platform_lib_reldir}/${_device_basename}.bc")
+
+  # Note this bc file as being part of the bundle that must be included with
+  # the compiler dylib.
+  set_property(GLOBAL APPEND PROPERTY IREE_COMPILER_DYLIB_RELPATHS "${_device_bc_relpath}")
+
+  list(APPEND _all_device_bc_deps "${_device_bc_path}")
+  list(APPEND _all_device_bc_files "${IREE_COMPILER_DYLIB_DIR}/${_device_bc_relpath}")
+
+  # Copy to lib/ tree.
+  list(APPEND _all_device_bc_copy_commands
+    COMMAND ${CMAKE_COMMAND} -E copy 
+      "${_device_output_path}" 
+      "${IREE_COMPILER_DYLIB_DIR}/${_device_bc_relpath}"
+  )
+endforeach()
+
+# Generate a custom target with all file level dependencies and commands to
+# copy to our build tree locations.
+# Our GenDeviceLibs target depends on all of the defined device lib targets.
+add_custom_command(
+  OUTPUT ${_all_device_bc_files}
+  DEPENDS ${_all_device_bc_deps}
+  POST_BUILD
+    ${_all_device_bc_copy_commands}
+)
+add_custom_target(iree_compiler_Dialect_HAL_Target_ROCM_GenDeviceLibs
+  DEPENDS ${_all_device_bc_files}
+)
+
+# Ensure that the device libs are built when the compiler dylib is built.
+set_property(GLOBAL APPEND PROPERTY IREE_COMPILER_DYLIB_DEPENDS 
+  iree_compiler_Dialect_HAL_Target_ROCM_GenDeviceLibs)
+
+# Install.
+install(FILES ${_all_device_bc_files} 
+  DESTINATION "${IREE_COMPILER_DYLIB_INSTALL_PREFIX}/${_platform_lib_reldir}")

compiler/src/iree/compiler/Dialect/HAL/Target/ROCM/ROCMTarget.cpp

@@ -12,11 +12,14 @@
 #include "iree/compiler/Codegen/LLVMGPU/Passes.h"
 #include "iree/compiler/Dialect/HAL/Target/TargetRegistry.h"
 #include "iree/compiler/Utils/FlatbufferUtils.h"
+#include "iree/compiler/Utils/ToolUtils.h"
 #include "iree/schemas/rocm_executable_def_builder.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/Module.h"
 #include "llvm/MC/TargetRegistry.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/Path.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Target/TargetMachine.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
@@ -37,10 +40,9 @@ static llvm::cl::opt<bool>
                  llvm::cl::desc("Whether to try Linking to AMD Bitcodes"),
                  llvm::cl::init(false));
 
-static llvm::cl::opt<std::string>
-    clROCMBitcodeDir("iree-rocm-bc-dir",
-                     llvm::cl::desc("Directory of ROCM Bitcode"),
-                     llvm::cl::init("/opt/rocm/amdgcn/bitcode"));
+static llvm::cl::opt<std::string> clROCMBitcodeDir(
+    "iree-rocm-bc-dir", llvm::cl::desc("Directory of ROCM Bitcode"),
+    llvm::cl::init(mlir::iree_compiler::findPlatformLibDirectory("rocm")));
 
 namespace mlir {
 namespace iree_compiler {
@@ -205,9 +207,15 @@ class ROCMTargetBackend final : public TargetBackend {
     iree_hal_rocm_ExecutableDef_start_as_root(builder);
 
     // Link module to Device Library
+    std::string rocmBitcodeDir = clROCMBitcodeDir;
     if (clROCMLinkBC) {
-      linkROCDLIfNecessary(llvmModule.get(), clROCMTargetChip,
-                           clROCMBitcodeDir);
+      if (clROCMBitcodeDir.empty()) {
+        return variantOp.emitError()
+               << "cannot find ROCM bitcode files. Check your installation "
+                  "consistency and in the worst case, set --iree-rocm-bc-dir= "
+                  "to an explicit location on your system.";
+      }
+      linkROCDLIfNecessary(llvmModule.get(), clROCMTargetChip, rocmBitcodeDir);
     }
 
     // Serialize hsaco kernel into the binary that we will embed in the

compiler/src/iree/compiler/Utils/ToolUtils.cpp

@@ -262,5 +262,22 @@ std::string findTool(std::string toolName) {
   return findTool(toolNames);
 }
 
+std::string findPlatformLibDirectory(StringRef platformName) {
+  std::string dylibPath = getCurrentDylibPath();
+  if (dylibPath.empty())
+    return {};
+
+  SmallString<256> path(dylibPath);
+  llvm::sys::path::remove_filename(path);
+  llvm::sys::path::append(path, "iree_platform_libs", platformName);
+  if (!llvm::sys::fs::is_directory(path))
+    return {};
+  llvm::sys::fs::make_absolute(path);
+  (void)llvm::sys::path::remove_dots(path, /*remove_dot_dot=*/true);
+
+  std::string pathStr(path);
+  return pathStr;
+}
+
 } // namespace iree_compiler
 } // namespace mlir

compiler/src/iree/compiler/Utils/ToolUtils.h

@@ -40,6 +40,20 @@ std::string findToolInEnvironment(SmallVector<std::string> toolNames);
 std::string findTool(SmallVector<std::string> toolNames);
 std::string findTool(std::string toolName);
 
+// Finds a bundled directory containing platform libraries for the given
+// platform name, returning an empty string if not found. We store bundled
+// platform libraries in a directory like:
+//   iree_platform_libs/{platformName}
+// adjacent to the shared library hosting the compiler (i.e. this entry
+// point). On a Posix system, this will typically be in a "lib" dir and
+// on Windows, it will be adjacent to executables (i.e. a "bin" or "tools"
+// dir). Note that if installed to a system library directory on a Posix
+// system, this would be under something like:
+//   /usr/lib/iree_platform_libs/{platformName}
+// This is not atypical to how other dependencies are located in a qualified
+// lib directory.
+std::string findPlatformLibDirectory(StringRef platformName);
+
 } // namespace iree_compiler
 } // namespace mlir