[llvm] [MLGO] Model selection for models lowered through EmitC (PR #212650)
Bhavesh M via llvm-commits
llvm-commits at lists.llvm.org
Tue Jul 28 18:43:49 PDT 2026
https://github.com/beamandala updated https://github.com/llvm/llvm-project/pull/212650
>From 306ba3ce3c6820ff36b5275338fac38bf969970b 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 1/4] 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 f3586c66cb056..0e1b02ff82dde 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 201f0ee86d7ff..d63972305fb95 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 b349c81f53f03..6d45a2e8de1d4 100644
--- a/llvm/test/lit.cfg.py
+++ b/llvm/test/lit.cfg.py
@@ -580,6 +580,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 64679c2f64034..eb0f938fa50ab 100644
--- a/llvm/test/lit.site.cfg.py.in
+++ b/llvm/test/lit.site.cfg.py.in
@@ -56,6 +56,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 14cab02135aa7278cfdd4a8a04d3fc148ad55815 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 2/4] 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 23dc6fbd6e500..aa322c74b5a17 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");
@@ -1020,7 +1063,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 *
@@ -1035,6 +1084,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 29c7515c6ba8a9a512d59e0d652d330bed46b4f1 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 3/4] 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 aa322c74b5a17..1d77f9068af37 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
{
@@ -1066,7 +1069,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;
@@ -1085,7 +1089,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 5dba7f026008afe1718ce18b8c531dc5c61d3de1 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 4/4] Address review comments
---
llvm/include/llvm/Analysis/EmitCModelRunner.h | 29 +++++++------------
llvm/lib/Analysis/CMakeLists.txt | 2 +-
llvm/lib/Analysis/MLInlineAdvisor.cpp | 28 +++++++++---------
3 files changed, 25 insertions(+), 34 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 0e1b02ff82dde..e66acb20e2131 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..38ac5c56c2d12 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,
+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
{
More information about the llvm-commits
mailing list