Roll-up of changes needed to support the nvgpu out of tree project. (#…

…12888)

* Adds CMake scoped IREE_PACKAGE_ROOT_DIR and IREE_PACKAGE_ROOT_PREFIX
to replace hard-coded path to namespace logic in iree_package_ns (and
uses within IREE/removes the special casing).
* Adds support for `BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_ABOVE_THIS_LINE
` to bazel_to_cmake. Been carrying this patch for a while and ended up
needing it.
* Further generalizes bazel_to_cmake target resolution so that it is
customizable as needed out of tree.
* Moves the `iree::runtime::src::defs` target to `iree::defs` and puts
it in the right place in the tree to avoid special casing.
* Ditto but for `iree::compiler::src::defs`
* Adds a bit more logging to `_DEBUG_IREE_PACKAGE_NAME` mode.
* Makes iree_tablegen_library consult a scoped
`IREE_COMPILER_TABLEGEN_INCLUDE_DIRS` var for additional include
directories (makes it possible to use out of tree).
* Adds `NVPTXDesc` and `NVPTXInfo` targets to HAL_Target_CUDA. No idea
why this was triggering for me but was getting undefined deps. Must have
been coming in elsewhere in a more full featured build.
* Fixes iree-opt initialization sequence with respect to command line
options. Also fixed the test which should have been verifying this.
* Fixed pytype issue in bazel_to_cmake that could theoretically happen.

Fixes build issues related to the out of tree build for
openxla/community#71

CMakeLists.txt

@@ -512,6 +512,11 @@ include(iree_static_linker_test)
 include(iree_fetch_artifact)
 include(iree_run_module_test)
 
+# Default any sub-tree which doesn't provide its own package namespacing
+# to derive it relative to this directory and prefixed with iree/.
+set(IREE_PACKAGE_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}")
+set(IREE_PACKAGE_ROOT_PREFIX "iree/")
+
 set(CMAKE_POSITION_INDEPENDENT_CODE TRUE)
 
 #-------------------------------------------------------------------------------

build_tools/bazel_to_cmake/bazel_to_cmake.py

@@ -62,7 +62,8 @@
     r"bazel[\s_]*to[\s_]*cmake[\s_]*:?[\s_]*do[\s_]*not[\s_]*edit",
     flags=re.IGNORECASE)
 
-PRESERVE_TAG = "### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###"
+PRESERVE_ABOVE_TAG = "### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_ABOVE_THIS_LINE ###"
+PRESERVE_BELOW_TAG = "### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###"
 REPO_CFG_FILE = ".bazel_to_cmake.cfg.py"
 REPO_CFG_MODULE_NAME = "bazel_to_cmake_repo_config"
 
@@ -140,12 +141,17 @@ def setup_environment():
   # Dynamically load the config file as a module.
   orig_dont_write_bytecode = sys.dont_write_bytecode
   sys.dont_write_bytecode = True  # Don't generate __pycache__ dir
-  spec = importlib.util.spec_from_file_location(
-      REPO_CFG_MODULE_NAME, os.path.join(repo_root, REPO_CFG_FILE))
-  repo_cfg = importlib.util.module_from_spec(spec)
-  sys.modules[REPO_CFG_MODULE_NAME] = repo_cfg
-  spec.loader.exec_module(repo_cfg)
-  sys.dont_write_bytecode = orig_dont_write_bytecode
+  repo_cfg_path = os.path.join(repo_root, REPO_CFG_FILE)
+  spec = importlib.util.spec_from_file_location(REPO_CFG_MODULE_NAME,
+                                                repo_cfg_path)
+  if spec and spec.loader:
+    repo_cfg = importlib.util.module_from_spec(spec)
+    sys.modules[REPO_CFG_MODULE_NAME] = repo_cfg
+    spec.loader.exec_module(repo_cfg)
+    sys.dont_write_bytecode = orig_dont_write_bytecode
+  else:
+    print(f"INTERNAL ERROR: Could not evaluate {repo_cfg_path} as module")
+    sys.exit(1)
 
 
 def repo_relpath(path):
@@ -229,26 +235,34 @@ def convert_directory(directory_path, write_files, allow_partial_conversion,
   ] + ["#" * 80])
 
   old_lines = []
-  preserved_footer_lines = ["\n" + PRESERVE_TAG + "\n"]
+  possible_preserved_header_lines = []
+  preserved_footer_lines = ["\n" + PRESERVE_BELOW_TAG + "\n"]
 
   # Read CMakeLists.txt and check if it has the auto-generated header.
