[llvm] [MLGO] Model selection for models lowered through EmitC (PR #212650)
Bhavesh M via llvm-commits
llvm-commits at lists.llvm.org
Thu Aug 13 14:08:52 PDT 2026
https://github.com/beamandala updated https://github.com/llvm/llvm-project/pull/212650
>From 48f91f07a8def5c25061e46eb90865dec179078e Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Mon, 27 Jul 2026 12:34:38 -0700
Subject: [PATCH 01/11] Model selection
---
llvm/CMakeLists.txt | 19 ++++++
llvm/cmake/modules/MLGOCompile.cmake | 99 ++++++++++++++++++++++++++++
llvm/lib/Analysis/CMakeLists.txt | 13 +++-
llvm/lib/CodeGen/CMakeLists.txt | 13 +++-
llvm/test/CMakeLists.txt | 1 +
llvm/test/lit.cfg.py | 3 +
llvm/test/lit.site.cfg.py.in | 1 +
7 files changed, 147 insertions(+), 2 deletions(-)
create mode 100644 llvm/cmake/modules/MLGOCompile.cmake
diff --git a/llvm/CMakeLists.txt b/llvm/CMakeLists.txt
index a9d9f1d11ab44..8490c70617a7c 100644
--- a/llvm/CMakeLists.txt
+++ b/llvm/CMakeLists.txt
@@ -1226,6 +1226,21 @@ set( CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${LLVM_LIBRARY_DIR} )
# For up-to-date instructions for installing the TFLite dependency, refer to
# the bot setup script: https://github.com/google/ml-compiler-opt/blob/main/buildbot/buildbot_init.sh
set(LLVM_HAVE_TFLITE "" CACHE BOOL "Use tflite")
+
+set(LLVM_MLGO_MODELS "" CACHE STRING "List of AOT MLGO models")
+set(LLVM_MLGO_MLIR_OPT "mlir-opt" CACHE STRING "Path or system binary name for mlir-opt")
+set(LLVM_MLGO_MLIR_TRANSLATE "mlir-translate" CACHE STRING "Path or system binary name for mlir-translate")
+
+if (NOT LLVM_MLGO_MODELS STREQUAL "")
+ set(LLVM_HAVE_EMITC_COMPILE "ON" CACHE BOOL "MLGO AOT models compiled with EmitC are available" FORCE)
+else()
+ set(LLVM_HAVE_EMITC_COMPILE "OFF" CACHE BOOL "MLGO AOT models compiled with EmitC are available" FORCE)
+endif()
+
+if (LLVM_HAVE_EMITC_COMPILE AND LLVM_HAVE_TFLITE)
+ message(FATAL_ERROR "Only one of LLVM_HAVE_TFLITE and LLVM_HAVE_EMITC_COMPILE can be enabled.")
+endif()
+
if (LLVM_HAVE_TFLITE)
find_package(tensorflow-lite REQUIRED)
endif()
@@ -1253,6 +1268,10 @@ if (NOT TENSORFLOW_AOT_PATH STREQUAL "")
install(TARGETS tf_xla_runtime EXPORT LLVMDevelopmentExports
ARCHIVE DESTINATION lib${LLVM_LIBDIR_SUFFIX} COMPONENT tf_xla_runtime)
set_property(GLOBAL APPEND PROPERTY LLVM_EXPORTS tf_xla_runtime)
+
+ if (LLVM_HAVE_TF_AOT AND LLVM_HAVE_EMITC_COMPILE)
+ message(FATAL_ERROR "Only one of LLVM_HAVE_TF_AOT and LLVM_HAVE_EMITC_COMPILE can be enabled.")
+ endif()
# Once we add more modules, we should handle this more automatically.
if (DEFINED LLVM_OVERRIDE_MODEL_HEADER_INLINERSIZEMODEL)
set(LLVM_INLINER_MODEL_PATH "none")
diff --git a/llvm/cmake/modules/MLGOCompile.cmake b/llvm/cmake/modules/MLGOCompile.cmake
new file mode 100644
index 0000000000000..c401f87bf23f7
--- /dev/null
+++ b/llvm/cmake/modules/MLGOCompile.cmake
@@ -0,0 +1,99 @@
+# Compile MLGO models expressed as MLIR to C++ headers via the EmitC pipeline.
+#
+# Each entry of ${models} has the form
+# "cli-flag,path/to/model.mlir,type". Entries whose type matches
+# ${target_type} are lowered with ${mlir_opt} and translated with
+# ${mlir_translate} to a C++ header defining Model<N>, placed in
+# ${LLVM_INCLUDE_DIR}/${include_subdir}. In the same directory, generate
+# ${def_file}, with one MLGO_MODEL(ClassName, "cli-flag") entry per model, and
+# ${umbrella_header}, including all the generated model headers.
+# Append the target driving generation to MLDeps in the caller's scope, and
+# define LLVM_HAVE_EMITC_COMPILE_<TARGET_TYPE>.
+function(mlgo_compile_models models mlir_opt mlir_translate target_type
+ include_subdir def_file umbrella_header)
+ if ("${mlir_opt}" MATCHES "/")
+ get_filename_component(mlir_opt "${mlir_opt}" ABSOLUTE BASE_DIR "${CMAKE_BINARY_DIR}")
+ endif()
+ if ("${mlir_translate}" MATCHES "/")
+ get_filename_component(mlir_translate "${mlir_translate}" ABSOLUTE BASE_DIR "${CMAKE_BINARY_DIR}")
+ endif()
+
+ set(DEF_CONTENT "/* Auto-generated by CMake */\n")
+ set(DEF_CONTENT "${DEF_CONTENT}#ifndef MLGO_MODEL\n#define MLGO_MODEL(CLASS_NAME, CLI_FLAG)\n#endif\n\n")
+
+ set(HEADERS_CONTENT "/* Auto-generated by CMake */\nnamespace llvm {\n")
+ set(MLGO_GEN_TARGETS "")
+ set(MODEL_INDEX 1)
+
+ foreach(MODEL_INFO IN LISTS models)
+ # Parse comma-separated fields: cli-flag,path/to/model.mlir,type
+ string(REPLACE "," ";" MODEL_FIELDS "${MODEL_INFO}")
+ list(GET MODEL_FIELDS 0 CLI_FLAG)
+ list(GET MODEL_FIELDS 1 MODEL_PATH)
+ list(GET MODEL_FIELDS 2 MODEL_TYPE)
+
+ # Only process models matching target_type
+ if (NOT MODEL_TYPE STREQUAL target_type)
+ continue()
+ endif()
+
+ set(CLASS_NAME "Model${MODEL_INDEX}")
+ math(EXPR MODEL_INDEX "${MODEL_INDEX} + 1")
+
+ if (NOT IS_ABSOLUTE "${MODEL_PATH}")
+ get_filename_component(MODEL_PATH "${MODEL_PATH}" ABSOLUTE BASE_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
+ endif()
+
+ set(EMITC_MLIR "${CMAKE_CURRENT_BINARY_DIR}/${CLASS_NAME}_emitc.mlir")
+ set(HEADER_FILE "${LLVM_INCLUDE_DIR}/${include_subdir}/${CLASS_NAME}.h")
+
+ # 1. Run MLIR pipeline to compile MODEL_PATH to EmitC MLIR
+ add_custom_command(OUTPUT ${EMITC_MLIR}
+ COMMAND ${mlir_opt} "--pass-pipeline=builtin.module(func.func(tosa-to-linalg-named,tosa-to-linalg,tosa-to-arith,tosa-to-tensor),symbol-privatize,mlgo-scalarize-single-element-tensor-return,one-shot-bufferize{bufferize-function-boundaries=true function-boundary-type-conversion=identity-layout-map buffer-alignment=0},buffer-results-to-out-params{hoist-static-allocs=true},func.func(promote-buffers-to-stack),buffer-deallocation-pipeline,func.func(convert-linalg-to-loops),expand-strided-metadata,canonicalize,memref-elide-reinterpret-cast,convert-to-emitc,wrap-emitc-func-in-class{class-name-format=${CLASS_NAME}},math-expand-ops{ops=rsqrt},arith-expand,convert-math-to-emitc,convert-arith-to-emitc)"
+ ${MODEL_PATH} -o ${EMITC_MLIR}
+ DEPENDS ${MODEL_PATH}
+ VERBATIM
+ )
+
+ # 2. Translate EmitC MLIR to C++ header
+ add_custom_command(OUTPUT ${HEADER_FILE}
+ COMMAND ${mlir_translate} -mlir-to-cpp ${EMITC_MLIR} -o ${HEADER_FILE} --mlir-print-stacktrace-on-diagnostic
+ DEPENDS ${EMITC_MLIR}
+ VERBATIM
+ )
+
+ # Set properties so CMake knows these are generated during the build
+ set_source_files_properties(${EMITC_MLIR} PROPERTIES GENERATED 1)
+ set_source_files_properties(${HEADER_FILE} PROPERTIES GENERATED 1)
+
+ # Custom target to force generation of this header
+ add_custom_target(mlgo_model_gen_${CLASS_NAME} DEPENDS ${HEADER_FILE})
+ list(APPEND MLGO_GEN_TARGETS mlgo_model_gen_${CLASS_NAME})
+
+ # Append the model metadata to the .def file
+ string(APPEND DEF_CONTENT "MLGO_MODEL(${CLASS_NAME}, \"${CLI_FLAG}\")\n")
+
+ # Append to the master header content
+ string(APPEND HEADERS_CONTENT "namespace ${CLASS_NAME}_ns {\n#include \"${include_subdir}/${CLASS_NAME}.h\"\n}\nusing ${CLASS_NAME}_ns::${CLASS_NAME};\n")
+ endforeach()
+
+ string(APPEND DEF_CONTENT "\n#undef MLGO_MODEL\n")
+ string(APPEND HEADERS_CONTENT "}\n")
+
+ # Stage the generated files, then only update the copies under
+ # LLVM_INCLUDE_DIR when their content changed, so that reconfiguring does not
+ # trigger rebuilds of their consumers.
+ file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/${def_file}.tmp" "${DEF_CONTENT}")
+ configure_file("${CMAKE_CURRENT_BINARY_DIR}/${def_file}.tmp"
+ "${LLVM_INCLUDE_DIR}/${include_subdir}/${def_file}" COPYONLY)
+ file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/${umbrella_header}.tmp" "${HEADERS_CONTENT}")
+ configure_file("${CMAKE_CURRENT_BINARY_DIR}/${umbrella_header}.tmp"
+ "${LLVM_INCLUDE_DIR}/${include_subdir}/${umbrella_header}" COPYONLY)
+
+ if (MLGO_GEN_TARGETS)
+ add_custom_target(mlgo_models_gen_${target_type} DEPENDS ${MLGO_GEN_TARGETS})
+ set(MLDeps ${MLDeps} mlgo_models_gen_${target_type} PARENT_SCOPE)
+ string(TOUPPER ${target_type} target_type_allcaps)
+ add_compile_definitions(LLVM_HAVE_EMITC_COMPILE_${target_type_allcaps})
+ endif()
+endfunction()
diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt
index aecea1bb92a1f..ed7fc6e7e26c0 100644
--- a/llvm/lib/Analysis/CMakeLists.txt
+++ b/llvm/lib/Analysis/CMakeLists.txt
@@ -1,4 +1,15 @@
-if (DEFINED LLVM_HAVE_TF_AOT OR LLVM_HAVE_TFLITE)
+if (LLVM_HAVE_EMITC_COMPILE)
+ include(MLGOCompile)
+ mlgo_compile_models(
+ "${LLVM_MLGO_MODELS}"
+ "${LLVM_MLGO_MLIR_OPT}"
+ "${LLVM_MLGO_MLIR_TRANSLATE}"
+ inliner
+ llvm/Analysis
+ MLGOModels.def
+ InlinerSizeModelMulti.h
+ )
+elseif (DEFINED LLVM_HAVE_TF_AOT OR LLVM_HAVE_TFLITE)
include(TensorFlowCompile)
set(LLVM_INLINER_MODEL_PATH_DEFAULT "models/inliner-Oz")
diff --git a/llvm/lib/CodeGen/CMakeLists.txt b/llvm/lib/CodeGen/CMakeLists.txt
index c572128b023c1..1c913c9c2f8ff 100644
--- a/llvm/lib/CodeGen/CMakeLists.txt
+++ b/llvm/lib/CodeGen/CMakeLists.txt
@@ -1,4 +1,15 @@
-if (DEFINED LLVM_HAVE_TF_AOT OR LLVM_HAVE_TFLITE)
+if (LLVM_HAVE_EMITC_COMPILE)
+ include(MLGOCompile)
+ mlgo_compile_models(
+ "${LLVM_MLGO_MODELS}"
+ "${LLVM_MLGO_MLIR_OPT}"
+ "${LLVM_MLGO_MLIR_TRANSLATE}"
+ regalloc
+ llvm/CodeGen
+ RegAllocEvictModels.def
+ RegAllocEvictModelMulti.h
+ )
+elseif (DEFINED LLVM_HAVE_TF_AOT OR LLVM_HAVE_TFLITE)
include(TensorFlowCompile)
set(LLVM_RAEVICT_MODEL_PATH_DEFAULT "models/regalloc-eviction")
diff --git a/llvm/test/CMakeLists.txt b/llvm/test/CMakeLists.txt
index 2a00a429e2626..90d3cd494565a 100644
--- a/llvm/test/CMakeLists.txt
+++ b/llvm/test/CMakeLists.txt
@@ -19,6 +19,7 @@ llvm_canonicalize_cmake_booleans(
LLVM_EXAMPLEIRTRANSFORMS_LINK_INTO_TOOLS
LLVM_HAVE_TF_AOT
LLVM_HAVE_TFLITE
+ LLVM_HAVE_EMITC_COMPILE
LLVM_ENABLE_PROFCHECK
LLVM_INLINER_MODEL_AUTOGENERATED
LLVM_RAEVICT_MODEL_AUTOGENERATED
diff --git a/llvm/test/lit.cfg.py b/llvm/test/lit.cfg.py
index 6f04a855d8cb8..3973b971149df 100644
--- a/llvm/test/lit.cfg.py
+++ b/llvm/test/lit.cfg.py
@@ -584,6 +584,9 @@ def enable_ptxas(ptxas_executable):
if config.have_tf_aot:
config.available_features.add("have_tf_aot")
+if getattr(config, "have_emitc_compile", False):
+ config.available_features.add("have_emitc_compile")
+
if getattr(config, "have_opencsd", False):
config.available_features.add("opencsd")
diff --git a/llvm/test/lit.site.cfg.py.in b/llvm/test/lit.site.cfg.py.in
index 130d9011df619..24279c94c2ae8 100644
--- a/llvm/test/lit.site.cfg.py.in
+++ b/llvm/test/lit.site.cfg.py.in
@@ -57,6 +57,7 @@ config.has_plugins = @LLVM_ENABLE_PLUGINS@
config.linked_bye_extension = @LLVM_BYE_LINK_INTO_TOOLS@
config.linked_exampleirtransforms_extension = @LLVM_EXAMPLEIRTRANSFORMS_LINK_INTO_TOOLS@
config.have_tf_aot = @LLVM_HAVE_TF_AOT@
+config.have_emitc_compile = @LLVM_HAVE_EMITC_COMPILE@
config.have_tflite = @LLVM_HAVE_TFLITE@
config.enable_profcheck = @LLVM_ENABLE_PROFCHECK@
config.llvm_inliner_model_autogenerated = @LLVM_INLINER_MODEL_AUTOGENERATED@
>From 4e31b259912736dd29f8e52b0a1cdbbfbbfeb555 Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Tue, 28 Jul 2026 16:11:29 -0700
Subject: [PATCH 02/11] EmitCModelRunner
---
llvm/cmake/modules/MLGOCompile.cmake | 26 ++++++-
llvm/include/llvm/Analysis/EmitCModelRunner.h | 76 +++++++++++++++++++
llvm/lib/Analysis/MLInlineAdvisor.cpp | 58 ++++++++++++--
llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp | 62 +++++++++++++--
4 files changed, 210 insertions(+), 12 deletions(-)
create mode 100644 llvm/include/llvm/Analysis/EmitCModelRunner.h
diff --git a/llvm/cmake/modules/MLGOCompile.cmake b/llvm/cmake/modules/MLGOCompile.cmake
index c401f87bf23f7..b4e340087bd62 100644
--- a/llvm/cmake/modules/MLGOCompile.cmake
+++ b/llvm/cmake/modules/MLGOCompile.cmake
@@ -47,9 +47,33 @@ function(mlgo_compile_models models mlir_opt mlir_translate target_type
set(EMITC_MLIR "${CMAKE_CURRENT_BINARY_DIR}/${CLASS_NAME}_emitc.mlir")
set(HEADER_FILE "${LLVM_INCLUDE_DIR}/${include_subdir}/${CLASS_NAME}.h")
+ # Pass pipeline to lower MLIR models to EmitC dialect
+ set(EMITC_PASSES
+ "func.func(tosa-to-linalg-named,tosa-to-linalg,tosa-to-arith,tosa-to-tensor)"
+ "symbol-privatize"
+ "scalarize-single-element-tensor-return"
+ "one-shot-bufferize{bufferize-function-boundaries=true function-boundary-type-conversion=identity-layout-map buffer-alignment=0}"
+ "buffer-results-to-out-params{hoist-static-allocs=true}"
+ "func.func(promote-buffers-to-stack)"
+ "buffer-deallocation-pipeline"
+ "func.func(convert-linalg-to-loops)"
+ "expand-strided-metadata"
+ "canonicalize"
+ "memref-elide-reinterpret-cast"
+ "convert-to-emitc"
+ "wrap-emitc-func-in-class{class-name-format=${CLASS_NAME}}"
+ "mlgo-add-reflection-map{included-field-attrs=tf_saved_model.index_path}"
+ "math-expand-ops{ops=rsqrt}"
+ "arith-expand"
+ "convert-math-to-emitc"
+ "convert-arith-to-emitc"
+ )
+ string(JOIN "," PASS_PIPELINE ${EMITC_PASSES})
+ set(PASS_PIPELINE "builtin.module(${PASS_PIPELINE})")
+
# 1. Run MLIR pipeline to compile MODEL_PATH to EmitC MLIR
add_custom_command(OUTPUT ${EMITC_MLIR}
- COMMAND ${mlir_opt} "--pass-pipeline=builtin.module(func.func(tosa-to-linalg-named,tosa-to-linalg,tosa-to-arith,tosa-to-tensor),symbol-privatize,mlgo-scalarize-single-element-tensor-return,one-shot-bufferize{bufferize-function-boundaries=true function-boundary-type-conversion=identity-layout-map buffer-alignment=0},buffer-results-to-out-params{hoist-static-allocs=true},func.func(promote-buffers-to-stack),buffer-deallocation-pipeline,func.func(convert-linalg-to-loops),expand-strided-metadata,canonicalize,memref-elide-reinterpret-cast,convert-to-emitc,wrap-emitc-func-in-class{class-name-format=${CLASS_NAME}},math-expand-ops{ops=rsqrt},arith-expand,convert-math-to-emitc,convert-arith-to-emitc)"
+ COMMAND ${mlir_opt} "--pass-pipeline=${PASS_PIPELINE}"
${MODEL_PATH} -o ${EMITC_MLIR}
DEPENDS ${MODEL_PATH}
VERBATIM
diff --git a/llvm/include/llvm/Analysis/EmitCModelRunner.h b/llvm/include/llvm/Analysis/EmitCModelRunner.h
new file mode 100644
index 0000000000000..6c36f930aeff0
--- /dev/null
+++ b/llvm/include/llvm/Analysis/EmitCModelRunner.h
@@ -0,0 +1,76 @@
+//===- EmitCModelRunner.h - Fast, precompiled model runner ---------------===//
+//
+// 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 a model runner wrapping an EmitC compiled ML model.
+// Only inference is supported.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_ANALYSIS_EMITCMODELRUNNER_H
+#define LLVM_ANALYSIS_EMITCMODELRUNNER_H
+
+#include "llvm/Analysis/MLModelRunner.h"
+#include "llvm/Analysis/TensorSpec.h"
+
+#include <memory>
+#include <type_traits>
+
+namespace llvm {
+
+template <class TGen>
+class EmitCModelRunner final : public MLModelRunner {
+public:
+ template <class FType>
+ EmitCModelRunner(LLVMContext &Ctx, const FType &InputSpec,
+ std::unique_ptr<TGen> Model = std::make_unique<TGen>())
+ : MLModelRunner(Ctx, MLModelRunner::Kind::Release, InputSpec.size()),
+ CompiledModel(std::move(Model)) {
+ assert(CompiledModel && "The CompiledModel should be valid");
+ for (size_t I = 0; I < InputSpec.size(); ++I)
+ populateTensor(I, InputSpec[I]);
+ }
+
+ ~EmitCModelRunner() override = default;
+
+ static bool classof(const MLModelRunner *R) {
+ return R->getKind() == MLModelRunner::Kind::Release;
+ }
+
+protected:
+ void *evaluateUntyped() override { return evaluateImpl(); }
+
+private:
+ void populateTensor(size_t Pos, const TensorSpec &Spec) {
+ void *Buffer = nullptr;
+ auto It = CompiledModel->reflectionMap.find(Spec.name());
+ if (It != CompiledModel->reflectionMap.end())
+ Buffer = static_cast<void *>(It->second);
+ setUpBufferForTensor(Pos, Spec, Buffer);
+ }
+
+ using ResultType = decltype(std::declval<TGen>()());
+
+ template <typename R = ResultType>
+ std::enable_if_t<!std::is_void_v<R>, void *> evaluateImpl() {
+ Result = (*CompiledModel)();
+ return &Result;
+ }
+
+ template <typename R = ResultType>
+ std::enable_if_t<std::is_void_v<R>, void *> evaluateImpl() {
+ (*CompiledModel)();
+ return nullptr;
+ }
+
+ std::conditional_t<std::is_void_v<ResultType>, char, ResultType> Result{};
+ std::unique_ptr<TGen> CompiledModel;
+};
+
+} // namespace llvm
+
+#endif // LLVM_ANALYSIS_EMITCMODELRUNNER_H
diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp
index 9a5ae2ae26799..7245a65141782 100644
--- a/llvm/lib/Analysis/MLInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp
@@ -72,20 +72,66 @@ using CompiledModelType = llvm::InlinerSizeModel;
using CompiledModelType = NoopSavedModelImpl;
#endif
+#if defined(LLVM_HAVE_EMITC_COMPILE_INLINER)
+#include "llvm/Analysis/EmitCModelRunner.h"
+#include "llvm/Analysis/InlinerSizeModelMulti.h"
+
+enum class MLGOModelChoice {
+ Default,
+#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) CLASS_NAME,
+#include "llvm/Analysis/MLGOModels.def"
+};
+
+static llvm::cl::opt<MLGOModelChoice> SelectedMLGOModel(
+ "mlgo-model",
+ llvm::cl::desc("Select the MLGO model to execute:"),
+ llvm::cl::init(MLGOModelChoice::Default),
+ llvm::cl::values(
+ clEnumValN(MLGOModelChoice::Default, "default", "Use standard heuristic")
+#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) \
+ , clEnumValN(MLGOModelChoice::CLASS_NAME, CLI_FLAG, "Use the " CLI_FLAG " MLGO model")
+#include "llvm/Analysis/MLGOModels.def"
+ )
+);
+
+static std::unique_ptr<MLModelRunner> createMLGOModelRunner(LLVMContext &Ctx, const std::vector<TensorSpec> &InputFeatures) {
+ switch (SelectedMLGOModel) {
+ case MLGOModelChoice::Default:
+ return nullptr;
+#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) \
+ case MLGOModelChoice::CLASS_NAME: \
+ return std::make_unique<EmitCModelRunner<CLASS_NAME>>(Ctx, InputFeatures);
+#include "llvm/Analysis/MLGOModels.def"
+ }
+ llvm_unreachable("Unknown MLGO model type!");
+}
+#endif
+
std::unique_ptr<InlineAdvisor>
llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM,
std::function<bool(CallBase &)> GetDefaultAdvice) {
if (!llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() &&
- InteractiveChannelBaseName.empty())
+ InteractiveChannelBaseName.empty()
+#if defined(LLVM_HAVE_EMITC_COMPILE_INLINER)
+ && SelectedMLGOModel == MLGOModelChoice::Default
+#endif
+ )
return nullptr;
auto RunnerFactory = [&](const std::vector<TensorSpec> &InputFeatures)
-> std::unique_ptr<MLModelRunner> {
std::unique_ptr<MLModelRunner> AOTRunner;
- if (InteractiveChannelBaseName.empty())
- AOTRunner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
- M.getContext(), InputFeatures, DecisionName,
- EmbeddedModelRunnerOptions().setModelSelector(ModelSelector));
- else {
+ if (InteractiveChannelBaseName.empty()) {
+#if defined(LLVM_HAVE_EMITC_COMPILE_INLINER)
+ if (SelectedMLGOModel != MLGOModelChoice::Default) {
+ AOTRunner = createMLGOModelRunner(M.getContext(), InputFeatures);
+ } else
+#endif
+ {
+ AOTRunner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
+ M.getContext(), InputFeatures, DecisionName,
+ EmbeddedModelRunnerOptions().setModelSelector(ModelSelector));
+ }
+ } else {
AOTRunner = std::make_unique<InteractiveModelRunner>(
M.getContext(), InputFeatures, InlineDecisionSpec,
InteractiveChannelBaseName + ".out",
diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index 54f1be77971d5..b4a55dbc79d7b 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -55,6 +55,41 @@ using CompiledModelType = RegAllocEvictModel;
using CompiledModelType = NoopSavedModelImpl;
#endif
+#if defined(LLVM_HAVE_EMITC_COMPILE_REGALLOC)
+#include "llvm/Analysis/EmitCModelRunner.h"
+#include "llvm/CodeGen/RegAllocEvictModelMulti.h"
+
+enum class MLGORegAllocModelChoice {
+ Default,
+#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) CLASS_NAME,
+#include "llvm/CodeGen/RegAllocEvictModels.def"
+};
+
+static llvm::cl::opt<MLGORegAllocModelChoice> SelectedMLGORegAllocModel(
+ "regalloc-mlgo-model",
+ llvm::cl::desc("Select the MLGO model to execute for register allocation:"),
+ llvm::cl::init(MLGORegAllocModelChoice::Default),
+ llvm::cl::values(
+ clEnumValN(MLGORegAllocModelChoice::Default, "default", "Use standard heuristic")
+#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) \
+ , clEnumValN(MLGORegAllocModelChoice::CLASS_NAME, CLI_FLAG, "Use the " CLI_FLAG " MLGO model")
+#include "llvm/CodeGen/RegAllocEvictModels.def"
+ )
+);
+
+static std::unique_ptr<MLModelRunner> createMLGORegAllocModelRunner(LLVMContext &Ctx, const std::vector<TensorSpec> &InputFeatures) {
+ switch (SelectedMLGORegAllocModel) {
+ case MLGORegAllocModelChoice::Default:
+ return nullptr;
+#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) \
+ case MLGORegAllocModelChoice::CLASS_NAME: \
+ return std::make_unique<EmitCModelRunner<CLASS_NAME>>(Ctx, InputFeatures);
+#include "llvm/CodeGen/RegAllocEvictModels.def"
+ }
+ llvm_unreachable("Unknown MLGO model type!");
+}
+#endif
+
static cl::opt<std::string> InteractiveChannelBaseName(
"regalloc-evict-interactive-channel-base", cl::Hidden,
cl::desc(
@@ -368,14 +403,22 @@ class ReleaseModeEvictionAdvisorProvider final
getAdvisor(const MachineFunction &MF, const RAGreedy &RA,
MachineBlockFrequencyInfo *MBFI, MachineLoopInfo *Loops) override {
if (!Runner) {
- if (InteractiveChannelBaseName.empty())
- Runner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
- MF.getFunction().getContext(), InputFeatures, DecisionName);
- else
+ if (InteractiveChannelBaseName.empty()) {
+#if defined(LLVM_HAVE_EMITC_COMPILE_REGALLOC)
+ if (SelectedMLGORegAllocModel != MLGORegAllocModelChoice::Default) {
+ Runner = createMLGORegAllocModelRunner(MF.getFunction().getContext(), InputFeatures);
+ } else
+#endif
+ {
+ Runner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
+ MF.getFunction().getContext(), InputFeatures, DecisionName);
+ }
+ } else {
Runner = std::make_unique<InteractiveModelRunner>(
MF.getFunction().getContext(), InputFeatures, DecisionSpec,
InteractiveChannelBaseName + ".out",
InteractiveChannelBaseName + ".in");
+ }
}
assert(MBFI && Loops &&
"Invalid provider state: must have analysis available");
@@ -1018,7 +1061,13 @@ bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) {
RegAllocEvictionAdvisorProvider *
llvm::createReleaseModeAdvisorProvider(LLVMContext &Ctx) {
- return new ReleaseModeEvictionAdvisorProvider(Ctx);
+ return llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
+ !InteractiveChannelBaseName.empty()
+#if defined(LLVM_HAVE_EMITC_COMPILE_REGALLOC)
+ || SelectedMLGORegAllocModel != MLGORegAllocModelChoice::Default
+#endif
+ ? new ReleaseModeEvictionAdvisorProvider(Ctx)
+ : nullptr;
}
RegAllocEvictionAdvisorProvider *
@@ -1033,6 +1082,9 @@ RegAllocEvictionAdvisorAnalysisLegacy *
llvm::createReleaseModeAdvisorAnalysisLegacy() {
return llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
!InteractiveChannelBaseName.empty()
+#if defined(LLVM_HAVE_EMITC_COMPILE_REGALLOC)
+ || SelectedMLGORegAllocModel != MLGORegAllocModelChoice::Default
+#endif
? new ReleaseModeEvictionAdvisorAnalysisLegacy()
: nullptr;
}
>From 531bb6f32b07bbe86f03241ad54456fde5d160e0 Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Tue, 28 Jul 2026 16:15:05 -0700
Subject: [PATCH 03/11] Formatting
---
llvm/include/llvm/Analysis/EmitCModelRunner.h | 3 +-
llvm/lib/Analysis/MLInlineAdvisor.cpp | 31 ++++++++--------
llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp | 35 +++++++++++--------
3 files changed, 37 insertions(+), 32 deletions(-)
diff --git a/llvm/include/llvm/Analysis/EmitCModelRunner.h b/llvm/include/llvm/Analysis/EmitCModelRunner.h
index 6c36f930aeff0..f60ca034d232d 100644
--- a/llvm/include/llvm/Analysis/EmitCModelRunner.h
+++ b/llvm/include/llvm/Analysis/EmitCModelRunner.h
@@ -22,8 +22,7 @@
namespace llvm {
-template <class TGen>
-class EmitCModelRunner final : public MLModelRunner {
+template <class TGen> class EmitCModelRunner final : public MLModelRunner {
public:
template <class FType>
EmitCModelRunner(LLVMContext &Ctx, const FType &InputSpec,
diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp
index 7245a65141782..e04805ff135af 100644
--- a/llvm/lib/Analysis/MLInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp
@@ -83,24 +83,25 @@ enum class MLGOModelChoice {
};
static llvm::cl::opt<MLGOModelChoice> SelectedMLGOModel(
- "mlgo-model",
- llvm::cl::desc("Select the MLGO model to execute:"),
+ "mlgo-model", llvm::cl::desc("Select the MLGO model to execute:"),
llvm::cl::init(MLGOModelChoice::Default),
- llvm::cl::values(
- clEnumValN(MLGOModelChoice::Default, "default", "Use standard heuristic")
-#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) \
- , clEnumValN(MLGOModelChoice::CLASS_NAME, CLI_FLAG, "Use the " CLI_FLAG " MLGO model")
+ llvm::cl::values(clEnumValN(MLGOModelChoice::Default, "default",
+ "Use standard heuristic")
+#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) \
+ , clEnumValN(MLGOModelChoice::CLASS_NAME, CLI_FLAG, \
+ "Use the " CLI_FLAG " MLGO model")
#include "llvm/Analysis/MLGOModels.def"
- )
-);
+ ));
-static std::unique_ptr<MLModelRunner> createMLGOModelRunner(LLVMContext &Ctx, const std::vector<TensorSpec> &InputFeatures) {
+static std::unique_ptr<MLModelRunner>
+createMLGOModelRunner(LLVMContext &Ctx,
+ const std::vector<TensorSpec> &InputFeatures) {
switch (SelectedMLGOModel) {
- case MLGOModelChoice::Default:
- return nullptr;
-#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) \
- case MLGOModelChoice::CLASS_NAME: \
- return std::make_unique<EmitCModelRunner<CLASS_NAME>>(Ctx, InputFeatures);
+ case MLGOModelChoice::Default:
+ return nullptr;
+#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) \
+ case MLGOModelChoice::CLASS_NAME: \
+ return std::make_unique<EmitCModelRunner<CLASS_NAME>>(Ctx, InputFeatures);
#include "llvm/Analysis/MLGOModels.def"
}
llvm_unreachable("Unknown MLGO model type!");
@@ -115,7 +116,7 @@ llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM,
#if defined(LLVM_HAVE_EMITC_COMPILE_INLINER)
&& SelectedMLGOModel == MLGOModelChoice::Default
#endif
- )
+ )
return nullptr;
auto RunnerFactory = [&](const std::vector<TensorSpec> &InputFeatures)
-> std::unique_ptr<MLModelRunner> {
diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index b4a55dbc79d7b..9366b81f37924 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -69,21 +69,23 @@ static llvm::cl::opt<MLGORegAllocModelChoice> SelectedMLGORegAllocModel(
"regalloc-mlgo-model",
llvm::cl::desc("Select the MLGO model to execute for register allocation:"),
llvm::cl::init(MLGORegAllocModelChoice::Default),
- llvm::cl::values(
- clEnumValN(MLGORegAllocModelChoice::Default, "default", "Use standard heuristic")
-#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) \
- , clEnumValN(MLGORegAllocModelChoice::CLASS_NAME, CLI_FLAG, "Use the " CLI_FLAG " MLGO model")
+ llvm::cl::values(clEnumValN(MLGORegAllocModelChoice::Default, "default",
+ "Use standard heuristic")
+#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) \
+ , clEnumValN(MLGORegAllocModelChoice::CLASS_NAME, CLI_FLAG, \
+ "Use the " CLI_FLAG " MLGO model")
#include "llvm/CodeGen/RegAllocEvictModels.def"
- )
-);
+ ));
-static std::unique_ptr<MLModelRunner> createMLGORegAllocModelRunner(LLVMContext &Ctx, const std::vector<TensorSpec> &InputFeatures) {
+static std::unique_ptr<MLModelRunner>
+createMLGORegAllocModelRunner(LLVMContext &Ctx,
+ const std::vector<TensorSpec> &InputFeatures) {
switch (SelectedMLGORegAllocModel) {
- case MLGORegAllocModelChoice::Default:
- return nullptr;
-#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) \
- case MLGORegAllocModelChoice::CLASS_NAME: \
- return std::make_unique<EmitCModelRunner<CLASS_NAME>>(Ctx, InputFeatures);
+ case MLGORegAllocModelChoice::Default:
+ return nullptr;
+#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) \
+ case MLGORegAllocModelChoice::CLASS_NAME: \
+ return std::make_unique<EmitCModelRunner<CLASS_NAME>>(Ctx, InputFeatures);
#include "llvm/CodeGen/RegAllocEvictModels.def"
}
llvm_unreachable("Unknown MLGO model type!");
@@ -406,7 +408,8 @@ class ReleaseModeEvictionAdvisorProvider final
if (InteractiveChannelBaseName.empty()) {
#if defined(LLVM_HAVE_EMITC_COMPILE_REGALLOC)
if (SelectedMLGORegAllocModel != MLGORegAllocModelChoice::Default) {
- Runner = createMLGORegAllocModelRunner(MF.getFunction().getContext(), InputFeatures);
+ Runner = createMLGORegAllocModelRunner(MF.getFunction().getContext(),
+ InputFeatures);
} else
#endif
{
@@ -1064,7 +1067,8 @@ llvm::createReleaseModeAdvisorProvider(LLVMContext &Ctx) {
return llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
!InteractiveChannelBaseName.empty()
#if defined(LLVM_HAVE_EMITC_COMPILE_REGALLOC)
- || SelectedMLGORegAllocModel != MLGORegAllocModelChoice::Default
+ ||
+ SelectedMLGORegAllocModel != MLGORegAllocModelChoice::Default
#endif
? new ReleaseModeEvictionAdvisorProvider(Ctx)
: nullptr;
@@ -1083,7 +1087,8 @@ llvm::createReleaseModeAdvisorAnalysisLegacy() {
return llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
!InteractiveChannelBaseName.empty()
#if defined(LLVM_HAVE_EMITC_COMPILE_REGALLOC)
- || SelectedMLGORegAllocModel != MLGORegAllocModelChoice::Default
+ ||
+ SelectedMLGORegAllocModel != MLGORegAllocModelChoice::Default
#endif
? new ReleaseModeEvictionAdvisorAnalysisLegacy()
: nullptr;
>From 7498ffea13b02027cf6f78c99b7430508b9be3d4 Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Tue, 28 Jul 2026 18:43:32 -0700
Subject: [PATCH 04/11] Address review comments
---
llvm/include/llvm/Analysis/EmitCModelRunner.h | 29 +++++++-----------
llvm/lib/Analysis/CMakeLists.txt | 2 +-
llvm/lib/Analysis/MLInlineAdvisor.cpp | 30 +++++++++----------
3 files changed, 26 insertions(+), 35 deletions(-)
diff --git a/llvm/include/llvm/Analysis/EmitCModelRunner.h b/llvm/include/llvm/Analysis/EmitCModelRunner.h
index f60ca034d232d..2517f070d22e3 100644
--- a/llvm/include/llvm/Analysis/EmitCModelRunner.h
+++ b/llvm/include/llvm/Analysis/EmitCModelRunner.h
@@ -17,7 +17,6 @@
#include "llvm/Analysis/MLModelRunner.h"
#include "llvm/Analysis/TensorSpec.h"
-#include <memory>
#include <type_traits>
namespace llvm {
@@ -25,11 +24,8 @@ namespace llvm {
template <class TGen> class EmitCModelRunner final : public MLModelRunner {
public:
template <class FType>
- EmitCModelRunner(LLVMContext &Ctx, const FType &InputSpec,
- std::unique_ptr<TGen> Model = std::make_unique<TGen>())
- : MLModelRunner(Ctx, MLModelRunner::Kind::Release, InputSpec.size()),
- CompiledModel(std::move(Model)) {
- assert(CompiledModel && "The CompiledModel should be valid");
+ EmitCModelRunner(LLVMContext &Ctx, const FType &InputSpec)
+ : MLModelRunner(Ctx, MLModelRunner::Kind::Release, InputSpec.size()) {
for (size_t I = 0; I < InputSpec.size(); ++I)
populateTensor(I, InputSpec[I]);
}
@@ -46,28 +42,23 @@ template <class TGen> class EmitCModelRunner final : public MLModelRunner {
private:
void populateTensor(size_t Pos, const TensorSpec &Spec) {
void *Buffer = nullptr;
- auto It = CompiledModel->reflectionMap.find(Spec.name());
- if (It != CompiledModel->reflectionMap.end())
+ auto It = CompiledModel.reflectionMap.find(Spec.name());
+ if (It != CompiledModel.reflectionMap.end())
Buffer = static_cast<void *>(It->second);
setUpBufferForTensor(Pos, Spec, Buffer);
}
using ResultType = decltype(std::declval<TGen>()());
+ static_assert(!std::is_void_v<ResultType>,
+ "EmitCModelRunner models must return a non-void result.");
- template <typename R = ResultType>
- std::enable_if_t<!std::is_void_v<R>, void *> evaluateImpl() {
- Result = (*CompiledModel)();
+ void *evaluateImpl() {
+ Result = CompiledModel();
return &Result;
}
- template <typename R = ResultType>
- std::enable_if_t<std::is_void_v<R>, void *> evaluateImpl() {
- (*CompiledModel)();
- return nullptr;
- }
-
- std::conditional_t<std::is_void_v<ResultType>, char, ResultType> Result{};
- std::unique_ptr<TGen> CompiledModel;
+ ResultType Result{};
+ TGen CompiledModel{};
};
} // namespace llvm
diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt
index ed7fc6e7e26c0..408d27474007b 100644
--- a/llvm/lib/Analysis/CMakeLists.txt
+++ b/llvm/lib/Analysis/CMakeLists.txt
@@ -6,7 +6,7 @@ if (LLVM_HAVE_EMITC_COMPILE)
"${LLVM_MLGO_MLIR_TRANSLATE}"
inliner
llvm/Analysis
- MLGOModels.def
+ InlinerModels.def
InlinerSizeModelMulti.h
)
elseif (DEFINED LLVM_HAVE_TF_AOT OR LLVM_HAVE_TFLITE)
diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp
index e04805ff135af..379d62b7ea25e 100644
--- a/llvm/lib/Analysis/MLInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp
@@ -76,33 +76,33 @@ using CompiledModelType = NoopSavedModelImpl;
#include "llvm/Analysis/EmitCModelRunner.h"
#include "llvm/Analysis/InlinerSizeModelMulti.h"
-enum class MLGOModelChoice {
+enum class EmitCModelChoice {
Default,
#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) CLASS_NAME,
-#include "llvm/Analysis/MLGOModels.def"
+#include "llvm/Analysis/InlinerModels.def"
};
-static llvm::cl::opt<MLGOModelChoice> SelectedMLGOModel(
+static llvm::cl::opt<EmitCModelChoice> SelectedMLGOModel(
"mlgo-model", llvm::cl::desc("Select the MLGO model to execute:"),
- llvm::cl::init(MLGOModelChoice::Default),
- llvm::cl::values(clEnumValN(MLGOModelChoice::Default, "default",
+ llvm::cl::init(EmitCModelChoice::Default),
+ llvm::cl::values(clEnumValN(EmitCModelChoice::Default, "default",
"Use standard heuristic")
#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) \
- , clEnumValN(MLGOModelChoice::CLASS_NAME, CLI_FLAG, \
+ , clEnumValN(EmitCModelChoice::CLASS_NAME, CLI_FLAG, \
"Use the " CLI_FLAG " MLGO model")
-#include "llvm/Analysis/MLGOModels.def"
+#include "llvm/Analysis/InlinerModels.def"
));
static std::unique_ptr<MLModelRunner>
-createMLGOModelRunner(LLVMContext &Ctx,
- const std::vector<TensorSpec> &InputFeatures) {
+createEmitCModelRunner(LLVMContext &Ctx,
+ const std::vector<TensorSpec> &InputFeatures) {
switch (SelectedMLGOModel) {
- case MLGOModelChoice::Default:
+ case EmitCModelChoice::Default:
return nullptr;
#define MLGO_MODEL(CLASS_NAME, CLI_FLAG) \
- case MLGOModelChoice::CLASS_NAME: \
+ case EmitCModelChoice::CLASS_NAME: \
return std::make_unique<EmitCModelRunner<CLASS_NAME>>(Ctx, InputFeatures);
-#include "llvm/Analysis/MLGOModels.def"
+#include "llvm/Analysis/InlinerModels.def"
}
llvm_unreachable("Unknown MLGO model type!");
}
@@ -114,7 +114,7 @@ llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM,
if (!llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() &&
InteractiveChannelBaseName.empty()
#if defined(LLVM_HAVE_EMITC_COMPILE_INLINER)
- && SelectedMLGOModel == MLGOModelChoice::Default
+ && SelectedMLGOModel == EmitCModelChoice::Default
#endif
)
return nullptr;
@@ -123,8 +123,8 @@ llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM,
std::unique_ptr<MLModelRunner> AOTRunner;
if (InteractiveChannelBaseName.empty()) {
#if defined(LLVM_HAVE_EMITC_COMPILE_INLINER)
- if (SelectedMLGOModel != MLGOModelChoice::Default) {
- AOTRunner = createMLGOModelRunner(M.getContext(), InputFeatures);
+ if (SelectedMLGOModel != EmitCModelChoice::Default) {
+ AOTRunner = createEmitCModelRunner(M.getContext(), InputFeatures);
} else
#endif
{
>From af063d5f3eaee02a58012b8090c9848a436841c8 Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Wed, 29 Jul 2026 14:00:12 -0700
Subject: [PATCH 05/11] Address review comments (1)
---
llvm/CMakeLists.txt | 2 +-
.../{MLGOCompile.cmake => MLGOLower.cmake} | 49 +++++++++++--------
llvm/include/llvm/Analysis/EmitCModelRunner.h | 4 +-
llvm/lib/Analysis/CMakeLists.txt | 7 ++-
llvm/lib/Analysis/MLInlineAdvisor.cpp | 8 +--
llvm/lib/CodeGen/CMakeLists.txt | 7 ++-
llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp | 10 ++--
7 files changed, 47 insertions(+), 40 deletions(-)
rename llvm/cmake/modules/{MLGOCompile.cmake => MLGOLower.cmake} (73%)
diff --git a/llvm/CMakeLists.txt b/llvm/CMakeLists.txt
index 8490c70617a7c..3a548cbdbd1cc 100644
--- a/llvm/CMakeLists.txt
+++ b/llvm/CMakeLists.txt
@@ -1227,7 +1227,7 @@ set( CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${LLVM_LIBRARY_DIR} )
# the bot setup script: https://github.com/google/ml-compiler-opt/blob/main/buildbot/buildbot_init.sh
set(LLVM_HAVE_TFLITE "" CACHE BOOL "Use tflite")
-set(LLVM_MLGO_MODELS "" CACHE STRING "List of AOT MLGO models")
+set(LLVM_MLGO_MODELS "" CACHE STRING "List of MLGO models")
set(LLVM_MLGO_MLIR_OPT "mlir-opt" CACHE STRING "Path or system binary name for mlir-opt")
set(LLVM_MLGO_MLIR_TRANSLATE "mlir-translate" CACHE STRING "Path or system binary name for mlir-translate")
diff --git a/llvm/cmake/modules/MLGOCompile.cmake b/llvm/cmake/modules/MLGOLower.cmake
similarity index 73%
rename from llvm/cmake/modules/MLGOCompile.cmake
rename to llvm/cmake/modules/MLGOLower.cmake
index b4e340087bd62..47ffc972265b8 100644
--- a/llvm/cmake/modules/MLGOCompile.cmake
+++ b/llvm/cmake/modules/MLGOLower.cmake
@@ -1,25 +1,34 @@
-# Compile MLGO models expressed as MLIR to C++ headers via the EmitC pipeline.
+# Lower MLGO models to C++ headers via the MLIR-based pipeline.
#
-# Each entry of ${models} has the form
-# "cli-flag,path/to/model.mlir,type". Entries whose type matches
-# ${target_type} are lowered with ${mlir_opt} and translated with
-# ${mlir_translate} to a C++ header defining Model<N>, placed in
-# ${LLVM_INCLUDE_DIR}/${include_subdir}. In the same directory, generate
-# ${def_file}, with one MLGO_MODEL(ClassName, "cli-flag") entry per model, and
-# ${umbrella_header}, including all the generated model headers.
-# Append the target driving generation to MLDeps in the caller's scope, and
-# define LLVM_HAVE_EMITC_COMPILE_<TARGET_TYPE>.
-function(mlgo_compile_models models mlir_opt mlir_translate target_type
- include_subdir def_file umbrella_header)
- if ("${mlir_opt}" MATCHES "/")
- get_filename_component(mlir_opt "${mlir_opt}" ABSOLUTE BASE_DIR "${CMAKE_BINARY_DIR}")
+# Each entry in ${models} has the form "cli-flag,path/to/model.mlir,type".
+# For entries whose type matches ${target_type}, this function:
+# 1. Lowers model MLIR with ${mlir_opt} and translates it with ${mlir_translate}
+# to a C++ header defining Model<N> in ${LLVM_INCLUDE_DIR}/${include_subdir}.
+# 2. Generates ${generated_file_basename}.def in ${LLVM_INCLUDE_DIR}/${include_subdir}
+# with one MLGO_MODEL(ClassName, "cli-flag") entry per model.
+# 3. Generates ${generated_file_basename}.h in ${LLVM_INCLUDE_DIR}/${include_subdir}
+# including all generated model headers.
+# 4. Appends the target driving generation to MLDeps in the caller's scope.
+# 5. Defines LLVM_HAVE_MLIR_LOWERING_<TARGET_TYPE>.
+function(mlgo_lower_models models mlir_opt mlir_translate target_type
+ include_subdir generated_file_basename)
+ set(def_file "${generated_file_basename}.def")
+ set(umbrella_header "${generated_file_basename}.h")
+
+ if(NOT IS_ABSOLUTE "${mlir_opt}")
+ cmake_path(ABSOLUTE_PATH mlir_opt BASE_DIRECTORY "${CMAKE_BINARY_DIR}")
endif()
- if ("${mlir_translate}" MATCHES "/")
- get_filename_component(mlir_translate "${mlir_translate}" ABSOLUTE BASE_DIR "${CMAKE_BINARY_DIR}")
+
+ if(NOT IS_ABSOLUTE "${mlir_translate}")
+ cmake_path(ABSOLUTE_PATH mlir_translate BASE_DIRECTORY "${CMAKE_BINARY_DIR}")
endif()
- set(DEF_CONTENT "/* Auto-generated by CMake */\n")
- set(DEF_CONTENT "${DEF_CONTENT}#ifndef MLGO_MODEL\n#define MLGO_MODEL(CLASS_NAME, CLI_FLAG)\n#endif\n\n")
+ string(APPEND DEF_CONTENT
+ "/* Auto-generated by CMake */\n"
+ "#ifndef MLGO_MODEL\n"
+ "# define MLGO_MODEL(CLASS_NAME, CLI_FLAG)\n"
+ "#endif\n\n"
+ )
set(HEADERS_CONTENT "/* Auto-generated by CMake */\nnamespace llvm {\n")
set(MLGO_GEN_TARGETS "")
@@ -71,7 +80,7 @@ function(mlgo_compile_models models mlir_opt mlir_translate target_type
string(JOIN "," PASS_PIPELINE ${EMITC_PASSES})
set(PASS_PIPELINE "builtin.module(${PASS_PIPELINE})")
- # 1. Run MLIR pipeline to compile MODEL_PATH to EmitC MLIR
+ # 1. Run MLIR pipeline to lower MODEL_PATH to EmitC MLIR
add_custom_command(OUTPUT ${EMITC_MLIR}
COMMAND ${mlir_opt} "--pass-pipeline=${PASS_PIPELINE}"
${MODEL_PATH} -o ${EMITC_MLIR}
@@ -118,6 +127,6 @@ function(mlgo_compile_models models mlir_opt mlir_translate target_type
add_custom_target(mlgo_models_gen_${target_type} DEPENDS ${MLGO_GEN_TARGETS})
set(MLDeps ${MLDeps} mlgo_models_gen_${target_type} PARENT_SCOPE)
string(TOUPPER ${target_type} target_type_allcaps)
- add_compile_definitions(LLVM_HAVE_EMITC_COMPILE_${target_type_allcaps})
+ add_compile_definitions(LLVM_HAVE_MLIR_LOWERING_${target_type_allcaps})
endif()
endfunction()
diff --git a/llvm/include/llvm/Analysis/EmitCModelRunner.h b/llvm/include/llvm/Analysis/EmitCModelRunner.h
index 2517f070d22e3..cb14553aaa0e8 100644
--- a/llvm/include/llvm/Analysis/EmitCModelRunner.h
+++ b/llvm/include/llvm/Analysis/EmitCModelRunner.h
@@ -26,8 +26,8 @@ template <class TGen> class EmitCModelRunner final : public MLModelRunner {
template <class FType>
EmitCModelRunner(LLVMContext &Ctx, const FType &InputSpec)
: MLModelRunner(Ctx, MLModelRunner::Kind::Release, InputSpec.size()) {
- for (size_t I = 0; I < InputSpec.size(); ++I)
- populateTensor(I, InputSpec[I]);
+ for (auto [I, Spec] : llvm::enumerate(InputSpec))
+ populateTensor(I, Spec);
}
~EmitCModelRunner() override = default;
diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt
index 408d27474007b..99c189562727f 100644
--- a/llvm/lib/Analysis/CMakeLists.txt
+++ b/llvm/lib/Analysis/CMakeLists.txt
@@ -1,13 +1,12 @@
if (LLVM_HAVE_EMITC_COMPILE)
- include(MLGOCompile)
- mlgo_compile_models(
+ include(MLGOLower)
+ mlgo_lower_models(
"${LLVM_MLGO_MODELS}"
"${LLVM_MLGO_MLIR_OPT}"
"${LLVM_MLGO_MLIR_TRANSLATE}"
inliner
llvm/Analysis
- InlinerModels.def
- InlinerSizeModelMulti.h
+ InlinerModels
)
elseif (DEFINED LLVM_HAVE_TF_AOT OR LLVM_HAVE_TFLITE)
include(TensorFlowCompile)
diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp
index 379d62b7ea25e..a4e3a167a8cb6 100644
--- a/llvm/lib/Analysis/MLInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp
@@ -72,9 +72,9 @@ using CompiledModelType = llvm::InlinerSizeModel;
using CompiledModelType = NoopSavedModelImpl;
#endif
-#if defined(LLVM_HAVE_EMITC_COMPILE_INLINER)
+#if defined(LLVM_HAVE_MLIR_LOWERING_INLINER)
#include "llvm/Analysis/EmitCModelRunner.h"
-#include "llvm/Analysis/InlinerSizeModelMulti.h"
+#include "llvm/Analysis/InlinerModels.h"
enum class EmitCModelChoice {
Default,
@@ -113,7 +113,7 @@ llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM,
std::function<bool(CallBase &)> GetDefaultAdvice) {
if (!llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() &&
InteractiveChannelBaseName.empty()
-#if defined(LLVM_HAVE_EMITC_COMPILE_INLINER)
+#if defined(LLVM_HAVE_MLIR_LOWERING_INLINER)
&& SelectedMLGOModel == EmitCModelChoice::Default
#endif
)
@@ -122,7 +122,7 @@ llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM,
-> std::unique_ptr<MLModelRunner> {
std::unique_ptr<MLModelRunner> AOTRunner;
if (InteractiveChannelBaseName.empty()) {
-#if defined(LLVM_HAVE_EMITC_COMPILE_INLINER)
+#if defined(LLVM_HAVE_MLIR_LOWERING_INLINER)
if (SelectedMLGOModel != EmitCModelChoice::Default) {
AOTRunner = createEmitCModelRunner(M.getContext(), InputFeatures);
} else
diff --git a/llvm/lib/CodeGen/CMakeLists.txt b/llvm/lib/CodeGen/CMakeLists.txt
index 1c913c9c2f8ff..4f88a4306e609 100644
--- a/llvm/lib/CodeGen/CMakeLists.txt
+++ b/llvm/lib/CodeGen/CMakeLists.txt
@@ -1,13 +1,12 @@
if (LLVM_HAVE_EMITC_COMPILE)
- include(MLGOCompile)
- mlgo_compile_models(
+ include(MLGOLower)
+ mlgo_lower_models(
"${LLVM_MLGO_MODELS}"
"${LLVM_MLGO_MLIR_OPT}"
"${LLVM_MLGO_MLIR_TRANSLATE}"
regalloc
llvm/CodeGen
- RegAllocEvictModels.def
- RegAllocEvictModelMulti.h
+ RegAllocEvictModels
)
elseif (DEFINED LLVM_HAVE_TF_AOT OR LLVM_HAVE_TFLITE)
include(TensorFlowCompile)
diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index 9366b81f37924..be02b96e8d921 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -55,9 +55,9 @@ using CompiledModelType = RegAllocEvictModel;
using CompiledModelType = NoopSavedModelImpl;
#endif
-#if defined(LLVM_HAVE_EMITC_COMPILE_REGALLOC)
+#if defined(LLVM_HAVE_MLIR_LOWERING_REGALLOC)
#include "llvm/Analysis/EmitCModelRunner.h"
-#include "llvm/CodeGen/RegAllocEvictModelMulti.h"
+#include "llvm/CodeGen/RegAllocEvictModels.h"
enum class MLGORegAllocModelChoice {
Default,
@@ -406,7 +406,7 @@ class ReleaseModeEvictionAdvisorProvider final
MachineBlockFrequencyInfo *MBFI, MachineLoopInfo *Loops) override {
if (!Runner) {
if (InteractiveChannelBaseName.empty()) {
-#if defined(LLVM_HAVE_EMITC_COMPILE_REGALLOC)
+#if defined(LLVM_HAVE_MLIR_LOWERING_REGALLOC)
if (SelectedMLGORegAllocModel != MLGORegAllocModelChoice::Default) {
Runner = createMLGORegAllocModelRunner(MF.getFunction().getContext(),
InputFeatures);
@@ -1066,7 +1066,7 @@ RegAllocEvictionAdvisorProvider *
llvm::createReleaseModeAdvisorProvider(LLVMContext &Ctx) {
return llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
!InteractiveChannelBaseName.empty()
-#if defined(LLVM_HAVE_EMITC_COMPILE_REGALLOC)
+#if defined(LLVM_HAVE_MLIR_LOWERING_REGALLOC)
||
SelectedMLGORegAllocModel != MLGORegAllocModelChoice::Default
#endif
@@ -1086,7 +1086,7 @@ RegAllocEvictionAdvisorAnalysisLegacy *
llvm::createReleaseModeAdvisorAnalysisLegacy() {
return llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
!InteractiveChannelBaseName.empty()
-#if defined(LLVM_HAVE_EMITC_COMPILE_REGALLOC)
+#if defined(LLVM_HAVE_MLIR_LOWERING_REGALLOC)
||
SelectedMLGORegAllocModel != MLGORegAllocModelChoice::Default
#endif
>From 767ab0719d2e5883a06fedb66e98641d4b7b4889 Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Wed, 29 Jul 2026 16:22:55 -0700
Subject: [PATCH 06/11] Address review comments (2)
---
llvm/CMakeLists.txt | 5 +--
llvm/cmake/modules/MLGOLower.cmake | 8 ++--
llvm/include/llvm/Analysis/EmitCModelRunner.h | 14 +++----
llvm/lib/Analysis/MLInlineAdvisor.cpp | 38 ++++++++++---------
4 files changed, 35 insertions(+), 30 deletions(-)
diff --git a/llvm/CMakeLists.txt b/llvm/CMakeLists.txt
index 3a548cbdbd1cc..660befb697a52 100644
--- a/llvm/CMakeLists.txt
+++ b/llvm/CMakeLists.txt
@@ -1231,10 +1231,9 @@ set(LLVM_MLGO_MODELS "" CACHE STRING "List of MLGO models")
set(LLVM_MLGO_MLIR_OPT "mlir-opt" CACHE STRING "Path or system binary name for mlir-opt")
set(LLVM_MLGO_MLIR_TRANSLATE "mlir-translate" CACHE STRING "Path or system binary name for mlir-translate")
+set(LLVM_HAVE_EMITC_COMPILE OFF CACHE BOOL "MLGO models to be compiled with EmitC are available")
if (NOT LLVM_MLGO_MODELS STREQUAL "")
- set(LLVM_HAVE_EMITC_COMPILE "ON" CACHE BOOL "MLGO AOT models compiled with EmitC are available" FORCE)
-else()
- set(LLVM_HAVE_EMITC_COMPILE "OFF" CACHE BOOL "MLGO AOT models compiled with EmitC are available" FORCE)
+ set(LLVM_HAVE_EMITC_COMPILE ON CACHE BOOL "MLGO models to be compiled with EmitC are available" FORCE)
endif()
if (LLVM_HAVE_EMITC_COMPILE AND LLVM_HAVE_TFLITE)
diff --git a/llvm/cmake/modules/MLGOLower.cmake b/llvm/cmake/modules/MLGOLower.cmake
index 47ffc972265b8..82e32ea5c3457 100644
--- a/llvm/cmake/modules/MLGOLower.cmake
+++ b/llvm/cmake/modules/MLGOLower.cmake
@@ -35,7 +35,8 @@ function(mlgo_lower_models models mlir_opt mlir_translate target_type
set(MODEL_INDEX 1)
foreach(MODEL_INFO IN LISTS models)
- # Parse comma-separated fields: cli-flag,path/to/model.mlir,type
+ # Parse comma-separated fields (from the LLVM_MLGO_MODELS flag)
+ # cli-flag,path/to/model.mlir,type
string(REPLACE "," ";" MODEL_FIELDS "${MODEL_INFO}")
list(GET MODEL_FIELDS 0 CLI_FLAG)
list(GET MODEL_FIELDS 1 MODEL_PATH)
@@ -57,6 +58,7 @@ function(mlgo_lower_models models mlir_opt mlir_translate target_type
set(HEADER_FILE "${LLVM_INCLUDE_DIR}/${include_subdir}/${CLASS_NAME}.h")
# Pass pipeline to lower MLIR models to EmitC dialect
+ # TODO: Simplify with builtin pipeline for translation.
set(EMITC_PASSES
"func.func(tosa-to-linalg-named,tosa-to-linalg,tosa-to-arith,tosa-to-tensor)"
"symbol-privatize"
@@ -107,11 +109,11 @@ function(mlgo_lower_models models mlir_opt mlir_translate target_type
string(APPEND DEF_CONTENT "MLGO_MODEL(${CLASS_NAME}, \"${CLI_FLAG}\")\n")
# Append to the master header content
- string(APPEND HEADERS_CONTENT "namespace ${CLASS_NAME}_ns {\n#include \"${include_subdir}/${CLASS_NAME}.h\"\n}\nusing ${CLASS_NAME}_ns::${CLASS_NAME};\n")
+ string(APPEND HEADERS_CONTENT "namespace ${CLASS_NAME}_ns {\n#include \"${include_subdir}/${CLASS_NAME}.h\"\n} // namespace ${CLASS_NAME}_ns\nusing ${CLASS_NAME}_ns::${CLASS_NAME};\n")
endforeach()
string(APPEND DEF_CONTENT "\n#undef MLGO_MODEL\n")
- string(APPEND HEADERS_CONTENT "}\n")
+ string(APPEND HEADERS_CONTENT "} // namespace llvm\n")
# Stage the generated files, then only update the copies under
# LLVM_INCLUDE_DIR when their content changed, so that reconfiguring does not
diff --git a/llvm/include/llvm/Analysis/EmitCModelRunner.h b/llvm/include/llvm/Analysis/EmitCModelRunner.h
index cb14553aaa0e8..21edb09450de6 100644
--- a/llvm/include/llvm/Analysis/EmitCModelRunner.h
+++ b/llvm/include/llvm/Analysis/EmitCModelRunner.h
@@ -1,14 +1,14 @@
-//===- EmitCModelRunner.h - Fast, precompiled model runner ---------------===//
+//===----------------------------------------------------------------------===//
//
// 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 a model runner wrapping an EmitC compiled ML model.
-// Only inference is supported.
-//
+///
+/// \file
+/// This file implements a model runner wrapping an EmitC compiled ML model.
+///
//===----------------------------------------------------------------------===//
#ifndef LLVM_ANALYSIS_EMITCMODELRUNNER_H
@@ -57,8 +57,8 @@ template <class TGen> class EmitCModelRunner final : public MLModelRunner {
return &Result;
}
- ResultType Result{};
- TGen CompiledModel{};
+ ResultType Result = {};
+ TGen CompiledModel = {};
};
} // namespace llvm
diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp
index a4e3a167a8cb6..8ff00b2bb6646 100644
--- a/llvm/lib/Analysis/MLInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp
@@ -73,6 +73,7 @@ using CompiledModelType = NoopSavedModelImpl;
#endif
#if defined(LLVM_HAVE_MLIR_LOWERING_INLINER)
+constexpr bool HaveMLIRLoweringInliner = true;
#include "llvm/Analysis/EmitCModelRunner.h"
#include "llvm/Analysis/InlinerModels.h"
@@ -106,39 +107,42 @@ createEmitCModelRunner(LLVMContext &Ctx,
}
llvm_unreachable("Unknown MLGO model type!");
}
+#else
+constexpr bool HaveMLIRLoweringInliner = false;
+enum class EmitCModelChoice { Default };
+static const EmitCModelChoice SelectedMLGOModel = EmitCModelChoice::Default;
+inline std::unique_ptr<MLModelRunner>
+createEmitCModelRunner(LLVMContext &, const std::vector<TensorSpec> &) {
+ return nullptr;
+}
#endif
std::unique_ptr<InlineAdvisor>
llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM,
std::function<bool(CallBase &)> GetDefaultAdvice) {
if (!llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() &&
- InteractiveChannelBaseName.empty()
-#if defined(LLVM_HAVE_MLIR_LOWERING_INLINER)
- && SelectedMLGOModel == EmitCModelChoice::Default
-#endif
- )
+ InteractiveChannelBaseName.empty() &&
+ SelectedMLGOModel == EmitCModelChoice::Default)
return nullptr;
auto RunnerFactory = [&](const std::vector<TensorSpec> &InputFeatures)
-> std::unique_ptr<MLModelRunner> {
- std::unique_ptr<MLModelRunner> AOTRunner;
+ std::unique_ptr<MLModelRunner> ModelRunner;
if (InteractiveChannelBaseName.empty()) {
-#if defined(LLVM_HAVE_MLIR_LOWERING_INLINER)
- if (SelectedMLGOModel != EmitCModelChoice::Default) {
- AOTRunner = createEmitCModelRunner(M.getContext(), InputFeatures);
- } else
-#endif
- {
- AOTRunner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
- M.getContext(), InputFeatures, DecisionName,
- EmbeddedModelRunnerOptions().setModelSelector(ModelSelector));
+ if constexpr (HaveMLIRLoweringInliner) {
+ ModelRunner = createEmitCModelRunner(M.getContext(), InputFeatures);
+ } else {
+ ModelRunner =
+ std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
+ M.getContext(), InputFeatures, DecisionName,
+ EmbeddedModelRunnerOptions().setModelSelector(ModelSelector));
}
} else {
- AOTRunner = std::make_unique<InteractiveModelRunner>(
+ ModelRunner = std::make_unique<InteractiveModelRunner>(
M.getContext(), InputFeatures, InlineDecisionSpec,
InteractiveChannelBaseName + ".out",
InteractiveChannelBaseName + ".in");
}
- return AOTRunner;
+ return ModelRunner;
};
return std::make_unique<MLInlineAdvisor>(M, MAM, RunnerFactory,
GetDefaultAdvice);
>From 5ffc10580af361bbf5ea769efcc002dd444ad7b4 Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Fri, 31 Jul 2026 12:29:21 -0700
Subject: [PATCH 07/11] Create helper functions for shared logic in inline and
regalloc advisors
---
llvm/include/llvm/Analysis/Utils/MLGOUtils.h | 82 ++++++++++++++++++++
llvm/lib/Analysis/MLInlineAdvisor.cpp | 31 +++-----
llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp | 51 +++++-------
3 files changed, 111 insertions(+), 53 deletions(-)
create mode 100644 llvm/include/llvm/Analysis/Utils/MLGOUtils.h
diff --git a/llvm/include/llvm/Analysis/Utils/MLGOUtils.h b/llvm/include/llvm/Analysis/Utils/MLGOUtils.h
new file mode 100644
index 0000000000000..6bd69d26ae8a5
--- /dev/null
+++ b/llvm/include/llvm/Analysis/Utils/MLGOUtils.h
@@ -0,0 +1,82 @@
+//===- MLGOUtils.h - Utilities for MLGO Release Mode ------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// This file provides helper functions for creating MLModelRunners and checking
+/// model validity in release mode.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_ANALYSIS_UTILS_MLGOUTILS_H
+#define LLVM_ANALYSIS_UTILS_MLGOUTILS_H
+
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Analysis/InteractiveModelRunner.h"
+#include "llvm/Analysis/MLModelRunner.h"
+#include "llvm/Analysis/ReleaseModeModelRunner.h"
+#include "llvm/Analysis/TensorSpec.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/Support/CommandLine.h"
+#include <memory>
+#include <string>
+#include <vector>
+
+namespace llvm {
+
+/// Helper to check if a release-mode ML advisor has a valid model to execute.
+/// Overload for cl::opt<EnumType>.
+template <class CompiledModelType, class EnumType, bool ExternalStorage,
+ class ParserClass>
+bool isReleaseModelValid(
+ StringRef InteractiveChannelBaseName,
+ const cl::opt<EnumType, ExternalStorage, ParserClass> &SelectedModel,
+ EnumType DefaultModelVal = EnumType::Default) {
+ return isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
+ !InteractiveChannelBaseName.empty() ||
+ SelectedModel != DefaultModelVal;
+}
+
+/// Helper to check if a release-mode ML advisor has a valid model to execute.
+/// Overload for plain EnumType.
+template <class CompiledModelType, class EnumType>
+bool isReleaseModelValid(StringRef InteractiveChannelBaseName,
+ EnumType SelectedModel,
+ EnumType DefaultModelVal = EnumType::Default) {
+ return isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
+ !InteractiveChannelBaseName.empty() ||
+ SelectedModel != DefaultModelVal;
+}
+
+/// Helper to construct the appropriate MLModelRunner in release mode:
+/// 1. InteractiveModelRunner if an interactive channel is specified.
+/// 2. EmitCModelRunner if MLIR lowering is enabled.
+/// 3. ReleaseModeModelRunner<CompiledModelType> otherwise.
+template <class CompiledModelType, bool HaveMLIRLowering, class CreateEmitCFunc>
+std::unique_ptr<MLModelRunner> createReleaseModeModelRunner(
+ LLVMContext &Ctx, const std::vector<TensorSpec> &InputFeatures,
+ StringRef DecisionName, const std::string &InteractiveChannelBaseName,
+ const TensorSpec &InteractiveDecisionSpec,
+ CreateEmitCFunc &&CreateEmitCModelRunner,
+ const EmbeddedModelRunnerOptions &Options = {}) {
+ if (!InteractiveChannelBaseName.empty()) {
+ return std::make_unique<InteractiveModelRunner>(
+ Ctx, InputFeatures, InteractiveDecisionSpec,
+ InteractiveChannelBaseName + ".out",
+ InteractiveChannelBaseName + ".in");
+ }
+ if constexpr (HaveMLIRLowering) {
+ return CreateEmitCModelRunner(Ctx, InputFeatures);
+ } else {
+ return std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
+ Ctx, InputFeatures, DecisionName, Options);
+ }
+}
+
+} // namespace llvm
+
+#endif // LLVM_ANALYSIS_UTILS_MLGOUTILS_H
diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp
index 8ff00b2bb6646..fd0e4b152bcc2 100644
--- a/llvm/lib/Analysis/MLInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp
@@ -19,7 +19,6 @@
#include "llvm/Analysis/FunctionPropertiesAnalysis.h"
#include "llvm/Analysis/InlineCost.h"
#include "llvm/Analysis/InlineModelFeatureMaps.h"
-#include "llvm/Analysis/InteractiveModelRunner.h"
#include "llvm/Analysis/LazyCallGraph.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Analysis/MLModelRunner.h"
@@ -28,6 +27,7 @@
#include "llvm/Analysis/ReleaseModeModelRunner.h"
#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/Analysis/TensorSpec.h"
+#include "llvm/Analysis/Utils/MLGOUtils.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/Module.h"
@@ -111,7 +111,7 @@ createEmitCModelRunner(LLVMContext &Ctx,
constexpr bool HaveMLIRLoweringInliner = false;
enum class EmitCModelChoice { Default };
static const EmitCModelChoice SelectedMLGOModel = EmitCModelChoice::Default;
-inline std::unique_ptr<MLModelRunner>
+static inline std::unique_ptr<MLModelRunner>
createEmitCModelRunner(LLVMContext &, const std::vector<TensorSpec> &) {
return nullptr;
}
@@ -120,29 +120,16 @@ createEmitCModelRunner(LLVMContext &, const std::vector<TensorSpec> &) {
std::unique_ptr<InlineAdvisor>
llvm::getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM,
std::function<bool(CallBase &)> GetDefaultAdvice) {
- if (!llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() &&
- InteractiveChannelBaseName.empty() &&
- SelectedMLGOModel == EmitCModelChoice::Default)
+ if (!isReleaseModelValid<CompiledModelType>(InteractiveChannelBaseName,
+ SelectedMLGOModel))
return nullptr;
auto RunnerFactory = [&](const std::vector<TensorSpec> &InputFeatures)
-> std::unique_ptr<MLModelRunner> {
- std::unique_ptr<MLModelRunner> ModelRunner;
- if (InteractiveChannelBaseName.empty()) {
- if constexpr (HaveMLIRLoweringInliner) {
- ModelRunner = createEmitCModelRunner(M.getContext(), InputFeatures);
- } else {
- ModelRunner =
- std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
- M.getContext(), InputFeatures, DecisionName,
- EmbeddedModelRunnerOptions().setModelSelector(ModelSelector));
- }
- } else {
- ModelRunner = std::make_unique<InteractiveModelRunner>(
- M.getContext(), InputFeatures, InlineDecisionSpec,
- InteractiveChannelBaseName + ".out",
- InteractiveChannelBaseName + ".in");
- }
- return ModelRunner;
+ return createReleaseModeModelRunner<CompiledModelType,
+ HaveMLIRLoweringInliner>(
+ M.getContext(), InputFeatures, DecisionName, InteractiveChannelBaseName,
+ InlineDecisionSpec, createEmitCModelRunner,
+ EmbeddedModelRunnerOptions().setModelSelector(ModelSelector));
};
return std::make_unique<MLInlineAdvisor>(M, MAM, RunnerFactory,
GetDefaultAdvice);
diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index be02b96e8d921..0ed413a7707ad 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -12,7 +12,6 @@
#include "AllocationOrder.h"
#include "RegAllocGreedy.h"
-#include "llvm/Analysis/InteractiveModelRunner.h"
#include "llvm/Analysis/MLModelRunner.h"
#include "llvm/Analysis/TensorSpec.h"
#include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
@@ -23,6 +22,7 @@
#endif
#include "MLRegAllocEvictAdvisor.h"
#include "llvm/Analysis/ReleaseModeModelRunner.h"
+#include "llvm/Analysis/Utils/MLGOUtils.h"
#include "llvm/CodeGen/CalcSpillWeights.h"
#include "llvm/CodeGen/LiveRegMatrix.h"
#include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
@@ -35,7 +35,6 @@
#include "llvm/IR/Module.h"
#include "llvm/InitializePasses.h"
#include "llvm/Pass.h"
-#include "llvm/PassRegistry.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/ErrorHandling.h"
@@ -56,6 +55,7 @@ using CompiledModelType = NoopSavedModelImpl;
#endif
#if defined(LLVM_HAVE_MLIR_LOWERING_REGALLOC)
+constexpr bool HaveMLIRLoweringRegAlloc = true;
#include "llvm/Analysis/EmitCModelRunner.h"
#include "llvm/CodeGen/RegAllocEvictModels.h"
@@ -90,6 +90,15 @@ createMLGORegAllocModelRunner(LLVMContext &Ctx,
}
llvm_unreachable("Unknown MLGO model type!");
}
+#else
+constexpr bool HaveMLIRLoweringRegAlloc = false;
+enum class MLGORegAllocModelChoice { Default };
+static const MLGORegAllocModelChoice SelectedMLGORegAllocModel =
+ MLGORegAllocModelChoice::Default;
+static inline std::unique_ptr<MLModelRunner>
+createMLGORegAllocModelRunner(LLVMContext &, const std::vector<TensorSpec> &) {
+ return nullptr;
+}
#endif
static cl::opt<std::string> InteractiveChannelBaseName(
@@ -405,23 +414,11 @@ class ReleaseModeEvictionAdvisorProvider final
getAdvisor(const MachineFunction &MF, const RAGreedy &RA,
MachineBlockFrequencyInfo *MBFI, MachineLoopInfo *Loops) override {
if (!Runner) {
- if (InteractiveChannelBaseName.empty()) {
-#if defined(LLVM_HAVE_MLIR_LOWERING_REGALLOC)
- if (SelectedMLGORegAllocModel != MLGORegAllocModelChoice::Default) {
- Runner = createMLGORegAllocModelRunner(MF.getFunction().getContext(),
- InputFeatures);
- } else
-#endif
- {
- Runner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>(
- MF.getFunction().getContext(), InputFeatures, DecisionName);
- }
- } else {
- Runner = std::make_unique<InteractiveModelRunner>(
- MF.getFunction().getContext(), InputFeatures, DecisionSpec,
- InteractiveChannelBaseName + ".out",
- InteractiveChannelBaseName + ".in");
- }
+ Runner = createReleaseModeModelRunner<CompiledModelType,
+ HaveMLIRLoweringRegAlloc>(
+ MF.getFunction().getContext(), InputFeatures, DecisionName,
+ InteractiveChannelBaseName, DecisionSpec,
+ createMLGORegAllocModelRunner);
}
assert(MBFI && Loops &&
"Invalid provider state: must have analysis available");
@@ -1064,12 +1061,8 @@ bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) {
RegAllocEvictionAdvisorProvider *
llvm::createReleaseModeAdvisorProvider(LLVMContext &Ctx) {
- return llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
- !InteractiveChannelBaseName.empty()
-#if defined(LLVM_HAVE_MLIR_LOWERING_REGALLOC)
- ||
- SelectedMLGORegAllocModel != MLGORegAllocModelChoice::Default
-#endif
+ return isReleaseModelValid<CompiledModelType>(InteractiveChannelBaseName,
+ SelectedMLGORegAllocModel)
? new ReleaseModeEvictionAdvisorProvider(Ctx)
: nullptr;
}
@@ -1084,12 +1077,8 @@ llvm::createDevelopmentModeAdvisorProvider(LLVMContext &Ctx) {
RegAllocEvictionAdvisorAnalysisLegacy *
llvm::createReleaseModeAdvisorAnalysisLegacy() {
- return llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
- !InteractiveChannelBaseName.empty()
-#if defined(LLVM_HAVE_MLIR_LOWERING_REGALLOC)
- ||
- SelectedMLGORegAllocModel != MLGORegAllocModelChoice::Default
-#endif
+ return isReleaseModelValid<CompiledModelType>(InteractiveChannelBaseName,
+ SelectedMLGORegAllocModel)
? new ReleaseModeEvictionAdvisorAnalysisLegacy()
: nullptr;
}
>From 9ff564ea56103c147bb0fea4e7d9f3a94f5b2b5f Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Mon, 3 Aug 2026 12:19:47 -0700
Subject: [PATCH 08/11] rename LLVM_HAVE_EMITC_COMPILE and default it to off
and prevent the user from setting the variable
---
llvm/CMakeLists.txt | 12 ++++++------
llvm/lib/Analysis/CMakeLists.txt | 2 +-
llvm/lib/CodeGen/CMakeLists.txt | 2 +-
llvm/test/CMakeLists.txt | 2 +-
llvm/test/lit.cfg.py | 4 ++--
llvm/test/lit.site.cfg.py.in | 2 +-
6 files changed, 12 insertions(+), 12 deletions(-)
diff --git a/llvm/CMakeLists.txt b/llvm/CMakeLists.txt
index 660befb697a52..ea749e23dd9b6 100644
--- a/llvm/CMakeLists.txt
+++ b/llvm/CMakeLists.txt
@@ -1231,13 +1231,13 @@ set(LLVM_MLGO_MODELS "" CACHE STRING "List of MLGO models")
set(LLVM_MLGO_MLIR_OPT "mlir-opt" CACHE STRING "Path or system binary name for mlir-opt")
set(LLVM_MLGO_MLIR_TRANSLATE "mlir-translate" CACHE STRING "Path or system binary name for mlir-translate")
-set(LLVM_HAVE_EMITC_COMPILE OFF CACHE BOOL "MLGO models to be compiled with EmitC are available")
+set(LLVM_HAVE_EMITC_LOWERING OFF)
if (NOT LLVM_MLGO_MODELS STREQUAL "")
- set(LLVM_HAVE_EMITC_COMPILE ON CACHE BOOL "MLGO models to be compiled with EmitC are available" FORCE)
+ set(LLVM_HAVE_EMITC_LOWERING ON)
endif()
-if (LLVM_HAVE_EMITC_COMPILE AND LLVM_HAVE_TFLITE)
- message(FATAL_ERROR "Only one of LLVM_HAVE_TFLITE and LLVM_HAVE_EMITC_COMPILE can be enabled.")
+if (LLVM_HAVE_EMITC_LOWERING AND LLVM_HAVE_TFLITE)
+ message(FATAL_ERROR "Only one of LLVM_HAVE_TFLITE and LLVM_HAVE_EMITC_LOWERING can be enabled.")
endif()
if (LLVM_HAVE_TFLITE)
@@ -1268,8 +1268,8 @@ 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)
- if (LLVM_HAVE_TF_AOT AND LLVM_HAVE_EMITC_COMPILE)
- message(FATAL_ERROR "Only one of LLVM_HAVE_TF_AOT and LLVM_HAVE_EMITC_COMPILE can be enabled.")
+ if (LLVM_HAVE_TF_AOT AND LLVM_HAVE_EMITC_LOWERING)
+ message(FATAL_ERROR "Only one of LLVM_HAVE_TF_AOT and LLVM_HAVE_EMITC_LOWERING can be enabled.")
endif()
# Once we add more modules, we should handle this more automatically.
if (DEFINED LLVM_OVERRIDE_MODEL_HEADER_INLINERSIZEMODEL)
diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt
index 99c189562727f..5aac720edc6b7 100644
--- a/llvm/lib/Analysis/CMakeLists.txt
+++ b/llvm/lib/Analysis/CMakeLists.txt
@@ -1,4 +1,4 @@
-if (LLVM_HAVE_EMITC_COMPILE)
+if (LLVM_HAVE_EMITC_LOWERING)
include(MLGOLower)
mlgo_lower_models(
"${LLVM_MLGO_MODELS}"
diff --git a/llvm/lib/CodeGen/CMakeLists.txt b/llvm/lib/CodeGen/CMakeLists.txt
index 4f88a4306e609..b3b6ed71874f8 100644
--- a/llvm/lib/CodeGen/CMakeLists.txt
+++ b/llvm/lib/CodeGen/CMakeLists.txt
@@ -1,4 +1,4 @@
-if (LLVM_HAVE_EMITC_COMPILE)
+if (LLVM_HAVE_EMITC_LOWERING)
include(MLGOLower)
mlgo_lower_models(
"${LLVM_MLGO_MODELS}"
diff --git a/llvm/test/CMakeLists.txt b/llvm/test/CMakeLists.txt
index 90d3cd494565a..a859812b71a45 100644
--- a/llvm/test/CMakeLists.txt
+++ b/llvm/test/CMakeLists.txt
@@ -19,7 +19,7 @@ llvm_canonicalize_cmake_booleans(
LLVM_EXAMPLEIRTRANSFORMS_LINK_INTO_TOOLS
LLVM_HAVE_TF_AOT
LLVM_HAVE_TFLITE
- LLVM_HAVE_EMITC_COMPILE
+ LLVM_HAVE_EMITC_LOWERING
LLVM_ENABLE_PROFCHECK
LLVM_INLINER_MODEL_AUTOGENERATED
LLVM_RAEVICT_MODEL_AUTOGENERATED
diff --git a/llvm/test/lit.cfg.py b/llvm/test/lit.cfg.py
index 3973b971149df..8b6d6ba311a5e 100644
--- a/llvm/test/lit.cfg.py
+++ b/llvm/test/lit.cfg.py
@@ -584,8 +584,8 @@ def enable_ptxas(ptxas_executable):
if config.have_tf_aot:
config.available_features.add("have_tf_aot")
-if getattr(config, "have_emitc_compile", False):
- config.available_features.add("have_emitc_compile")
+if getattr(config, "have_emitc_lowering", False):
+ config.available_features.add("have_emitc_lowering")
if getattr(config, "have_opencsd", False):
config.available_features.add("opencsd")
diff --git a/llvm/test/lit.site.cfg.py.in b/llvm/test/lit.site.cfg.py.in
index 24279c94c2ae8..2c8bd9687d4d8 100644
--- a/llvm/test/lit.site.cfg.py.in
+++ b/llvm/test/lit.site.cfg.py.in
@@ -57,7 +57,7 @@ config.has_plugins = @LLVM_ENABLE_PLUGINS@
config.linked_bye_extension = @LLVM_BYE_LINK_INTO_TOOLS@
config.linked_exampleirtransforms_extension = @LLVM_EXAMPLEIRTRANSFORMS_LINK_INTO_TOOLS@
config.have_tf_aot = @LLVM_HAVE_TF_AOT@
-config.have_emitc_compile = @LLVM_HAVE_EMITC_COMPILE@
+config.have_emitc_lowering = @LLVM_HAVE_EMITC_LOWERING@
config.have_tflite = @LLVM_HAVE_TFLITE@
config.enable_profcheck = @LLVM_ENABLE_PROFCHECK@
config.llvm_inliner_model_autogenerated = @LLVM_INLINER_MODEL_AUTOGENERATED@
>From 2918c7ba76c25c5a225aa39f1e43558b6a8306b7 Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Thu, 6 Aug 2026 12:36:11 -0700
Subject: [PATCH 09/11] Update CMAKE variable naming to not expose
implementation detail
---
llvm/CMakeLists.txt | 12 ++++++------
llvm/cmake/modules/MLGOLower.cmake | 4 ++--
llvm/lib/Analysis/CMakeLists.txt | 2 +-
llvm/lib/CodeGen/CMakeLists.txt | 2 +-
llvm/test/CMakeLists.txt | 2 +-
llvm/test/lit.cfg.py | 4 ++--
llvm/test/lit.site.cfg.py.in | 2 +-
7 files changed, 14 insertions(+), 14 deletions(-)
diff --git a/llvm/CMakeLists.txt b/llvm/CMakeLists.txt
index ea749e23dd9b6..b0a40dbdfd423 100644
--- a/llvm/CMakeLists.txt
+++ b/llvm/CMakeLists.txt
@@ -1231,13 +1231,13 @@ set(LLVM_MLGO_MODELS "" CACHE STRING "List of MLGO models")
set(LLVM_MLGO_MLIR_OPT "mlir-opt" CACHE STRING "Path or system binary name for mlir-opt")
set(LLVM_MLGO_MLIR_TRANSLATE "mlir-translate" CACHE STRING "Path or system binary name for mlir-translate")
-set(LLVM_HAVE_EMITC_LOWERING OFF)
+set(LLVM_HAVE_MLIR_LOWERING OFF)
if (NOT LLVM_MLGO_MODELS STREQUAL "")
- set(LLVM_HAVE_EMITC_LOWERING ON)
+ set(LLVM_HAVE_MLIR_LOWERING ON)
endif()
-if (LLVM_HAVE_EMITC_LOWERING AND LLVM_HAVE_TFLITE)
- message(FATAL_ERROR "Only one of LLVM_HAVE_TFLITE and LLVM_HAVE_EMITC_LOWERING can be enabled.")
+if (LLVM_HAVE_MLIR_LOWERING AND LLVM_HAVE_TFLITE)
+ message(FATAL_ERROR "Only one of LLVM_HAVE_TFLITE and LLVM_HAVE_MLIR_LOWERING can be enabled.")
endif()
if (LLVM_HAVE_TFLITE)
@@ -1268,8 +1268,8 @@ 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)
- if (LLVM_HAVE_TF_AOT AND LLVM_HAVE_EMITC_LOWERING)
- message(FATAL_ERROR "Only one of LLVM_HAVE_TF_AOT and LLVM_HAVE_EMITC_LOWERING can be enabled.")
+ if (LLVM_HAVE_TF_AOT AND LLVM_HAVE_MLIR_LOWERING)
+ message(FATAL_ERROR "Only one of LLVM_HAVE_TF_AOT and LLVM_HAVE_MLIR_LOWERING can be enabled.")
endif()
# Once we add more modules, we should handle this more automatically.
if (DEFINED LLVM_OVERRIDE_MODEL_HEADER_INLINERSIZEMODEL)
diff --git a/llvm/cmake/modules/MLGOLower.cmake b/llvm/cmake/modules/MLGOLower.cmake
index 82e32ea5c3457..28c18eb4defb3 100644
--- a/llvm/cmake/modules/MLGOLower.cmake
+++ b/llvm/cmake/modules/MLGOLower.cmake
@@ -59,7 +59,7 @@ function(mlgo_lower_models models mlir_opt mlir_translate target_type
# Pass pipeline to lower MLIR models to EmitC dialect
# TODO: Simplify with builtin pipeline for translation.
- set(EMITC_PASSES
+ set(MLIR_PASSES
"func.func(tosa-to-linalg-named,tosa-to-linalg,tosa-to-arith,tosa-to-tensor)"
"symbol-privatize"
"scalarize-single-element-tensor-return"
@@ -79,7 +79,7 @@ function(mlgo_lower_models models mlir_opt mlir_translate target_type
"convert-math-to-emitc"
"convert-arith-to-emitc"
)
- string(JOIN "," PASS_PIPELINE ${EMITC_PASSES})
+ string(JOIN "," PASS_PIPELINE ${MLIR_PASSES})
set(PASS_PIPELINE "builtin.module(${PASS_PIPELINE})")
# 1. Run MLIR pipeline to lower MODEL_PATH to EmitC MLIR
diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt
index 5aac720edc6b7..e248f3916b476 100644
--- a/llvm/lib/Analysis/CMakeLists.txt
+++ b/llvm/lib/Analysis/CMakeLists.txt
@@ -1,4 +1,4 @@
-if (LLVM_HAVE_EMITC_LOWERING)
+if (LLVM_HAVE_MLIR_LOWERING)
include(MLGOLower)
mlgo_lower_models(
"${LLVM_MLGO_MODELS}"
diff --git a/llvm/lib/CodeGen/CMakeLists.txt b/llvm/lib/CodeGen/CMakeLists.txt
index b3b6ed71874f8..1b675d9b1322c 100644
--- a/llvm/lib/CodeGen/CMakeLists.txt
+++ b/llvm/lib/CodeGen/CMakeLists.txt
@@ -1,4 +1,4 @@
-if (LLVM_HAVE_EMITC_LOWERING)
+if (LLVM_HAVE_MLIR_LOWERING)
include(MLGOLower)
mlgo_lower_models(
"${LLVM_MLGO_MODELS}"
diff --git a/llvm/test/CMakeLists.txt b/llvm/test/CMakeLists.txt
index a859812b71a45..a2269597fe5e6 100644
--- a/llvm/test/CMakeLists.txt
+++ b/llvm/test/CMakeLists.txt
@@ -19,7 +19,7 @@ llvm_canonicalize_cmake_booleans(
LLVM_EXAMPLEIRTRANSFORMS_LINK_INTO_TOOLS
LLVM_HAVE_TF_AOT
LLVM_HAVE_TFLITE
- LLVM_HAVE_EMITC_LOWERING
+ LLVM_HAVE_MLIR_LOWERING
LLVM_ENABLE_PROFCHECK
LLVM_INLINER_MODEL_AUTOGENERATED
LLVM_RAEVICT_MODEL_AUTOGENERATED
diff --git a/llvm/test/lit.cfg.py b/llvm/test/lit.cfg.py
index 8b6d6ba311a5e..d60690dd76514 100644
--- a/llvm/test/lit.cfg.py
+++ b/llvm/test/lit.cfg.py
@@ -584,8 +584,8 @@ def enable_ptxas(ptxas_executable):
if config.have_tf_aot:
config.available_features.add("have_tf_aot")
-if getattr(config, "have_emitc_lowering", False):
- config.available_features.add("have_emitc_lowering")
+if getattr(config, "have_mlir_lowering", False):
+ config.available_features.add("have_mlir_lowering")
if getattr(config, "have_opencsd", False):
config.available_features.add("opencsd")
diff --git a/llvm/test/lit.site.cfg.py.in b/llvm/test/lit.site.cfg.py.in
index 2c8bd9687d4d8..eb9934f85ad42 100644
--- a/llvm/test/lit.site.cfg.py.in
+++ b/llvm/test/lit.site.cfg.py.in
@@ -57,7 +57,7 @@ config.has_plugins = @LLVM_ENABLE_PLUGINS@
config.linked_bye_extension = @LLVM_BYE_LINK_INTO_TOOLS@
config.linked_exampleirtransforms_extension = @LLVM_EXAMPLEIRTRANSFORMS_LINK_INTO_TOOLS@
config.have_tf_aot = @LLVM_HAVE_TF_AOT@
-config.have_emitc_lowering = @LLVM_HAVE_EMITC_LOWERING@
+config.have_mlir_lowering = @LLVM_HAVE_MLIR_LOWERING@
config.have_tflite = @LLVM_HAVE_TFLITE@
config.enable_profcheck = @LLVM_ENABLE_PROFCHECK@
config.llvm_inliner_model_autogenerated = @LLVM_INLINER_MODEL_AUTOGENERATED@
>From fc5ccceecf1adbb07d9a8c38dd81481bcf7d914a Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Thu, 6 Aug 2026 13:32:46 -0700
Subject: [PATCH 10/11] EmitCModelRunner tests
---
llvm/unittests/Analysis/MLModelRunnerTest.cpp | 57 +++++++++++++++++++
1 file changed, 57 insertions(+)
diff --git a/llvm/unittests/Analysis/MLModelRunnerTest.cpp b/llvm/unittests/Analysis/MLModelRunnerTest.cpp
index 62f90699c8718..5179d673cfa97 100644
--- a/llvm/unittests/Analysis/MLModelRunnerTest.cpp
+++ b/llvm/unittests/Analysis/MLModelRunnerTest.cpp
@@ -8,6 +8,7 @@
#include "llvm/Analysis/MLModelRunner.h"
#include "llvm/ADT/StringExtras.h"
+#include "llvm/Analysis/EmitCModelRunner.h"
#include "llvm/Analysis/InteractiveModelRunner.h"
#include "llvm/Analysis/NoInferenceModelRunner.h"
#include "llvm/Analysis/ReleaseModeModelRunner.h"
@@ -122,6 +123,18 @@ class ComposedAOTModel final {
void Run() { getModel()->Run(); }
};
+class MockEmitCModel final {
+ int64_t A = 0;
+ int64_t B = 0;
+
+public:
+ std::map<std::string, void *> reflectionMap;
+
+ MockEmitCModel() : reflectionMap{{"a", &A}, {"b", &B}} {}
+
+ int64_t operator()() { return A - B; }
+};
+
static EmbeddedModelRunnerOptions makeOptions() {
EmbeddedModelRunnerOptions Opts;
Opts.setFeedPrefix("prefix_");
@@ -262,6 +275,50 @@ TEST(ReleaseModelRunner, ModelSelector) {
// expect the model implementation to fail at a point.
}
+TEST(EmitCModelRunner, NormalUse) {
+ LLVMContext Ctx;
+ std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),
+ TensorSpec::createSpec<int64_t>("b", {1})};
+ auto Evaluator =
+ std::make_unique<EmitCModelRunner<MockEmitCModel>>(Ctx, Inputs);
+ EXPECT_TRUE(EmitCModelRunner<MockEmitCModel>::classof(Evaluator.get()));
+ EXPECT_EQ(Evaluator->getKind(), MLModelRunner::Kind::Release);
+ *Evaluator->getTensor<int64_t>(0) = 10;
+ *Evaluator->getTensor<int64_t>(1) = 3;
+ EXPECT_EQ(Evaluator->evaluate<int64_t>(), 7);
+ EXPECT_EQ(*Evaluator->getTensor<int64_t>(0), 10);
+ EXPECT_EQ(*Evaluator->getTensor<int64_t>(1), 3);
+}
+
+TEST(EmitCModelRunner, ExtraFeatures) {
+ LLVMContext Ctx;
+ std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),
+ TensorSpec::createSpec<int64_t>("b", {1}),
+ TensorSpec::createSpec<int64_t>("c", {1})};
+ auto Evaluator =
+ std::make_unique<EmitCModelRunner<MockEmitCModel>>(Ctx, Inputs);
+ *Evaluator->getTensor<int64_t>(0) = 10;
+ *Evaluator->getTensor<int64_t>(1) = 3;
+ *Evaluator->getTensor<int64_t>(2) = -5;
+ EXPECT_EQ(Evaluator->evaluate<int64_t>(), 7);
+ EXPECT_EQ(*Evaluator->getTensor<int64_t>(0), 10);
+ EXPECT_EQ(*Evaluator->getTensor<int64_t>(1), 3);
+ EXPECT_EQ(*Evaluator->getTensor<int64_t>(2), -5);
+}
+
+TEST(EmitCModelRunner, ExtraFeaturesOutOfOrder) {
+ LLVMContext Ctx;
+ std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("b", {1}),
+ TensorSpec::createSpec<int64_t>("a", {1})};
+ auto Evaluator =
+ std::make_unique<EmitCModelRunner<MockEmitCModel>>(Ctx, Inputs);
+ *Evaluator->getTensor<int64_t>(0) = 3; // b
+ *Evaluator->getTensor<int64_t>(1) = 10; // a
+ EXPECT_EQ(Evaluator->evaluate<int64_t>(), 7); // a - b = 10 - 3 = 7
+ EXPECT_EQ(*Evaluator->getTensor<int64_t>(0), 3);
+ EXPECT_EQ(*Evaluator->getTensor<int64_t>(1), 10);
+}
+
#if defined(LLVM_ON_UNIX)
TEST(InteractiveModelRunner, Evaluation) {
LLVMContext Ctx;
>From 488501df758a3edf27179157c7a7b5317e8164e7 Mon Sep 17 00:00:00 2001
From: Bhavesh Mandalapu <bmandalapu at google.com>
Date: Thu, 13 Aug 2026 14:08:31 -0700
Subject: [PATCH 11/11] Update tests
---
llvm/cmake/modules/MLGOLower.cmake | 6 +-
.../MLRegAlloc/default-eviction-advisor.ll | 7 +-
llvm/test/CodeGen/MLRegAlloc/rel-codepath.ll | 6 +-
.../Transforms/Inline/ML/bounds-checks.ll | 6 +-
.../Inline/ML/coro-split-func-levels.ll | 4 +-
llvm/test/Transforms/Inline/ML/dead-callee.ll | 4 +-
.../ML/enable-inline-advisor-printing-ml.ll | 6 +-
llvm/test/Transforms/Inline/ML/fpi-update.ll | 4 +-
.../Inline/ML/ml-test-release-mode.ll | 4 +-
llvm/test/Transforms/Inline/ML/recursive.ll | 4 +-
.../Inline/ML/scc-dead-accounting.ll | 4 +-
.../Transforms/Inline/ML/skip-unreachable.ll | 4 +-
.../ML/state-accounting-skip-non-cold.ll | 6 +-
.../Inline/inlining-advisor-default.ll | 9 +-
llvm/unittests/Analysis/CMakeLists.txt | 1 +
llvm/unittests/Analysis/MLGOUtilsTest.cpp | 177 ++++++++++++++++++
16 files changed, 218 insertions(+), 34 deletions(-)
create mode 100644 llvm/unittests/Analysis/MLGOUtilsTest.cpp
diff --git a/llvm/cmake/modules/MLGOLower.cmake b/llvm/cmake/modules/MLGOLower.cmake
index 28c18eb4defb3..54118bea5df90 100644
--- a/llvm/cmake/modules/MLGOLower.cmake
+++ b/llvm/cmake/modules/MLGOLower.cmake
@@ -47,7 +47,7 @@ function(mlgo_lower_models models mlir_opt mlir_translate target_type
continue()
endif()
- set(CLASS_NAME "Model${MODEL_INDEX}")
+ set(CLASS_NAME "${target_type}_Model${MODEL_INDEX}")
math(EXPR MODEL_INDEX "${MODEL_INDEX} + 1")
if (NOT IS_ABSOLUTE "${MODEL_PATH}")
@@ -102,8 +102,8 @@ function(mlgo_lower_models models mlir_opt mlir_translate target_type
set_source_files_properties(${HEADER_FILE} PROPERTIES GENERATED 1)
# Custom target to force generation of this header
- add_custom_target(mlgo_model_gen_${CLASS_NAME} DEPENDS ${HEADER_FILE})
- list(APPEND MLGO_GEN_TARGETS mlgo_model_gen_${CLASS_NAME})
+ add_custom_target(mlgo_model_gen_${MODEL_TYPE}_${CLASS_NAME} DEPENDS ${HEADER_FILE})
+ list(APPEND MLGO_GEN_TARGETS mlgo_model_gen_${MODEL_TYPE}_${CLASS_NAME})
# Append the model metadata to the .def file
string(APPEND DEF_CONTENT "MLGO_MODEL(${CLASS_NAME}, \"${CLI_FLAG}\")\n")
diff --git a/llvm/test/CodeGen/MLRegAlloc/default-eviction-advisor.ll b/llvm/test/CodeGen/MLRegAlloc/default-eviction-advisor.ll
index 881a80c41361c..f9509e57e1ff1 100644
--- a/llvm/test/CodeGen/MLRegAlloc/default-eviction-advisor.ll
+++ b/llvm/test/CodeGen/MLRegAlloc/default-eviction-advisor.ll
@@ -1,11 +1,13 @@
-; Check that, in the absence of dependencies, we emit an error message when
-; trying to use ML-driven advisor.
+; Check that, in the absence of dependencies or a selected model, we emit an
+; error message when trying to use ML-driven advisor.
; REQUIRES: !have_tf_aot
; REQUIRES: !have_tflite
; REQUIRES: default_triple
; RUN: not llc -O2 -regalloc-enable-advisor=development < %s 2>&1 | FileCheck %s
; RUN: not llc -O2 -regalloc-enable-advisor=release < %s 2>&1 | FileCheck %s
; RUN: llc -O2 -regalloc-enable-advisor=default < %s 2>&1 | FileCheck %s --check-prefix=DEFAULT
+; RUN: %if have_mlir_lowering %{ not llc -O2 -regalloc-enable-advisor=release -regalloc-mlgo-model=default < %s 2>&1 | FileCheck %s %}
+; RUN: %if have_mlir_lowering %{ not llc -O2 -regalloc-enable-advisor=release -regalloc-mlgo-model=invalid_model < %s 2>&1 | FileCheck %s --check-prefix=INVALID %}
; regalloc-enable-advisor is not enabled for NVPTX
; UNSUPPORTED: target=nvptx{{.*}}
@@ -18,3 +20,4 @@ define void @f2(i64 %lhs, i64 %rhs, ptr %addr) {
; CHECK: Requested regalloc eviction advisor analysis could not be created. Using default
; DEFAULT-NOT: Requested regalloc eviction advisor analysis could not be created. Using default
+; INVALID: {{.*}}llc{{.*}}: for the --regalloc-mlgo-model option: Cannot find option named 'invalid_model'!
diff --git a/llvm/test/CodeGen/MLRegAlloc/rel-codepath.ll b/llvm/test/CodeGen/MLRegAlloc/rel-codepath.ll
index 03866d0add9bc..d426ce6f8f3e3 100644
--- a/llvm/test/CodeGen/MLRegAlloc/rel-codepath.ll
+++ b/llvm/test/CodeGen/MLRegAlloc/rel-codepath.ll
@@ -1,6 +1,6 @@
-; REQUIRES: have_tf_aot
+; REQUIRES: have_tf_aot || have_mlir_lowering
; REQUIRES: x86_64-linux
-; REQUIRES: llvm_raevict_model_autogenerated
+; REQUIRES: llvm_raevict_model_autogenerated || have_mlir_lowering
;
; Check the code path for release mode is correctly taken. It is shared with
; development mode, and we separately test the internals of that (logged
@@ -10,7 +10,7 @@
; RUN: llc -mtriple=x86_64-linux-unknown -regalloc=greedy -regalloc-enable-advisor=default \
; RUN: %S/Inputs/input.ll -o %t.default
-; RUN: llc -mtriple=x86_64-linux-unknown -regalloc=greedy -regalloc-enable-advisor=release \
+; RUN: llc -mtriple=x86_64-linux-unknown -regalloc=greedy %if have_mlir_lowering %{ -regalloc-enable-advisor=release -regalloc-mlgo-model=regalloc %} %else %{ -regalloc-enable-advisor=release %} \
; RUN: %S/Inputs/input.ll -o %t.release
; RUN: not diff %t.release %t.default
diff --git a/llvm/test/Transforms/Inline/ML/bounds-checks.ll b/llvm/test/Transforms/Inline/ML/bounds-checks.ll
index 7270fa3a93b1e..f1e2911a2004c 100644
--- a/llvm/test/Transforms/Inline/ML/bounds-checks.ll
+++ b/llvm/test/Transforms/Inline/ML/bounds-checks.ll
@@ -2,9 +2,9 @@
; 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: 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
+; REQUIRES: llvm_inliner_model_autogenerated || have_mlir_lowering
+; RUN: opt -passes=scc-oz-module-inliner %if have_mlir_lowering %{ -enable-ml-inliner=release -mlgo-model=inliner %} %else %{ -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 %if have_mlir_lowering %{ -enable-ml-inliner=release -mlgo-model=inliner %} %else %{ -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"
diff --git a/llvm/test/Transforms/Inline/ML/coro-split-func-levels.ll b/llvm/test/Transforms/Inline/ML/coro-split-func-levels.ll
index 79e1ebeec17b9..b3660cb7c5646 100644
--- a/llvm/test/Transforms/Inline/ML/coro-split-func-levels.ll
+++ b/llvm/test/Transforms/Inline/ML/coro-split-func-levels.ll
@@ -1,6 +1,6 @@
-; REQUIRES: llvm_inliner_model_autogenerated
+; REQUIRES: llvm_inliner_model_autogenerated || have_mlir_lowering
; RUN: opt -S -passes='coro-early,scc-oz-module-inliner,print<inline-advisor>' \
-; RUN: -enable-ml-inliner=release -keep-inline-advisor-for-printing < %s
+; RUN: %if have_mlir_lowering %{ -enable-ml-inliner=release -mlgo-model=inliner %} %else %{ -enable-ml-inliner=release %} -keep-inline-advisor-for-printing < %s
define void @_Z5get_sv() presplitcoroutine {
%1 = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
diff --git a/llvm/test/Transforms/Inline/ML/dead-callee.ll b/llvm/test/Transforms/Inline/ML/dead-callee.ll
index a88655777e6f9..34d44eda1a625 100644
--- a/llvm/test/Transforms/Inline/ML/dead-callee.ll
+++ b/llvm/test/Transforms/Inline/ML/dead-callee.ll
@@ -1,5 +1,5 @@
-; REQUIRES: llvm_inliner_model_autogenerated
-; RUN: opt -passes=inliner-ml-advisor-release -S < %s | FileCheck %s
+; REQUIRES: llvm_inliner_model_autogenerated || have_mlir_lowering
+; RUN: opt %if have_mlir_lowering %{ -passes=inliner-ml-advisor-release -mlgo-model=inliner %} %else %{ -passes=inliner-ml-advisor-release %} -S < %s | FileCheck %s
; Check that our accounting works when a function in a non-trivial SCC is dead.
diff --git a/llvm/test/Transforms/Inline/ML/enable-inline-advisor-printing-ml.ll b/llvm/test/Transforms/Inline/ML/enable-inline-advisor-printing-ml.ll
index 376548b3f06fe..f72fb741cc037 100644
--- a/llvm/test/Transforms/Inline/ML/enable-inline-advisor-printing-ml.ll
+++ b/llvm/test/Transforms/Inline/ML/enable-inline-advisor-printing-ml.ll
@@ -1,10 +1,10 @@
-; REQUIRES: llvm_inliner_model_autogenerated
+; REQUIRES: llvm_inliner_model_autogenerated || have_mlir_lowering
-; RUN: opt -enable-ml-inliner=release -passes=scc-oz-module-inliner \
+; RUN: opt %if have_mlir_lowering %{ -enable-ml-inliner=release -mlgo-model=inliner %} %else %{ -enable-ml-inliner=release %} -passes=scc-oz-module-inliner \
; RUN: -keep-inline-advisor-for-printing -mandatory-inlining-first=1 \
; RUN: -enable-scc-inline-advisor-printing -S < %s 2>&1 | FileCheck %s
-; RUN: opt -enable-ml-inliner=release -passes=scc-oz-module-inliner \
+; RUN: opt %if have_mlir_lowering %{ -enable-ml-inliner=release -mlgo-model=inliner %} %else %{ -enable-ml-inliner=release %} -passes=scc-oz-module-inliner \
; RUN: -keep-inline-advisor-for-printing -mandatory-inlining-first=0 \
; RUN: -enable-scc-inline-advisor-printing -S < %s 2>&1 \
; RUN: | FileCheck %s --check-prefix=TWO
diff --git a/llvm/test/Transforms/Inline/ML/fpi-update.ll b/llvm/test/Transforms/Inline/ML/fpi-update.ll
index 9208ad70b9a35..2a3fa12616629 100644
--- a/llvm/test/Transforms/Inline/ML/fpi-update.ll
+++ b/llvm/test/Transforms/Inline/ML/fpi-update.ll
@@ -3,9 +3,9 @@ target triple = "x86_64-grtev4-linux-gnu"
; TODO: we could instantiate the MLInliner with a non-model generated evaluator
; and drop the requirement
-; REQUIRES: llvm_inliner_model_autogenerated
+; REQUIRES: llvm_inliner_model_autogenerated || have_mlir_lowering
-; RUN: opt -enable-ml-inliner=release -passes='scc-oz-module-inliner,print<inline-advisor>' \
+; RUN: opt %if have_mlir_lowering %{ -enable-ml-inliner=release -mlgo-model=inliner %} %else %{ -enable-ml-inliner=release %} -passes='scc-oz-module-inliner,print<inline-advisor>' \
; RUN: -keep-inline-advisor-for-printing -max-devirt-iterations=0 \
; RUN: -mandatory-inlining-first=0 -S < %s 2>&1 | FileCheck %s
diff --git a/llvm/test/Transforms/Inline/ML/ml-test-release-mode.ll b/llvm/test/Transforms/Inline/ML/ml-test-release-mode.ll
index 9cbbfdd124e1f..711652a197fbf 100644
--- a/llvm/test/Transforms/Inline/ML/ml-test-release-mode.ll
+++ b/llvm/test/Transforms/Inline/ML/ml-test-release-mode.ll
@@ -5,6 +5,6 @@
; This test uses Inputs/test-module.ll, as it will share it with a similar test
; for the 'development' mode.
;
-; 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
+; REQUIRES: llvm_inliner_model_autogenerated || have_mlir_lowering
+; RUN: opt -passes=scc-oz-module-inliner %if have_mlir_lowering %{ -enable-ml-inliner=release -mlgo-model=inliner %} %else %{ -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
diff --git a/llvm/test/Transforms/Inline/ML/recursive.ll b/llvm/test/Transforms/Inline/ML/recursive.ll
index 2d9240a12a713..db5b01d90bfbd 100644
--- a/llvm/test/Transforms/Inline/ML/recursive.ll
+++ b/llvm/test/Transforms/Inline/ML/recursive.ll
@@ -1,6 +1,6 @@
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals all --version 6
-; REQUIRES: llvm_inliner_model_autogenerated
-; RUN: opt -S %s -o - -passes='inliner-ml-advisor-release' -ml-inliner-skip-policy=if-caller-not-cold | FileCheck %s
+; REQUIRES: llvm_inliner_model_autogenerated || have_mlir_lowering
+; RUN: opt -S %s -o - %if have_mlir_lowering %{ -passes='inliner-ml-advisor-release' -mlgo-model=inliner %} %else %{ -passes='inliner-ml-advisor-release' %} -ml-inliner-skip-policy=if-caller-not-cold | FileCheck %s
target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128-Fn32"
target triple = "aarch64-unknown-linux-android29"
diff --git a/llvm/test/Transforms/Inline/ML/scc-dead-accounting.ll b/llvm/test/Transforms/Inline/ML/scc-dead-accounting.ll
index 41c0e4c313811..897475cd6dbae 100644
--- a/llvm/test/Transforms/Inline/ML/scc-dead-accounting.ll
+++ b/llvm/test/Transforms/Inline/ML/scc-dead-accounting.ll
@@ -9,9 +9,9 @@
; In this example if loop-unroll is ran after a mandatory inlining CGSCC pass,
; edges would increase but wouldn't be tracked
-; REQUIRES: llvm_inliner_model_autogenerated
+; REQUIRES: llvm_inliner_model_autogenerated || have_mlir_lowering
-; RUN: opt -enable-ml-inliner=release -passes=inliner-ml-advisor-release \
+; RUN: opt %if have_mlir_lowering %{ -enable-ml-inliner=release -mlgo-model=inliner %} %else %{ -enable-ml-inliner=release %} -passes=inliner-ml-advisor-release \
; RUN: -keep-inline-advisor-for-printing \
; RUN: -enable-scc-inline-advisor-printing -S < %s 2>&1 | FileCheck %s
diff --git a/llvm/test/Transforms/Inline/ML/skip-unreachable.ll b/llvm/test/Transforms/Inline/ML/skip-unreachable.ll
index 47a75c3ad898b..2f548edd65dd1 100644
--- a/llvm/test/Transforms/Inline/ML/skip-unreachable.ll
+++ b/llvm/test/Transforms/Inline/ML/skip-unreachable.ll
@@ -1,7 +1,7 @@
; Test skipping inlining when the callsite is unreachable (for both mandatory
; and non-mandatory cases)
-; REQUIRES: llvm_inliner_model_autogenerated
-; RUN: opt -passes=inliner-ml-advisor-release -S < %s | FileCheck %s --check-prefix=CHECK
+; REQUIRES: llvm_inliner_model_autogenerated || have_mlir_lowering
+; RUN: opt %if have_mlir_lowering %{ -passes=inliner-ml-advisor-release -mlgo-model=inliner %} %else %{ -passes=inliner-ml-advisor-release %} -S < %s | FileCheck %s --check-prefix=CHECK
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"
diff --git a/llvm/test/Transforms/Inline/ML/state-accounting-skip-non-cold.ll b/llvm/test/Transforms/Inline/ML/state-accounting-skip-non-cold.ll
index 0887f5e29187d..e75523cf90b58 100644
--- a/llvm/test/Transforms/Inline/ML/state-accounting-skip-non-cold.ll
+++ b/llvm/test/Transforms/Inline/ML/state-accounting-skip-non-cold.ll
@@ -1,7 +1,7 @@
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
-; REQUIRES: llvm_inliner_model_autogenerated && asserts
-; RUN: opt -passes='default<O3>' -enable-ml-inliner=release -ml-inliner-skip-policy=if-caller-not-cold -S %s -o - | FileCheck %s
-; RUN: opt -passes='default<O3>' -ml-inliner-stop-immediately -enable-ml-inliner=release -ml-inliner-skip-policy=if-caller-not-cold -S %s -o - | FileCheck %s
+; REQUIRES: (llvm_inliner_model_autogenerated || have_mlir_lowering) && asserts
+; RUN: opt -passes='default<O3>' %if have_mlir_lowering %{ -enable-ml-inliner=release -mlgo-model=inliner %} %else %{ -enable-ml-inliner=release %} -ml-inliner-skip-policy=if-caller-not-cold -S %s -o - | FileCheck %s
+; RUN: opt -passes='default<O3>' -ml-inliner-stop-immediately %if have_mlir_lowering %{ -enable-ml-inliner=release -mlgo-model=inliner %} %else %{ -enable-ml-inliner=release %} -ml-inliner-skip-policy=if-caller-not-cold -S %s -o - | FileCheck %s
declare ptr @f()
diff --git a/llvm/test/Transforms/Inline/inlining-advisor-default.ll b/llvm/test/Transforms/Inline/inlining-advisor-default.ll
index 502a16280192d..87ddc286092d4 100644
--- a/llvm/test/Transforms/Inline/inlining-advisor-default.ll
+++ b/llvm/test/Transforms/Inline/inlining-advisor-default.ll
@@ -1,10 +1,13 @@
-; Check that, in the absence of dependencies, we emit an error message when
-; trying to use ML-driven inlining.
+; Check that, in the absence of dependencies or a selected model, we emit an
+; error message when trying to use ML-driven inlining.
; REQUIRES: !have_tf_aot
; REQUIRES: !have_tflite
; RUN: not opt -passes=scc-oz-module-inliner -enable-ml-inliner=development -S < %s 2>&1 | FileCheck %s
; RUN: not opt -passes=scc-oz-module-inliner -enable-ml-inliner=release -S < %s 2>&1 | FileCheck %s
+; RUN: %if have_mlir_lowering %{ not opt -passes=scc-oz-module-inliner -enable-ml-inliner=release -mlgo-model=default -S < %s 2>&1 | FileCheck %s %}
+; RUN: %if have_mlir_lowering %{ not opt -passes=scc-oz-module-inliner -enable-ml-inliner=release -mlgo-model=invalid_model -S < %s 2>&1 | FileCheck %s --check-prefix=INVALID %}
declare i64 @f1()
-; CHECK: Could not setup Inlining Advisor for the requested mode and/or options
\ No newline at end of file
+; CHECK: Could not setup Inlining Advisor for the requested mode and/or options
+; INVALID: {{.*}}opt{{.*}}: for the --mlgo-model option: Cannot find option named 'invalid_model'!
\ No newline at end of file
diff --git a/llvm/unittests/Analysis/CMakeLists.txt b/llvm/unittests/Analysis/CMakeLists.txt
index b1f8cb918d590..9ca57024a6134 100644
--- a/llvm/unittests/Analysis/CMakeLists.txt
+++ b/llvm/unittests/Analysis/CMakeLists.txt
@@ -43,6 +43,7 @@ set(ANALYSIS_TEST_SOURCES
MemoryBuiltinsTest.cpp
MemoryProfileInfoTest.cpp
MemorySSATest.cpp
+ MLGOUtilsTest.cpp
MLModelRunnerTest.cpp
PhiValuesTest.cpp
PluginInlineAdvisorAnalysisTest.cpp
diff --git a/llvm/unittests/Analysis/MLGOUtilsTest.cpp b/llvm/unittests/Analysis/MLGOUtilsTest.cpp
new file mode 100644
index 0000000000000..1afa584df92dd
--- /dev/null
+++ b/llvm/unittests/Analysis/MLGOUtilsTest.cpp
@@ -0,0 +1,177 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Analysis/Utils/MLGOUtils.h"
+#include "llvm/Analysis/EmitCModelRunner.h"
+#include "llvm/Analysis/ReleaseModeModelRunner.h"
+#include "llvm/Analysis/TensorSpec.h"
+#include "llvm/IR/LLVMContext.h"
+#include "gtest/gtest.h"
+#include <map>
+#include <string>
+
+using namespace llvm;
+
+namespace {
+
+class MockAOTModel final {
+ int64_t A = 0;
+ int64_t B = 0;
+ int64_t R = 0;
+
+public:
+ MockAOTModel() = default;
+ int LookupArgIndex(const std::string &Name) {
+ if (Name == "feed_a")
+ return 0;
+ if (Name == "feed_b")
+ return 1;
+ return -1;
+ }
+ int LookupResultIndex(const std::string &) { return 0; }
+ void Run() { R = A + B; }
+ void *result_data(int RIndex) { return (RIndex == 0) ? &R : nullptr; }
+ void *arg_data(int Index) {
+ switch (Index) {
+ case 0:
+ return &A;
+ case 1:
+ return &B;
+ default:
+ return nullptr;
+ }
+ }
+};
+
+class MockEmitCModel1 final {
+ int64_t A = 0;
+ int64_t B = 0;
+
+public:
+ std::map<std::string, void *> reflectionMap;
+
+ MockEmitCModel1() : reflectionMap{{"a", &A}, {"b", &B}} {}
+
+ int64_t operator()() { return A - B; }
+};
+
+class MockEmitCModel2 final {
+ int64_t A = 0;
+ int64_t B = 0;
+
+public:
+ std::map<std::string, void *> reflectionMap;
+
+ MockEmitCModel2() : reflectionMap{{"a", &A}, {"b", &B}} {}
+
+ int64_t operator()() { return A + B; }
+};
+
+enum class TestModelChoice { Default, Model1, Model2 };
+
+TEST(MLGOUtilsTest, IsReleaseModelValid) {
+ // With NoopSavedModelImpl (no embedded AOT model):
+ // 1. Default model choice without interactive channel -> invalid
+ EXPECT_FALSE(
+ (isReleaseModelValid<NoopSavedModelImpl>("", TestModelChoice::Default)));
+ // 2. Selected model choice -> valid
+ EXPECT_TRUE(
+ (isReleaseModelValid<NoopSavedModelImpl>("", TestModelChoice::Model1)));
+ EXPECT_TRUE(
+ (isReleaseModelValid<NoopSavedModelImpl>("", TestModelChoice::Model2)));
+ // 3. Interactive channel specified -> valid regardless of model choice
+ EXPECT_TRUE((isReleaseModelValid<NoopSavedModelImpl>(
+ "channel", TestModelChoice::Default)));
+
+ // With a custom default model value:
+ EXPECT_FALSE((isReleaseModelValid<NoopSavedModelImpl>(
+ "", TestModelChoice::Model1, /*DefaultModelVal=*/TestModelChoice::Model1)));
+
+ // With a valid embedded AOT model (MockAOTModel):
+ EXPECT_TRUE(
+ (isReleaseModelValid<MockAOTModel>("", TestModelChoice::Default)));
+
+ // Overload with cl::opt<EnumType>:
+ cl::opt<TestModelChoice> OptChoice("test-mlgo-utils-choice",
+ cl::init(TestModelChoice::Default));
+ EXPECT_FALSE((isReleaseModelValid<NoopSavedModelImpl>("", OptChoice)));
+ OptChoice = TestModelChoice::Model1;
+ EXPECT_TRUE((isReleaseModelValid<NoopSavedModelImpl>("", OptChoice)));
+ OptChoice = TestModelChoice::Default;
+ EXPECT_TRUE((isReleaseModelValid<NoopSavedModelImpl>("channel", OptChoice)));
+ EXPECT_TRUE((isReleaseModelValid<MockAOTModel>("", OptChoice)));
+}
+
+TEST(MLGOUtilsTest, CreateReleaseModeModelRunnerModelSelection) {
+ LLVMContext Ctx;
+ std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),
+ TensorSpec::createSpec<int64_t>("b", {1})};
+ TensorSpec OutputSpec = TensorSpec::createSpec<int64_t>("result", {1});
+
+ auto CreateRunnerForChoice =
+ [&](TestModelChoice Choice) -> std::unique_ptr<MLModelRunner> {
+ auto Factory = [&](LLVMContext &C, const std::vector<TensorSpec> &Specs)
+ -> std::unique_ptr<MLModelRunner> {
+ switch (Choice) {
+ case TestModelChoice::Default:
+ return nullptr;
+ case TestModelChoice::Model1:
+ return std::make_unique<EmitCModelRunner<MockEmitCModel1>>(C, Specs);
+ case TestModelChoice::Model2:
+ return std::make_unique<EmitCModelRunner<MockEmitCModel2>>(C, Specs);
+ }
+ llvm_unreachable("unknown model choice");
+ };
+ return createReleaseModeModelRunner<NoopSavedModelImpl,
+ /*HaveMLIRLowering=*/true>(
+ Ctx, Inputs, "decision", "", OutputSpec, Factory);
+ };
+
+ // Default choice produces no runner
+ EXPECT_EQ(CreateRunnerForChoice(TestModelChoice::Default), nullptr);
+
+ // Model1 choice produces Model1 (A - B)
+ auto Runner1 = CreateRunnerForChoice(TestModelChoice::Model1);
+ ASSERT_NE(Runner1, nullptr);
+ EXPECT_TRUE(EmitCModelRunner<MockEmitCModel1>::classof(Runner1.get()));
+ *Runner1->getTensor<int64_t>(0) = 10;
+ *Runner1->getTensor<int64_t>(1) = 3;
+ EXPECT_EQ(Runner1->evaluate<int64_t>(), 7);
+
+ // Model2 choice produces Model2 (A + B)
+ auto Runner2 = CreateRunnerForChoice(TestModelChoice::Model2);
+ ASSERT_NE(Runner2, nullptr);
+ EXPECT_TRUE(EmitCModelRunner<MockEmitCModel2>::classof(Runner2.get()));
+ *Runner2->getTensor<int64_t>(0) = 10;
+ *Runner2->getTensor<int64_t>(1) = 3;
+ EXPECT_EQ(Runner2->evaluate<int64_t>(), 13);
+}
+
+TEST(MLGOUtilsTest, CreateReleaseModeModelRunnerAOTFallback) {
+ LLVMContext Ctx;
+ std::vector<TensorSpec> Inputs{TensorSpec::createSpec<int64_t>("a", {1}),
+ TensorSpec::createSpec<int64_t>("b", {1})};
+ TensorSpec OutputSpec = TensorSpec::createSpec<int64_t>("result", {1});
+
+ auto DummyEmitCFactory = [](LLVMContext &, const std::vector<TensorSpec> &)
+ -> std::unique_ptr<MLModelRunner> {
+ llvm_unreachable(
+ "EmitC factory should not be called when HaveMLIRLowering=false");
+ };
+
+ auto Runner =
+ createReleaseModeModelRunner<MockAOTModel, /*HaveMLIRLowering=*/false>(
+ Ctx, Inputs, "result", "", OutputSpec, DummyEmitCFactory);
+ ASSERT_NE(Runner, nullptr);
+ EXPECT_TRUE(ReleaseModeModelRunner<MockAOTModel>::classof(Runner.get()));
+ *Runner->getTensor<int64_t>(0) = 10;
+ *Runner->getTensor<int64_t>(1) = 3;
+ EXPECT_EQ(Runner->evaluate<int64_t>(), 13);
+}
+
+} // namespace
More information about the llvm-commits
mailing list