[MLGO] Improved support for AOT cross-targeting scenarios

The tensorflow AOT compiler can cross-target, but it can't run on (for example) arm64. We added earlier support where the AOT-ed header and object would be built on a separate builder and then passed at build time to a build host where the AOT compiler can't run, but clang can be otherwise built. To simplify such scenarios given we now support more than one AOT-able case (regalloc and inliner), we make the AOT scenario centered on whether files are generated, case by case (this includes the "passed from a different builder" scenario). This means we shouldn't need an 'umbrella' LLVM_HAVE_TF_AOT, in favor of case by case control. A builder can opt out of an AOT case by passing that case's model path as `none`. Note that the overrides still take precedence. This patch controls conditional compilation with case-specific flags, which can be enabled locally, for the component where those are available. We still keep an overall flag for some tests. The 'development/training' mode is unchanged, because there the model is passed from the command line and interpreted. Differential Revision: https://reviews.llvm.org/D117752
2022-01-19 21:19:53 -08:00
parent 81d35f27dd
commit f29256a64a
12 changed files with 81 additions and 80 deletions
--- a/llvm/CMakeLists.txt
+++ b/llvm/CMakeLists.txt
@@ -891,13 +891,17 @@ if (NOT TENSORFLOW_AOT_PATH STREQUAL "")
    ARCHIVE DESTINATION lib${LLVM_LIBDIR_SUFFIX} COMPONENT tf_xla_runtime)
  set_property(GLOBAL APPEND PROPERTY LLVM_EXPORTS tf_xla_runtime)
  # Once we add more modules, we should handle this more automatically.
