[llvm] [Draft] Integrate EmitC-translated MLGO models (PR #209829)
ioana ghiban via llvm-commits
llvm-commits at lists.llvm.org
Thu Jul 23 07:19:08 PDT 2026
https://github.com/ioghiban updated https://github.com/llvm/llvm-project/pull/209829
>From f1714dbc43381a6f02ef39cdc6adf703d0238245 Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Mon, 13 Jul 2026 17:08:13 +0200
Subject: [PATCH 1/7] Integrate EmitC translated model
---
.../llvm/Analysis/EmitCInlinerSizeModel.h | 82 ++++++
.../llvm/CodeGen/EmitCRegAllocEvictModel.h | 98 +++++++
llvm/lib/Analysis/CMakeLists.txt | 51 ++--
llvm/lib/Analysis/EmitCInlinerSizeModel.cpp | 250 ++++++++++++++++++
llvm/lib/Analysis/MLInlineAdvisor.cpp | 5 +-
llvm/lib/CodeGen/CMakeLists.txt | 17 +-
llvm/lib/CodeGen/EmitCRegAllocEvictModel.cpp | 209 +++++++++++++++
llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp | 7 +-
.../MLRegAlloc/default-eviction-advisor.ll | 1 +
llvm/test/lit.cfg.py | 3 +
llvm/test/lit.site.cfg.py.in | 1 +
11 files changed, 700 insertions(+), 24 deletions(-)
create mode 100644 llvm/include/llvm/Analysis/EmitCInlinerSizeModel.h
create mode 100644 llvm/include/llvm/CodeGen/EmitCRegAllocEvictModel.h
create mode 100644 llvm/lib/Analysis/EmitCInlinerSizeModel.cpp
create mode 100644 llvm/lib/CodeGen/EmitCRegAllocEvictModel.cpp
diff --git a/llvm/include/llvm/Analysis/EmitCInlinerSizeModel.h b/llvm/include/llvm/Analysis/EmitCInlinerSizeModel.h
new file mode 100644
index 0000000000000..63e298e4ce9fa
--- /dev/null
+++ b/llvm/include/llvm/Analysis/EmitCInlinerSizeModel.h
@@ -0,0 +1,82 @@
+//===- EmitCInlinerSizeModel.h - EmitC inliner model wrapper ----*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+/// Declares the wrapper around the EmitC-translated MLGO inliner model.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_ANALYSIS_EMITCINLINERSIZEMODEL_H
+#define LLVM_LIB_ANALYSIS_EMITCINLINERSIZEMODEL_H
+
+#include <array>
+#include <cstdint>
+#include <string>
+
+namespace llvm {
+
+class EmitCInlinerSizeModel final {
+public:
+ int LookupArgIndex(const std::string &Name);
+ int LookupResultIndex(const std::string &Name);
+ void *arg_data(int Index);
+ void *result_data(int Index);
+ void Run();
+
+private:
+ enum ArgIndex : int {
+ DeadBlocks = 0,
+ CaseClusterPenalty,
+ SroaSavings,
+ JumpTablePenalty,
+ CallsiteHeight,
+ CalleeBasicBlockCount,
+ CallArgumentSetup,
+ LoweredCallArgSetup,
+ SimplifiedInstructions,
+ NrCtantParams,
+ IsMultipleBlocks,
+ LoadElimination,
+ EdgeCount,
+ CallerUsers,
+ CallerConditionallyExecutedBlocks,
+ ConstantOffsetPtrArgs,
+ CallsiteCost,
+ CallerBasicBlockCount,
+ LoadRelativeIntrinsic,
+ IndirectCallPenalty,
+ CostEstimate,
+ Threshold,
+ NestedInlineCostEstimate,
+ UnsimplifiedCommonInstructions,
+ SroaLosses,
+ NumLoops,
+ SwitchPenalty,
+ CalleeUsers,
+ NodeCount,
+ ConstantArgs,
+ LastCallToStaticBonus,
+ ColdCCPenalty,
+ CalleeConditionallyExecutedBlocks,
+ CallPenalty,
+ NestedInlines,
+
+ NumArgs
+ };
+
+ std::array<std::array<int64_t, 1>, NumArgs> Inputs{};
+ std::array<int64_t, 1> Result{};
+
+ std::array<int64_t, 1> DummyInliningDefault{};
+ std::array<int32_t, 1> DummyStepType{};
+ std::array<float, 1> DummyDiscount{};
+ std::array<float, 1> DummyReward{};
+};
+
+} // namespace llvm
+
+#endif
diff --git a/llvm/include/llvm/CodeGen/EmitCRegAllocEvictModel.h b/llvm/include/llvm/CodeGen/EmitCRegAllocEvictModel.h
new file mode 100644
index 0000000000000..41a19ec62522e
--- /dev/null
+++ b/llvm/include/llvm/CodeGen/EmitCRegAllocEvictModel.h
@@ -0,0 +1,98 @@
+//===- EmitCRegAllocEvictModel.h - EmitC regalloc model wrapper -*- 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
+/// Declares the wrapper around the EmitC-translated MLGO regalloc eviction
+/// model.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGEN_EMITCREGALLOCEVICTMODEL_H
+#define LLVM_CODEGEN_EMITCREGALLOCEVICTMODEL_H
+
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
+#include <string>
+
+namespace llvm {
+
+class EmitCRegAllocEvictModel final {
+public:
+ int LookupArgIndex(const std::string &Name);
+ int LookupResultIndex(const std::string &Name);
+ void *arg_data(int Index);
+ void *result_data(int Index);
+ void Run();
+
+private:
+ static constexpr std::size_t InterferenceCount = 33;
+
+ using F32InterferenceTensor = float[1][InterferenceCount];
+ using I64InterferenceTensor = int64_t[1][InterferenceCount];
+ using F32Scalar = float[1];
+ using I32Scalar = int32_t[1];
+ using I64Scalar = int64_t[1];
+
+ enum ArgIndex : int {
+ Mask = 0,
+ IsFree,
+ NrUrgent,
+ NrBrokenHints,
+ IsHint,
+ IsLocal,
+ NrRematerializable,
+ NrDefsAndUses,
+ WeighedReadsByMax,
+ WeighedWritesByMax,
+ WeighedReadWritesByMax,
+ WeighedIndvarsByMax,
+ HintWeightsByMax,
+ StartBBFreqByMax,
+ EndBBFreqByMax,
+ HottestBBFreqByMax,
+ LiverangeSize,
+ UseDefDensity,
+ MaxStage,
+ MinStage,
+ Progress,
+
+ NumArgs
+ };
+
+ I64InterferenceTensor MaskInput{};
+ I64InterferenceTensor IsFreeInput{};
+ F32InterferenceTensor NrUrgentInput{};
+ F32InterferenceTensor NrBrokenHintsInput{};
+ I64InterferenceTensor IsHintInput{};
+ I64InterferenceTensor IsLocalInput{};
+ F32InterferenceTensor NrRematerializableInput{};
+ F32InterferenceTensor NrDefsAndUsesInput{};
+ F32InterferenceTensor WeighedReadsByMaxInput{};
+ F32InterferenceTensor WeighedWritesByMaxInput{};
+ F32InterferenceTensor WeighedReadWritesByMaxInput{};
+ F32InterferenceTensor WeighedIndvarsByMaxInput{};
+ F32InterferenceTensor HintWeightsByMaxInput{};
+ F32InterferenceTensor StartBBFreqByMaxInput{};
+ F32InterferenceTensor EndBBFreqByMaxInput{};
+ F32InterferenceTensor HottestBBFreqByMaxInput{};
+ F32InterferenceTensor LiverangeSizeInput{};
+ F32InterferenceTensor UseDefDensityInput{};
+ I64InterferenceTensor MaxStageInput{};
+ I64InterferenceTensor MinStageInput{};
+ F32Scalar ProgressInput{};
+
+ I32Scalar DummyStepType{};
+ F32Scalar DummyDiscount{};
+ F32Scalar DummyReward{};
+ I64Scalar Result{};
+};
+
+} // namespace llvm
+
+#endif
diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt
index f3586c66cb056..474d96df9b619 100644
--- a/llvm/lib/Analysis/CMakeLists.txt
+++ b/llvm/lib/Analysis/CMakeLists.txt
@@ -1,26 +1,31 @@
-if (DEFINED LLVM_HAVE_TF_AOT OR LLVM_HAVE_TFLITE)
- include(TensorFlowCompile)
- set(LLVM_INLINER_MODEL_PATH_DEFAULT "models/inliner-Oz")
+option(LLVM_USE_EMITC_INLINER_MODEL
+ "Use the in-tree EmitC inliner model instead of TensorFlow AOT"
+ OFF)
- set(LLVM_INLINER_MODEL_CURRENT_URL "<UNSPECIFIED>" CACHE STRING "URL to download the LLVM inliner model")
+if (LLVM_HAVE_TFLITE)
+ list(APPEND MLLinkDeps
+ tensorflow-lite::tensorflow-lite)
+endif()
- if (DEFINED LLVM_HAVE_TF_AOT)
- tf_find_and_compile(
- ${LLVM_INLINER_MODEL_PATH}
- ${LLVM_INLINER_MODEL_CURRENT_URL}
- ${LLVM_INLINER_MODEL_PATH_DEFAULT}
- "models/gen-inline-oz-test-model.py"
- serve
- action
- InlinerSizeModel
- llvm::InlinerSizeModel
- )
- endif()
+if (LLVM_USE_EMITC_INLINER_MODEL)
+ list(APPEND LLVM_COMPILE_DEFINITIONS
+ LLVM_HAVE_EMITC_INLINERSIZEMODEL)
+elseif (DEFINED LLVM_HAVE_TF_AOT)
+ include(TensorFlowCompile)
+ set(LLVM_INLINER_MODEL_PATH_DEFAULT "models/inliner-Oz")
+ set(LLVM_INLINER_MODEL_CURRENT_URL "<UNSPECIFIED>" CACHE STRING
+ "URL to download the LLVM inliner model")
- if (LLVM_HAVE_TFLITE)
- list(APPEND MLLinkDeps
- tensorflow-lite::tensorflow-lite)
- endif()
+ tf_find_and_compile(
+ ${LLVM_INLINER_MODEL_PATH}
+ ${LLVM_INLINER_MODEL_CURRENT_URL}
+ ${LLVM_INLINER_MODEL_PATH_DEFAULT}
+ "models/gen-inline-oz-test-model.py"
+ serve
+ action
+ InlinerSizeModel
+ llvm::InlinerSizeModel
+ )
endif()
# The implementation of ConstantFolding.cpp relies on the use of math functions
@@ -77,6 +82,7 @@ add_llvm_component_library(LLVMAnalysis
DominanceFrontier.cpp
DXILResource.cpp
DXILMetadataAnalysis.cpp
+ EmitCInlinerSizeModel.cpp
EphemeralValuesCache.cpp
FloatingPointPredicateUtils.cpp
FunctionPropertiesAnalysis.cpp
@@ -187,6 +193,11 @@ add_llvm_component_library(LLVMAnalysis
TargetParser
)
+if (LLVM_USE_EMITC_INLINER_MODEL)
+ target_compile_definitions(LLVMAnalysis
+ PRIVATE LLVM_HAVE_EMITC_INLINERSIZEMODEL)
+endif()
+
include(CheckCXXSymbolExists)
check_cxx_symbol_exists(logf128 math.h HAS_LOGF128)
if(HAS_LOGF128)
diff --git a/llvm/lib/Analysis/EmitCInlinerSizeModel.cpp b/llvm/lib/Analysis/EmitCInlinerSizeModel.cpp
new file mode 100644
index 0000000000000..a8e0094737bd6
--- /dev/null
+++ b/llvm/lib/Analysis/EmitCInlinerSizeModel.cpp
@@ -0,0 +1,250 @@
+//===- EmitCInlinerSizeModel.cpp - EmitC inliner model wrapper ------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+/// This file implements the wrapper around the EmitC-translated MLGO inliner
+/// model.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Analysis/EmitCInlinerSizeModel.h"
+
+#include "llvm/ADT/StringSwitch.h"
+#include "llvm/Support/ErrorHandling.h"
+
+#include <math.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <type_traits>
+
+#if defined(__clang__)
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wmissing-braces"
+#endif
+
+namespace llvm::emitc_inliner_model {
+#define main action
+#include "llvm/Analysis/EmitCInlinerSizeModel.inc"
+#undef main
+} // namespace llvm::emitc_inliner_model
+
+#if defined(__clang__)
+#pragma clang diagnostic pop
+#endif
+
+using namespace llvm;
+
+namespace {
+template <typename T> inline constexpr bool AlwaysFalse = false;
+using I64Ptr = int64_t *;
+using I32Ptr = int32_t *;
+using F32Ptr = float *;
+
+using InlinerProductionActionTy =
+ int64_t (*)(I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
+ I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
+ I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
+ I64Ptr, I64Ptr, I64Ptr, I32Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
+ I64Ptr, I64Ptr, I64Ptr, I64Ptr, F32Ptr, I64Ptr, F32Ptr);
+using InlinerMockActionTy = int64_t (*)(I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
+ I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
+ I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
+ I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
+ I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
+ I64Ptr, I64Ptr, I32Ptr, I64Ptr, I64Ptr,
+ I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
+ I64Ptr, I64Ptr, F32Ptr, I64Ptr, F32Ptr);
+
+struct InlinerRunInputs {
+ I64Ptr deadBlocks;
+ I64Ptr caseClusterPenalty;
+ I64Ptr sroaSavings;
+ I64Ptr jumpTablePenalty;
+ I64Ptr callsiteHeight;
+ I64Ptr calleeBasicBlockCount;
+ I64Ptr callArgumentSetup;
+ I64Ptr loweredCallArgSetup;
+ I64Ptr simplifiedInstructions;
+ I64Ptr nrCtantParams;
+ I64Ptr isMultipleBlocks;
+ I64Ptr loadElimination;
+ I64Ptr edgeCount;
+ I64Ptr callerUsers;
+ I64Ptr callerConditionallyExecutedBlocks;
+ I64Ptr constantOffsetPtrArgs;
+ I64Ptr callsiteCost;
+ I64Ptr callerBasicBlockCount;
+ I64Ptr loadRelativeIntrinsic;
+ I64Ptr indirectCallPenalty;
+ I64Ptr costEstimate;
+ I64Ptr threshold;
+ I64Ptr nestedInlineCostEstimate;
+ I64Ptr unsimplifiedCommonInstructions;
+ I64Ptr sroaLosses;
+ I64Ptr numLoops;
+ I64Ptr switchPenalty;
+ I64Ptr calleeUsers;
+ I64Ptr nodeCount;
+ I64Ptr constantArgs;
+ I64Ptr lastCallToStaticBonus;
+ I64Ptr coldCCPenalty;
+ I64Ptr calleeConditionallyExecutedBlocks;
+ I64Ptr callPenalty;
+ I64Ptr nestedInlines;
+ I32Ptr dummyStepType;
+ F32Ptr dummyDiscount;
+ F32Ptr dummyReward;
+ I64Ptr dummyInliningDefault;
+};
+
+template <typename ActionTy>
+int64_t runEmitCInlinerAction(const InlinerRunInputs &I) {
+ if constexpr (std::is_same_v<ActionTy, InlinerProductionActionTy>) {
+ return static_cast<ActionTy>(emitc_inliner_model::action)(
+ I.callsiteCost, I.isMultipleBlocks, I.callerConditionallyExecutedBlocks,
+ I.dummyInliningDefault, I.coldCCPenalty,
+ I.calleeConditionallyExecutedBlocks, I.calleeUsers,
+ I.calleeBasicBlockCount, I.nrCtantParams, I.loadRelativeIntrinsic,
+ I.jumpTablePenalty, I.unsimplifiedCommonInstructions,
+ I.indirectCallPenalty, I.loadElimination, I.callPenalty, I.costEstimate,
+ I.caseClusterPenalty, I.nodeCount, I.callArgumentSetup, I.sroaSavings,
+ I.loweredCallArgSetup, I.threshold, I.deadBlocks, I.constantArgs,
+ I.sroaLosses, I.simplifiedInstructions, I.numLoops, I.dummyStepType,
+ I.edgeCount, I.nestedInlines, I.callerBasicBlockCount,
+ I.lastCallToStaticBonus, I.nestedInlineCostEstimate, I.callsiteHeight,
+ I.constantOffsetPtrArgs, I.switchPenalty, I.dummyDiscount,
+ I.callerUsers, I.dummyReward);
+ } else if constexpr (std::is_same_v<ActionTy, InlinerMockActionTy>) {
+ return static_cast<ActionTy>(emitc_inliner_model::action)(
+ I.callerBasicBlockCount, I.callerConditionallyExecutedBlocks,
+ I.callerUsers, I.calleeBasicBlockCount,
+ I.calleeConditionallyExecutedBlocks, I.calleeUsers, I.nrCtantParams,
+ I.nodeCount, I.edgeCount, I.callsiteHeight, I.costEstimate,
+ I.sroaSavings, I.sroaLosses, I.loadElimination, I.callPenalty,
+ I.callArgumentSetup, I.loadRelativeIntrinsic, I.loweredCallArgSetup,
+ I.indirectCallPenalty, I.jumpTablePenalty, I.caseClusterPenalty,
+ I.switchPenalty, I.unsimplifiedCommonInstructions, I.numLoops,
+ I.deadBlocks, I.simplifiedInstructions, I.constantArgs, I.dummyStepType,
+ I.constantOffsetPtrArgs, I.callsiteCost, I.coldCCPenalty,
+ I.lastCallToStaticBonus, I.isMultipleBlocks, I.nestedInlines,
+ I.nestedInlineCostEstimate, I.threshold, I.dummyInliningDefault,
+ I.dummyDiscount, I.callerUsers, I.dummyReward);
+ } else {
+ static_assert(AlwaysFalse<ActionTy>,
+ "Unsupported EmitC inliner model signature");
+ }
+}
+} // namespace
+
+int EmitCInlinerSizeModel::LookupArgIndex(const std::string &Name) {
+ return StringSwitch<int>(Name)
+ .Case("feed_dead_blocks", DeadBlocks)
+ .Case("feed_case_cluster_penalty", CaseClusterPenalty)
+ .Case("feed_sroa_savings", SroaSavings)
+ .Case("feed_jump_table_penalty", JumpTablePenalty)
+ .Case("feed_callsite_height", CallsiteHeight)
+ .Case("feed_callee_basic_block_count", CalleeBasicBlockCount)
+ .Case("feed_call_argument_setup", CallArgumentSetup)
+ .Case("feed_lowered_call_arg_setup", LoweredCallArgSetup)
+ .Case("feed_simplified_instructions", SimplifiedInstructions)
+ .Case("feed_nr_ctant_params", NrCtantParams)
+ .Case("feed_is_multiple_blocks", IsMultipleBlocks)
+ .Case("feed_load_elimination", LoadElimination)
+ .Case("feed_edge_count", EdgeCount)
+ .Case("feed_caller_users", CallerUsers)
+ .Case("feed_caller_conditionally_executed_blocks",
+ CallerConditionallyExecutedBlocks)
+ .Case("feed_constant_offset_ptr_args", ConstantOffsetPtrArgs)
+ .Case("feed_callsite_cost", CallsiteCost)
+ .Case("feed_caller_basic_block_count", CallerBasicBlockCount)
+ .Case("feed_load_relative_intrinsic", LoadRelativeIntrinsic)
+ .Case("feed_indirect_call_penalty", IndirectCallPenalty)
+ .Case("feed_cost_estimate", CostEstimate)
+ .Case("feed_threshold", Threshold)
+ .Case("feed_nested_inline_cost_estimate", NestedInlineCostEstimate)
+ .Case("feed_unsimplified_common_instructions",
+ UnsimplifiedCommonInstructions)
+ .Case("feed_sroa_losses", SroaLosses)
+ .Case("feed_num_loops", NumLoops)
+ .Case("feed_switch_penalty", SwitchPenalty)
+ .Case("feed_callee_users", CalleeUsers)
+ .Case("feed_node_count", NodeCount)
+ .Case("feed_constant_args", ConstantArgs)
+ .Case("feed_last_call_to_static_bonus", LastCallToStaticBonus)
+ .Case("feed_cold_cc_penalty", ColdCCPenalty)
+ .Case("feed_callee_conditionally_executed_blocks",
+ CalleeConditionallyExecutedBlocks)
+ .Case("feed_call_penalty", CallPenalty)
+ .Case("feed_nested_inlines", NestedInlines)
+ .Default(-1);
+}
+
+int EmitCInlinerSizeModel::LookupResultIndex(const std::string &Name) {
+ return Name == "fetch_inlining_decision" ? 0 : -1;
+}
+
+void *EmitCInlinerSizeModel::arg_data(int Index) {
+ if (Index < 0 || Index >= NumArgs)
+ llvm_unreachable("invalid EmitC inliner input index");
+ return Inputs[Index].data();
+}
+
+void *EmitCInlinerSizeModel::result_data(int Index) {
+ if (Index != 0)
+ llvm_unreachable("invalid EmitC inliner result index");
+ return Result.data();
+}
+
+void EmitCInlinerSizeModel::Run() {
+ using ActionTy = decltype(&emitc_inliner_model::action);
+ InlinerRunInputs I{};
+ I.deadBlocks = Inputs[DeadBlocks].data();
+ I.caseClusterPenalty = Inputs[CaseClusterPenalty].data();
+ I.sroaSavings = Inputs[SroaSavings].data();
+ I.jumpTablePenalty = Inputs[JumpTablePenalty].data();
+ I.callsiteHeight = Inputs[CallsiteHeight].data();
+ I.calleeBasicBlockCount = Inputs[CalleeBasicBlockCount].data();
+ I.callArgumentSetup = Inputs[CallArgumentSetup].data();
+ I.loweredCallArgSetup = Inputs[LoweredCallArgSetup].data();
+ I.simplifiedInstructions = Inputs[SimplifiedInstructions].data();
+ I.nrCtantParams = Inputs[NrCtantParams].data();
+ I.isMultipleBlocks = Inputs[IsMultipleBlocks].data();
+ I.loadElimination = Inputs[LoadElimination].data();
+ I.edgeCount = Inputs[EdgeCount].data();
+ I.callerUsers = Inputs[CallerUsers].data();
+ I.callerConditionallyExecutedBlocks =
+ Inputs[CallerConditionallyExecutedBlocks].data();
+ I.constantOffsetPtrArgs = Inputs[ConstantOffsetPtrArgs].data();
+ I.callsiteCost = Inputs[CallsiteCost].data();
+ I.callerBasicBlockCount = Inputs[CallerBasicBlockCount].data();
+ I.loadRelativeIntrinsic = Inputs[LoadRelativeIntrinsic].data();
+ I.indirectCallPenalty = Inputs[IndirectCallPenalty].data();
+ I.costEstimate = Inputs[CostEstimate].data();
+ I.threshold = Inputs[Threshold].data();
+ I.nestedInlineCostEstimate = Inputs[NestedInlineCostEstimate].data();
+ I.unsimplifiedCommonInstructions =
+ Inputs[UnsimplifiedCommonInstructions].data();
+ I.sroaLosses = Inputs[SroaLosses].data();
+ I.numLoops = Inputs[NumLoops].data();
+ I.switchPenalty = Inputs[SwitchPenalty].data();
+ I.calleeUsers = Inputs[CalleeUsers].data();
+ I.nodeCount = Inputs[NodeCount].data();
+ I.constantArgs = Inputs[ConstantArgs].data();
+ I.lastCallToStaticBonus = Inputs[LastCallToStaticBonus].data();
+ I.coldCCPenalty = Inputs[ColdCCPenalty].data();
+ I.calleeConditionallyExecutedBlocks =
+ Inputs[CalleeConditionallyExecutedBlocks].data();
+ I.callPenalty = Inputs[CallPenalty].data();
+ I.nestedInlines = Inputs[NestedInlines].data();
+ I.dummyStepType = DummyStepType.data();
+ I.dummyDiscount = DummyDiscount.data();
+ I.dummyReward = DummyReward.data();
+ I.dummyInliningDefault = DummyInliningDefault.data();
+ Result[0] = runEmitCInlinerAction<ActionTy>(I);
+}
diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp
index 9a5ae2ae26799..e698f066b9336 100644
--- a/llvm/lib/Analysis/MLInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp
@@ -64,7 +64,10 @@ static cl::opt<std::string> ModelSelector("ml-inliner-model-selector",
static cl::opt<bool> StopImmediatelyForTest("ml-inliner-stop-immediately",
cl::Hidden);
-#if defined(LLVM_HAVE_TF_AOT_INLINERSIZEMODEL)
+#if defined(LLVM_HAVE_EMITC_INLINERSIZEMODEL)
+#include "llvm/Analysis/EmitCInlinerSizeModel.h"
+using CompiledModelType = llvm::EmitCInlinerSizeModel;
+#elif defined(LLVM_HAVE_TF_AOT_INLINERSIZEMODEL)
// codegen-ed file
#include "InlinerSizeModel.h" // NOLINT
using CompiledModelType = llvm::InlinerSizeModel;
diff --git a/llvm/lib/CodeGen/CMakeLists.txt b/llvm/lib/CodeGen/CMakeLists.txt
index c572128b023c1..bf05ca6d7cb19 100644
--- a/llvm/lib/CodeGen/CMakeLists.txt
+++ b/llvm/lib/CodeGen/CMakeLists.txt
@@ -1,10 +1,19 @@
+option(LLVM_USE_EMITC_REGALLOC_EVICT_MODEL
+ "Use the in-tree EmitC regalloc eviction model instead of TensorFlow AOT"
+ OFF)
+
+if (LLVM_USE_EMITC_REGALLOC_EVICT_MODEL)
+ list(APPEND LLVM_COMPILE_DEFINITIONS
+ LLVM_HAVE_EMITC_REGALLOCEVICTMODEL)
+endif()
+
if (DEFINED LLVM_HAVE_TF_AOT OR LLVM_HAVE_TFLITE)
include(TensorFlowCompile)
set(LLVM_RAEVICT_MODEL_PATH_DEFAULT "models/regalloc-eviction")
set(LLVM_RAEVICT_MODEL_CURRENT_URL "<UNSPECIFIED>" CACHE STRING "URL to download the LLVM register allocator eviction model")
- if (DEFINED LLVM_HAVE_TF_AOT)
+ if (DEFINED LLVM_HAVE_TF_AOT AND NOT LLVM_USE_EMITC_REGALLOC_EVICT_MODEL)
tf_find_and_compile(
${LLVM_RAEVICT_MODEL_PATH}
${LLVM_RAEVICT_MODEL_CURRENT_URL}
@@ -53,6 +62,7 @@ add_llvm_component_library(LLVMCodeGen
DroppedVariableStatsMIR.cpp
DwarfEHPrepare.cpp
EarlyIfConversion.cpp
+ EmitCRegAllocEvictModel.cpp
EdgeBundles.cpp
EHContGuardTargets.cpp
ExecutionDomainFix.cpp
@@ -295,6 +305,11 @@ add_llvm_component_library(LLVMCodeGen
TransformUtils
)
+if (LLVM_USE_EMITC_REGALLOC_EVICT_MODEL)
+ target_compile_definitions(LLVMCodeGen
+ PRIVATE LLVM_HAVE_EMITC_REGALLOCEVICTMODEL)
+endif()
+
add_subdirectory(SelectionDAG)
add_subdirectory(AsmPrinter)
add_subdirectory(MIRParser)
diff --git a/llvm/lib/CodeGen/EmitCRegAllocEvictModel.cpp b/llvm/lib/CodeGen/EmitCRegAllocEvictModel.cpp
new file mode 100644
index 0000000000000..120268ec1f389
--- /dev/null
+++ b/llvm/lib/CodeGen/EmitCRegAllocEvictModel.cpp
@@ -0,0 +1,209 @@
+//===- EmitCRegAllocEvictModel.cpp - EmitC regalloc model wrapper ---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+/// This file implements the wrapper around the EmitC-translated MLGO
+/// regalloc eviction model.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/EmitCRegAllocEvictModel.h"
+
+#include "llvm/ADT/StringSwitch.h"
+#include "llvm/Support/ErrorHandling.h"
+
+#include <stddef.h>
+#include <stdint.h>
+#include <type_traits>
+
+namespace llvm::emitc_regalloc_evict_model {
+#define main action
+#include "llvm/CodeGen/EmitCRegAllocEvictModel.inc"
+#undef main
+} // namespace llvm::emitc_regalloc_evict_model
+
+using namespace llvm;
+
+namespace {
+template <typename T> inline constexpr bool AlwaysFalse = false;
+constexpr std::size_t EmitCRegAllocInterferenceCount = 33;
+using F32TensorPtr = float (*)[EmitCRegAllocInterferenceCount];
+using I64TensorPtr = int64_t (*)[EmitCRegAllocInterferenceCount];
+using F32ScalarPtr = float *;
+using I32ScalarPtr = int32_t *;
+using I64ScalarPtr = int64_t *;
+
+using RegAllocProductionActionTy = int64_t (*)(
+ F32TensorPtr, F32TensorPtr, I64TensorPtr, F32TensorPtr, F32TensorPtr,
+ F32TensorPtr, F32ScalarPtr, F32TensorPtr, F32TensorPtr, F32TensorPtr,
+ I64TensorPtr, I64TensorPtr, F32TensorPtr, F32TensorPtr, I32ScalarPtr,
+ F32TensorPtr, I64TensorPtr, F32TensorPtr, I64TensorPtr, F32TensorPtr,
+ F32ScalarPtr, F32TensorPtr, I64TensorPtr, F32ScalarPtr);
+using RegAllocMaskOnlyActionTy = int64_t (*)(I64ScalarPtr);
+
+struct RegAllocRunInputs {
+ F32TensorPtr liverangeSize;
+ F32TensorPtr hintWeightsByMax;
+ I64TensorPtr isFree;
+ F32TensorPtr weighedReadsByMax;
+ F32TensorPtr weighedReadWritesByMax;
+ F32TensorPtr nrBrokenHints;
+ F32ScalarPtr progress;
+ F32TensorPtr hottestBBFreqByMax;
+ F32TensorPtr useDefDensity;
+ F32TensorPtr startBBFreqByMax;
+ I64TensorPtr maxStage;
+ I64TensorPtr isHint;
+ F32TensorPtr nrRematerializable;
+ F32TensorPtr weighedWritesByMax;
+ I32ScalarPtr dummyStepType;
+ F32TensorPtr nrUrgent;
+ I64TensorPtr mask;
+ F32TensorPtr nrDefsAndUses;
+ I64TensorPtr isLocal;
+ F32TensorPtr endBBFreqByMax;
+ F32ScalarPtr dummyDiscount;
+ F32TensorPtr weighedIndvarsByMax;
+ I64TensorPtr minStage;
+ F32ScalarPtr dummyReward;
+ I64ScalarPtr maskFlat;
+};
+
+template <typename ActionTy>
+int64_t runEmitCRegAllocAction(const RegAllocRunInputs &I) {
+ if constexpr (std::is_same_v<ActionTy, RegAllocProductionActionTy>) {
+ return static_cast<ActionTy>(emitc_regalloc_evict_model::action)(
+ I.liverangeSize, I.hintWeightsByMax, I.isFree, I.weighedReadsByMax,
+ I.weighedReadWritesByMax, I.nrBrokenHints, I.progress,
+ I.hottestBBFreqByMax, I.useDefDensity, I.startBBFreqByMax, I.maxStage,
+ I.isHint, I.nrRematerializable, I.weighedWritesByMax, I.dummyStepType,
+ I.nrUrgent, I.mask, I.nrDefsAndUses, I.isLocal, I.endBBFreqByMax,
+ I.dummyDiscount, I.weighedIndvarsByMax, I.minStage, I.dummyReward);
+ } else if constexpr (std::is_same_v<ActionTy, RegAllocMaskOnlyActionTy>) {
+ return static_cast<ActionTy>(emitc_regalloc_evict_model::action)(
+ I.maskFlat);
+ } else {
+ static_assert(AlwaysFalse<ActionTy>,
+ "Unsupported EmitC regalloc eviction model signature");
+ }
+}
+} // namespace
+
+int EmitCRegAllocEvictModel::LookupArgIndex(const std::string &Name) {
+ return StringSwitch<int>(Name)
+ .Case("feed_mask", Mask)
+ .Case("feed_is_free", IsFree)
+ .Case("feed_nr_urgent", NrUrgent)
+ .Case("feed_nr_broken_hints", NrBrokenHints)
+ .Case("feed_is_hint", IsHint)
+ .Case("feed_is_local", IsLocal)
+ .Case("feed_nr_rematerializable", NrRematerializable)
+ .Case("feed_nr_defs_and_uses", NrDefsAndUses)
+ .Case("feed_weighed_reads_by_max", WeighedReadsByMax)
+ .Case("feed_weighed_writes_by_max", WeighedWritesByMax)
+ .Case("feed_weighed_read_writes_by_max", WeighedReadWritesByMax)
+ .Case("feed_weighed_indvars_by_max", WeighedIndvarsByMax)
+ .Case("feed_hint_weights_by_max", HintWeightsByMax)
+ .Case("feed_start_bb_freq_by_max", StartBBFreqByMax)
+ .Case("feed_end_bb_freq_by_max", EndBBFreqByMax)
+ .Case("feed_hottest_bb_freq_by_max", HottestBBFreqByMax)
+ .Case("feed_liverange_size", LiverangeSize)
+ .Case("feed_use_def_density", UseDefDensity)
+ .Case("feed_max_stage", MaxStage)
+ .Case("feed_min_stage", MinStage)
+ .Case("feed_progress", Progress)
+ .Default(-1);
+}
+
+int EmitCRegAllocEvictModel::LookupResultIndex(const std::string &Name) {
+ return Name == "fetch_index_to_evict" ? 0 : -1;
+}
+
+void *EmitCRegAllocEvictModel::arg_data(int Index) {
+ switch (Index) {
+ case Mask:
+ return MaskInput;
+ case IsFree:
+ return IsFreeInput;
+ case NrUrgent:
+ return NrUrgentInput;
+ case NrBrokenHints:
+ return NrBrokenHintsInput;
+ case IsHint:
+ return IsHintInput;
+ case IsLocal:
+ return IsLocalInput;
+ case NrRematerializable:
+ return NrRematerializableInput;
+ case NrDefsAndUses:
+ return NrDefsAndUsesInput;
+ case WeighedReadsByMax:
+ return WeighedReadsByMaxInput;
+ case WeighedWritesByMax:
+ return WeighedWritesByMaxInput;
+ case WeighedReadWritesByMax:
+ return WeighedReadWritesByMaxInput;
+ case WeighedIndvarsByMax:
+ return WeighedIndvarsByMaxInput;
+ case HintWeightsByMax:
+ return HintWeightsByMaxInput;
+ case StartBBFreqByMax:
+ return StartBBFreqByMaxInput;
+ case EndBBFreqByMax:
+ return EndBBFreqByMaxInput;
+ case HottestBBFreqByMax:
+ return HottestBBFreqByMaxInput;
+ case LiverangeSize:
+ return LiverangeSizeInput;
+ case UseDefDensity:
+ return UseDefDensityInput;
+ case MaxStage:
+ return MaxStageInput;
+ case MinStage:
+ return MinStageInput;
+ case Progress:
+ return ProgressInput;
+ }
+ llvm_unreachable("invalid EmitC regalloc eviction input index");
+}
+
+void *EmitCRegAllocEvictModel::result_data(int Index) {
+ if (Index != 0)
+ llvm_unreachable("invalid EmitC regalloc eviction result index");
+ return Result;
+}
+
+void EmitCRegAllocEvictModel::Run() {
+ using ActionTy = decltype(&emitc_regalloc_evict_model::action);
+ RegAllocRunInputs I{};
+ I.liverangeSize = LiverangeSizeInput;
+ I.hintWeightsByMax = HintWeightsByMaxInput;
+ I.isFree = IsFreeInput;
+ I.weighedReadsByMax = WeighedReadsByMaxInput;
+ I.weighedReadWritesByMax = WeighedReadWritesByMaxInput;
+ I.nrBrokenHints = NrBrokenHintsInput;
+ I.progress = ProgressInput;
+ I.hottestBBFreqByMax = HottestBBFreqByMaxInput;
+ I.useDefDensity = UseDefDensityInput;
+ I.startBBFreqByMax = StartBBFreqByMaxInput;
+ I.maxStage = MaxStageInput;
+ I.isHint = IsHintInput;
+ I.nrRematerializable = NrRematerializableInput;
+ I.weighedWritesByMax = WeighedWritesByMaxInput;
+ I.dummyStepType = DummyStepType;
+ I.nrUrgent = NrUrgentInput;
+ I.mask = MaskInput;
+ I.nrDefsAndUses = NrDefsAndUsesInput;
+ I.isLocal = IsLocalInput;
+ I.endBBFreqByMax = EndBBFreqByMaxInput;
+ I.dummyDiscount = DummyDiscount;
+ I.weighedIndvarsByMax = WeighedIndvarsByMaxInput;
+ I.minStage = MinStageInput;
+ I.dummyReward = DummyReward;
+ I.maskFlat = MaskInput[0];
+ Result[0] = runEmitCRegAllocAction<ActionTy>(I);
+}
diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index 23dc6fbd6e500..03ba2da4e945e 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -47,8 +47,11 @@ using namespace llvm;
#define DEBUG_TYPE "ml-regalloc"
-// Generated header in release (AOT) mode
-#if defined(LLVM_HAVE_TF_AOT_REGALLOCEVICTMODEL)
+// Generated header in release (AOT / EmitC) mode
+#if defined(LLVM_HAVE_EMITC_REGALLOCEVICTMODEL)
+#include "llvm/CodeGen/EmitCRegAllocEvictModel.h"
+using CompiledModelType = llvm::EmitCRegAllocEvictModel;
+#elif defined(LLVM_HAVE_TF_AOT_REGALLOCEVICTMODEL)
#include "RegAllocEvictModel.h"
using CompiledModelType = RegAllocEvictModel;
#else
diff --git a/llvm/test/CodeGen/MLRegAlloc/default-eviction-advisor.ll b/llvm/test/CodeGen/MLRegAlloc/default-eviction-advisor.ll
index 881a80c41361c..fc9d244a617a9 100644
--- a/llvm/test/CodeGen/MLRegAlloc/default-eviction-advisor.ll
+++ b/llvm/test/CodeGen/MLRegAlloc/default-eviction-advisor.ll
@@ -2,6 +2,7 @@
; trying to use ML-driven advisor.
; REQUIRES: !have_tf_aot
; REQUIRES: !have_tflite
+; REQUIRES: !have_emitc_raevict_model
; 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
diff --git a/llvm/test/lit.cfg.py b/llvm/test/lit.cfg.py
index cd028963dd59e..796d258515e5b 100644
--- a/llvm/test/lit.cfg.py
+++ b/llvm/test/lit.cfg.py
@@ -597,6 +597,9 @@ def enable_ptxas(ptxas_executable):
if config.have_tflite:
config.available_features.add("have_tflite")
+if config.have_emitc_raevict_model:
+ config.available_features.add("have_emitc_raevict_model")
+
if config.llvm_inliner_model_autogenerated:
config.available_features.add("llvm_inliner_model_autogenerated")
diff --git a/llvm/test/lit.site.cfg.py.in b/llvm/test/lit.site.cfg.py.in
index 64679c2f64034..8c60b0d1a5d1b 100644
--- a/llvm/test/lit.site.cfg.py.in
+++ b/llvm/test/lit.site.cfg.py.in
@@ -57,6 +57,7 @@ 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_tflite = @LLVM_HAVE_TFLITE@
+config.have_emitc_raevict_model = @LLVM_USE_EMITC_REGALLOC_EVICT_MODEL@
config.enable_profcheck = @LLVM_ENABLE_PROFCHECK@
config.llvm_inliner_model_autogenerated = @LLVM_INLINER_MODEL_AUTOGENERATED@
config.llvm_raevict_model_autogenerated = @LLVM_RAEVICT_MODEL_AUTOGENERATED@
>From 33bdfa86d8a408260351bb220b2fa0944549f75d Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Thu, 16 Jul 2026 11:26:09 +0200
Subject: [PATCH 2/7] Mark model wrapper implementations optional
---
llvm/lib/Analysis/CMakeLists.txt | 13 ++++++++++---
llvm/lib/CodeGen/CMakeLists.txt | 14 +++++++++++---
2 files changed, 21 insertions(+), 6 deletions(-)
diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt
index 474d96df9b619..50f10c58e6892 100644
--- a/llvm/lib/Analysis/CMakeLists.txt
+++ b/llvm/lib/Analysis/CMakeLists.txt
@@ -2,14 +2,21 @@ option(LLVM_USE_EMITC_INLINER_MODEL
"Use the in-tree EmitC inliner model instead of TensorFlow AOT"
OFF)
+# The EmitC wrapper is only built when the in-tree EmitC model is enabled.
+# Mark it optional so LLVM's source audit does not require it in default builds.
+list(APPEND LLVM_OPTIONAL_SOURCES
+ EmitCInlinerSizeModel.cpp)
+
if (LLVM_HAVE_TFLITE)
list(APPEND MLLinkDeps
tensorflow-lite::tensorflow-lite)
endif()
+# Only compile the EmitC inliner wrapper when the EmitC-backed release model
+# is selected. Otherwise the generated .inc file may not exist.
if (LLVM_USE_EMITC_INLINER_MODEL)
- list(APPEND LLVM_COMPILE_DEFINITIONS
- LLVM_HAVE_EMITC_INLINERSIZEMODEL)
+ list(APPEND LLVMAnalysisOptionalSources
+ EmitCInlinerSizeModel.cpp)
elseif (DEFINED LLVM_HAVE_TF_AOT)
include(TensorFlowCompile)
set(LLVM_INLINER_MODEL_PATH_DEFAULT "models/inliner-Oz")
@@ -82,7 +89,6 @@ add_llvm_component_library(LLVMAnalysis
DominanceFrontier.cpp
DXILResource.cpp
DXILMetadataAnalysis.cpp
- EmitCInlinerSizeModel.cpp
EphemeralValuesCache.cpp
FloatingPointPredicateUtils.cpp
FunctionPropertiesAnalysis.cpp
@@ -170,6 +176,7 @@ add_llvm_component_library(LLVMAnalysis
ValueLatticeUtils.cpp
ValueTracking.cpp
VectorUtils.cpp
+ ${LLVMAnalysisOptionalSources}
${GeneratedMLSources}
ADDITIONAL_HEADER_DIRS
diff --git a/llvm/lib/CodeGen/CMakeLists.txt b/llvm/lib/CodeGen/CMakeLists.txt
index bf05ca6d7cb19..72496671b3d68 100644
--- a/llvm/lib/CodeGen/CMakeLists.txt
+++ b/llvm/lib/CodeGen/CMakeLists.txt
@@ -2,9 +2,17 @@ option(LLVM_USE_EMITC_REGALLOC_EVICT_MODEL
"Use the in-tree EmitC regalloc eviction model instead of TensorFlow AOT"
OFF)
+# The EmitC wrapper is only built when the in-tree EmitC regalloc model is
+# enabled. Mark it optional so LLVM's source audit accepts its absence from the
+# default target source list.
+list(APPEND LLVM_OPTIONAL_SOURCES
+ EmitCRegAllocEvictModel.cpp)
+
+# Only compile the EmitC regalloc wrapper when the EmitC-backed release model
+# is selected. Otherwise the generated .inc file may not exist.
if (LLVM_USE_EMITC_REGALLOC_EVICT_MODEL)
- list(APPEND LLVM_COMPILE_DEFINITIONS
- LLVM_HAVE_EMITC_REGALLOCEVICTMODEL)
+ list(APPEND LLVMCodeGenOptionalSources
+ EmitCRegAllocEvictModel.cpp)
endif()
if (DEFINED LLVM_HAVE_TF_AOT OR LLVM_HAVE_TFLITE)
@@ -62,7 +70,6 @@ add_llvm_component_library(LLVMCodeGen
DroppedVariableStatsMIR.cpp
DwarfEHPrepare.cpp
EarlyIfConversion.cpp
- EmitCRegAllocEvictModel.cpp
EdgeBundles.cpp
EHContGuardTargets.cpp
ExecutionDomainFix.cpp
@@ -268,6 +275,7 @@ add_llvm_component_library(LLVMCodeGen
WindowsSecureHotPatching.cpp
WinEHPrepare.cpp
XRayInstrumentation.cpp
+ ${LLVMCodeGenOptionalSources}
${GeneratedMLSources}
LiveDebugValues/LiveDebugValues.cpp
>From 4482a777d430d7757b1c4484224992ac8894261e Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Thu, 16 Jul 2026 11:58:45 +0200
Subject: [PATCH 3/7] Refine comments and naming
---
llvm/lib/Analysis/CMakeLists.txt | 12 +++++-------
llvm/lib/CodeGen/CMakeLists.txt | 13 +++++--------
2 files changed, 10 insertions(+), 15 deletions(-)
diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt
index 50f10c58e6892..323d9cd2e6df2 100644
--- a/llvm/lib/Analysis/CMakeLists.txt
+++ b/llvm/lib/Analysis/CMakeLists.txt
@@ -1,9 +1,8 @@
option(LLVM_USE_EMITC_INLINER_MODEL
- "Use the in-tree EmitC inliner model instead of TensorFlow AOT"
+ "Use the fully in-tree inliner model integration instead of TensorFlow AOT"
OFF)
-# The EmitC wrapper is only built when the in-tree EmitC model is enabled.
-# Mark it optional so LLVM's source audit does not require it in default builds.
+# Marked optional so LLVM's source audit does not require it in default builds.
list(APPEND LLVM_OPTIONAL_SOURCES
EmitCInlinerSizeModel.cpp)
@@ -12,10 +11,9 @@ if (LLVM_HAVE_TFLITE)
tensorflow-lite::tensorflow-lite)
endif()
-# Only compile the EmitC inliner wrapper when the EmitC-backed release model
-# is selected. Otherwise the generated .inc file may not exist.
+# The inlining model wrapper is only built when the fully in-tree model integration is enabled.
if (LLVM_USE_EMITC_INLINER_MODEL)
- list(APPEND LLVMAnalysisOptionalSources
+ list(APPEND AnalysisInTreeModelSources
EmitCInlinerSizeModel.cpp)
elseif (DEFINED LLVM_HAVE_TF_AOT)
include(TensorFlowCompile)
@@ -176,7 +174,7 @@ add_llvm_component_library(LLVMAnalysis
ValueLatticeUtils.cpp
ValueTracking.cpp
VectorUtils.cpp
- ${LLVMAnalysisOptionalSources}
+ ${AnalysisInTreeModelSources}
${GeneratedMLSources}
ADDITIONAL_HEADER_DIRS
diff --git a/llvm/lib/CodeGen/CMakeLists.txt b/llvm/lib/CodeGen/CMakeLists.txt
index 72496671b3d68..a32300db1093b 100644
--- a/llvm/lib/CodeGen/CMakeLists.txt
+++ b/llvm/lib/CodeGen/CMakeLists.txt
@@ -1,17 +1,14 @@
option(LLVM_USE_EMITC_REGALLOC_EVICT_MODEL
- "Use the in-tree EmitC regalloc eviction model instead of TensorFlow AOT"
+ "Use the fully in-tree regalloc eviction model integration instead of TensorFlow AOT"
OFF)
-# The EmitC wrapper is only built when the in-tree EmitC regalloc model is
-# enabled. Mark it optional so LLVM's source audit accepts its absence from the
-# default target source list.
+# Marked optional so LLVM's source audit does not require it in default builds.
list(APPEND LLVM_OPTIONAL_SOURCES
EmitCRegAllocEvictModel.cpp)
-# Only compile the EmitC regalloc wrapper when the EmitC-backed release model
-# is selected. Otherwise the generated .inc file may not exist.
+# The regalloc model wrapper is only built when the fully in-tree model integration is enabled.
if (LLVM_USE_EMITC_REGALLOC_EVICT_MODEL)
- list(APPEND LLVMCodeGenOptionalSources
+ list(APPEND CodeGenInTreeModelSources
EmitCRegAllocEvictModel.cpp)
endif()
@@ -275,7 +272,7 @@ add_llvm_component_library(LLVMCodeGen
WindowsSecureHotPatching.cpp
WinEHPrepare.cpp
XRayInstrumentation.cpp
- ${LLVMCodeGenOptionalSources}
+ ${CodeGenInTreeModelSources}
${GeneratedMLSources}
LiveDebugValues/LiveDebugValues.cpp
>From 9e879d2020ff4499a6445e49fde94e27e32bd567 Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Tue, 21 Jul 2026 19:59:39 +0200
Subject: [PATCH 4/7] Add comments, rename and structure for clarity
---
llvm/CMakeLists.txt | 45 ++++-
llvm/cmake/modules/TensorFlowCompile.cmake | 188 +++++++++++++++---
...inerSizeModel.h => MLIRInlinerSizeModel.h} | 16 +-
...cEvictModel.h => MLIRRegAllocEvictModel.h} | 20 +-
llvm/lib/Analysis/CMakeLists.txt | 59 +++---
...SizeModel.cpp => MLIRInlinerSizeModel.cpp} | 48 +++--
llvm/lib/Analysis/MLInlineAdvisor.cpp | 6 +-
llvm/lib/CodeGen/CMakeLists.txt | 45 +++--
...ctModel.cpp => MLIRRegAllocEvictModel.cpp} | 64 +++---
llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp | 8 +-
.../MLRegAlloc/default-eviction-advisor.ll | 2 +-
llvm/test/lit.cfg.py | 4 +-
llvm/test/lit.site.cfg.py.in | 2 +-
13 files changed, 352 insertions(+), 155 deletions(-)
rename llvm/include/llvm/Analysis/{EmitCInlinerSizeModel.h => MLIRInlinerSizeModel.h} (74%)
rename llvm/include/llvm/CodeGen/{EmitCRegAllocEvictModel.h => MLIRRegAllocEvictModel.h} (77%)
rename llvm/lib/Analysis/{EmitCInlinerSizeModel.cpp => MLIRInlinerSizeModel.cpp} (85%)
rename llvm/lib/CodeGen/{EmitCRegAllocEvictModel.cpp => MLIRRegAllocEvictModel.cpp} (72%)
diff --git a/llvm/CMakeLists.txt b/llvm/CMakeLists.txt
index 3f8d47a8997fa..20ac3ff0ff3a8 100644
--- a/llvm/CMakeLists.txt
+++ b/llvm/CMakeLists.txt
@@ -144,6 +144,13 @@ set(LLVM_ALL_PROJECTS "bolt;clang;clang-tools-extra;cross-project-tests;lld;lldb
set(LLVM_EXTRA_PROJECTS "flang" "libc" "compiler-rt")
# List of all known projects in the mono repo
set(LLVM_KNOWN_PROJECTS "${LLVM_ALL_PROJECTS};${LLVM_EXTRA_PROJECTS}")
+option(LLVM_USE_MLIR_FOR_MLGO
+ "Compile release-mode MLGO models with the in-tree MLIR-based flow"
+ OFF)
+set(LLVM_MLIR_OPT_PATH "" CACHE FILEPATH
+ "Path to the mlir-opt executable used by the MLIR-based MLGO model flow")
+set(LLVM_MLIR_TRANSLATE_PATH "" CACHE FILEPATH
+ "Path to the mlir-translate executable used by the MLIR-based MLGO model flow")
set(LLVM_ENABLE_PROJECTS "" CACHE STRING
"Semicolon-separated list of projects to build (${LLVM_KNOWN_PROJECTS}), or \"all\".")
# Make sure expansion happens first to not handle "all" in rest of the checks.
@@ -1248,8 +1255,24 @@ set(LLVM_ENABLE_PROFCHECK OFF CACHE BOOL "Enable profile checking in test tools"
#
set(TENSORFLOW_AOT_PATH "" CACHE PATH "Path to TensorFlow pip install dir")
+set(LLVM_HAVE_TF_AOT OFF CACHE BOOL "Tensorflow AOT available" FORCE)
+set(LLVM_BUILD_INLINERSIZEMODEL OFF)
+set(LLVM_BUILD_REGALLOCEVICTMODEL OFF)
+
+if (LLVM_USE_MLIR_FOR_MLGO AND NOT TENSORFLOW_AOT_PATH STREQUAL "")
+ message(FATAL_ERROR
+ "LLVM_USE_MLIR_FOR_MLGO and TENSORFLOW_AOT_PATH are mutually exclusive. "
+ "Select exactly one MLGO model compilation flow.")
+endif()
+
+# LLVM currently supports two release-mode MLGO compilation flows:
+# TensorFlow AOT, which emits a compiled serving interface directly, and the
+# in-tree MLIR flow, which lowers the model to generated C++ and serves it
+# through LLVM-owned wrappers. Both flows feed the same per-model selection
+# logic below so lib/Analysis and lib/CodeGen only need to check whether a
+# given model should be built.
if (NOT TENSORFLOW_AOT_PATH STREQUAL "")
- set(LLVM_HAVE_TF_AOT "ON" CACHE BOOL "Tensorflow AOT available")
+ set(LLVM_HAVE_TF_AOT ON CACHE BOOL "Tensorflow AOT available" FORCE)
set(TENSORFLOW_AOT_COMPILER
"${TENSORFLOW_AOT_PATH}/../../../../bin/saved_model_cli"
CACHE PATH "Path to the Tensorflow AOT compiler")
@@ -1261,24 +1284,36 @@ 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)
+endif()
+
+if (NOT TENSORFLOW_AOT_PATH STREQUAL "" OR LLVM_USE_MLIR_FOR_MLGO)
+ # Resolve user intent once for each release-mode model: disable it, use
+ # provided override artifacts, or autogenerate a test model to compile.
# Once we add more modules, we should handle this more automatically.
if (DEFINED LLVM_OVERRIDE_MODEL_HEADER_INLINERSIZEMODEL)
set(LLVM_INLINER_MODEL_PATH "none")
elseif(NOT DEFINED LLVM_INLINER_MODEL_PATH
- OR "${LLVM_INLINER_MODEL_PATH}" STREQUAL ""
- OR "${LLVM_INLINER_MODEL_PATH}" STREQUAL "autogenerate")
+ OR LLVM_INLINER_MODEL_PATH STREQUAL ""
+ OR LLVM_INLINER_MODEL_PATH STREQUAL "autogenerate")
set(LLVM_INLINER_MODEL_PATH "autogenerate")
set(LLVM_INLINER_MODEL_AUTOGENERATED 1)
endif()
+
if (DEFINED LLVM_OVERRIDE_MODEL_HEADER_REGALLOCEVICTMODEL)
set(LLVM_RAEVICT_MODEL_PATH "none")
elseif(NOT DEFINED LLVM_RAEVICT_MODEL_PATH
- OR "${LLVM_RAEVICT_MODEL_PATH}" STREQUAL ""
- OR "${LLVM_RAEVICT_MODEL_PATH}" STREQUAL "autogenerate")
+ OR LLVM_RAEVICT_MODEL_PATH STREQUAL ""
+ OR LLVM_RAEVICT_MODEL_PATH STREQUAL "autogenerate")
set(LLVM_RAEVICT_MODEL_PATH "autogenerate")
set(LLVM_RAEVICT_MODEL_AUTOGENERATED 1)
endif()
+ if (NOT LLVM_INLINER_MODEL_PATH STREQUAL "none")
+ set(LLVM_BUILD_INLINERSIZEMODEL ON)
+ endif()
+ if (NOT LLVM_RAEVICT_MODEL_PATH STREQUAL "none")
+ set(LLVM_BUILD_REGALLOCEVICTMODEL ON)
+ endif()
endif()
# Configure the LLVM configuration header files.
diff --git a/llvm/cmake/modules/TensorFlowCompile.cmake b/llvm/cmake/modules/TensorFlowCompile.cmake
index c4dae39f37e8c..19652a556c5ba 100644
--- a/llvm/cmake/modules/TensorFlowCompile.cmake
+++ b/llvm/cmake/modules/TensorFlowCompile.cmake
@@ -1,4 +1,4 @@
-function(tf_get_absolute_path path base final_path)
+function(mlgo_get_absolute_path path base final_path)
if (IS_ABSOLUTE ${path})
set(${final_path} ${path} PARENT_SCOPE)
else()
@@ -6,7 +6,7 @@ function(tf_get_absolute_path path base final_path)
endif()
endfunction()
-function(tf_get_model model final_path)
+function(mlgo_get_model model final_path)
string(FIND ${model} "http:" pos_http)
string(FIND ${model} "https:" pos_https)
if (${pos_http} EQUAL 0 OR ${pos_https} EQUAL 0)
@@ -21,15 +21,15 @@ function(tf_get_model model final_path)
DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/${fname}_model)
set(${final_path} ${CMAKE_CURRENT_BINARY_DIR}/${fname}_model/model PARENT_SCOPE)
else()
- tf_get_absolute_path(${model} ${CMAKE_CURRENT_BINARY_DIR} model_path)
+ mlgo_get_absolute_path(${model} ${CMAKE_CURRENT_BINARY_DIR} model_path)
set(${final_path} ${model_path} PARENT_SCOPE)
endif()
endfunction()
# Generate a mock model for tests.
function(generate_mock_model generator output)
- tf_get_absolute_path(${generator} ${CMAKE_CURRENT_SOURCE_DIR} generator_absolute_path)
- tf_get_absolute_path(${output} ${CMAKE_CURRENT_BINARY_DIR} output_absolute_path)
+ mlgo_get_absolute_path(${generator} ${CMAKE_CURRENT_SOURCE_DIR} generator_absolute_path)
+ mlgo_get_absolute_path(${output} ${CMAKE_CURRENT_BINARY_DIR} output_absolute_path)
message(WARNING "Autogenerated mock models should not be used in production builds.")
execute_process(COMMAND ${Python3_EXECUTABLE}
${generator_absolute_path}
@@ -38,6 +38,41 @@ function(generate_mock_model generator output)
)
endfunction()
+# Shared MLGO model build helpers. Both release-mode flows start from the same
+# model selection and discovery steps. The TensorFlow AOT path consumes a
+# SavedModel and asks TensorFlow to emit a compiled serving interface, while
+# the MLIR path converts the SavedModel through TFLite and TOSA, lowers it with
+# an MLIR pass pipeline, and emits C++ that LLVM wraps locally.
+
+function(mlgo_resolve_model model default_url default_path
+ test_model_generator model_label should_skip final_path)
+ if (${model} STREQUAL "none")
+ message(STATUS "Will skip enabling mlgo for ${model_label}")
+ set(${should_skip} TRUE PARENT_SCOPE)
+ return()
+ endif()
+
+ if (${model} STREQUAL "download")
+ # Crash if the user wants to download a model but a URL is not configured.
+ if (${default_url} STREQUAL "<UNSPECIFIED>")
+ message(FATAL_ERROR "Model path was set to 'download' but there is no"
+ " model url currently specified in cmake. You can generate a model"
+ " using, for example, the tools at http://github.com/google/ml-compiler-opt."
+ " Some reference models are also periodically released there.")
+ endif()
+ set(model ${default_url})
+ endif()
+
+ if (${model} STREQUAL "autogenerate")
+ set(model ${default_path}-autogenerated)
+ generate_mock_model(${test_model_generator} ${model})
+ endif()
+
+ mlgo_get_model(${model} model_input)
+ set(${should_skip} FALSE PARENT_SCOPE)
+ set(${final_path} ${model_input} PARENT_SCOPE)
+endfunction()
+
# Run the tensorflow compiler (saved_model_cli) on the saved model in the
# ${model} directory, looking for the ${tag_set} tag set, and the SignatureDef
# ${signature_def_key}.
@@ -45,7 +80,7 @@ endfunction()
# ${CMAKE_CURRENT_BINARY_DIR}. The generated header will define a C++ class
# called ${cpp_class} - which may be a namespace-qualified class name.
function(tf_compile model tag_set signature_def_key fname cpp_class hdr_file obj_file)
- tf_get_absolute_path(${model} ${CMAKE_CURRENT_BINARY_DIR} LLVM_ML_MODELS_ABSOLUTE)
+ mlgo_get_absolute_path(${model} ${CMAKE_CURRENT_BINARY_DIR} LLVM_ML_MODELS_ABSOLUTE)
message("Using model at " ${LLVM_ML_MODELS_ABSOLUTE})
add_custom_command(OUTPUT ${obj_file} ${hdr_file}
COMMAND ${TENSORFLOW_AOT_COMPILER} aot_compile_cpu
@@ -79,35 +114,21 @@ function(tf_find_and_compile model default_url default_path test_model_generator
set(override_object ${LLVM_OVERRIDE_MODEL_OBJECT_${fname_allcaps}})
# If the user specified overrides, that indicates intent to use AOT and we
# don't care what the model path is
- if (EXISTS "${override_header}" AND EXISTS "${override_object}")
- configure_file(${override_header} ${hdr_file} COPYONLY)
- configure_file(${override_object} ${obj_file} COPYONLY)
- message(STATUS "Using provided header " ${hdr_file} " and object " ${obj_file} "
- files for model " ${fname})
- set(GENERATED_OBJS ${GENERATED_OBJS} ${obj_file})
- set(GENERATED_HEADERS ${GENERATED_HEADERS} ${hdr_file})
- elseif("${model}" STREQUAL "none")
- message(STATUS "Will skip enabling mlgo for ${fname}")
- return()
- else()
- if ("${model}" STREQUAL "download")
- # Crash if the user wants to download a model but a URL is set to "TO_BE_UPDATED"
- if ("${default_url}" STREQUAL "<UNSPECIFIED>")
- message(FATAL_ERROR "Model path was set to 'download' but there is no"
- " model url currently specified in cmake. You can generate a model"
- " using, for example, the tools at http://github.com/google/ml-compiler-opt."
- " Some reference models are also periodically released there.")
- endif()
-
- set(model ${default_url})
+ if (override_header AND override_object)
+ if (EXISTS ${override_header} AND EXISTS ${override_object})
+ configure_file(${override_header} ${hdr_file} COPYONLY)
+ configure_file(${override_object} ${obj_file} COPYONLY)
+ message(STATUS "Using provided header " ${hdr_file} " and object " ${obj_file} "
+ files for model " ${fname})
+ set(GENERATED_OBJS ${GENERATED_OBJS} ${obj_file})
+ set(GENERATED_HEADERS ${GENERATED_HEADERS} ${hdr_file})
endif()
-
- if ("${model}" STREQUAL "autogenerate")
- set(model ${default_path}-autogenerated)
- generate_mock_model(${test_model_generator} ${model})
+ else()
+ mlgo_resolve_model(${model} ${default_url} ${default_path}
+ ${test_model_generator} ${fname} should_skip LLVM_ML_MODELS_ABSOLUTE)
+ if (should_skip)
+ return()
endif()
-
- tf_get_model(${model} LLVM_ML_MODELS_ABSOLUTE)
tf_compile(${LLVM_ML_MODELS_ABSOLUTE} ${tag_set} ${signature_def_key} ${fname} ${cpp_class} ${hdr_file} ${obj_file})
endif()
@@ -116,3 +137,104 @@ function(tf_find_and_compile model default_url default_path test_model_generator
set(MLLinkDeps ${MLLinkDeps} tf_xla_runtime PARENT_SCOPE)
add_compile_definitions(LLVM_HAVE_TF_AOT_${fname_allcaps})
endfunction()
+
+set(LLVM_MLGO_MLIR_PASS_PIPELINE
+ "builtin.module(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},buffer-deallocation-pipeline,func.func(convert-linalg-to-loops),expand-strided-metadata,canonicalize,memref-elide-reinterpret-cast,convert-to-emitc,math-expand-ops{ops=rsqrt},arith-expand,convert-math-to-emitc,convert-arith-to-emitc)"
+ CACHE STRING
+ "MLIR pass pipeline used to lower MLGO TOSA models to EmitC.")
+
+# Lower an MLGO model with the in-tree MLIR flow and emit the generated model
+# body as a header under llvm/include. The generated header is not a complete
+# serving interface by itself; LLVM-side wrappers in lib/Analysis and
+# lib/CodeGen provide the stable API expected by ReleaseModeModelRunner.
+function(mlir_find_and_compile model default_url default_path
+ test_model_generator header_relative_path)
+ mlgo_resolve_model(${model} ${default_url} ${default_path}
+ ${test_model_generator} ${header_relative_path} should_skip model_input)
+ if (should_skip)
+ return()
+ endif()
+
+ set(mlir_opt_path ${LLVM_MLIR_OPT_PATH})
+ if (NOT mlir_opt_path)
+ find_program(mlir_opt_path
+ NAMES mlir-opt
+ DOC "Path to the mlir-opt executable used by the MLIR-based MLGO model flow")
+ endif()
+ if (NOT mlir_opt_path)
+ message(FATAL_ERROR
+ "LLVM_USE_MLIR_FOR_MLGO requires 'mlir-opt'. Set LLVM_MLIR_OPT_PATH "
+ "or make 'mlir-opt' available on PATH.")
+ endif()
+ get_filename_component(mlir_opt_path ${mlir_opt_path} ABSOLUTE
+ BASE_DIR ${CMAKE_BINARY_DIR})
+
+ set(mlir_translate_path ${LLVM_MLIR_TRANSLATE_PATH})
+ if (NOT mlir_translate_path)
+ find_program(mlir_translate_path
+ NAMES mlir-translate
+ DOC "Path to the mlir-translate executable used by the MLIR-based MLGO model flow")
+ endif()
+ if (NOT mlir_translate_path)
+ message(FATAL_ERROR
+ "LLVM_USE_MLIR_FOR_MLGO requires 'mlir-translate'. Set "
+ "LLVM_MLIR_TRANSLATE_PATH or make 'mlir-translate' available on PATH.")
+ endif()
+ get_filename_component(mlir_translate_path ${mlir_translate_path} ABSOLUTE
+ BASE_DIR ${CMAKE_BINARY_DIR})
+
+ set(tosa_converter_path ${LLVM_TFLITE_TOSA_CONVERTER})
+ if (NOT tosa_converter_path)
+ find_program(tosa_converter_path
+ NAMES tosa-converter-for-tflite
+ DOC "Path to the tosa-converter-for-tflite executable")
+ endif()
+ if (NOT tosa_converter_path)
+ message(FATAL_ERROR
+ "LLVM_USE_MLIR_FOR_MLGO requires 'tosa-converter-for-tflite' to be "
+ "available on PATH.")
+ endif()
+ get_filename_component(tosa_converter_path ${tosa_converter_path} ABSOLUTE
+ BASE_DIR ${CMAKE_BINARY_DIR})
+
+ set(generated_header ${LLVM_INCLUDE_DIR}/${header_relative_path})
+ get_filename_component(generated_header_dir ${generated_header} DIRECTORY)
+ get_filename_component(generated_header_stem ${generated_header} NAME_WE)
+ set(tflite_dir
+ ${CMAKE_CURRENT_BINARY_DIR}/${generated_header_stem}-tflite)
+ set(tflite_model ${tflite_dir}/model.tflite)
+ set(tosa_mlir
+ ${CMAKE_CURRENT_BINARY_DIR}/${generated_header_stem}-tosa.mlir)
+ set(lowered_mlir
+ ${CMAKE_CURRENT_BINARY_DIR}/${generated_header_stem}-emitc.mlir)
+ add_custom_command(
+ OUTPUT ${generated_header}
+ BYPRODUCTS ${tflite_model} ${tosa_mlir} ${lowered_mlir}
+ COMMAND ${CMAKE_COMMAND} -E make_directory ${generated_header_dir}
+ COMMAND ${Python3_EXECUTABLE}
+ ${LLVM_MAIN_SRC_DIR}/lib/Analysis/models/saved-model-to-tflite.py
+ ${model_input}
+ ${tflite_dir}
+ COMMAND ${tosa_converter_path}
+ ${tflite_model}
+ --text
+ -o ${tosa_mlir}
+ COMMAND ${mlir_opt_path}
+ --pass-pipeline=${LLVM_MLGO_MLIR_PASS_PIPELINE}
+ ${tosa_mlir}
+ -o ${lowered_mlir}
+ COMMAND ${mlir_translate_path}
+ -mlir-to-cpp
+ ${lowered_mlir}
+ -o ${generated_header}
+ DEPENDS
+ ${model_input}/saved_model.pb
+ ${LLVM_MAIN_SRC_DIR}/lib/Analysis/models/saved-model-to-tflite.py
+ ${mlir_opt_path}
+ ${mlir_translate_path}
+ VERBATIM)
+
+ set_source_files_properties(${generated_header} PROPERTIES GENERATED 1)
+ set(GeneratedMLSources ${GeneratedMLSources} ${generated_header}
+ PARENT_SCOPE)
+endfunction()
diff --git a/llvm/include/llvm/Analysis/EmitCInlinerSizeModel.h b/llvm/include/llvm/Analysis/MLIRInlinerSizeModel.h
similarity index 74%
rename from llvm/include/llvm/Analysis/EmitCInlinerSizeModel.h
rename to llvm/include/llvm/Analysis/MLIRInlinerSizeModel.h
index 63e298e4ce9fa..5e16ddbff6432 100644
--- a/llvm/include/llvm/Analysis/EmitCInlinerSizeModel.h
+++ b/llvm/include/llvm/Analysis/MLIRInlinerSizeModel.h
@@ -1,4 +1,4 @@
-//===- EmitCInlinerSizeModel.h - EmitC inliner model wrapper ----*- C++ -*-===//
+//===- MLIRInlinerSizeModel.h - MLIR inliner model wrapper ----*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
@@ -6,12 +6,18 @@
//
//===----------------------------------------------------------------------===//
//
-/// Declares the wrapper around the EmitC-translated MLGO inliner model.
+/// Wraps the MLIR-translated MLGO inliner model in the interface expected by
+/// ReleaseModeModelRunner.
+///
+/// The generated `.inc` file only contains lowered model code. This wrapper
+/// owns the named inliner tensors, keeps their layout stable across generated
+/// model variants, and exposes the same serving API that the rest of LLVM
+/// already uses for release-mode MLGO models.
//
//===----------------------------------------------------------------------===//
-#ifndef LLVM_LIB_ANALYSIS_EMITCINLINERSIZEMODEL_H
-#define LLVM_LIB_ANALYSIS_EMITCINLINERSIZEMODEL_H
+#ifndef LLVM_LIB_ANALYSIS_MLIRINLINERSIZEMODEL_H
+#define LLVM_LIB_ANALYSIS_MLIRINLINERSIZEMODEL_H
#include <array>
#include <cstdint>
@@ -19,7 +25,7 @@
namespace llvm {
-class EmitCInlinerSizeModel final {
+class MLIRInlinerSizeModel final {
public:
int LookupArgIndex(const std::string &Name);
int LookupResultIndex(const std::string &Name);
diff --git a/llvm/include/llvm/CodeGen/EmitCRegAllocEvictModel.h b/llvm/include/llvm/CodeGen/MLIRRegAllocEvictModel.h
similarity index 77%
rename from llvm/include/llvm/CodeGen/EmitCRegAllocEvictModel.h
rename to llvm/include/llvm/CodeGen/MLIRRegAllocEvictModel.h
index 41a19ec62522e..d507d76599270 100644
--- a/llvm/include/llvm/CodeGen/EmitCRegAllocEvictModel.h
+++ b/llvm/include/llvm/CodeGen/MLIRRegAllocEvictModel.h
@@ -1,4 +1,4 @@
-//===- EmitCRegAllocEvictModel.h - EmitC regalloc model wrapper -*- C++ -*-===//
+//===- MLIRRegAllocEvictModel.h - MLIR regalloc model wrapper -*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
@@ -6,14 +6,18 @@
//
//===----------------------------------------------------------------------===//
//
-/// \file
-/// Declares the wrapper around the EmitC-translated MLGO regalloc eviction
-/// model.
+/// Wraps the MLIR-translated MLGO regalloc eviction model in the interface
+/// expected by ReleaseModeModelRunner.
+///
+/// The generated `.inc` file only contains lowered model code. This wrapper
+/// owns the named regalloc tensors, keeps their layout stable across generated
+/// model variants, and exposes the same serving API that the rest of LLVM
+/// already uses for release-mode MLGO models.
//
//===----------------------------------------------------------------------===//
-#ifndef LLVM_CODEGEN_EMITCREGALLOCEVICTMODEL_H
-#define LLVM_CODEGEN_EMITCREGALLOCEVICTMODEL_H
+#ifndef LLVM_CODEGEN_MLIRREGALLOCEVICTMODEL_H
+#define LLVM_CODEGEN_MLIRREGALLOCEVICTMODEL_H
#include <cmath>
#include <cstddef>
@@ -22,7 +26,7 @@
namespace llvm {
-class EmitCRegAllocEvictModel final {
+class MLIRRegAllocEvictModel final {
public:
int LookupArgIndex(const std::string &Name);
int LookupResultIndex(const std::string &Name);
@@ -87,6 +91,8 @@ class EmitCRegAllocEvictModel final {
I64InterferenceTensor MinStageInput{};
F32Scalar ProgressInput{};
+ // These scalars remain part of the serving API even when a translated model
+ // does not make semantic use of all of them.
I32Scalar DummyStepType{};
F32Scalar DummyDiscount{};
F32Scalar DummyReward{};
diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt
index 323d9cd2e6df2..aa84630c1e562 100644
--- a/llvm/lib/Analysis/CMakeLists.txt
+++ b/llvm/lib/Analysis/CMakeLists.txt
@@ -1,36 +1,43 @@
-option(LLVM_USE_EMITC_INLINER_MODEL
- "Use the fully in-tree inliner model integration instead of TensorFlow AOT"
- OFF)
-
# Marked optional so LLVM's source audit does not require it in default builds.
list(APPEND LLVM_OPTIONAL_SOURCES
- EmitCInlinerSizeModel.cpp)
+ MLIRInlinerSizeModel.cpp)
if (LLVM_HAVE_TFLITE)
list(APPEND MLLinkDeps
tensorflow-lite::tensorflow-lite)
endif()
-# The inlining model wrapper is only built when the fully in-tree model integration is enabled.
-if (LLVM_USE_EMITC_INLINER_MODEL)
- list(APPEND AnalysisInTreeModelSources
- EmitCInlinerSizeModel.cpp)
-elseif (DEFINED LLVM_HAVE_TF_AOT)
- include(TensorFlowCompile)
- set(LLVM_INLINER_MODEL_PATH_DEFAULT "models/inliner-Oz")
- set(LLVM_INLINER_MODEL_CURRENT_URL "<UNSPECIFIED>" CACHE STRING
- "URL to download the LLVM inliner model")
+set(LLVM_INLINER_MODEL_PATH_DEFAULT "models/inliner-Oz")
+set(LLVM_INLINER_MODEL_CURRENT_URL "<UNSPECIFIED>" CACHE STRING
+ "URL to download the LLVM inliner model")
- tf_find_and_compile(
- ${LLVM_INLINER_MODEL_PATH}
- ${LLVM_INLINER_MODEL_CURRENT_URL}
- ${LLVM_INLINER_MODEL_PATH_DEFAULT}
- "models/gen-inline-oz-test-model.py"
- serve
- action
- InlinerSizeModel
- llvm::InlinerSizeModel
- )
+# Release-mode inlining can be backed either by TensorFlow AOT artifacts or by
+# the in-tree MLIR flow. In both cases LLVM builds a thin wrapper so
+# MLInlineAdvisor continues to talk to one serving interface.
+if (LLVM_BUILD_INLINERSIZEMODEL)
+ include(TensorFlowCompile)
+ if (LLVM_USE_MLIR_FOR_MLGO)
+ list(APPEND AnalysisInTreeModelSources
+ MLIRInlinerSizeModel.cpp)
+ mlir_find_and_compile(
+ ${LLVM_INLINER_MODEL_PATH}
+ ${LLVM_INLINER_MODEL_CURRENT_URL}
+ ${LLVM_INLINER_MODEL_PATH_DEFAULT}
+ "models/gen-inline-oz-test-model.py"
+ "llvm/Analysis/MLIRInlinerSizeModel.inc"
+ )
+ else()
+ tf_find_and_compile(
+ ${LLVM_INLINER_MODEL_PATH}
+ ${LLVM_INLINER_MODEL_CURRENT_URL}
+ ${LLVM_INLINER_MODEL_PATH_DEFAULT}
+ "models/gen-inline-oz-test-model.py"
+ serve
+ action
+ InlinerSizeModel
+ llvm::InlinerSizeModel
+ )
+ endif()
endif()
# The implementation of ConstantFolding.cpp relies on the use of math functions
@@ -198,9 +205,9 @@ add_llvm_component_library(LLVMAnalysis
TargetParser
)
-if (LLVM_USE_EMITC_INLINER_MODEL)
+if (LLVM_BUILD_INLINERSIZEMODEL AND LLVM_USE_MLIR_FOR_MLGO)
target_compile_definitions(LLVMAnalysis
- PRIVATE LLVM_HAVE_EMITC_INLINERSIZEMODEL)
+ PRIVATE LLVM_HAVE_MLIR_INLINERSIZEMODEL)
endif()
include(CheckCXXSymbolExists)
diff --git a/llvm/lib/Analysis/EmitCInlinerSizeModel.cpp b/llvm/lib/Analysis/MLIRInlinerSizeModel.cpp
similarity index 85%
rename from llvm/lib/Analysis/EmitCInlinerSizeModel.cpp
rename to llvm/lib/Analysis/MLIRInlinerSizeModel.cpp
index a8e0094737bd6..1baff8c1b6d23 100644
--- a/llvm/lib/Analysis/EmitCInlinerSizeModel.cpp
+++ b/llvm/lib/Analysis/MLIRInlinerSizeModel.cpp
@@ -1,4 +1,4 @@
-//===- EmitCInlinerSizeModel.cpp - EmitC inliner model wrapper ------------===//
+//===- MLIRInlinerSizeModel.cpp - MLIR inliner model wrapper ------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
@@ -6,12 +6,12 @@
//
//===----------------------------------------------------------------------===//
//
-/// This file implements the wrapper around the EmitC-translated MLGO inliner
+/// This file implements the wrapper around the MLIR-translated MLGO inliner
/// model.
//
//===----------------------------------------------------------------------===//
-#include "llvm/Analysis/EmitCInlinerSizeModel.h"
+#include "llvm/Analysis/MLIRInlinerSizeModel.h"
#include "llvm/ADT/StringSwitch.h"
#include "llvm/Support/ErrorHandling.h"
@@ -28,11 +28,15 @@
#pragma clang diagnostic ignored "-Wmissing-braces"
#endif
-namespace llvm::emitc_inliner_model {
+namespace llvm::mlir_inliner_model {
+// Mock models generated for tests use `main` as entry point, while
+// pretrained models lowered use `action`. Rename
+// locally so the wrapper can always dispatch through one symbol without
+// rewriting the generated `.inc` file.
#define main action
-#include "llvm/Analysis/EmitCInlinerSizeModel.inc"
+#include "llvm/Analysis/MLIRInlinerSizeModel.inc"
#undef main
-} // namespace llvm::emitc_inliner_model
+} // namespace llvm::mlir_inliner_model
#if defined(__clang__)
#pragma clang diagnostic pop
@@ -103,10 +107,14 @@ struct InlinerRunInputs {
I64Ptr dummyInliningDefault;
};
+// The MLIR flow currently supports both the production inlining models API
+// and the simplified mock-model API used by tests. Keep that API
+// variance local to the wrapper so the rest of LLVM only sees one
+// ReleaseModeModelRunner-compatible object.
template <typename ActionTy>
-int64_t runEmitCInlinerAction(const InlinerRunInputs &I) {
+int64_t runMLIRInlinerAction(const InlinerRunInputs &I) {
if constexpr (std::is_same_v<ActionTy, InlinerProductionActionTy>) {
- return static_cast<ActionTy>(emitc_inliner_model::action)(
+ return static_cast<ActionTy>(mlir_inliner_model::action)(
I.callsiteCost, I.isMultipleBlocks, I.callerConditionallyExecutedBlocks,
I.dummyInliningDefault, I.coldCCPenalty,
I.calleeConditionallyExecutedBlocks, I.calleeUsers,
@@ -121,7 +129,7 @@ int64_t runEmitCInlinerAction(const InlinerRunInputs &I) {
I.constantOffsetPtrArgs, I.switchPenalty, I.dummyDiscount,
I.callerUsers, I.dummyReward);
} else if constexpr (std::is_same_v<ActionTy, InlinerMockActionTy>) {
- return static_cast<ActionTy>(emitc_inliner_model::action)(
+ return static_cast<ActionTy>(mlir_inliner_model::action)(
I.callerBasicBlockCount, I.callerConditionallyExecutedBlocks,
I.callerUsers, I.calleeBasicBlockCount,
I.calleeConditionallyExecutedBlocks, I.calleeUsers, I.nrCtantParams,
@@ -137,12 +145,12 @@ int64_t runEmitCInlinerAction(const InlinerRunInputs &I) {
I.dummyDiscount, I.callerUsers, I.dummyReward);
} else {
static_assert(AlwaysFalse<ActionTy>,
- "Unsupported EmitC inliner model signature");
+ "Unsupported MLIR inliner model signature");
}
}
} // namespace
-int EmitCInlinerSizeModel::LookupArgIndex(const std::string &Name) {
+int MLIRInlinerSizeModel::LookupArgIndex(const std::string &Name) {
return StringSwitch<int>(Name)
.Case("feed_dead_blocks", DeadBlocks)
.Case("feed_case_cluster_penalty", CaseClusterPenalty)
@@ -185,24 +193,26 @@ int EmitCInlinerSizeModel::LookupArgIndex(const std::string &Name) {
.Default(-1);
}
-int EmitCInlinerSizeModel::LookupResultIndex(const std::string &Name) {
+int MLIRInlinerSizeModel::LookupResultIndex(const std::string &Name) {
return Name == "fetch_inlining_decision" ? 0 : -1;
}
-void *EmitCInlinerSizeModel::arg_data(int Index) {
+void *MLIRInlinerSizeModel::arg_data(int Index) {
if (Index < 0 || Index >= NumArgs)
- llvm_unreachable("invalid EmitC inliner input index");
+ llvm_unreachable("invalid MLIR inliner input index");
return Inputs[Index].data();
}
-void *EmitCInlinerSizeModel::result_data(int Index) {
+void *MLIRInlinerSizeModel::result_data(int Index) {
if (Index != 0)
- llvm_unreachable("invalid EmitC inliner result index");
+ llvm_unreachable("invalid MLIR inliner result index");
return Result.data();
}
-void EmitCInlinerSizeModel::Run() {
- using ActionTy = decltype(&emitc_inliner_model::action);
+void MLIRInlinerSizeModel::Run() {
+ using ActionTy = decltype(&mlir_inliner_model::action);
+ // Gather the stored named tensors once, then dispatch through the adapter
+ // selected from the generated `action` signature.
InlinerRunInputs I{};
I.deadBlocks = Inputs[DeadBlocks].data();
I.caseClusterPenalty = Inputs[CaseClusterPenalty].data();
@@ -246,5 +256,5 @@ void EmitCInlinerSizeModel::Run() {
I.dummyDiscount = DummyDiscount.data();
I.dummyReward = DummyReward.data();
I.dummyInliningDefault = DummyInliningDefault.data();
- Result[0] = runEmitCInlinerAction<ActionTy>(I);
+ Result[0] = runMLIRInlinerAction<ActionTy>(I);
}
diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp
index e698f066b9336..cdc28a3a7a72b 100644
--- a/llvm/lib/Analysis/MLInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/MLInlineAdvisor.cpp
@@ -64,9 +64,9 @@ static cl::opt<std::string> ModelSelector("ml-inliner-model-selector",
static cl::opt<bool> StopImmediatelyForTest("ml-inliner-stop-immediately",
cl::Hidden);
-#if defined(LLVM_HAVE_EMITC_INLINERSIZEMODEL)
-#include "llvm/Analysis/EmitCInlinerSizeModel.h"
-using CompiledModelType = llvm::EmitCInlinerSizeModel;
+#if defined(LLVM_HAVE_MLIR_INLINERSIZEMODEL)
+#include "llvm/Analysis/MLIRInlinerSizeModel.h"
+using CompiledModelType = llvm::MLIRInlinerSizeModel;
#elif defined(LLVM_HAVE_TF_AOT_INLINERSIZEMODEL)
// codegen-ed file
#include "InlinerSizeModel.h" // NOLINT
diff --git a/llvm/lib/CodeGen/CMakeLists.txt b/llvm/lib/CodeGen/CMakeLists.txt
index a32300db1093b..e90f50b36c8e2 100644
--- a/llvm/lib/CodeGen/CMakeLists.txt
+++ b/llvm/lib/CodeGen/CMakeLists.txt
@@ -1,24 +1,27 @@
-option(LLVM_USE_EMITC_REGALLOC_EVICT_MODEL
- "Use the fully in-tree regalloc eviction model integration instead of TensorFlow AOT"
- OFF)
-
# Marked optional so LLVM's source audit does not require it in default builds.
list(APPEND LLVM_OPTIONAL_SOURCES
- EmitCRegAllocEvictModel.cpp)
+ MLIRRegAllocEvictModel.cpp)
-# The regalloc model wrapper is only built when the fully in-tree model integration is enabled.
-if (LLVM_USE_EMITC_REGALLOC_EVICT_MODEL)
- list(APPEND CodeGenInTreeModelSources
- EmitCRegAllocEvictModel.cpp)
-endif()
+set(LLVM_RAEVICT_MODEL_PATH_DEFAULT "models/regalloc-eviction")
+set(LLVM_RAEVICT_MODEL_CURRENT_URL "<UNSPECIFIED>" CACHE STRING
+ "URL to download the LLVM register allocator eviction model")
-if (DEFINED LLVM_HAVE_TF_AOT OR LLVM_HAVE_TFLITE)
+# Release-mode regalloc eviction can be backed either by TensorFlow AOT
+# artifacts or by the in-tree MLIR flow. In both cases LLVM builds a thin
+# wrapper so MLRegAllocEvictAdvisor continues to talk to one serving interface.
+if (LLVM_BUILD_REGALLOCEVICTMODEL)
include(TensorFlowCompile)
- set(LLVM_RAEVICT_MODEL_PATH_DEFAULT "models/regalloc-eviction")
-
- set(LLVM_RAEVICT_MODEL_CURRENT_URL "<UNSPECIFIED>" CACHE STRING "URL to download the LLVM register allocator eviction model")
-
- if (DEFINED LLVM_HAVE_TF_AOT AND NOT LLVM_USE_EMITC_REGALLOC_EVICT_MODEL)
+ if (LLVM_USE_MLIR_FOR_MLGO)
+ list(APPEND CodeGenInTreeModelSources
+ MLIRRegAllocEvictModel.cpp)
+ mlir_find_and_compile(
+ ${LLVM_RAEVICT_MODEL_PATH}
+ ${LLVM_RAEVICT_MODEL_CURRENT_URL}
+ ${LLVM_RAEVICT_MODEL_PATH_DEFAULT}
+ "../Analysis/models/gen-regalloc-eviction-test-model.py"
+ "llvm/CodeGen/MLIRRegAllocEvictModel.inc"
+ )
+ else()
tf_find_and_compile(
${LLVM_RAEVICT_MODEL_PATH}
${LLVM_RAEVICT_MODEL_CURRENT_URL}
@@ -30,10 +33,10 @@ if (DEFINED LLVM_HAVE_TF_AOT OR LLVM_HAVE_TFLITE)
llvm::RegAllocEvictModel
)
endif()
+endif()
- if (LLVM_HAVE_TFLITE)
- list(APPEND MLLinkDeps ${tensorflow_c_api} ${tensorflow_fx})
- endif()
+if (LLVM_HAVE_TFLITE)
+ list(APPEND MLLinkDeps ${tensorflow_c_api} ${tensorflow_fx})
endif()
add_llvm_component_library(LLVMCodeGen
@@ -310,9 +313,9 @@ add_llvm_component_library(LLVMCodeGen
TransformUtils
)
-if (LLVM_USE_EMITC_REGALLOC_EVICT_MODEL)
+if (LLVM_BUILD_REGALLOCEVICTMODEL AND LLVM_USE_MLIR_FOR_MLGO)
target_compile_definitions(LLVMCodeGen
- PRIVATE LLVM_HAVE_EMITC_REGALLOCEVICTMODEL)
+ PRIVATE LLVM_HAVE_MLIR_REGALLOCEVICTMODEL)
endif()
add_subdirectory(SelectionDAG)
diff --git a/llvm/lib/CodeGen/EmitCRegAllocEvictModel.cpp b/llvm/lib/CodeGen/MLIRRegAllocEvictModel.cpp
similarity index 72%
rename from llvm/lib/CodeGen/EmitCRegAllocEvictModel.cpp
rename to llvm/lib/CodeGen/MLIRRegAllocEvictModel.cpp
index 120268ec1f389..14ecae10b29f1 100644
--- a/llvm/lib/CodeGen/EmitCRegAllocEvictModel.cpp
+++ b/llvm/lib/CodeGen/MLIRRegAllocEvictModel.cpp
@@ -1,4 +1,4 @@
-//===- EmitCRegAllocEvictModel.cpp - EmitC regalloc model wrapper ---------===//
+//===- MLIRRegAllocEvictModel.cpp - MLIR regalloc model wrapper ---------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
@@ -6,12 +6,12 @@
//
//===----------------------------------------------------------------------===//
//
-/// This file implements the wrapper around the EmitC-translated MLGO
+/// This file implements the wrapper around the MLIR-translated MLGO
/// regalloc eviction model.
//
//===----------------------------------------------------------------------===//
-#include "llvm/CodeGen/EmitCRegAllocEvictModel.h"
+#include "llvm/CodeGen/MLIRRegAllocEvictModel.h"
#include "llvm/ADT/StringSwitch.h"
#include "llvm/Support/ErrorHandling.h"
@@ -20,22 +20,25 @@
#include <stdint.h>
#include <type_traits>
-namespace llvm::emitc_regalloc_evict_model {
+namespace llvm::mlir_regalloc_evict_model {
+// Mock models generated for tests use `main` as entry point, while
+// pretrained models lowered use `action`. Rename
+// locally so the wrapper can always dispatch through one symbol without
+// rewriting the generated `.inc` file.
#define main action
-#include "llvm/CodeGen/EmitCRegAllocEvictModel.inc"
+#include "llvm/CodeGen/MLIRRegAllocEvictModel.inc"
#undef main
-} // namespace llvm::emitc_regalloc_evict_model
+} // namespace llvm::mlir_regalloc_evict_model
using namespace llvm;
namespace {
template <typename T> inline constexpr bool AlwaysFalse = false;
-constexpr std::size_t EmitCRegAllocInterferenceCount = 33;
-using F32TensorPtr = float (*)[EmitCRegAllocInterferenceCount];
-using I64TensorPtr = int64_t (*)[EmitCRegAllocInterferenceCount];
+constexpr std::size_t MLIRRegAllocInterferenceCount = 33;
+using F32TensorPtr = float (*)[MLIRRegAllocInterferenceCount];
+using I64TensorPtr = int64_t (*)[MLIRRegAllocInterferenceCount];
using F32ScalarPtr = float *;
using I32ScalarPtr = int32_t *;
-using I64ScalarPtr = int64_t *;
using RegAllocProductionActionTy = int64_t (*)(
F32TensorPtr, F32TensorPtr, I64TensorPtr, F32TensorPtr, F32TensorPtr,
@@ -43,7 +46,8 @@ using RegAllocProductionActionTy = int64_t (*)(
I64TensorPtr, I64TensorPtr, F32TensorPtr, F32TensorPtr, I32ScalarPtr,
F32TensorPtr, I64TensorPtr, F32TensorPtr, I64TensorPtr, F32TensorPtr,
F32ScalarPtr, F32TensorPtr, I64TensorPtr, F32ScalarPtr);
-using RegAllocMaskOnlyActionTy = int64_t (*)(I64ScalarPtr);
+using RegAllocMaskOnlyActionTy = int64_t (*)(I64TensorPtr);
+using RegAllocFlatMaskActionTy = int64_t (*)(int64_t *);
struct RegAllocRunInputs {
F32TensorPtr liverangeSize;
@@ -70,30 +74,33 @@ struct RegAllocRunInputs {
F32TensorPtr weighedIndvarsByMax;
I64TensorPtr minStage;
F32ScalarPtr dummyReward;
- I64ScalarPtr maskFlat;
};
+// The MLIR flow currently supports both the production regalloc models API
+// and the simplified mock-model API used by tests. Keep that API
+// variance local to the wrapper so the rest of LLVM only sees one
+// ReleaseModeModelRunner-compatible object.
template <typename ActionTy>
-int64_t runEmitCRegAllocAction(const RegAllocRunInputs &I) {
+int64_t runMLIRRegAllocAction(const RegAllocRunInputs &I) {
if constexpr (std::is_same_v<ActionTy, RegAllocProductionActionTy>) {
- return static_cast<ActionTy>(emitc_regalloc_evict_model::action)(
+ return static_cast<ActionTy>(mlir_regalloc_evict_model::action)(
I.liverangeSize, I.hintWeightsByMax, I.isFree, I.weighedReadsByMax,
I.weighedReadWritesByMax, I.nrBrokenHints, I.progress,
I.hottestBBFreqByMax, I.useDefDensity, I.startBBFreqByMax, I.maxStage,
I.isHint, I.nrRematerializable, I.weighedWritesByMax, I.dummyStepType,
I.nrUrgent, I.mask, I.nrDefsAndUses, I.isLocal, I.endBBFreqByMax,
I.dummyDiscount, I.weighedIndvarsByMax, I.minStage, I.dummyReward);
- } else if constexpr (std::is_same_v<ActionTy, RegAllocMaskOnlyActionTy>) {
- return static_cast<ActionTy>(emitc_regalloc_evict_model::action)(
- I.maskFlat);
+ } else if constexpr (std::is_same_v<ActionTy, RegAllocFlatMaskActionTy>) {
+ return static_cast<ActionTy>(mlir_regalloc_evict_model::action)(
+ &(*I.mask)[0]);
} else {
static_assert(AlwaysFalse<ActionTy>,
- "Unsupported EmitC regalloc eviction model signature");
+ "Unsupported MLIR regalloc eviction model signature");
}
}
} // namespace
-int EmitCRegAllocEvictModel::LookupArgIndex(const std::string &Name) {
+int MLIRRegAllocEvictModel::LookupArgIndex(const std::string &Name) {
return StringSwitch<int>(Name)
.Case("feed_mask", Mask)
.Case("feed_is_free", IsFree)
@@ -119,11 +126,11 @@ int EmitCRegAllocEvictModel::LookupArgIndex(const std::string &Name) {
.Default(-1);
}
-int EmitCRegAllocEvictModel::LookupResultIndex(const std::string &Name) {
+int MLIRRegAllocEvictModel::LookupResultIndex(const std::string &Name) {
return Name == "fetch_index_to_evict" ? 0 : -1;
}
-void *EmitCRegAllocEvictModel::arg_data(int Index) {
+void *MLIRRegAllocEvictModel::arg_data(int Index) {
switch (Index) {
case Mask:
return MaskInput;
@@ -168,17 +175,19 @@ void *EmitCRegAllocEvictModel::arg_data(int Index) {
case Progress:
return ProgressInput;
}
- llvm_unreachable("invalid EmitC regalloc eviction input index");
+ llvm_unreachable("invalid MLIR regalloc eviction input index");
}
-void *EmitCRegAllocEvictModel::result_data(int Index) {
+void *MLIRRegAllocEvictModel::result_data(int Index) {
if (Index != 0)
- llvm_unreachable("invalid EmitC regalloc eviction result index");
+ llvm_unreachable("invalid MLIR regalloc eviction result index");
return Result;
}
-void EmitCRegAllocEvictModel::Run() {
- using ActionTy = decltype(&emitc_regalloc_evict_model::action);
+void MLIRRegAllocEvictModel::Run() {
+ using ActionTy = decltype(&mlir_regalloc_evict_model::action);
+ // Gather the stored named tensors once, then dispatch through the adapter
+ // selected from the generated `action` signature.
RegAllocRunInputs I{};
I.liverangeSize = LiverangeSizeInput;
I.hintWeightsByMax = HintWeightsByMaxInput;
@@ -204,6 +213,5 @@ void EmitCRegAllocEvictModel::Run() {
I.weighedIndvarsByMax = WeighedIndvarsByMaxInput;
I.minStage = MinStageInput;
I.dummyReward = DummyReward;
- I.maskFlat = MaskInput[0];
- Result[0] = runEmitCRegAllocAction<ActionTy>(I);
+ Result[0] = runMLIRRegAllocAction<ActionTy>(I);
}
diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index 03ba2da4e945e..7bb051498c8a2 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -47,10 +47,10 @@ using namespace llvm;
#define DEBUG_TYPE "ml-regalloc"
-// Generated header in release (AOT / EmitC) mode
-#if defined(LLVM_HAVE_EMITC_REGALLOCEVICTMODEL)
-#include "llvm/CodeGen/EmitCRegAllocEvictModel.h"
-using CompiledModelType = llvm::EmitCRegAllocEvictModel;
+// Generated header in release (AOT / MLIR) mode
+#if defined(LLVM_HAVE_MLIR_REGALLOCEVICTMODEL)
+#include "llvm/CodeGen/MLIRRegAllocEvictModel.h"
+using CompiledModelType = llvm::MLIRRegAllocEvictModel;
#elif defined(LLVM_HAVE_TF_AOT_REGALLOCEVICTMODEL)
#include "RegAllocEvictModel.h"
using CompiledModelType = RegAllocEvictModel;
diff --git a/llvm/test/CodeGen/MLRegAlloc/default-eviction-advisor.ll b/llvm/test/CodeGen/MLRegAlloc/default-eviction-advisor.ll
index fc9d244a617a9..6675411f04572 100644
--- a/llvm/test/CodeGen/MLRegAlloc/default-eviction-advisor.ll
+++ b/llvm/test/CodeGen/MLRegAlloc/default-eviction-advisor.ll
@@ -2,7 +2,7 @@
; trying to use ML-driven advisor.
; REQUIRES: !have_tf_aot
; REQUIRES: !have_tflite
-; REQUIRES: !have_emitc_raevict_model
+; REQUIRES: !have_mlir_mlgo
; 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
diff --git a/llvm/test/lit.cfg.py b/llvm/test/lit.cfg.py
index 796d258515e5b..096d1a4d5b69d 100644
--- a/llvm/test/lit.cfg.py
+++ b/llvm/test/lit.cfg.py
@@ -597,8 +597,8 @@ def enable_ptxas(ptxas_executable):
if config.have_tflite:
config.available_features.add("have_tflite")
-if config.have_emitc_raevict_model:
- config.available_features.add("have_emitc_raevict_model")
+if config.have_mlir_mlgo:
+ config.available_features.add("have_mlir_mlgo")
if config.llvm_inliner_model_autogenerated:
config.available_features.add("llvm_inliner_model_autogenerated")
diff --git a/llvm/test/lit.site.cfg.py.in b/llvm/test/lit.site.cfg.py.in
index 8c60b0d1a5d1b..0a7579a1a0c66 100644
--- a/llvm/test/lit.site.cfg.py.in
+++ b/llvm/test/lit.site.cfg.py.in
@@ -57,7 +57,7 @@ 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_tflite = @LLVM_HAVE_TFLITE@
-config.have_emitc_raevict_model = @LLVM_USE_EMITC_REGALLOC_EVICT_MODEL@
+config.have_mlir_mlgo = @LLVM_USE_MLIR_FOR_MLGO@
config.enable_profcheck = @LLVM_ENABLE_PROFCHECK@
config.llvm_inliner_model_autogenerated = @LLVM_INLINER_MODEL_AUTOGENERATED@
config.llvm_raevict_model_autogenerated = @LLVM_RAEVICT_MODEL_AUTOGENERATED@
>From 4b15bb0dfd0e96f2757c296c7202a11327e1fa3e Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Tue, 21 Jul 2026 21:55:28 +0200
Subject: [PATCH 5/7] Bind MLIR-generated model inputs by name
---
llvm/cmake/modules/TensorFlowCompile.cmake | 2 +-
.../llvm/Analysis/MLIRInlinerSizeModel.h | 60 +---
.../llvm/CodeGen/MLIRRegAllocEvictModel.h | 80 +----
llvm/lib/Analysis/MLIRInlinerSizeModel.cpp | 288 ++++++------------
llvm/lib/CodeGen/MLIRRegAllocEvictModel.cpp | 251 +++++----------
5 files changed, 185 insertions(+), 496 deletions(-)
diff --git a/llvm/cmake/modules/TensorFlowCompile.cmake b/llvm/cmake/modules/TensorFlowCompile.cmake
index 19652a556c5ba..727bcd7c2d7c6 100644
--- a/llvm/cmake/modules/TensorFlowCompile.cmake
+++ b/llvm/cmake/modules/TensorFlowCompile.cmake
@@ -139,7 +139,7 @@ function(tf_find_and_compile model default_url default_path test_model_generator
endfunction()
set(LLVM_MLGO_MLIR_PASS_PIPELINE
- "builtin.module(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},buffer-deallocation-pipeline,func.func(convert-linalg-to-loops),expand-strided-metadata,canonicalize,memref-elide-reinterpret-cast,convert-to-emitc,math-expand-ops{ops=rsqrt},arith-expand,convert-math-to-emitc,convert-arith-to-emitc)"
+ "builtin.module(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},buffer-deallocation-pipeline,func.func(convert-linalg-to-loops),expand-strided-metadata,canonicalize,memref-elide-reinterpret-cast,convert-to-emitc,math-expand-ops{ops=rsqrt},arith-expand,convert-math-to-emitc,convert-arith-to-emitc,wrap-emitc-func-in-class,mlgo-add-reflection-map{included-field-attrs=tf_saved_model.index_path})"
CACHE STRING
"MLIR pass pipeline used to lower MLGO TOSA models to EmitC.")
diff --git a/llvm/include/llvm/Analysis/MLIRInlinerSizeModel.h b/llvm/include/llvm/Analysis/MLIRInlinerSizeModel.h
index 5e16ddbff6432..1877abb388df7 100644
--- a/llvm/include/llvm/Analysis/MLIRInlinerSizeModel.h
+++ b/llvm/include/llvm/Analysis/MLIRInlinerSizeModel.h
@@ -9,10 +9,10 @@
/// Wraps the MLIR-translated MLGO inliner model in the interface expected by
/// ReleaseModeModelRunner.
///
-/// The generated `.inc` file only contains lowered model code. This wrapper
-/// owns the named inliner tensors, keeps their layout stable across generated
-/// model variants, and exposes the same serving API that the rest of LLVM
-/// already uses for release-mode MLGO models.
+/// The generated `.inc` file only contains the lowered model class. This
+/// wrapper adapts that name-based surface to the index-based
+/// ReleaseModeModelRunner contract that the rest of LLVM already uses for
+/// release-mode MLGO models.
//
//===----------------------------------------------------------------------===//
@@ -21,12 +21,16 @@
#include <array>
#include <cstdint>
+#include <memory>
#include <string>
namespace llvm {
class MLIRInlinerSizeModel final {
public:
+ MLIRInlinerSizeModel();
+ ~MLIRInlinerSizeModel();
+
int LookupArgIndex(const std::string &Name);
int LookupResultIndex(const std::string &Name);
void *arg_data(int Index);
@@ -34,53 +38,9 @@ class MLIRInlinerSizeModel final {
void Run();
private:
- enum ArgIndex : int {
- DeadBlocks = 0,
- CaseClusterPenalty,
- SroaSavings,
- JumpTablePenalty,
- CallsiteHeight,
- CalleeBasicBlockCount,
- CallArgumentSetup,
- LoweredCallArgSetup,
- SimplifiedInstructions,
- NrCtantParams,
- IsMultipleBlocks,
- LoadElimination,
- EdgeCount,
- CallerUsers,
- CallerConditionallyExecutedBlocks,
- ConstantOffsetPtrArgs,
- CallsiteCost,
- CallerBasicBlockCount,
- LoadRelativeIntrinsic,
- IndirectCallPenalty,
- CostEstimate,
- Threshold,
- NestedInlineCostEstimate,
- UnsimplifiedCommonInstructions,
- SroaLosses,
- NumLoops,
- SwitchPenalty,
- CalleeUsers,
- NodeCount,
- ConstantArgs,
- LastCallToStaticBonus,
- ColdCCPenalty,
- CalleeConditionallyExecutedBlocks,
- CallPenalty,
- NestedInlines,
-
- NumArgs
- };
-
- std::array<std::array<int64_t, 1>, NumArgs> Inputs{};
+ struct Impl;
+ std::unique_ptr<Impl> Model;
std::array<int64_t, 1> Result{};
-
- std::array<int64_t, 1> DummyInliningDefault{};
- std::array<int32_t, 1> DummyStepType{};
- std::array<float, 1> DummyDiscount{};
- std::array<float, 1> DummyReward{};
};
} // namespace llvm
diff --git a/llvm/include/llvm/CodeGen/MLIRRegAllocEvictModel.h b/llvm/include/llvm/CodeGen/MLIRRegAllocEvictModel.h
index d507d76599270..e5a4539aab3ad 100644
--- a/llvm/include/llvm/CodeGen/MLIRRegAllocEvictModel.h
+++ b/llvm/include/llvm/CodeGen/MLIRRegAllocEvictModel.h
@@ -9,25 +9,28 @@
/// Wraps the MLIR-translated MLGO regalloc eviction model in the interface
/// expected by ReleaseModeModelRunner.
///
-/// The generated `.inc` file only contains lowered model code. This wrapper
-/// owns the named regalloc tensors, keeps their layout stable across generated
-/// model variants, and exposes the same serving API that the rest of LLVM
-/// already uses for release-mode MLGO models.
+/// The generated `.inc` file only contains the lowered model class. This
+/// wrapper adapts that name-based surface to the index-based
+/// ReleaseModeModelRunner contract that the rest of LLVM already uses for
+/// release-mode MLGO models.
//
//===----------------------------------------------------------------------===//
#ifndef LLVM_CODEGEN_MLIRREGALLOCEVICTMODEL_H
#define LLVM_CODEGEN_MLIRREGALLOCEVICTMODEL_H
-#include <cmath>
-#include <cstddef>
+#include <array>
#include <cstdint>
+#include <memory>
#include <string>
namespace llvm {
class MLIRRegAllocEvictModel final {
public:
+ MLIRRegAllocEvictModel();
+ ~MLIRRegAllocEvictModel();
+
int LookupArgIndex(const std::string &Name);
int LookupResultIndex(const std::string &Name);
void *arg_data(int Index);
@@ -35,68 +38,9 @@ class MLIRRegAllocEvictModel final {
void Run();
private:
- static constexpr std::size_t InterferenceCount = 33;
-
- using F32InterferenceTensor = float[1][InterferenceCount];
- using I64InterferenceTensor = int64_t[1][InterferenceCount];
- using F32Scalar = float[1];
- using I32Scalar = int32_t[1];
- using I64Scalar = int64_t[1];
-
- enum ArgIndex : int {
- Mask = 0,
- IsFree,
- NrUrgent,
- NrBrokenHints,
- IsHint,
- IsLocal,
- NrRematerializable,
- NrDefsAndUses,
- WeighedReadsByMax,
- WeighedWritesByMax,
- WeighedReadWritesByMax,
- WeighedIndvarsByMax,
- HintWeightsByMax,
- StartBBFreqByMax,
- EndBBFreqByMax,
- HottestBBFreqByMax,
- LiverangeSize,
- UseDefDensity,
- MaxStage,
- MinStage,
- Progress,
-
- NumArgs
- };
-
- I64InterferenceTensor MaskInput{};
- I64InterferenceTensor IsFreeInput{};
- F32InterferenceTensor NrUrgentInput{};
- F32InterferenceTensor NrBrokenHintsInput{};
- I64InterferenceTensor IsHintInput{};
- I64InterferenceTensor IsLocalInput{};
- F32InterferenceTensor NrRematerializableInput{};
- F32InterferenceTensor NrDefsAndUsesInput{};
- F32InterferenceTensor WeighedReadsByMaxInput{};
- F32InterferenceTensor WeighedWritesByMaxInput{};
- F32InterferenceTensor WeighedReadWritesByMaxInput{};
- F32InterferenceTensor WeighedIndvarsByMaxInput{};
- F32InterferenceTensor HintWeightsByMaxInput{};
- F32InterferenceTensor StartBBFreqByMaxInput{};
- F32InterferenceTensor EndBBFreqByMaxInput{};
- F32InterferenceTensor HottestBBFreqByMaxInput{};
- F32InterferenceTensor LiverangeSizeInput{};
- F32InterferenceTensor UseDefDensityInput{};
- I64InterferenceTensor MaxStageInput{};
- I64InterferenceTensor MinStageInput{};
- F32Scalar ProgressInput{};
-
- // These scalars remain part of the serving API even when a translated model
- // does not make semantic use of all of them.
- I32Scalar DummyStepType{};
- F32Scalar DummyDiscount{};
- F32Scalar DummyReward{};
- I64Scalar Result{};
+ struct Impl;
+ std::unique_ptr<Impl> Model;
+ std::array<int64_t, 1> Result{};
};
} // namespace llvm
diff --git a/llvm/lib/Analysis/MLIRInlinerSizeModel.cpp b/llvm/lib/Analysis/MLIRInlinerSizeModel.cpp
index 1baff8c1b6d23..73b85e879f0c4 100644
--- a/llvm/lib/Analysis/MLIRInlinerSizeModel.cpp
+++ b/llvm/lib/Analysis/MLIRInlinerSizeModel.cpp
@@ -13,184 +13,116 @@
#include "llvm/Analysis/MLIRInlinerSizeModel.h"
-#include "llvm/ADT/StringSwitch.h"
#include "llvm/Support/ErrorHandling.h"
#include <math.h>
+#include <memory>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
-#include <type_traits>
+#include <string>
#if defined(__clang__)
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wmissing-braces"
#endif
-namespace llvm::mlir_inliner_model {
-// Mock models generated for tests use `main` as entry point, while
-// pretrained models lowered use `action`. Rename
-// locally so the wrapper can always dispatch through one symbol without
-// rewriting the generated `.inc` file.
-#define main action
+// Generated models may name the wrapper class either `mainClass` or
+// `actionClass`. Normalize both spellings to a TU-local name so the inliner and
+// regalloc wrappers do not emit conflicting global symbols.
+#define actionClass MLIRInlinerGeneratedModel
+#define mainClass MLIRInlinerGeneratedModel
#include "llvm/Analysis/MLIRInlinerSizeModel.inc"
-#undef main
-} // namespace llvm::mlir_inliner_model
+#undef actionClass
+#undef mainClass
#if defined(__clang__)
#pragma clang diagnostic pop
#endif
-using namespace llvm;
+namespace llvm {
namespace {
-template <typename T> inline constexpr bool AlwaysFalse = false;
-using I64Ptr = int64_t *;
-using I32Ptr = int32_t *;
-using F32Ptr = float *;
-using InlinerProductionActionTy =
- int64_t (*)(I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
- I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
- I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
- I64Ptr, I64Ptr, I64Ptr, I32Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
- I64Ptr, I64Ptr, I64Ptr, I64Ptr, F32Ptr, I64Ptr, F32Ptr);
-using InlinerMockActionTy = int64_t (*)(I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
- I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
- I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
- I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
- I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
- I64Ptr, I64Ptr, I32Ptr, I64Ptr, I64Ptr,
- I64Ptr, I64Ptr, I64Ptr, I64Ptr, I64Ptr,
- I64Ptr, I64Ptr, F32Ptr, I64Ptr, F32Ptr);
+using GeneratedInlinerModel = MLIRInlinerGeneratedModel;
-struct InlinerRunInputs {
- I64Ptr deadBlocks;
- I64Ptr caseClusterPenalty;
- I64Ptr sroaSavings;
- I64Ptr jumpTablePenalty;
- I64Ptr callsiteHeight;
- I64Ptr calleeBasicBlockCount;
- I64Ptr callArgumentSetup;
- I64Ptr loweredCallArgSetup;
- I64Ptr simplifiedInstructions;
- I64Ptr nrCtantParams;
- I64Ptr isMultipleBlocks;
- I64Ptr loadElimination;
- I64Ptr edgeCount;
- I64Ptr callerUsers;
- I64Ptr callerConditionallyExecutedBlocks;
- I64Ptr constantOffsetPtrArgs;
- I64Ptr callsiteCost;
- I64Ptr callerBasicBlockCount;
- I64Ptr loadRelativeIntrinsic;
- I64Ptr indirectCallPenalty;
- I64Ptr costEstimate;
- I64Ptr threshold;
- I64Ptr nestedInlineCostEstimate;
- I64Ptr unsimplifiedCommonInstructions;
- I64Ptr sroaLosses;
- I64Ptr numLoops;
- I64Ptr switchPenalty;
- I64Ptr calleeUsers;
- I64Ptr nodeCount;
- I64Ptr constantArgs;
- I64Ptr lastCallToStaticBonus;
- I64Ptr coldCCPenalty;
- I64Ptr calleeConditionallyExecutedBlocks;
- I64Ptr callPenalty;
- I64Ptr nestedInlines;
- I32Ptr dummyStepType;
- F32Ptr dummyDiscount;
- F32Ptr dummyReward;
- I64Ptr dummyInliningDefault;
+struct NamedArg {
+ const char *FeedName;
+ const char *ModelName;
};
-// The MLIR flow currently supports both the production inlining models API
-// and the simplified mock-model API used by tests. Keep that API
-// variance local to the wrapper so the rest of LLVM only sees one
-// ReleaseModeModelRunner-compatible object.
-template <typename ActionTy>
-int64_t runMLIRInlinerAction(const InlinerRunInputs &I) {
- if constexpr (std::is_same_v<ActionTy, InlinerProductionActionTy>) {
- return static_cast<ActionTy>(mlir_inliner_model::action)(
- I.callsiteCost, I.isMultipleBlocks, I.callerConditionallyExecutedBlocks,
- I.dummyInliningDefault, I.coldCCPenalty,
- I.calleeConditionallyExecutedBlocks, I.calleeUsers,
- I.calleeBasicBlockCount, I.nrCtantParams, I.loadRelativeIntrinsic,
- I.jumpTablePenalty, I.unsimplifiedCommonInstructions,
- I.indirectCallPenalty, I.loadElimination, I.callPenalty, I.costEstimate,
- I.caseClusterPenalty, I.nodeCount, I.callArgumentSetup, I.sroaSavings,
- I.loweredCallArgSetup, I.threshold, I.deadBlocks, I.constantArgs,
- I.sroaLosses, I.simplifiedInstructions, I.numLoops, I.dummyStepType,
- I.edgeCount, I.nestedInlines, I.callerBasicBlockCount,
- I.lastCallToStaticBonus, I.nestedInlineCostEstimate, I.callsiteHeight,
- I.constantOffsetPtrArgs, I.switchPenalty, I.dummyDiscount,
- I.callerUsers, I.dummyReward);
- } else if constexpr (std::is_same_v<ActionTy, InlinerMockActionTy>) {
- return static_cast<ActionTy>(mlir_inliner_model::action)(
- I.callerBasicBlockCount, I.callerConditionallyExecutedBlocks,
- I.callerUsers, I.calleeBasicBlockCount,
- I.calleeConditionallyExecutedBlocks, I.calleeUsers, I.nrCtantParams,
- I.nodeCount, I.edgeCount, I.callsiteHeight, I.costEstimate,
- I.sroaSavings, I.sroaLosses, I.loadElimination, I.callPenalty,
- I.callArgumentSetup, I.loadRelativeIntrinsic, I.loweredCallArgSetup,
- I.indirectCallPenalty, I.jumpTablePenalty, I.caseClusterPenalty,
- I.switchPenalty, I.unsimplifiedCommonInstructions, I.numLoops,
- I.deadBlocks, I.simplifiedInstructions, I.constantArgs, I.dummyStepType,
- I.constantOffsetPtrArgs, I.callsiteCost, I.coldCCPenalty,
- I.lastCallToStaticBonus, I.isMultipleBlocks, I.nestedInlines,
- I.nestedInlineCostEstimate, I.threshold, I.dummyInliningDefault,
- I.dummyDiscount, I.callerUsers, I.dummyReward);
- } else {
- static_assert(AlwaysFalse<ActionTy>,
- "Unsupported MLIR inliner model signature");
- }
-}
+constexpr NamedArg InlinerArgs[] = {
+ {"feed_dead_blocks", "dead_blocks"},
+ {"feed_case_cluster_penalty", "case_cluster_penalty"},
+ {"feed_sroa_savings", "sroa_savings"},
+ {"feed_jump_table_penalty", "jump_table_penalty"},
+ {"feed_callsite_height", "callsite_height"},
+ {"feed_callee_basic_block_count", "callee_basic_block_count"},
+ {"feed_call_argument_setup", "call_argument_setup"},
+ {"feed_lowered_call_arg_setup", "lowered_call_arg_setup"},
+ {"feed_simplified_instructions", "simplified_instructions"},
+ {"feed_nr_ctant_params", "nr_ctant_params"},
+ {"feed_is_multiple_blocks", "is_multiple_blocks"},
+ {"feed_load_elimination", "load_elimination"},
+ {"feed_edge_count", "edge_count"},
+ {"feed_caller_users", "caller_users"},
+ {"feed_caller_conditionally_executed_blocks",
+ "caller_conditionally_executed_blocks"},
+ {"feed_constant_offset_ptr_args", "constant_offset_ptr_args"},
+ {"feed_callsite_cost", "callsite_cost"},
+ {"feed_caller_basic_block_count", "caller_basic_block_count"},
+ {"feed_load_relative_intrinsic", "load_relative_intrinsic"},
+ {"feed_indirect_call_penalty", "indirect_call_penalty"},
+ {"feed_cost_estimate", "cost_estimate"},
+ {"feed_threshold", "threshold"},
+ {"feed_nested_inline_cost_estimate", "nested_inline_cost_estimate"},
+ {"feed_unsimplified_common_instructions",
+ "unsimplified_common_instructions"},
+ {"feed_sroa_losses", "sroa_losses"},
+ {"feed_num_loops", "num_loops"},
+ {"feed_switch_penalty", "switch_penalty"},
+ {"feed_callee_users", "callee_users"},
+ {"feed_node_count", "node_count"},
+ {"feed_constant_args", "constant_args"},
+ {"feed_last_call_to_static_bonus", "last_call_to_static_bonus"},
+ {"feed_cold_cc_penalty", "cold_cc_penalty"},
+ {"feed_callee_conditionally_executed_blocks",
+ "callee_conditionally_executed_blocks"},
+ {"feed_call_penalty", "call_penalty"},
+ {"feed_nested_inlines", "nested_inlines"},
+};
+
+constexpr size_t NumInlinerArgs = sizeof(InlinerArgs) / sizeof(*InlinerArgs);
+
+static_assert(NumInlinerArgs == 35,
+ "Unexpected number of inliner model inputs");
+
} // namespace
+struct MLIRInlinerSizeModel::Impl {
+ GeneratedInlinerModel Model{};
+
+ bool hasBufferForName(const char *Name) const {
+ return Model.reflectionMap.find(Name) != Model.reflectionMap.end();
+ }
+
+ void *getBufferForName(const char *Name) {
+ return Model.getBufferForName(std::string(Name));
+ }
+};
+
+MLIRInlinerSizeModel::MLIRInlinerSizeModel()
+ : Model(std::make_unique<Impl>()) {}
+
+MLIRInlinerSizeModel::~MLIRInlinerSizeModel() = default;
+
int MLIRInlinerSizeModel::LookupArgIndex(const std::string &Name) {
- return StringSwitch<int>(Name)
- .Case("feed_dead_blocks", DeadBlocks)
- .Case("feed_case_cluster_penalty", CaseClusterPenalty)
- .Case("feed_sroa_savings", SroaSavings)
- .Case("feed_jump_table_penalty", JumpTablePenalty)
- .Case("feed_callsite_height", CallsiteHeight)
- .Case("feed_callee_basic_block_count", CalleeBasicBlockCount)
- .Case("feed_call_argument_setup", CallArgumentSetup)
- .Case("feed_lowered_call_arg_setup", LoweredCallArgSetup)
- .Case("feed_simplified_instructions", SimplifiedInstructions)
- .Case("feed_nr_ctant_params", NrCtantParams)
- .Case("feed_is_multiple_blocks", IsMultipleBlocks)
- .Case("feed_load_elimination", LoadElimination)
- .Case("feed_edge_count", EdgeCount)
- .Case("feed_caller_users", CallerUsers)
- .Case("feed_caller_conditionally_executed_blocks",
- CallerConditionallyExecutedBlocks)
- .Case("feed_constant_offset_ptr_args", ConstantOffsetPtrArgs)
- .Case("feed_callsite_cost", CallsiteCost)
- .Case("feed_caller_basic_block_count", CallerBasicBlockCount)
- .Case("feed_load_relative_intrinsic", LoadRelativeIntrinsic)
- .Case("feed_indirect_call_penalty", IndirectCallPenalty)
- .Case("feed_cost_estimate", CostEstimate)
- .Case("feed_threshold", Threshold)
- .Case("feed_nested_inline_cost_estimate", NestedInlineCostEstimate)
- .Case("feed_unsimplified_common_instructions",
- UnsimplifiedCommonInstructions)
- .Case("feed_sroa_losses", SroaLosses)
- .Case("feed_num_loops", NumLoops)
- .Case("feed_switch_penalty", SwitchPenalty)
- .Case("feed_callee_users", CalleeUsers)
- .Case("feed_node_count", NodeCount)
- .Case("feed_constant_args", ConstantArgs)
- .Case("feed_last_call_to_static_bonus", LastCallToStaticBonus)
- .Case("feed_cold_cc_penalty", ColdCCPenalty)
- .Case("feed_callee_conditionally_executed_blocks",
- CalleeConditionallyExecutedBlocks)
- .Case("feed_call_penalty", CallPenalty)
- .Case("feed_nested_inlines", NestedInlines)
- .Default(-1);
+ for (size_t I = 0; I < NumInlinerArgs; ++I)
+ if (Name == InlinerArgs[I].FeedName &&
+ Model->hasBufferForName(InlinerArgs[I].ModelName))
+ return static_cast<int>(I);
+ return -1;
}
int MLIRInlinerSizeModel::LookupResultIndex(const std::string &Name) {
@@ -198,9 +130,9 @@ int MLIRInlinerSizeModel::LookupResultIndex(const std::string &Name) {
}
void *MLIRInlinerSizeModel::arg_data(int Index) {
- if (Index < 0 || Index >= NumArgs)
+ if (Index < 0 || static_cast<size_t>(Index) >= NumInlinerArgs)
llvm_unreachable("invalid MLIR inliner input index");
- return Inputs[Index].data();
+ return Model->getBufferForName(InlinerArgs[Index].ModelName);
}
void *MLIRInlinerSizeModel::result_data(int Index) {
@@ -209,52 +141,6 @@ void *MLIRInlinerSizeModel::result_data(int Index) {
return Result.data();
}
-void MLIRInlinerSizeModel::Run() {
- using ActionTy = decltype(&mlir_inliner_model::action);
- // Gather the stored named tensors once, then dispatch through the adapter
- // selected from the generated `action` signature.
- InlinerRunInputs I{};
- I.deadBlocks = Inputs[DeadBlocks].data();
- I.caseClusterPenalty = Inputs[CaseClusterPenalty].data();
- I.sroaSavings = Inputs[SroaSavings].data();
- I.jumpTablePenalty = Inputs[JumpTablePenalty].data();
- I.callsiteHeight = Inputs[CallsiteHeight].data();
- I.calleeBasicBlockCount = Inputs[CalleeBasicBlockCount].data();
- I.callArgumentSetup = Inputs[CallArgumentSetup].data();
- I.loweredCallArgSetup = Inputs[LoweredCallArgSetup].data();
- I.simplifiedInstructions = Inputs[SimplifiedInstructions].data();
- I.nrCtantParams = Inputs[NrCtantParams].data();
- I.isMultipleBlocks = Inputs[IsMultipleBlocks].data();
- I.loadElimination = Inputs[LoadElimination].data();
- I.edgeCount = Inputs[EdgeCount].data();
- I.callerUsers = Inputs[CallerUsers].data();
- I.callerConditionallyExecutedBlocks =
- Inputs[CallerConditionallyExecutedBlocks].data();
- I.constantOffsetPtrArgs = Inputs[ConstantOffsetPtrArgs].data();
- I.callsiteCost = Inputs[CallsiteCost].data();
- I.callerBasicBlockCount = Inputs[CallerBasicBlockCount].data();
- I.loadRelativeIntrinsic = Inputs[LoadRelativeIntrinsic].data();
- I.indirectCallPenalty = Inputs[IndirectCallPenalty].data();
- I.costEstimate = Inputs[CostEstimate].data();
- I.threshold = Inputs[Threshold].data();
- I.nestedInlineCostEstimate = Inputs[NestedInlineCostEstimate].data();
- I.unsimplifiedCommonInstructions =
- Inputs[UnsimplifiedCommonInstructions].data();
- I.sroaLosses = Inputs[SroaLosses].data();
- I.numLoops = Inputs[NumLoops].data();
- I.switchPenalty = Inputs[SwitchPenalty].data();
- I.calleeUsers = Inputs[CalleeUsers].data();
- I.nodeCount = Inputs[NodeCount].data();
- I.constantArgs = Inputs[ConstantArgs].data();
- I.lastCallToStaticBonus = Inputs[LastCallToStaticBonus].data();
- I.coldCCPenalty = Inputs[ColdCCPenalty].data();
- I.calleeConditionallyExecutedBlocks =
- Inputs[CalleeConditionallyExecutedBlocks].data();
- I.callPenalty = Inputs[CallPenalty].data();
- I.nestedInlines = Inputs[NestedInlines].data();
- I.dummyStepType = DummyStepType.data();
- I.dummyDiscount = DummyDiscount.data();
- I.dummyReward = DummyReward.data();
- I.dummyInliningDefault = DummyInliningDefault.data();
- Result[0] = runMLIRInlinerAction<ActionTy>(I);
-}
+void MLIRInlinerSizeModel::Run() { Result[0] = (Model->Model)(); }
+
+} // namespace llvm
diff --git a/llvm/lib/CodeGen/MLIRRegAllocEvictModel.cpp b/llvm/lib/CodeGen/MLIRRegAllocEvictModel.cpp
index 14ecae10b29f1..b833e3254c80e 100644
--- a/llvm/lib/CodeGen/MLIRRegAllocEvictModel.cpp
+++ b/llvm/lib/CodeGen/MLIRRegAllocEvictModel.cpp
@@ -13,117 +13,86 @@
#include "llvm/CodeGen/MLIRRegAllocEvictModel.h"
-#include "llvm/ADT/StringSwitch.h"
#include "llvm/Support/ErrorHandling.h"
-#include <stddef.h>
-#include <stdint.h>
-#include <type_traits>
+#include <cstddef>
+#include <memory>
+#include <string>
-namespace llvm::mlir_regalloc_evict_model {
-// Mock models generated for tests use `main` as entry point, while
-// pretrained models lowered use `action`. Rename
-// locally so the wrapper can always dispatch through one symbol without
-// rewriting the generated `.inc` file.
-#define main action
+// Generated models may name the wrapper class either `mainClass` or
+// `actionClass`. Normalize both spellings to a TU-local name so the inliner and
+// regalloc wrappers do not emit conflicting global symbols.
+#define actionClass MLIRRegAllocGeneratedModel
+#define mainClass MLIRRegAllocGeneratedModel
#include "llvm/CodeGen/MLIRRegAllocEvictModel.inc"
-#undef main
-} // namespace llvm::mlir_regalloc_evict_model
+#undef actionClass
+#undef mainClass
-using namespace llvm;
+namespace llvm {
namespace {
-template <typename T> inline constexpr bool AlwaysFalse = false;
-constexpr std::size_t MLIRRegAllocInterferenceCount = 33;
-using F32TensorPtr = float (*)[MLIRRegAllocInterferenceCount];
-using I64TensorPtr = int64_t (*)[MLIRRegAllocInterferenceCount];
-using F32ScalarPtr = float *;
-using I32ScalarPtr = int32_t *;
-
-using RegAllocProductionActionTy = int64_t (*)(
- F32TensorPtr, F32TensorPtr, I64TensorPtr, F32TensorPtr, F32TensorPtr,
- F32TensorPtr, F32ScalarPtr, F32TensorPtr, F32TensorPtr, F32TensorPtr,
- I64TensorPtr, I64TensorPtr, F32TensorPtr, F32TensorPtr, I32ScalarPtr,
- F32TensorPtr, I64TensorPtr, F32TensorPtr, I64TensorPtr, F32TensorPtr,
- F32ScalarPtr, F32TensorPtr, I64TensorPtr, F32ScalarPtr);
-using RegAllocMaskOnlyActionTy = int64_t (*)(I64TensorPtr);
-using RegAllocFlatMaskActionTy = int64_t (*)(int64_t *);
-
-struct RegAllocRunInputs {
- F32TensorPtr liverangeSize;
- F32TensorPtr hintWeightsByMax;
- I64TensorPtr isFree;
- F32TensorPtr weighedReadsByMax;
- F32TensorPtr weighedReadWritesByMax;
- F32TensorPtr nrBrokenHints;
- F32ScalarPtr progress;
- F32TensorPtr hottestBBFreqByMax;
- F32TensorPtr useDefDensity;
- F32TensorPtr startBBFreqByMax;
- I64TensorPtr maxStage;
- I64TensorPtr isHint;
- F32TensorPtr nrRematerializable;
- F32TensorPtr weighedWritesByMax;
- I32ScalarPtr dummyStepType;
- F32TensorPtr nrUrgent;
- I64TensorPtr mask;
- F32TensorPtr nrDefsAndUses;
- I64TensorPtr isLocal;
- F32TensorPtr endBBFreqByMax;
- F32ScalarPtr dummyDiscount;
- F32TensorPtr weighedIndvarsByMax;
- I64TensorPtr minStage;
- F32ScalarPtr dummyReward;
+
+using GeneratedRegAllocModel = MLIRRegAllocGeneratedModel;
+
+struct NamedArg {
+ const char *FeedName;
+ const char *ModelName;
};
-// The MLIR flow currently supports both the production regalloc models API
-// and the simplified mock-model API used by tests. Keep that API
-// variance local to the wrapper so the rest of LLVM only sees one
-// ReleaseModeModelRunner-compatible object.
-template <typename ActionTy>
-int64_t runMLIRRegAllocAction(const RegAllocRunInputs &I) {
- if constexpr (std::is_same_v<ActionTy, RegAllocProductionActionTy>) {
- return static_cast<ActionTy>(mlir_regalloc_evict_model::action)(
- I.liverangeSize, I.hintWeightsByMax, I.isFree, I.weighedReadsByMax,
- I.weighedReadWritesByMax, I.nrBrokenHints, I.progress,
- I.hottestBBFreqByMax, I.useDefDensity, I.startBBFreqByMax, I.maxStage,
- I.isHint, I.nrRematerializable, I.weighedWritesByMax, I.dummyStepType,
- I.nrUrgent, I.mask, I.nrDefsAndUses, I.isLocal, I.endBBFreqByMax,
- I.dummyDiscount, I.weighedIndvarsByMax, I.minStage, I.dummyReward);
- } else if constexpr (std::is_same_v<ActionTy, RegAllocFlatMaskActionTy>) {
- return static_cast<ActionTy>(mlir_regalloc_evict_model::action)(
- &(*I.mask)[0]);
- } else {
- static_assert(AlwaysFalse<ActionTy>,
- "Unsupported MLIR regalloc eviction model signature");
- }
-}
+constexpr NamedArg RegAllocArgs[] = {
+ {"feed_mask", "mask"},
+ {"feed_is_free", "is_free"},
+ {"feed_nr_urgent", "nr_urgent"},
+ {"feed_nr_broken_hints", "nr_broken_hints"},
+ {"feed_is_hint", "is_hint"},
+ {"feed_is_local", "is_local"},
+ {"feed_nr_rematerializable", "nr_rematerializable"},
+ {"feed_nr_defs_and_uses", "nr_defs_and_uses"},
+ {"feed_weighed_reads_by_max", "weighed_reads_by_max"},
+ {"feed_weighed_writes_by_max", "weighed_writes_by_max"},
+ {"feed_weighed_read_writes_by_max", "weighed_read_writes_by_max"},
+ {"feed_weighed_indvars_by_max", "weighed_indvars_by_max"},
+ {"feed_hint_weights_by_max", "hint_weights_by_max"},
+ {"feed_start_bb_freq_by_max", "start_bb_freq_by_max"},
+ {"feed_end_bb_freq_by_max", "end_bb_freq_by_max"},
+ {"feed_hottest_bb_freq_by_max", "hottest_bb_freq_by_max"},
+ {"feed_liverange_size", "liverange_size"},
+ {"feed_use_def_density", "use_def_density"},
+ {"feed_max_stage", "max_stage"},
+ {"feed_min_stage", "min_stage"},
+ {"feed_progress", "progress"},
+};
+
+constexpr size_t NumRegAllocArgs = sizeof(RegAllocArgs) / sizeof(*RegAllocArgs);
+
+static_assert(NumRegAllocArgs == 21,
+ "Unexpected number of regalloc model inputs");
+
} // namespace
+struct MLIRRegAllocEvictModel::Impl {
+ GeneratedRegAllocModel Model{};
+
+ bool hasBufferForName(const char *Name) const {
+ return Model.reflectionMap.find(Name) != Model.reflectionMap.end();
+ }
+
+ void *getBufferForName(const char *Name) {
+ return Model.getBufferForName(std::string(Name));
+ }
+};
+
+MLIRRegAllocEvictModel::MLIRRegAllocEvictModel()
+ : Model(std::make_unique<Impl>()) {}
+
+MLIRRegAllocEvictModel::~MLIRRegAllocEvictModel() = default;
+
int MLIRRegAllocEvictModel::LookupArgIndex(const std::string &Name) {
- return StringSwitch<int>(Name)
- .Case("feed_mask", Mask)
- .Case("feed_is_free", IsFree)
- .Case("feed_nr_urgent", NrUrgent)
- .Case("feed_nr_broken_hints", NrBrokenHints)
- .Case("feed_is_hint", IsHint)
- .Case("feed_is_local", IsLocal)
- .Case("feed_nr_rematerializable", NrRematerializable)
- .Case("feed_nr_defs_and_uses", NrDefsAndUses)
- .Case("feed_weighed_reads_by_max", WeighedReadsByMax)
- .Case("feed_weighed_writes_by_max", WeighedWritesByMax)
- .Case("feed_weighed_read_writes_by_max", WeighedReadWritesByMax)
- .Case("feed_weighed_indvars_by_max", WeighedIndvarsByMax)
- .Case("feed_hint_weights_by_max", HintWeightsByMax)
- .Case("feed_start_bb_freq_by_max", StartBBFreqByMax)
- .Case("feed_end_bb_freq_by_max", EndBBFreqByMax)
- .Case("feed_hottest_bb_freq_by_max", HottestBBFreqByMax)
- .Case("feed_liverange_size", LiverangeSize)
- .Case("feed_use_def_density", UseDefDensity)
- .Case("feed_max_stage", MaxStage)
- .Case("feed_min_stage", MinStage)
- .Case("feed_progress", Progress)
- .Default(-1);
+ for (size_t I = 0; I < NumRegAllocArgs; ++I)
+ if (Name == RegAllocArgs[I].FeedName &&
+ Model->hasBufferForName(RegAllocArgs[I].ModelName))
+ return static_cast<int>(I);
+ return -1;
}
int MLIRRegAllocEvictModel::LookupResultIndex(const std::string &Name) {
@@ -131,87 +100,17 @@ int MLIRRegAllocEvictModel::LookupResultIndex(const std::string &Name) {
}
void *MLIRRegAllocEvictModel::arg_data(int Index) {
- switch (Index) {
- case Mask:
- return MaskInput;
- case IsFree:
- return IsFreeInput;
- case NrUrgent:
- return NrUrgentInput;
- case NrBrokenHints:
- return NrBrokenHintsInput;
- case IsHint:
- return IsHintInput;
- case IsLocal:
- return IsLocalInput;
- case NrRematerializable:
- return NrRematerializableInput;
- case NrDefsAndUses:
- return NrDefsAndUsesInput;
- case WeighedReadsByMax:
- return WeighedReadsByMaxInput;
- case WeighedWritesByMax:
- return WeighedWritesByMaxInput;
- case WeighedReadWritesByMax:
- return WeighedReadWritesByMaxInput;
- case WeighedIndvarsByMax:
- return WeighedIndvarsByMaxInput;
- case HintWeightsByMax:
- return HintWeightsByMaxInput;
- case StartBBFreqByMax:
- return StartBBFreqByMaxInput;
- case EndBBFreqByMax:
- return EndBBFreqByMaxInput;
- case HottestBBFreqByMax:
- return HottestBBFreqByMaxInput;
- case LiverangeSize:
- return LiverangeSizeInput;
- case UseDefDensity:
- return UseDefDensityInput;
- case MaxStage:
- return MaxStageInput;
- case MinStage:
- return MinStageInput;
- case Progress:
- return ProgressInput;
- }
- llvm_unreachable("invalid MLIR regalloc eviction input index");
+ if (Index < 0 || static_cast<size_t>(Index) >= NumRegAllocArgs)
+ llvm_unreachable("invalid MLIR regalloc eviction input index");
+ return Model->getBufferForName(RegAllocArgs[Index].ModelName);
}
void *MLIRRegAllocEvictModel::result_data(int Index) {
if (Index != 0)
llvm_unreachable("invalid MLIR regalloc eviction result index");
- return Result;
+ return Result.data();
}
-void MLIRRegAllocEvictModel::Run() {
- using ActionTy = decltype(&mlir_regalloc_evict_model::action);
- // Gather the stored named tensors once, then dispatch through the adapter
- // selected from the generated `action` signature.
- RegAllocRunInputs I{};
- I.liverangeSize = LiverangeSizeInput;
- I.hintWeightsByMax = HintWeightsByMaxInput;
- I.isFree = IsFreeInput;
- I.weighedReadsByMax = WeighedReadsByMaxInput;
- I.weighedReadWritesByMax = WeighedReadWritesByMaxInput;
- I.nrBrokenHints = NrBrokenHintsInput;
- I.progress = ProgressInput;
- I.hottestBBFreqByMax = HottestBBFreqByMaxInput;
- I.useDefDensity = UseDefDensityInput;
- I.startBBFreqByMax = StartBBFreqByMaxInput;
- I.maxStage = MaxStageInput;
- I.isHint = IsHintInput;
- I.nrRematerializable = NrRematerializableInput;
- I.weighedWritesByMax = WeighedWritesByMaxInput;
- I.dummyStepType = DummyStepType;
- I.nrUrgent = NrUrgentInput;
- I.mask = MaskInput;
- I.nrDefsAndUses = NrDefsAndUsesInput;
- I.isLocal = IsLocalInput;
- I.endBBFreqByMax = EndBBFreqByMaxInput;
- I.dummyDiscount = DummyDiscount;
- I.weighedIndvarsByMax = WeighedIndvarsByMaxInput;
- I.minStage = MinStageInput;
- I.dummyReward = DummyReward;
- Result[0] = runMLIRRegAllocAction<ActionTy>(I);
-}
+void MLIRRegAllocEvictModel::Run() { Result[0] = (Model->Model)(); }
+
+} // namespace llvm
>From e041140d30dd253ab2a5e45be3cb69af06252d53 Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Tue, 21 Jul 2026 22:27:30 +0200
Subject: [PATCH 6/7] Fix tests
---
llvm/test/CodeGen/MLRegAlloc/default-priority-advisor.ll | 1 +
llvm/test/Transforms/Inline/inlining-advisor-default.ll | 1 +
llvm/test/lit.site.cfg.py.in | 6 +++---
3 files changed, 5 insertions(+), 3 deletions(-)
diff --git a/llvm/test/CodeGen/MLRegAlloc/default-priority-advisor.ll b/llvm/test/CodeGen/MLRegAlloc/default-priority-advisor.ll
index e2908259facf9..f47218ad99fb1 100644
--- a/llvm/test/CodeGen/MLRegAlloc/default-priority-advisor.ll
+++ b/llvm/test/CodeGen/MLRegAlloc/default-priority-advisor.ll
@@ -2,6 +2,7 @@
; trying to use ML-driven advisor.
; REQUIRES: !have_tf_aot
; REQUIRES: !have_tflite
+; REQUIRES: !have_mlir_mlgo
; REQUIRES: default_triple
; RUN: not llc -O2 -regalloc-enable-priority-advisor=development < %s 2>&1 | FileCheck %s
; RUN: not llc -O2 -regalloc-enable-priority-advisor=release < %s 2>&1 | FileCheck %s
diff --git a/llvm/test/Transforms/Inline/inlining-advisor-default.ll b/llvm/test/Transforms/Inline/inlining-advisor-default.ll
index 502a16280192d..73540730a6f22 100644
--- a/llvm/test/Transforms/Inline/inlining-advisor-default.ll
+++ b/llvm/test/Transforms/Inline/inlining-advisor-default.ll
@@ -2,6 +2,7 @@
; trying to use ML-driven inlining.
; REQUIRES: !have_tf_aot
; REQUIRES: !have_tflite
+; REQUIRES: !have_mlir_mlgo
; 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
diff --git a/llvm/test/lit.site.cfg.py.in b/llvm/test/lit.site.cfg.py.in
index 0a7579a1a0c66..cce11ba1575e5 100644
--- a/llvm/test/lit.site.cfg.py.in
+++ b/llvm/test/lit.site.cfg.py.in
@@ -55,9 +55,9 @@ config.libcxx_used = @LLVM_LIBCXX_USED@
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_tflite = @LLVM_HAVE_TFLITE@
-config.have_mlir_mlgo = @LLVM_USE_MLIR_FOR_MLGO@
+config.have_tf_aot = "@LLVM_HAVE_TF_AOT@" in ("1", "ON", "TRUE", "True")
+config.have_tflite = "@LLVM_HAVE_TFLITE@" in ("1", "ON", "TRUE", "True")
+config.have_mlir_mlgo = "@LLVM_USE_MLIR_FOR_MLGO@" in ("1", "ON", "TRUE", "True")
config.enable_profcheck = @LLVM_ENABLE_PROFCHECK@
config.llvm_inliner_model_autogenerated = @LLVM_INLINER_MODEL_AUTOGENERATED@
config.llvm_raevict_model_autogenerated = @LLVM_RAEVICT_MODEL_AUTOGENERATED@
>From be272a108b10afa782b3fab1681840817d8d2682 Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Thu, 23 Jul 2026 16:17:58 +0200
Subject: [PATCH 7/7] Require TOSA input for in-tree flow
---
llvm/cmake/modules/TensorFlowCompile.cmake | 277 ++++++++++++++-------
llvm/lib/Analysis/CMakeLists.txt | 3 +-
llvm/lib/CodeGen/CMakeLists.txt | 3 +-
3 files changed, 187 insertions(+), 96 deletions(-)
diff --git a/llvm/cmake/modules/TensorFlowCompile.cmake b/llvm/cmake/modules/TensorFlowCompile.cmake
index 727bcd7c2d7c6..8231b9e4a3471 100644
--- a/llvm/cmake/modules/TensorFlowCompile.cmake
+++ b/llvm/cmake/modules/TensorFlowCompile.cmake
@@ -1,4 +1,4 @@
-function(mlgo_get_absolute_path path base final_path)
+function(get_absolute_path path base final_path)
if (IS_ABSOLUTE ${path})
set(${final_path} ${path} PARENT_SCOPE)
else()
@@ -6,7 +6,7 @@ function(mlgo_get_absolute_path path base final_path)
endif()
endfunction()
-function(mlgo_get_model model final_path)
+function(tf_get_model model final_path)
string(FIND ${model} "http:" pos_http)
string(FIND ${model} "https:" pos_https)
if (${pos_http} EQUAL 0 OR ${pos_https} EQUAL 0)
@@ -21,15 +21,15 @@ function(mlgo_get_model model final_path)
DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/${fname}_model)
set(${final_path} ${CMAKE_CURRENT_BINARY_DIR}/${fname}_model/model PARENT_SCOPE)
else()
- mlgo_get_absolute_path(${model} ${CMAKE_CURRENT_BINARY_DIR} model_path)
+ get_absolute_path(${model} ${CMAKE_CURRENT_BINARY_DIR} model_path)
set(${final_path} ${model_path} PARENT_SCOPE)
endif()
endfunction()
# Generate a mock model for tests.
function(generate_mock_model generator output)
- mlgo_get_absolute_path(${generator} ${CMAKE_CURRENT_SOURCE_DIR} generator_absolute_path)
- mlgo_get_absolute_path(${output} ${CMAKE_CURRENT_BINARY_DIR} output_absolute_path)
+ get_absolute_path(${generator} ${CMAKE_CURRENT_SOURCE_DIR} generator_absolute_path)
+ get_absolute_path(${output} ${CMAKE_CURRENT_BINARY_DIR} output_absolute_path)
message(WARNING "Autogenerated mock models should not be used in production builds.")
execute_process(COMMAND ${Python3_EXECUTABLE}
${generator_absolute_path}
@@ -38,41 +38,6 @@ function(generate_mock_model generator output)
)
endfunction()
-# Shared MLGO model build helpers. Both release-mode flows start from the same
-# model selection and discovery steps. The TensorFlow AOT path consumes a
-# SavedModel and asks TensorFlow to emit a compiled serving interface, while
-# the MLIR path converts the SavedModel through TFLite and TOSA, lowers it with
-# an MLIR pass pipeline, and emits C++ that LLVM wraps locally.
-
-function(mlgo_resolve_model model default_url default_path
- test_model_generator model_label should_skip final_path)
- if (${model} STREQUAL "none")
- message(STATUS "Will skip enabling mlgo for ${model_label}")
- set(${should_skip} TRUE PARENT_SCOPE)
- return()
- endif()
-
- if (${model} STREQUAL "download")
- # Crash if the user wants to download a model but a URL is not configured.
- if (${default_url} STREQUAL "<UNSPECIFIED>")
- message(FATAL_ERROR "Model path was set to 'download' but there is no"
- " model url currently specified in cmake. You can generate a model"
- " using, for example, the tools at http://github.com/google/ml-compiler-opt."
- " Some reference models are also periodically released there.")
- endif()
- set(model ${default_url})
- endif()
-
- if (${model} STREQUAL "autogenerate")
- set(model ${default_path}-autogenerated)
- generate_mock_model(${test_model_generator} ${model})
- endif()
-
- mlgo_get_model(${model} model_input)
- set(${should_skip} FALSE PARENT_SCOPE)
- set(${final_path} ${model_input} PARENT_SCOPE)
-endfunction()
-
# Run the tensorflow compiler (saved_model_cli) on the saved model in the
# ${model} directory, looking for the ${tag_set} tag set, and the SignatureDef
# ${signature_def_key}.
@@ -80,7 +45,7 @@ endfunction()
# ${CMAKE_CURRENT_BINARY_DIR}. The generated header will define a C++ class
# called ${cpp_class} - which may be a namespace-qualified class name.
function(tf_compile model tag_set signature_def_key fname cpp_class hdr_file obj_file)
- mlgo_get_absolute_path(${model} ${CMAKE_CURRENT_BINARY_DIR} LLVM_ML_MODELS_ABSOLUTE)
+ get_absolute_path(${model} ${CMAKE_CURRENT_BINARY_DIR} LLVM_ML_MODELS_ABSOLUTE)
message("Using model at " ${LLVM_ML_MODELS_ABSOLUTE})
add_custom_command(OUTPUT ${obj_file} ${hdr_file}
COMMAND ${TENSORFLOW_AOT_COMPILER} aot_compile_cpu
@@ -114,21 +79,35 @@ function(tf_find_and_compile model default_url default_path test_model_generator
set(override_object ${LLVM_OVERRIDE_MODEL_OBJECT_${fname_allcaps}})
# If the user specified overrides, that indicates intent to use AOT and we
# don't care what the model path is
- if (override_header AND override_object)
- if (EXISTS ${override_header} AND EXISTS ${override_object})
- configure_file(${override_header} ${hdr_file} COPYONLY)
- configure_file(${override_object} ${obj_file} COPYONLY)
- message(STATUS "Using provided header " ${hdr_file} " and object " ${obj_file} "
- files for model " ${fname})
- set(GENERATED_OBJS ${GENERATED_OBJS} ${obj_file})
- set(GENERATED_HEADERS ${GENERATED_HEADERS} ${hdr_file})
- endif()
+ if (EXISTS "${override_header}" AND EXISTS "${override_object}")
+ configure_file(${override_header} ${hdr_file} COPYONLY)
+ configure_file(${override_object} ${obj_file} COPYONLY)
+ message(STATUS "Using provided header " ${hdr_file} " and object " ${obj_file} "
+ files for model " ${fname})
+ set(GENERATED_OBJS ${GENERATED_OBJS} ${obj_file})
+ set(GENERATED_HEADERS ${GENERATED_HEADERS} ${hdr_file})
+ elseif("${model}" STREQUAL "none")
+ message(STATUS "Will skip enabling mlgo for ${fname}")
+ return()
else()
- mlgo_resolve_model(${model} ${default_url} ${default_path}
- ${test_model_generator} ${fname} should_skip LLVM_ML_MODELS_ABSOLUTE)
- if (should_skip)
- return()
+ if ("${model}" STREQUAL "download")
+ # Crash if the user wants to download a model but a URL is set to "TO_BE_UPDATED"
+ if ("${default_url}" STREQUAL "<UNSPECIFIED>")
+ message(FATAL_ERROR "Model path was set to 'download' but there is no"
+ " model url currently specified in cmake. You can generate a model"
+ " using, for example, the tools at http://github.com/google/ml-compiler-opt."
+ " Some reference models are also periodically released there.")
+ endif()
+
+ set(model ${default_url})
endif()
+
+ if ("${model}" STREQUAL "autogenerate")
+ set(model ${default_path}-autogenerated)
+ generate_mock_model(${test_model_generator} ${model})
+ endif()
+
+ tf_get_model(${model} LLVM_ML_MODELS_ABSOLUTE)
tf_compile(${LLVM_ML_MODELS_ABSOLUTE} ${tag_set} ${signature_def_key} ${fname} ${cpp_class} ${hdr_file} ${obj_file})
endif()
@@ -139,18 +118,157 @@ function(tf_find_and_compile model default_url default_path test_model_generator
endfunction()
set(LLVM_MLGO_MLIR_PASS_PIPELINE
- "builtin.module(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},buffer-deallocation-pipeline,func.func(convert-linalg-to-loops),expand-strided-metadata,canonicalize,memref-elide-reinterpret-cast,convert-to-emitc,math-expand-ops{ops=rsqrt},arith-expand,convert-math-to-emitc,convert-arith-to-emitc,wrap-emitc-func-in-class,mlgo-add-reflection-map{included-field-attrs=tf_saved_model.index_path})"
+ "builtin.module( \
+ 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 \
+ }, \
+ buffer-deallocation-pipeline, \
+ func.func( \
+ convert-linalg-to-loops \
+ ), \
+ expand-strided-metadata, \
+ canonicalize, \
+ memref-elide-reinterpret-cast, \
+ convert-to-emitc, \
+ math-expand-ops{ \
+ ops=rsqrt \
+ }, \
+ arith-expand, \
+ convert-math-to-emitc, \
+ convert-arith-to-emitc, \
+ wrap-emitc-func-in-class, \
+ mlgo-add-reflection-map{ \
+ included-field-attrs=tf_saved_model.index_path \
+ })"
CACHE STRING
"MLIR pass pipeline used to lower MLGO TOSA models to EmitC.")
-# Lower an MLGO model with the in-tree MLIR flow and emit the generated model
-# body as a header under llvm/include. The generated header is not a complete
-# serving interface by itself; LLVM-side wrappers in lib/Analysis and
+# Resolve a TOSA MLIR model used by the MLIR-based MLGO flow.
+#
+# Unlike tf_get_model(), this function expects a single .mlir file rather
+# than a SavedModel directory or an archive containing a SavedModel.
+function(get_tosa_model model final_path)
+ string(FIND "${model}" "http:" pos_http)
+ string(FIND "${model}" "https:" pos_https)
+
+ if (${pos_http} EQUAL 0 OR ${pos_https} EQUAL 0)
+ string(FIND "${model}" "/" fname_start REVERSE)
+ math(EXPR fname_start "${fname_start}+1")
+ string(SUBSTRING "${model}" ${fname_start} -1 fname)
+
+ if ("${fname}" STREQUAL "")
+ set(fname "model_tosa.mlir")
+ endif()
+
+ set(downloaded_model
+ "${CMAKE_CURRENT_BINARY_DIR}/${fname}")
+ message(STATUS "Downloading TOSA model ${model}")
+
+ file(DOWNLOAD
+ "${model}"
+ "${downloaded_model}"
+ STATUS download_status)
+
+ list(GET download_status 0 download_error)
+ list(GET download_status 1 download_message)
+ if (NOT download_error EQUAL 0)
+ message(FATAL_ERROR
+ "Could not download TOSA model '${model}': ${download_message}")
+ endif()
+
+ set(${final_path} "${downloaded_model}" PARENT_SCOPE)
+ else()
+ get_absolute_path(
+ "${model}" "${CMAKE_CURRENT_BINARY_DIR}" model_path)
+ set(${final_path} "${model_path}" PARENT_SCOPE)
+ endif()
+endfunction()
+
+# Resolve a TOSA model for LLVM_USE_MLIR_FOR_MLGO.
+#
+# In this flow, "autogenerate" selects an in-tree TOSA integration-test
+# model. It does not invoke the TensorFlow mock-model generator.
+function(resolve_tosa_model model default_url test_model
+ model_label should_skip final_path)
+ if ("${model}" STREQUAL "none")
+ message(STATUS "Will skip enabling mlgo for ${model_label}")
+ set(${should_skip} TRUE PARENT_SCOPE)
+ return()
+ endif()
+
+ if ("${model}" STREQUAL "download")
+ if ("${default_url}" STREQUAL "<UNSPECIFIED>"
+ OR "${default_url}" STREQUAL "")
+ message(FATAL_ERROR
+ "Model path was set to 'download', but no TOSA model URL is "
+ "configured for ${model_label}.")
+ endif()
+ set(model "${default_url}")
+ endif()
+
+ if ("${model}" STREQUAL "autogenerate")
+ get_absolute_path(
+ "${test_model}" "${CMAKE_CURRENT_SOURCE_DIR}" model_input)
+ else()
+ get_tosa_model("${model}" model_input)
+ endif()
+
+ if (NOT EXISTS "${model_input}")
+ message(FATAL_ERROR
+ "TOSA model for ${model_label} does not exist: ${model_input}")
+ endif()
+
+ if (IS_DIRECTORY "${model_input}")
+ message(FATAL_ERROR
+ "Expected a TOSA MLIR file for ${model_label}, but got a directory: "
+ "${model_input}")
+ endif()
+
+ set(${should_skip} FALSE PARENT_SCOPE)
+ set(${final_path} "${model_input}" PARENT_SCOPE)
+endfunction()
+
+# Lower a TOSA MLIR model with the in-tree MLIR flow and emit the generated
+# model body as a header under llvm/include. The generated header is not a
+# complete serving interface by itself; LLVM-side wrappers in lib/Analysis and
# lib/CodeGen provide the stable API expected by ReleaseModeModelRunner.
function(mlir_find_and_compile model default_url default_path
- test_model_generator header_relative_path)
- mlgo_resolve_model(${model} ${default_url} ${default_path}
- ${test_model_generator} ${header_relative_path} should_skip model_input)
+ test_model_generator fname header_relative_path)
+ set(prefix ${CMAKE_CURRENT_BINARY_DIR}/${fname})
+ set(inc_file ${LLVM_INCLUDE_DIR}/${header_relative_path})
+ string(TOUPPER ${fname} fname_allcaps)
+ set(override_impl ${LLVM_OVERRIDE_MODEL_IMPLEMENTATION_${fname_allcaps}})
+ # If the user specified an override, that indicates intent to use the MLIR
+ # model and we do not care what the model path is. The supplied file must be
+ # suitable for textual inclusion by the LLVM-side model wrapper.
+ if (EXISTS "${override_impl}")
+ get_filename_component(inc_dir ${inc_file} DIRECTORY)
+ file(MAKE_DIRECTORY ${inc_dir})
+ configure_file(${override_impl} ${inc_file} COPYONLY)
+ message(STATUS
+ "Using provided implementation ${override_impl} for model ${fname}")
+ set(GENERATED_SOURCES ${GENERATED_SOURCES} ${inc_file})
+ elseif ("${model}" STREQUAL "none")
+ message(STATUS "Will skip enabling mlgo for ${fname}")
+ return()
+ else()
+ resolve_tosa_model(
+ "${model}" "${default_url}" "${test_model_generator}"
+ "${header_relative_path}" should_skip model_input)
+ endif()
if (should_skip)
return()
endif()
@@ -182,54 +300,25 @@ function(mlir_find_and_compile model default_url default_path
endif()
get_filename_component(mlir_translate_path ${mlir_translate_path} ABSOLUTE
BASE_DIR ${CMAKE_BINARY_DIR})
-
- set(tosa_converter_path ${LLVM_TFLITE_TOSA_CONVERTER})
- if (NOT tosa_converter_path)
- find_program(tosa_converter_path
- NAMES tosa-converter-for-tflite
- DOC "Path to the tosa-converter-for-tflite executable")
- endif()
- if (NOT tosa_converter_path)
- message(FATAL_ERROR
- "LLVM_USE_MLIR_FOR_MLGO requires 'tosa-converter-for-tflite' to be "
- "available on PATH.")
- endif()
- get_filename_component(tosa_converter_path ${tosa_converter_path} ABSOLUTE
- BASE_DIR ${CMAKE_BINARY_DIR})
-
set(generated_header ${LLVM_INCLUDE_DIR}/${header_relative_path})
get_filename_component(generated_header_dir ${generated_header} DIRECTORY)
get_filename_component(generated_header_stem ${generated_header} NAME_WE)
- set(tflite_dir
- ${CMAKE_CURRENT_BINARY_DIR}/${generated_header_stem}-tflite)
- set(tflite_model ${tflite_dir}/model.tflite)
- set(tosa_mlir
- ${CMAKE_CURRENT_BINARY_DIR}/${generated_header_stem}-tosa.mlir)
set(lowered_mlir
- ${CMAKE_CURRENT_BINARY_DIR}/${generated_header_stem}-emitc.mlir)
+ ${CMAKE_CURRENT_BINARY_DIR}/${generated_header_stem}_emitc.mlir)
add_custom_command(
OUTPUT ${generated_header}
- BYPRODUCTS ${tflite_model} ${tosa_mlir} ${lowered_mlir}
+ BYPRODUCTS ${lowered_mlir}
COMMAND ${CMAKE_COMMAND} -E make_directory ${generated_header_dir}
- COMMAND ${Python3_EXECUTABLE}
- ${LLVM_MAIN_SRC_DIR}/lib/Analysis/models/saved-model-to-tflite.py
- ${model_input}
- ${tflite_dir}
- COMMAND ${tosa_converter_path}
- ${tflite_model}
- --text
- -o ${tosa_mlir}
COMMAND ${mlir_opt_path}
--pass-pipeline=${LLVM_MLGO_MLIR_PASS_PIPELINE}
- ${tosa_mlir}
+ ${model_input}
-o ${lowered_mlir}
COMMAND ${mlir_translate_path}
-mlir-to-cpp
${lowered_mlir}
-o ${generated_header}
DEPENDS
- ${model_input}/saved_model.pb
- ${LLVM_MAIN_SRC_DIR}/lib/Analysis/models/saved-model-to-tflite.py
+ ${model_input}
${mlir_opt_path}
${mlir_translate_path}
VERBATIM)
diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt
index aa84630c1e562..c49a15339fbaa 100644
--- a/llvm/lib/Analysis/CMakeLists.txt
+++ b/llvm/lib/Analysis/CMakeLists.txt
@@ -23,7 +23,8 @@ if (LLVM_BUILD_INLINERSIZEMODEL)
${LLVM_INLINER_MODEL_PATH}
${LLVM_INLINER_MODEL_CURRENT_URL}
${LLVM_INLINER_MODEL_PATH_DEFAULT}
- "models/gen-inline-oz-test-model.py"
+ "${LLVM_MAIN_SRC_DIR}/../mlir/test/Integration/Dialect/EmitC/inline-oz-test-model-tosa.mlir"
+ InlinerSizeModel
"llvm/Analysis/MLIRInlinerSizeModel.inc"
)
else()
diff --git a/llvm/lib/CodeGen/CMakeLists.txt b/llvm/lib/CodeGen/CMakeLists.txt
index e90f50b36c8e2..0a80f69c53816 100644
--- a/llvm/lib/CodeGen/CMakeLists.txt
+++ b/llvm/lib/CodeGen/CMakeLists.txt
@@ -18,7 +18,8 @@ if (LLVM_BUILD_REGALLOCEVICTMODEL)
${LLVM_RAEVICT_MODEL_PATH}
${LLVM_RAEVICT_MODEL_CURRENT_URL}
${LLVM_RAEVICT_MODEL_PATH_DEFAULT}
- "../Analysis/models/gen-regalloc-eviction-test-model.py"
+ "${LLVM_MAIN_SRC_DIR}/../mlir/test/Integration/Dialect/EmitC/regalloc-eviction-test-model-tosa.mlir"
+ RegAllocEvictModel
"llvm/CodeGen/MLIRRegAllocEvictModel.inc"
)
else()
More information about the llvm-commits
mailing list