+  found_preserve_below_tag = False
+  found_preserve_above_tag = False
   if os.path.isfile(cmakelists_file_path):
     found_autogeneration_tag = False
-    found_preserve_tag = False
     with open(cmakelists_file_path) as f:
       old_lines = f.readlines()
 
     for line in old_lines:
+      if not found_preserve_above_tag:
+        possible_preserved_header_lines.append(line)
       if not found_autogeneration_tag and autogeneration_tag in line:
         found_autogeneration_tag = True
-      if not found_preserve_tag and PRESERVE_TAG in line:
-        found_preserve_tag = True
-      elif found_preserve_tag:
+      if not found_preserve_below_tag and PRESERVE_BELOW_TAG in line:
+        found_preserve_below_tag = True
+      elif not found_preserve_above_tag and PRESERVE_ABOVE_TAG in line:
+        found_preserve_above_tag = True
+      elif found_preserve_below_tag:
         preserved_footer_lines.append(line)
     if not found_autogeneration_tag:
       if verbosity >= 1:
         log(f"Skipped. Did not find autogeneration line.", indent=2)
       return Status.SKIPPED
+  preserved_header = ("".join(possible_preserved_header_lines)
+                      if found_preserve_above_tag else "")
   preserved_footer = "".join(preserved_footer_lines)
 
   # Read the Bazel BUILD file and interpret it.
@@ -276,7 +290,8 @@ def convert_directory(directory_path, write_files, allow_partial_conversion,
         f"Reason: `{type(e).__name__}: {e}`",
         indent=2)
     return Status.FAILED
-  converted_content = header + converted_build_file + preserved_footer
+  converted_content = (preserved_header + header + converted_build_file +
+                       preserved_footer)
   if write_files:
     with open(cmakelists_file_path, "wt") as cmakelists_file:
       cmakelists_file.write(converted_content)

build_tools/bazel_to_cmake/bazel_to_cmake_targets.py

@@ -58,6 +58,7 @@ def __init__(self, repo_map: Dict[str, str]):
         "@llvm-project//llvm:FileCheck": ["FileCheck"],
         "@llvm-project//llvm:not": ["not"],
         "@llvm-project//llvm:llvm-link": ["${IREE_LLVM_LINK_TARGET}"],
+        "@llvm-project//llvm:NVPTXUtilsAndDesc": ["LLVMNVPTXDesc",],
 
         # MLIR
         "@llvm-project//mlir:AllPassesAndDialects": ["MLIRAllDialects"],
