[MLGO] Improved support for AOT cross-targeting scenarios

[mtrofin]

Summary

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.

[phosek]

LGTM

llvm/CMakeLists.txt

896–897

I think this could be elif?

904–905

The same here?

llvm/cmake/modules/TensorFlowCompile.cmake

47–59

This is unrelated to this change, but since you're already modifying this function, I'd also consider renaming it to tf_compile to match other tf_* functions.

85

I'd probably use STATUS level for this.

90

I'd probably use STATUS level for this.

[mtrofin]

feedback & rebase

llvm/CMakeLists.txt

set(TENSORFLOW_AOT_COMPILER
     CACHE PATH "Path to the Tensorflow AOT compiler")
   include_directories(${TENSORFLOW_AOT_PATH}/include)
   add_subdirectory(${TENSORFLOW_AOT_PATH}/xla_aot_runtime_src
     ${CMAKE_ARCHIVE_OUTPUT_DIRECTORY}/tf_runtime)
   install(TARGETS tf_xla_runtime EXPORT LLVMExports
     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")
     set(LLVM_RAEVICT_MODEL_AUTOGENERATED 1)
   endif()
 
 endif()
 
 # Configure the three LLVM configuration header files.

llvm/cmake/modules/TensorFlowCompile.cmake

 endfunction()
 
 # Run the tensorflow compiler (saved_model_cli) on the saved model in the
 # ${model} directory, looking for the ${tag_set} tag set, and the SignatureDef
 # ${signature_def_key}.
 # 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)
 
-  # 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
     GENERATED 1 EXTERNAL_OBJECT 1)
 
   set(GENERATED_HEADERS ${GENERATED_HEADERS} ${hdr_file} PARENT_SCOPE)
   set_source_files_properties(${hdr_file} PROPERTIES
     GENERATED 1)
 
   endfunction()
 
 function(tf_find_and_compile model default_url default_path test_model_generator 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 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.")
+          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()
 
     if ("${model}" STREQUAL "autogenerate")
-    set(model ${default_path}-autogenerated)  
+      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)
+    tf_compile(${LLVM_ML_MODELS_ABSOLUTE} ${tag_set} ${signature_def_key} ${fname} ${cpp_class} ${hdr_file} ${obj_file})
+  endif()
 
   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()

llvm/include/llvm/Analysis/InlineAdvisor.h

class InlineAdvisorAnalysisPrinterPass
   raw_ostream &OS;
 
 public:
   explicit InlineAdvisorAnalysisPrinterPass(raw_ostream &OS) : OS(OS) {}
 
   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.
 
 /// Return the cost only if the inliner should attempt to inline at the given
 /// CallSite. If we return the cost, we will emit an optimisation remark later
 /// using that cost, so we won't do so from this function. Return None if
 /// inlining should not be attempted.

llvm/include/llvm/Config/llvm-config.h.cmake

 #cmakedefine LLVM_WITH_Z3 ${LLVM_WITH_Z3}
 
 /* Define if we have curl and want to use it */
 #cmakedefine LLVM_ENABLE_CURL ${LLVM_ENABLE_CURL}
 
 /* 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}
 
 /* Define to 1 to enable the experimental new pass manager by default */
 #cmakedefine01 LLVM_ENABLE_NEW_PASS_MANAGER
 
 /* Define if the xar_open() function is supported on this platform. */
 #cmakedefine LLVM_HAVE_LIBXAR ${LLVM_HAVE_LIBXAR}
 
 /* Define if building libLLVM shared library */
 #cmakedefine LLVM_BUILD_LLVM_DYLIB
 
 /* Define if building LLVM with BUILD_SHARED_LIBS */
 #cmakedefine LLVM_BUILD_SHARED_LIBS
 
 #endif

llvm/lib/Analysis/InlineAdvisor.cpp

 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/raw_ostream.h"
 
 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
 // if those would be more profitable and blocked inline steps.
 STATISTIC(NumCallerCallersAnalyzed, "Number of caller-callers analyzed");
 
 /// Flag to add inline messages as callsite attributes 'inline-remark'.
 static cl::opt<bool>

llvm/lib/Analysis/MLInlineAdvisor.cpp

 //===- MLInlineAdvisor.cpp - machine learned InlineAdvisor ----------------===//
 //
 // 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 file implements the interface between the inliner and a learned model.
 // It delegates model evaluation to either the AOT compiled model (the
 // '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) {
   auto AOTRunner =
       std::make_unique<ReleaseModeModelRunner<llvm::InlinerSizeModel>>(
           M.getContext(), FeatureNameMap, DecisionName);
   return std::make_unique<MLInlineAdvisor>(M, MAM, std::move(AOTRunner));
 }
 #endif
 
 #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 "
              "blocking any further inlining."),
     cl::init(2.0));
 
 // clang-format off
 const std::array<std::string, NumberOfFeatures> llvm::FeatureNameMap{

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

 #include "llvm/Target/TargetMachine.h"
 
 #include <array>
 #include <memory>
 
 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"
 #endif
 
 // Options that only make sense in development mode
 #ifdef LLVM_HAVE_TF_API
 static cl::opt<std::string> TrainingLog(

llvm/lib/CodeGen/RegAllocEvictionAdvisor.h

private:
   StringRef getPassName() const override;
   const AdvisorMode Mode;
 };
 
 /// Specialization for the API used by the analysis infrastructure to create
 /// 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
 class DefaultEvictionAdvisor : public RegAllocEvictionAdvisor {
 public:
   DefaultEvictionAdvisor(const MachineFunction &MF, const RAGreedy &RA)
       : RegAllocEvictionAdvisor(MF, RA) {}
 

llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp

 
 static cl::opt<bool> EnableLocalReassignment(
     "enable-local-reassign", cl::Hidden,
     cl::desc("Local reassignment can yield better allocation decisions, but "
              "may be compile time intensive"),
     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",
                 "Regalloc eviction policy", false, true)
 
 namespace {
 class DefaultEvictionAdvisorAnalysis final
     : public RegAllocEvictionAdvisorAnalysis {

llvm/test/CodeGen/MLRegalloc/dev-rel-equivalence.ll

 ; REQUIRES: have_tf_api
-; REQUIRES: have_tf_aot
 ; REQUIRES: llvm_raevict_model_autogenerated
 ; REQUIRES: x86_64-linux
 ;
 ; Check that the same model (==the autogenerated one) produces the same output
 ; regardless of how it's evaluated, which is different from the default
 ;
 ; RUN: llc -mtriple=x86_64-linux-unknown -regalloc=greedy -regalloc-enable-advisor=default \
 ; RUN:   %S/Inputs/input.ll -o %t.default

llvm/test/Transforms/Inline/ML/bounds-checks.ll

 ; Test behavior when inlining policy grows size out of control.
 ; 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
 
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-grtev4-linux-gnu"
 
 declare i64 @f1()

llvm/test/Transforms/Inline/ML/ml-test-release-mode.ll

 ; The default inliner doesn't elide @adder, it believes it's too costly to inline 
 ; adder into switcher. The ML inliner carries out that inlining, resulting in
 ; a smaller result (part of it is that adder gets elided).
 ;
 ; 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