-  if (NOT DEFINED LLVM_INLINER_MODEL_PATH
+  if (DEFINED LLVM_OVERRIDE_MODEL_HEADER_INLINERSIZEMODEL)
+    set(LLVM_INLINER_MODEL_PATH "none")
+  elseif(NOT DEFINED LLVM_INLINER_MODEL_PATH
      OR "${LLVM_INLINER_MODEL_PATH}" STREQUAL ""
      OR "${LLVM_INLINER_MODEL_PATH}" STREQUAL "autogenerate")
    set(LLVM_INLINER_MODEL_PATH "autogenerate")
    set(LLVM_INLINER_MODEL_AUTOGENERATED 1)
  endif()
-  if (NOT DEFINED LLVM_RAEVICT_MODEL_PATH
+  if (DEFINED LLVM_OVERRIDE_MODEL_HEADER_REGALLOCEVICTMODEL)
+    set(LLVM_RAEVICT_MODEL_PATH "none")
+  elseif(NOT DEFINED LLVM_RAEVICT_MODEL_PATH
      OR "${LLVM_RAEVICT_MODEL_PATH}" STREQUAL ""
      OR "${LLVM_RAEVICT_MODEL_PATH}" STREQUAL "autogenerate")
    set(LLVM_RAEVICT_MODEL_PATH "autogenerate")
--- a/llvm/cmake/modules/TensorFlowCompile.cmake
+++ b/llvm/cmake/modules/TensorFlowCompile.cmake
@@ -44,35 +44,21 @@ endfunction()
 # Produce a pair of files called ${fname}.h and  ${fname}.o in the
 # ${CMAKE_CURRENT_BINARY_DIR}. The generated header will define a C++ class
 # called ${cpp_class} - which may be a namespace-qualified class name.
-function(tfcompile model tag_set signature_def_key fname cpp_class)
-  set(prefix ${CMAKE_CURRENT_BINARY_DIR}/${fname})
-  set(obj_file ${prefix}.o)
-  set(hdr_file ${prefix}.h)
-  string(TOUPPER ${fname} fname_allcaps)
-  set(override_header ${LLVM_OVERRIDE_MODEL_HEADER_${fname_allcaps}})
-  set(override_object ${LLVM_OVERRIDE_MODEL_OBJECT_${fname_allcaps}})
-  if (EXISTS "${override_header}" AND EXISTS "${override_object}")
-    configure_file(${override_header} ${hdr_file} COPYONLY)
-    configure_file(${override_object} ${obj_file} COPYONLY)
-    message("Using provided header "
-      ${hdr_file} " and object "   ${obj_file}
-      " files for model " ${model})
-  else()
-    tf_get_absolute_path(${model} ${CMAKE_CURRENT_BINARY_DIR} LLVM_ML_MODELS_ABSOLUTE)
-    message("Using model at " ${LLVM_ML_MODELS_ABSOLUTE})
-    add_custom_command(OUTPUT ${obj_file} ${hdr_file}
-      COMMAND ${TENSORFLOW_AOT_COMPILER} aot_compile_cpu
-            --multithreading false
-            --dir ${LLVM_ML_MODELS_ABSOLUTE}
-            --tag_set ${tag_set}
-            --signature_def_key ${signature_def_key}
-            --output_prefix ${prefix}
-            --cpp_class ${cpp_class}
-            --target_triple ${LLVM_HOST_TRIPLE}
-    )
-  endif()
+function(tf_compile model tag_set signature_def_key fname cpp_class hdr_file obj_file)
+  tf_get_absolute_path(${model} ${CMAKE_CURRENT_BINARY_DIR} LLVM_ML_MODELS_ABSOLUTE)
+  message("Using model at " ${LLVM_ML_MODELS_ABSOLUTE})
+  add_custom_command(OUTPUT ${obj_file} ${hdr_file}
+    COMMAND ${TENSORFLOW_AOT_COMPILER} aot_compile_cpu
+          --multithreading false
+          --dir ${LLVM_ML_MODELS_ABSOLUTE}
+          --tag_set ${tag_set}
+          --signature_def_key ${signature_def_key}
+          --output_prefix ${prefix}
+          --cpp_class ${cpp_class}
+          --target_triple ${LLVM_HOST_TRIPLE}
+  )

-  # Aggregate the objects so that results of different tfcompile calls may be
+  # Aggregate the objects so that results of different tf_compile calls may be
  # grouped into one target.
  set(GENERATED_OBJS ${GENERATED_OBJS} ${obj_file} PARENT_SCOPE)
  set_source_files_properties(${obj_file} PROPERTIES
@@ -82,36 +68,50 @@ function(tfcompile model tag_set signature_def_key fname cpp_class)
  set_source_files_properties(${hdr_file} PROPERTIES
    GENERATED 1)

-endfunction()
+  endfunction()

 function(tf_find_and_compile model default_url default_path test_model_generator tag_set signature_def_key fname cpp_class)
-  if ("${model}" STREQUAL "download")
-    # Crash if the user wants to download a model but a URL is set to "TO_BE_UPDATED"
-    if ("${default_url}" STREQUAL "TO_BE_UPDATED")
-        message(FATAL_ERROR "Default URL was set to 'download' but there is no model url currently specified in cmake - likely, the model interface recently changed, and so there is not a released model available.")
+  set(prefix ${CMAKE_CURRENT_BINARY_DIR}/${fname})
+  set(obj_file ${prefix}.o)
+  set(hdr_file ${prefix}.h)
+  string(TOUPPER ${fname} fname_allcaps)
+  set(override_header ${LLVM_OVERRIDE_MODEL_HEADER_${fname_allcaps}})
+  set(override_object ${LLVM_OVERRIDE_MODEL_OBJECT_${fname_allcaps}})
+  # If the user specified overrides, that indicates intent to use AOT and we
+  # don't care what the model path is
+  if (EXISTS "${override_header}" AND EXISTS "${override_object}")
+    configure_file(${override_header} ${hdr_file} COPYONLY)
+    configure_file(${override_object} ${obj_file} COPYONLY)
+    message(STATUS "Using provided header " ${hdr_file} " and object " ${obj_file} "
+      files for model " ${fname})  
+    set(GENERATED_OBJS ${GENERATED_OBJS} ${obj_file})
+    set(GENERATED_HEADERS ${GENERATED_HEADERS} ${hdr_file})
+  elseif("${model}" STREQUAL "none")
+    message(STATUS "Will skip enabling mlgo for ${fname}")
+    return()
+  else()
+    if ("${model}" STREQUAL "download")
+      # Crash if the user wants to download a model but a URL is set to "TO_BE_UPDATED"
+      if ("${default_url}" STREQUAL "TO_BE_UPDATED")
+          message(FATAL_ERROR "Default URL was set to 'download' but there is no"
+          " model url currently specified in cmake - likely, the model interface"
+          " recently changed, and so there is not a released model available.")
+      endif()
+
+      set(model ${default_url})
    endif()

-    set(model ${default_url})
+    if ("${model}" STREQUAL "autogenerate")
+      set(model ${default_path}-autogenerated)
+      generate_mock_model(${test_model_generator} ${model})
+    endif()
+
+    tf_get_model(${model} LLVM_ML_MODELS_ABSOLUTE)
+    tf_compile(${LLVM_ML_MODELS_ABSOLUTE} ${tag_set} ${signature_def_key} ${fname} ${cpp_class} ${hdr_file} ${obj_file})
  endif()

-  if ("${model}" STREQUAL "autogenerate")
-    set(model ${default_path}-autogenerated)  
-    generate_mock_model(${test_model_generator} ${model})
-  endif()
-
-  tf_get_model(${model} LLVM_ML_MODELS_ABSOLUTE)
-  tfcompile(${LLVM_ML_MODELS_ABSOLUTE} ${tag_set} ${signature_def_key} ${fname} ${cpp_class})
-
-  set(GENERATED_OBJS ${GENERATED_OBJS} ${obj_file} PARENT_SCOPE)
-  set_source_files_properties(${obj_file} PROPERTIES
-    GENERATED 1 EXTERNAL_OBJECT 1)
-
-  set(GENERATED_HEADERS ${GENERATED_HEADERS} ${hdr_file} PARENT_SCOPE)
-  set_source_files_properties(${hdr_file} PROPERTIES
-    GENERATED 1)
-
  set(GeneratedMLSources ${GeneratedMLSources} ${GENERATED_HEADERS} PARENT_SCOPE)
  set(MLDeps ${MLDeps} tf_xla_runtime PARENT_SCOPE)
  set(MLLinkDeps ${MLLinkDeps} tf_xla_runtime ${GENERATED_OBJS} PARENT_SCOPE)
-
+  add_definitions(-DLLVM_HAVE_TF_AOT_${fname_allcaps})
 endfunction()
--- a/llvm/include/llvm/Analysis/InlineAdvisor.h
+++ b/llvm/include/llvm/Analysis/InlineAdvisor.h
@@ -246,16 +246,12 @@ public:
  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
 };

-#ifdef LLVM_HAVE_TF_AOT
 std::unique_ptr<InlineAdvisor>
 getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM);
-#endif

-#ifdef LLVM_HAVE_TF_API
 std::unique_ptr<InlineAdvisor>
 getDevelopmentModeAdvisor(Module &M, ModuleAnalysisManager &MAM,
                          std::function<bool(CallBase &)> GetDefaultAdvice);
-#endif

 // Default (manual policy) decision making helper APIs. Shared with the legacy
 // pass manager inliner.
--- a/llvm/include/llvm/Config/llvm-config.h.cmake
+++ b/llvm/include/llvm/Config/llvm-config.h.cmake
@@ -91,9 +91,6 @@
 /* Define if LLVM was built with a dependency to the libtensorflow dynamic library */
 #cmakedefine LLVM_HAVE_TF_API

-/* Define if LLVM was built with a dependency to the tensorflow compiler */
-#cmakedefine LLVM_HAVE_TF_AOT
-
 /* Define to 1 if you have the <sysexits.h> header file. */
 #cmakedefine HAVE_SYSEXITS_H ${HAVE_SYSEXITS_H}

--- a/llvm/lib/Analysis/InlineAdvisor.cpp
+++ b/llvm/lib/Analysis/InlineAdvisor.cpp
@@ -27,6 +27,9 @@

 using namespace llvm;
 #define DEBUG_TYPE "inline"
+#ifdef LLVM_HAVE_TF_AOT_INLINERSIZEMODEL
+#define LLVM_HAVE_TF_AOT
+#endif

 // This weirdly named statistic tracks the number of times that, when attempting
 // to inline a function A into B, we analyze the callers of B in order to see
--- a/llvm/lib/Analysis/MLInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp
@@ -11,36 +11,38 @@
 // 'release' mode) or a runtime-loaded model (the 'development' case).
 //
 //===----------------------------------------------------------------------===//
-#include "llvm/Config/config.h"
-#if defined(LLVM_HAVE_TF_AOT) || defined(LLVM_HAVE_TF_API)
-
-#include <limits>
-#include <unordered_map>
-#include <unordered_set>
-
+#include "llvm/Analysis/MLInlineAdvisor.h"
 #include "llvm/ADT/SCCIterator.h"
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/FunctionPropertiesAnalysis.h"
 #include "llvm/Analysis/InlineCost.h"
+#include "llvm/Analysis/InlineModelFeatureMaps.h"
 #include "llvm/Analysis/LazyCallGraph.h"
-#include "llvm/Analysis/MLInlineAdvisor.h"
 #include "llvm/Analysis/MLModelRunner.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
+#include "llvm/Analysis/ReleaseModeModelRunner.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/Config/config.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Path.h"

+#include <limits>
+#include <unordered_map>
+#include <unordered_set>
+
 using namespace llvm;

-#ifdef LLVM_HAVE_TF_AOT
-#include "llvm/Analysis/ReleaseModeModelRunner.h"
+#ifdef LLVM_HAVE_TF_AOT_INLINERSIZEMODEL
+#define LLVM_HAVE_TF_AOT
+#endif
+
+#if defined(LLVM_HAVE_TF_AOT)
 // codegen-ed file
 #include "InlinerSizeModel.h" // NOLINT
-#include "llvm/Analysis/InlineModelFeatureMaps.h"

 std::unique_ptr<InlineAdvisor>
 llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM) {
@@ -53,6 +55,8 @@ llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM) {

 #define DEBUG_TYPE "inline-ml"

+#if defined(LLVM_HAVE_TF_AOT) || defined(LLVM_HAVE_TF_API)
+
 static cl::opt<float> SizeIncreaseThreshold(
    "ml-advisor-size-increase-threshold", cl::Hidden,
    cl::desc("Maximum factor by which expected native size may increase before "
--- a/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp
@@ -42,7 +42,9 @@
 using namespace llvm;

 #define DEBUG_TYPE "ml-regalloc"
-
+#ifdef LLVM_HAVE_TF_AOT_REGALLOCEVICTMODEL
+#define LLVM_HAVE_TF_AOT
+#endif
 // Generated header in release (AOT) mode
 #if defined LLVM_HAVE_TF_AOT
 #include "RegallocEvictModel.h"
--- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
+++ b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
@@ -43,6 +43,9 @@ static cl::opt<bool> EnableLocalReassignment(
    cl::init(false));

 #define DEBUG_TYPE "regalloc"
+#ifdef LLVM_HAVE_TF_AOT_REGALLOCEVICTMODEL
+#define LLVM_HAVE_TF_AOT
+#endif

 char RegAllocEvictionAdvisorAnalysis::ID = 0;
 INITIALIZE_PASS(RegAllocEvictionAdvisorAnalysis, "regalloc-evict",
--- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h
+++ b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h
@@ -192,14 +192,9 @@ private:
 /// an instance of the eviction advisor.
 template <> Pass *callDefaultCtor<RegAllocEvictionAdvisorAnalysis>();

-// TODO(mtrofin): implement these.
-#ifdef LLVM_HAVE_TF_AOT
 RegAllocEvictionAdvisorAnalysis *createReleaseModeAdvisor();
-#endif

-#ifdef LLVM_HAVE_TF_API
 RegAllocEvictionAdvisorAnalysis *createDevelopmentModeAdvisor();
-#endif

 // TODO: move to RegAllocEvictionAdvisor.cpp when we move implementation
 // out of RegAllocGreedy.cpp
--- a/llvm/test/CodeGen/MLRegalloc/dev-rel-equivalence.ll
+++ b/llvm/test/CodeGen/MLRegalloc/dev-rel-equivalence.ll
@@ -1,5 +1,4 @@
 ; REQUIRES: have_tf_api
-; REQUIRES: have_tf_aot
 ; REQUIRES: llvm_raevict_model_autogenerated
 ; REQUIRES: x86_64-linux
 ;
--- a/llvm/test/Transforms/Inline/ML/bounds-checks.ll
+++ b/llvm/test/Transforms/Inline/ML/bounds-checks.ll
@@ -2,7 +2,6 @@
 ; In all cases, the end result is the same: mandatory inlinings must happen.
 ; However, when we discover we 'trip' over the artificially-low size increase
 ; factor, we don't inline anymore.
-; REQUIRES: have_tf_aot
 ; REQUIRES: llvm_inliner_model_autogenerated
 ; RUN: opt -passes=scc-oz-module-inliner -enable-ml-inliner=release -ml-advisor-size-increase-threshold=10.0 -S < %s 2>&1 | FileCheck %s --check-prefix=CHECK --check-prefix=NOBOUNDS
 ; RUN: opt -passes=scc-oz-module-inliner -enable-ml-inliner=release -ml-advisor-size-increase-threshold=1.0 -S < %s 2>&1 | FileCheck %s --check-prefix=CHECK --check-prefix=BOUNDS
--- a/llvm/test/Transforms/Inline/ML/ml-test-release-mode.ll
+++ b/llvm/test/Transforms/Inline/ML/ml-test-release-mode.ll
@@ -5,7 +5,6 @@
 ; This test uses Inputs/test-module.ll, as it will share it with a similar test
 ; for the 'development' mode.
 ;
-; REQUIRES: have_tf_aot
 ; REQUIRES: llvm_inliner_model_autogenerated
 ; RUN: opt -passes=scc-oz-module-inliner -enable-ml-inliner=release -S < %S/Inputs/test-module.ll 2>&1 | FileCheck %S/Inputs/test-module.ll --check-prefix=CHECK
 ; RUN: opt -passes=scc-oz-module-inliner -enable-ml-inliner=default -S < %S/Inputs/test-module.ll 2>&1 | FileCheck %S/Inputs/test-module.ll --check-prefix=DEFAULT