@@ -212,6 +213,7 @@ def convert_target(self, target):
     Raises:
       KeyError: No conversion was found for the target.
     """
+    iree_core_repo = self._repo_alias("@iree_core")
     if target in self._explicit_target_mapping:
       return self._explicit_target_mapping[target]
     if target.startswith("@llvm-project//llvm"):
@@ -220,9 +222,21 @@ def convert_target(self, target):
       return self._convert_mlir_target(target)
     if target.startswith("@iree_cuda//"):
       return self._convert_iree_cuda_target(target)
+    if target.startswith(f"{iree_core_repo}//"):
+      return self._convert_iree_core_target(target)
     if target.startswith("@"):
       raise KeyError(f"No conversion found for target '{target}'")
 
+    # Pass through package-relative targets
+    #   :target_name
+    #   file_name.txt
+    if target.startswith(":") or (":" not in target and
+                                  not target.startswith("/")):
+      return [self._convert_to_cmake_path(target)]
+
+    return self._convert_unmatched_target(target)
+
+  def _convert_iree_core_target(self, target):
     iree_core_repo = self._repo_alias("@iree_core")
     if target.startswith(
         f"{iree_core_repo}//llvm-external-projects/iree-dialects"):
@@ -251,12 +265,6 @@ def convert_target(self, target):
     if m:
       return [self._convert_to_cmake_path(m.group(1))]
 
-    # Pass through package-relative targets
-    #   :target_name
-    #   file_name.txt
-    if target.startswith(":") or ":" not in target:
-      return [self._convert_to_cmake_path(target)]
-
     return self._convert_unmatched_target(target)
 
   def _convert_unmatched_target(self, target: str) -> str:

build_tools/cmake/iree_c_module.cmake

@@ -101,7 +101,7 @@ function(iree_c_module)
       "${_TESTONLY_ARG}"
     DEPS
       # Include paths and options for the runtime sources.
-      iree::runtime::src::defs
+      iree::defs
   )
 
   if(_RULE_NO_RUNTIME)

build_tools/cmake/iree_cc_binary.cmake

@@ -68,12 +68,19 @@ function(iree_cc_binary)
     set(_NAME "${_PACKAGE_NAME}_${_RULE_NAME}")
   endif()
 
+  if(_DEBUG_IREE_PACKAGE_NAME)
+    message(STATUS "  : iree_cc_binary(${_NAME})")
+  endif()
+
   add_executable(${_NAME} "")
 
   if(NOT "${_PACKAGE_NS}" STREQUAL "")
     # Alias the iree_package_name binary to iree::package::name.
     # This lets us more clearly map to Bazel and makes it possible to
     # disambiguate the underscores in paths vs. the separators.
+    if(_DEBUG_IREE_PACKAGE_NAME)
+      message(STATUS "  + alias ${_PACKAGE_NS}::${_RULE_NAME}")
+    endif()
     add_executable(${_PACKAGE_NS}::${_RULE_NAME} ALIAS ${_NAME})
 
     # If the binary name matches the package then treat it as a default. For

build_tools/cmake/iree_cc_library.cmake

@@ -82,6 +82,10 @@ function(iree_cc_library)
   set(_NAME "${_PACKAGE_NAME}_${_RULE_NAME}")
   set(_OBJECTS_NAME ${_NAME}.objects)
 
+  if(_DEBUG_IREE_PACKAGE_NAME)
+    message(STATUS "  : iree_cc_library(${_NAME})")
+  endif()
+
   # Replace dependencies passed by ::name with iree::package::name
   list(TRANSFORM _RULE_DEPS REPLACE "^::" "${_PACKAGE_NS}::")
 
@@ -253,6 +257,9 @@ function(iree_cc_library)
   # Alias the iree_package_name library to iree::package::name.
   # This lets us more clearly map to Bazel and makes it possible to
   # disambiguate the underscores in paths vs. the separators.
+  if(_DEBUG_IREE_PACKAGE_NAME)
+    message(STATUS "  + alias ${_PACKAGE_NS}::${_RULE_NAME}")
+  endif()
   add_library(${_PACKAGE_NS}::${_RULE_NAME} ALIAS ${_NAME})
 
   if(NOT "${_PACKAGE_NS}" STREQUAL "")

build_tools/cmake/iree_cc_test.cmake

@@ -69,6 +69,10 @@ function(iree_cc_test)
   iree_package_ns(_PACKAGE_NS)
   set(_NAME "${_PACKAGE_NAME}_${_RULE_NAME}")
 
+  if(_DEBUG_IREE_PACKAGE_NAME)
+    message(STATUS "  : iree_cc_test(${_NAME})")
+    message(STATUS "  + alias ${_PACKAGE_NS}::${_RULE_NAME}")
+  endif()
   add_executable(${_NAME} "")
   # Alias the iree_package_name test binary to iree::package::name.
   # This lets us more clearly map to Bazel and makes it possible to

build_tools/cmake/iree_macros.cmake

@@ -131,40 +131,44 @@ endfunction()
 #   runtime/src/iree/base/CMakeLists.txt -> iree::base
 #   tests/e2e/CMakeLists.txt -> iree::tests::e2e
 function(iree_package_ns PACKAGE_NS)
-  # Get the relative path of the current dir (i.e. runtime/src/iree/vm).
-  string(REPLACE ${IREE_ROOT_DIR} "" _IREE_RELATIVE_PATH ${CMAKE_CURRENT_LIST_DIR})
-  string(SUBSTRING ${_IREE_RELATIVE_PATH} 1 -1 _IREE_RELATIVE_PATH)
-
-  if(NOT ${CMAKE_CURRENT_LIST_DIR} MATCHES "^${IREE_ROOT_DIR}/.*")
-    # Function is being called from outside IREE. Use the source-relative path.
-    # Please check the README.md to see the potential risk.
-    string(REPLACE ${PROJECT_SOURCE_DIR} "" _SOURCE_RELATIVE_PATH ${CMAKE_CURRENT_LIST_DIR})
-    string(SUBSTRING ${_SOURCE_RELATIVE_PATH} 1 -1 _SOURCE_RELATIVE_PATH)
-    set(_PACKAGE "${_SOURCE_RELATIVE_PATH}")
-
-  # If changing the directory/package mapping rules, please also implement
-  # the corresponding rule in:
-  #   build_tools/bazel_to_cmake/bazel_to_cmake_targets.py
-  # Some sub-trees form their own roots for package purposes. Rewrite them.
-  elseif(_IREE_RELATIVE_PATH MATCHES "^compiler/src/(.*)")
-    # compiler/src/iree/compiler -> iree/compiler
-    set(_PACKAGE "${CMAKE_MATCH_1}")
-  elseif(_IREE_RELATIVE_PATH MATCHES "^runtime/src/(.*)")
-    # runtime/src/iree/base -> iree/base
-    set(_PACKAGE "${CMAKE_MATCH_1}")
-  elseif(_IREE_RELATIVE_PATH MATCHES "^tools$")
-    # Special case for tools/ -> "" (empty string)
-    # For example, tools/iree-compile -> iree-compile (no namespace)
-    set(_PACKAGE "")
+  if(DEFINED IREE_PACKAGE_ROOT_DIR)
+    # If an enclosing package root dir is set, then the package is just the
+    # relative part after that.
+    cmake_path(RELATIVE_PATH CMAKE_CURRENT_LIST_DIR
+      BASE_DIRECTORY "${IREE_PACKAGE_ROOT_DIR}"
+      OUTPUT_VARIABLE _PACKAGE)
+    if(_PACKAGE STREQUAL ".")
+      set(_PACKAGE "")
+    endif()
+    if(IREE_PACKAGE_ROOT_PREFIX)
+      set(_PACKAGE "${IREE_PACKAGE_ROOT_PREFIX}${_PACKAGE}")
+    endif()
   else()
-    # Default to prefixing with iree/
-    set(_PACKAGE "iree/${_IREE_RELATIVE_PATH}")
+    # Get the relative path of the current dir (i.e. runtime/src/iree/vm).
+    string(REPLACE ${IREE_ROOT_DIR} "" _IREE_RELATIVE_PATH ${CMAKE_CURRENT_LIST_DIR})
+    string(SUBSTRING ${_IREE_RELATIVE_PATH} 1 -1 _IREE_RELATIVE_PATH)
+
+    if(NOT ${CMAKE_CURRENT_LIST_DIR} MATCHES "^${IREE_ROOT_DIR}/.*")
+      # Function is being called from outside IREE. Use the source-relative path.
+      # Please check the README.md to see the potential risk.
+      string(REPLACE ${PROJECT_SOURCE_DIR} "" _SOURCE_RELATIVE_PATH ${CMAKE_CURRENT_LIST_DIR})
+      string(SUBSTRING ${_SOURCE_RELATIVE_PATH} 1 -1 _SOURCE_RELATIVE_PATH)
+      set(_PACKAGE "${_SOURCE_RELATIVE_PATH}")
+
+    # If changing the directory/package mapping rules, please also implement
+    # the corresponding rule in:
+    #   build_tools/bazel_to_cmake/bazel_to_cmake_targets.py
+    # Some sub-trees form their own roots for package purposes. Rewrite them.
+    else()
+      message(SEND_ERROR "iree_package_ns(): Could not determine package for ${CMAKE_CURRENT_LIST_DIR}")
+      set(_PACKAGE "iree/unknown")
+    endif()
   endif()
 
   string(REPLACE "/" "::" _PACKAGE_NS "${_PACKAGE}")
 
   if(_DEBUG_IREE_PACKAGE_NAME)
-    message(STATUS "iree_package_ns(): map ${_IREE_RELATIVE_PATH} -> ${_PACKAGE_NS}")
+    message(STATUS "iree_package_ns(): map ${CMAKE_CURRENT_LIST_DIR} -> ${_PACKAGE_NS}")
   endif()
 
   set(${PACKAGE_NS} ${_PACKAGE_NS} PARENT_SCOPE)

build_tools/cmake/iree_tablegen_library.cmake

@@ -38,6 +38,9 @@ function(iree_tablegen_library)
     "${IREE_SOURCE_DIR}/compiler/src"
     "${IREE_BINARY_DIR}/compiler/src"
   )
+  if(DEFINED IREE_COMPILER_TABLEGEN_INCLUDE_DIRS)
+    list(APPEND _INCLUDE_DIRS ${IREE_COMPILER_TABLEGEN_INCLUDE_DIRS})
+  endif()
   list(APPEND _INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR})
   list(TRANSFORM _INCLUDE_DIRS PREPEND "-I")
   set(_OUTPUTS)

compiler/src/CMakeLists.txt

@@ -4,18 +4,12 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-iree_cc_library(
-  NAME
-    defs
-  INCLUDES
-    ${CMAKE_CURRENT_SOURCE_DIR}
-    ${CMAKE_CURRENT_BINARY_DIR}
-  PUBLIC
-)
+set(IREE_PACKAGE_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}")
+set(IREE_PACKAGE_ROOT_PREFIX "")
 
 # Configures all iree_cc_* targets to take this implicit dep,
 # which provides common includes and copts for the tree.
-set(IREE_IMPLICIT_DEFS_CC_DEPS iree::compiler::src::defs)
+set(IREE_IMPLICIT_DEFS_CC_DEPS iree::compiler::defs)
 
 # Force enable BUILD_SHARED_LIBS for the compiler if instructed.
 set(_IREE_ORIG_BUILD_SHARED_LIBS ${BUILD_SHARED_LIBS})

compiler/src/iree/compiler/API/Internal/IREEOptToolEntryPoint.cpp

@@ -66,6 +66,11 @@ static LogicalResult ireeOptMainFromCL(int argc, char **argv,
     llvm::errs() << "error: Failed to initialize IREE compiler plugins\n";
     return failure();
   }
+  // Initialization that applies to all global plugins must be done prior
+  // to CL parsing.
+  pluginManager.globalInitialize();
+  pluginManager.registerPasses();
+  pluginManager.registerGlobalDialects(registry);
   pluginManager.initializeCLI();
 
   // Parse pass names in main to ensure static initialization completed.
@@ -75,9 +80,6 @@ static LogicalResult ireeOptMainFromCL(int argc, char **argv,
   // The local binder is meant for overriding session-level options, but for
   // tools like this it is unused.
   auto localBinder = mlir::iree_compiler::OptionsBinder::local();
-  pluginManager.globalInitialize();
-  pluginManager.registerPasses();
-  pluginManager.registerGlobalDialects(registry);
   mlir::iree_compiler::PluginManagerSession pluginSession(
       pluginManager, localBinder, pluginManagerOptions);
   if (failed(pluginSession.initializePlugins())) return failure();

compiler/src/iree/compiler/CMakeLists.txt

@@ -1,3 +1,20 @@
+# Copyright 2023 The IREE Authors
+#
+# Licensed 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
+
+iree_cc_library(
+  NAME
+    defs
+  INCLUDES
+    ${CMAKE_CURRENT_SOURCE_DIR}/../..
+    ${CMAKE_CURRENT_BINARY_DIR}/../..
+  PUBLIC
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_ABOVE_THIS_LINE ###
+
 ################################################################################
 # Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from           #
 # iree/compiler/BUILD                                                          #

compiler/src/iree/compiler/Dialect/HAL/Target/CUDA/BUILD.bazel

@@ -47,6 +47,8 @@ iree_compiler_cc_library(
         "@llvm-project//llvm:Linker",
         "@llvm-project//llvm:MC",
         "@llvm-project//llvm:NVPTXCodeGen",
+        "@llvm-project//llvm:NVPTXInfo",
+        "@llvm-project//llvm:NVPTXUtilsAndDesc",
         "@llvm-project//llvm:Passes",
         "@llvm-project//llvm:Support",
         "@llvm-project//llvm:Target",

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

@@ -29,6 +29,8 @@ iree_cc_library(
     LLVMLinker
     LLVMMC
     LLVMNVPTXCodeGen
+    LLVMNVPTXDesc
+    LLVMNVPTXInfo
     LLVMPasses
     LLVMSupport
     LLVMTarget

runtime/src/CMakeLists.txt

@@ -4,17 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-iree_cc_library(
-  NAME
-    defs
-  INCLUDES
-    ${CMAKE_CURRENT_SOURCE_DIR}
-    ${CMAKE_CURRENT_BINARY_DIR}
-  PUBLIC
-)
-
-# Configures all iree_cc_* targets to take this implicit dep,
-# which provides common includes and copts for the tree.
-set(IREE_IMPLICIT_DEFS_CC_DEPS iree::runtime::src::defs)
+set(IREE_PACKAGE_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}")
+set(IREE_PACKAGE_ROOT_PREFIX "")
 
 add_subdirectory(iree)