diff --git a/.github/workflows/android.yml b/.github/workflows/android.yml index 64f5de7e78dcc..b3d6b895a97fb 100644 --- a/.github/workflows/android.yml +++ b/.github/workflows/android.yml @@ -79,8 +79,8 @@ jobs: run: | set -e -x BINARY_SIZE_THRESHOLD_ARGS="" - echo "Binary size threshold in bytes: 1595392" - BINARY_SIZE_THRESHOLD_ARGS="--threshold_size_in_bytes 1595392" + echo "Binary size threshold in bytes: 1596416" + BINARY_SIZE_THRESHOLD_ARGS="--threshold_size_in_bytes 1596416" # Ensure ANDROID_NDK_HOME is available and get its real path if [ -z "$ANDROID_NDK_HOME" ]; then diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h index 977fd1c54571b..bc7237ee570ff 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h @@ -99,6 +99,11 @@ static const char* const kOrtSessionOptionsEnableCastChainElimination = "optimiz // Its default value is "0". static const char* const kOrtSessionOptionsDisableAheadOfTimeFunctionInlining = "session.disable_aot_function_inlining"; +// Limits cumulative model-local function expansion across AOT and fallback inlining. +// Values must be positive decimal integers. Defaults are 1,000,000 nodes and 1 GiB of serialized node payload. +static const char* const kOrtSessionOptionsFunctionExpansionNodeLimit = "session.function_expansion_node_limit"; +static const char* const kOrtSessionOptionsFunctionExpansionByteLimit = "session.function_expansion_byte_limit"; + #ifdef ENABLE_TRAINING // Specifies a path of the file containing a list of memory optimization configurations. // The value should be a string indicating the file path of the config file. diff --git a/onnxruntime/core/framework/graph_partitioner.cc b/onnxruntime/core/framework/graph_partitioner.cc index 489c0c7dc9da9..821c9c1bc4e0b 100644 --- a/onnxruntime/core/framework/graph_partitioner.cc +++ b/onnxruntime/core/framework/graph_partitioner.cc @@ -1207,13 +1207,41 @@ static Status PartitionOnnxFormatModelImpl(Graph& graph, FuncManager& func_mgr, return Status::OK(); } +struct FunctionExpansionCost { + size_t node_count; + size_t proto_bytes; +}; + +enum class FunctionExpansionLimit { + kNone, + kNodes, + kProtoBytes, +}; + +static Status GetFunctionExpansionCost(const Node& node, FunctionExpansionCost& cost); + +static FunctionExpansionLimit TryChargeFunctionExpansion(const FunctionExpansionCost& cost, + size_t node_limit, + size_t& expanded_node_count, + size_t byte_limit, + size_t& expanded_proto_bytes); + // expand any nodes that have an ONNX function definition but no matching ORT kernel -static Status InlineNodes(Graph& graph, bool& modified_graph, LayeringIndex* layering_index) { +static Status InlineNodes(Graph& graph, + bool& modified_graph, + LayeringIndex* layering_index, + const logging::Logger& logger, + size_t expansion_node_limit, + size_t& expanded_node_count, + size_t expansion_byte_limit, + size_t& expanded_proto_bytes) { // recurse into nested graphs first so we process from bottom up for (auto& node : graph.Nodes()) { for (auto& entry : node.GetAttributeNameToMutableSubgraphMap()) { Graph* subgraph = entry.second; - ORT_RETURN_IF_ERROR(InlineNodes(*subgraph, modified_graph, layering_index)); + ORT_RETURN_IF_ERROR(InlineNodes(*subgraph, modified_graph, layering_index, logger, + expansion_node_limit, expanded_node_count, + expansion_byte_limit, expanded_proto_bytes)); } } @@ -1233,6 +1261,23 @@ static Status InlineNodes(Graph& graph, bool& modified_graph, LayeringIndex* lay InlinedVector new_node_indices; for (auto* node : nodes_to_inline) { + FunctionExpansionCost expansion_cost{}; + ORT_RETURN_IF_ERROR(GetFunctionExpansionCost(*node, expansion_cost)); + const auto limit_exceeded = TryChargeFunctionExpansion(expansion_cost, + expansion_node_limit, + expanded_node_count, + expansion_byte_limit, + expanded_proto_bytes); + if (limit_exceeded != FunctionExpansionLimit::kNone) { + const auto function_id = + function_utils::GetFunctionIdentifier(node->Domain(), node->OpType(), node->Overload()); + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "Function inlining exceeded the configured cumulative ", + limit_exceeded == FunctionExpansionLimit::kNodes ? "node" : "protobuf", + " expansion limit while expanding '", function_id, "'."); + } + // Check for an effective layering assignment: either from an explicit annotation // on the node, or from an inherited assignment via the LayeringIndex (e.g., a function // call node inside an annotated If/Loop subgraph that inherited its parent's rule). @@ -1277,6 +1322,108 @@ static Status InlineNodes(Graph& graph, bool& modified_graph, LayeringIndex* lay return Status::OK(); } +constexpr size_t kDefaultFunctionExpansionNodeLimit = 1'000'000; +constexpr size_t kDefaultFunctionExpansionByteLimit = 1024ULL * 1024ULL * 1024ULL; + +static size_t CountNodesIncludingSubgraphs(const ONNX_NAMESPACE::GraphProto& graph); + +static size_t CountNodesIncludingSubgraphs(const ONNX_NAMESPACE::AttributeProto& attribute) { + SafeInt node_count = 0; + if (attribute.has_g()) { + node_count += CountNodesIncludingSubgraphs(attribute.g()); + } + for (const auto& attribute_graph : attribute.graphs()) { + node_count += CountNodesIncludingSubgraphs(attribute_graph); + } + + return node_count; +} + +static size_t CountNodesIncludingSubgraphs(const ONNX_NAMESPACE::GraphProto& graph) { + SafeInt node_count = graph.node_size(); + for (const auto& node : graph.node()) { + for (const auto& attribute : node.attribute()) { + node_count += CountNodesIncludingSubgraphs(attribute); + } + } + + return node_count; +} + +static Status GetFunctionExpansionCost(const Node& node, FunctionExpansionCost& cost) { + if (const auto* function_body = node.GetFunctionBody()) { + const auto graph_proto = function_body->Body().ToGraphProto(); + SafeInt proto_bytes = 0; + for (const auto& function_node : graph_proto.node()) { + proto_bytes += function_node.ByteSizeLong(); + } + cost = {CountNodesIncludingSubgraphs(graph_proto), proto_bytes}; + return Status::OK(); + } + + ONNX_NAMESPACE::FunctionProto function_proto; + ORT_RETURN_IF_NOT(node.TryGetFunctionProto(function_proto), + "Unable to get function body for node '", node.Name(), "'."); + std::string accounting_prefix = "_inlfunc_" + node.OpType(); + accounting_prefix.append(32, '_'); + function_utils::Specialize(function_proto, node, accounting_prefix); + SafeInt node_count = function_proto.node_size(); + SafeInt proto_bytes = 0; + for (const auto& function_node : function_proto.node()) { + proto_bytes += function_node.ByteSizeLong(); + for (const auto& attribute : function_node.attribute()) { + node_count += CountNodesIncludingSubgraphs(attribute); + } + } + cost = {node_count, proto_bytes}; + return Status::OK(); +} + +static FunctionExpansionLimit TryChargeFunctionExpansion(const FunctionExpansionCost& cost, + size_t node_limit, + size_t& expanded_node_count, + size_t byte_limit, + size_t& expanded_proto_bytes) { + if (cost.node_count > node_limit - expanded_node_count) { + return FunctionExpansionLimit::kNodes; + } + if (cost.proto_bytes > byte_limit - expanded_proto_bytes) { + return FunctionExpansionLimit::kProtoBytes; + } + + expanded_node_count += cost.node_count; + expanded_proto_bytes += cost.proto_bytes; + return FunctionExpansionLimit::kNone; +} + +static Status InitializeFunctionExpansionLimits(const ConfigOptions& config_options, + bool& initialized, + size_t& node_limit, + size_t& byte_limit) { + if (initialized) { + return Status::OK(); + } + + const auto parse_limit = [&config_options](const char* config_key, + size_t default_value, + size_t& value) -> Status { + const auto config_value = config_options.GetConfigOrDefault(config_key, std::to_string(default_value)); + if (!TryParseStringWithClassicLocale(config_value, value) || value == 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Invalid positive integer value '", config_value, + "' for session configuration '", config_key, "'."); + } + return Status::OK(); + }; + + ORT_RETURN_IF_ERROR(parse_limit(kOrtSessionOptionsFunctionExpansionNodeLimit, + kDefaultFunctionExpansionNodeLimit, node_limit)); + ORT_RETURN_IF_ERROR(parse_limit(kOrtSessionOptionsFunctionExpansionByteLimit, + kDefaultFunctionExpansionByteLimit, byte_limit)); + initialized = true; + return Status::OK(); +} + static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_providers, const KernelRegistryManager& kernel_registry_mgr, Graph& graph, @@ -1284,7 +1431,11 @@ static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_provide const logging::Logger& logger, const CheckLoadCancellationFn& check_load_cancellation_fn, InlinedHashSet& not_inlined, - size_t& inlined_count) { + size_t& inlined_count, + size_t expansion_node_limit, + size_t& expanded_node_count, + size_t expansion_byte_limit, + size_t& expanded_proto_bytes) { // handle testing edge case where optimizers or constant lifting results in graph with no nodes. // doing it here saves all providers checking for this in GetCapability if (graph.NumberOfNodes() == 0) { @@ -1302,7 +1453,11 @@ static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_provide logger, check_load_cancellation_fn, not_inlined, - inlined_count)); + inlined_count, + expansion_node_limit, + expanded_node_count, + expansion_byte_limit, + expanded_proto_bytes)); } } @@ -1355,6 +1510,23 @@ static Status InlineFunctionsAOTImpl(const ExecutionProviders& execution_provide auto* node = graph.GetNode(node_index); if (node != nullptr) { if (claimed_by_ep.count(node_index) == 0) { + auto function_id = function_utils::GetFunctionIdentifier(node->Domain(), node->OpType(), node->Overload()); + FunctionExpansionCost expansion_cost{}; + ORT_RETURN_IF_ERROR(GetFunctionExpansionCost(*node, expansion_cost)); + const auto limit_exceeded = TryChargeFunctionExpansion(expansion_cost, + expansion_node_limit, + expanded_node_count, + expansion_byte_limit, + expanded_proto_bytes); + if (limit_exceeded != FunctionExpansionLimit::kNone) { + LOGS(logger, WARNING) << "AOT function inlining reached the cumulative " + << (limit_exceeded == FunctionExpansionLimit::kNodes ? "node" : "protobuf") + << " expansion limit. " + << "Retaining function call '" << function_id + << "' for execution-provider partitioning."; + ORT_IGNORE_RETURN_VALUE(not_inlined.insert(function_id)); + continue; + } ORT_RETURN_IF_ERROR(graph.InlineFunction(*node)); ++inlined_count; } else { @@ -1523,7 +1695,11 @@ static Status PartitionOnnxFormatModel(const PartitionParams& partition_params, KernelRegistryManager& kernel_registry_manager, const std::optional& acc_map, const GraphOptimizerRegistry& graph_optimizer_registry, - const logging::Logger& logger, bool disable_model_compile) { // Added arg + const logging::Logger& logger, bool disable_model_compile, + size_t expansion_node_limit, + size_t& expanded_node_count, + size_t expansion_byte_limit, + size_t& expanded_proto_bytes) { // Added arg bool modified_graph = false; auto& graph = partition_params.graph.get(); @@ -1570,7 +1746,9 @@ static Status PartitionOnnxFormatModel(const PartitionParams& partition_params, // expand any nodes that have an ONNX function definition but no matching ORT kernel. modified_graph = false; - ORT_RETURN_IF_ERROR(InlineNodes(graph, modified_graph, partition_params.layering_index)); + ORT_RETURN_IF_ERROR(InlineNodes(graph, modified_graph, partition_params.layering_index, logger, + expansion_node_limit, expanded_node_count, + expansion_byte_limit, expanded_proto_bytes)); // Resolve and rerun graph partitioning and inlining if there was a change if (modified_graph) { @@ -1730,6 +1908,7 @@ static Status PartitionOrtFormatModel(const PartitionParams& partition_params, Status GraphPartitioner::InlineFunctionsAOT(Model& model, const ExecutionProviders& execution_providers, const KernelRegistryManager& kernel_registry_manager, + const ConfigOptions& config_options, const logging::Logger& logger) const { const auto local_functions_num = model.GetModelLocalFunctionTemplates().size(); const bool is_there_local_functions = local_functions_num > 0; @@ -1740,6 +1919,11 @@ Status GraphPartitioner::InlineFunctionsAOT(Model& model, } auto check_load_cancellation_fn = [this]() -> bool { return IsLoadCancellationFlagSet(); }; + ORT_RETURN_IF_ERROR(InitializeFunctionExpansionLimits( + config_options, + function_expansion_limits_initialized_, + function_expansion_node_limit_, + function_expansion_byte_limit_)); auto& graph = model.MainGraph(); InlinedHashSet not_inlined; @@ -1752,7 +1936,11 @@ Status GraphPartitioner::InlineFunctionsAOT(Model& model, logger, check_load_cancellation_fn, not_inlined, - inlined_count)); + inlined_count, + function_expansion_node_limit_, + expanded_function_node_count_, + function_expansion_byte_limit_, + expanded_function_proto_bytes_)); if (inlined_count == 0) { break; @@ -1848,11 +2036,20 @@ Status GraphPartitioner::Partition(Graph& graph, FuncManager& func_mgr, // The map is empty if not created if not enabled std::optional ep_acc_map; ORT_RETURN_IF_ERROR(CreateAccountants(config_options, graph.ModelPath(), ep_acc_map)); + ORT_RETURN_IF_ERROR(InitializeFunctionExpansionLimits( + config_options, + function_expansion_limits_initialized_, + function_expansion_node_limit_, + function_expansion_byte_limit_)); bool disable_model_compile = config_options.GetConfigOrDefault(kOrtSessionOptionsDisableModelCompile, "0") == "1"; ORT_RETURN_IF_ERROR(PartitionOnnxFormatModel(partition_params, mode, providers_, kernel_registry_mgr_, ep_acc_map, *graph_optimizer_registry_, logger, - disable_model_compile)); // Pass param + disable_model_compile, + function_expansion_node_limit_, + expanded_function_node_count_, + function_expansion_byte_limit_, + expanded_function_proto_bytes_)); // Pass param if (ep_acc_map.has_value()) { for (const auto& [ep_type, accountant] : *ep_acc_map) { diff --git a/onnxruntime/core/framework/graph_partitioner.h b/onnxruntime/core/framework/graph_partitioner.h index 6cf547c31107a..ad91b23282cea 100644 --- a/onnxruntime/core/framework/graph_partitioner.h +++ b/onnxruntime/core/framework/graph_partitioner.h @@ -103,6 +103,7 @@ class GraphPartitioner { Status InlineFunctionsAOT(Model& model, const ExecutionProviders& execution_providers, const KernelRegistryManager& kernel_registry_manager, + const ConfigOptions& config_options, const logging::Logger& logger) const; #endif @@ -114,6 +115,14 @@ class GraphPartitioner { std::unique_ptr graph_optimizer_registry_; CheckLoadCancellationFn check_load_cancellation_fn_; OnPartitionAssignmentFunction on_partition_assignment_fn_; +#ifndef ORT_MINIMAL_BUILD + // Shared by AOT and fallback inlining so neither path can bypass the cumulative limit. + mutable bool function_expansion_limits_initialized_ = false; + mutable size_t function_expansion_node_limit_ = 0; + mutable size_t function_expansion_byte_limit_ = 0; + mutable size_t expanded_function_node_count_ = 0; + mutable size_t expanded_function_proto_bytes_ = 0; +#endif }; } // namespace onnxruntime diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 16202678e0153..6ced42d6fd229 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -1625,6 +1625,7 @@ common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph, bool ORT_RETURN_IF_ERROR_SESSIONID_(partitioner.InlineFunctionsAOT(*model_, execution_providers_, kernel_registry_manager_, + session_options_.config_options, *session_logger_)); } diff --git a/onnxruntime/test/framework/function_test.cc b/onnxruntime/test/framework/function_test.cc index ab3317885b5d5..ad6ccfbe38fc2 100644 --- a/onnxruntime/test/framework/function_test.cc +++ b/onnxruntime/test/framework/function_test.cc @@ -4,6 +4,7 @@ #include "gmock/gmock.h" #include "gtest/gtest.h" +#include #include #include "core/graph/onnx_protobuf.h" @@ -16,9 +17,13 @@ #include "core/graph/model.h" #include "core/graph/model_helpers.h" #include "core/providers/cpu/cpu_execution_provider.h" +#include "core/providers/partitioning_utils.h" +#include "core/session/environment.h" #include "core/session/inference_session.h" +#include "core/session/onnxruntime_session_options_config_keys.h" #include "test/common/tensor_op_test_utils.h" +#include "test/capturing_sink.h" #include "test/unittest_util/framework_test_utils.h" #include "test/internal_testing_ep/internal_testing_execution_provider.h" #include "test/test_environment.h" @@ -120,6 +125,270 @@ static Status LoadModel(const char* source) { return session_object.Load(sstr); } +static ONNX_NAMESPACE::ModelProto CreateFunctionExpansionModel(size_t body_node_count, size_t call_count) { + ONNX_NAMESPACE::ModelProto model; + model.set_ir_version(8); + auto* onnx_opset = model.add_opset_import(); + onnx_opset->set_domain(""); + onnx_opset->set_version(13); + auto* local_opset = model.add_opset_import(); + local_opset->set_domain("local"); + local_opset->set_version(1); + + auto set_float_value = [](ONNX_NAMESPACE::ValueInfoProto& value, const std::string& name) { + value.set_name(name); + auto* tensor_type = value.mutable_type()->mutable_tensor_type(); + tensor_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + tensor_type->mutable_shape()->add_dim()->set_dim_value(1); + }; + + auto* graph = model.mutable_graph(); + graph->set_name("function_expansion"); + set_float_value(*graph->add_input(), "input"); + + auto* function = model.add_functions(); + function->set_domain("local"); + function->set_name("FunctionToExpand"); + function->add_input("function_input"); + function->add_output("function_output"); + auto* function_opset = function->add_opset_import(); + function_opset->set_domain(""); + function_opset->set_version(13); + + std::string previous_value = "function_input"; + for (size_t i = 0; i < body_node_count; ++i) { + auto* node = function->add_node(); + node->set_op_type("Identity"); + node->add_input(previous_value); + previous_value = i + 1 == body_node_count ? "function_output" : "body_" + std::to_string(i); + node->add_output(previous_value); + } + + previous_value = "input"; + for (size_t i = 0; i < call_count; ++i) { + auto* node = graph->add_node(); + node->set_domain("local"); + node->set_op_type("FunctionToExpand"); + node->add_input(previous_value); + previous_value = "call_" + std::to_string(i); + node->add_output(previous_value); + } + set_float_value(*graph->add_output(), previous_value); + + return model; +} + +static ONNX_NAMESPACE::ModelProto CreateRecursiveFunctionExpansionModel(size_t branch_node_count, + size_t call_count) { + auto model = CreateFunctionExpansionModel(0, call_count); + model.set_doc_string(std::string(1024 * 1024, 'x')); + + auto* function = model.mutable_functions(0); + auto* condition = function->add_node(); + condition->set_op_type("Constant"); + condition->add_output("condition"); + auto* condition_value = condition->add_attribute(); + condition_value->set_name("value"); + condition_value->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_TENSOR); + condition_value->mutable_t()->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_BOOL); + condition_value->mutable_t()->add_int32_data(1); + + auto* if_node = function->add_node(); + if_node->set_op_type("If"); + if_node->add_input("condition"); + if_node->add_output("function_output"); + + for (const auto* attribute_name : {"then_branch", "else_branch"}) { + auto* attribute = if_node->add_attribute(); + attribute->set_name(attribute_name); + attribute->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_GRAPH); + auto* branch = attribute->mutable_g(); + branch->set_name(attribute_name); + + std::string previous_value = "function_input"; + for (size_t i = 0; i < branch_node_count; ++i) { + auto* node = branch->add_node(); + node->set_op_type("Identity"); + node->add_input(previous_value); + previous_value = "branch_" + std::to_string(i); + node->add_output(previous_value); + } + + auto* output = branch->add_output(); + output->set_name(previous_value); + auto* tensor_type = output->mutable_type()->mutable_tensor_type(); + tensor_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + tensor_type->mutable_shape()->add_dim()->set_dim_value(1); + } + + return model; +} + +struct FunctionExpansionTestOptions { + bool claim_first_function_call = false; + bool disable_aot_inlining = false; + std::optional node_limit; + std::optional byte_limit; +}; + +static Status InitializeFunctionExpansionModel( + ONNX_NAMESPACE::ModelProto model, + std::vector& log_messages, + const FunctionExpansionTestOptions& options = {}) { + std::string serialized_model; + ORT_RETURN_IF_NOT(model.SerializeToString(&serialized_model), + "Failed to serialize function expansion model."); + + SessionOptions session_options; + if (options.disable_aot_inlining) { + ORT_RETURN_IF_ERROR(session_options.config_options.AddConfigEntry( + kOrtSessionOptionsDisableAheadOfTimeFunctionInlining, "1")); + } + if (options.node_limit.has_value()) { + ORT_RETURN_IF_ERROR(session_options.config_options.AddConfigEntry( + kOrtSessionOptionsFunctionExpansionNodeLimit, + std::to_string(*options.node_limit).c_str())); + } + if (options.byte_limit.has_value()) { + ORT_RETURN_IF_ERROR(session_options.config_options.AddConfigEntry( + kOrtSessionOptionsFunctionExpansionByteLimit, + std::to_string(*options.byte_limit).c_str())); + } + + auto capturing_sink = std::make_unique(); + auto* capturing_sink_ptr = capturing_sink.get(); + auto logging_manager = std::make_unique( + std::move(capturing_sink), logging::Severity::kWARNING, false, + logging::LoggingManager::InstanceType::Temporal); + std::unique_ptr environment; + ORT_RETURN_IF_ERROR(Environment::Create(std::move(logging_manager), environment)); + + InferenceSession session{session_options, *environment}; + if (options.claim_first_function_call) { + class FirstFunctionCallExecutionProvider final + : public internal_testing_ep::InternalTestingExecutionProvider { + public: + FirstFunctionCallExecutionProvider() + : InternalTestingExecutionProvider({}, {}, DataLayout::NCHW) {} + + std::vector> GetCapability( + const GraphViewer& graph_view, + const IKernelLookup&, + const GraphOptimizerRegistry&, + IResourceAccountant*) const override { + for (const auto node_index : graph_view.GetNodesInTopologicalOrder()) { + const auto* node = graph_view.GetNode(node_index); + if (node != nullptr && node->CanBeInlined()) { + std::vector> capabilities; + capabilities.push_back(utils::MakeComputeCapability( + graph_view, std::vector{node}, + [node_index]() { return "FirstFunctionCall_" + std::to_string(node_index); }, + Type(), false)); + return capabilities; + } + } + + return {}; + } + }; + + ORT_RETURN_IF_ERROR(session.RegisterExecutionProvider( + std::make_unique())); + } + + std::istringstream stream(serialized_model); + ORT_RETURN_IF_ERROR(session.Load(stream)); + const auto status = session.Initialize(); + log_messages = capturing_sink_ptr->Messages(); + return status; +} + +TEST(FunctionTest, AotInliningLimitsFunctionExpansionByNodeCount) { + auto model = CreateRecursiveFunctionExpansionModel(20, 14); + std::vector log_messages; + const auto status = InitializeFunctionExpansionModel( + std::move(model), log_messages, + {.claim_first_function_call = false, + .disable_aot_inlining = false, + .node_limit = 500, + .byte_limit = std::nullopt}); + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("node expansion limit")); + EXPECT_THAT(log_messages, testing::Contains(testing::HasSubstr("node expansion limit"))); + EXPECT_THAT(log_messages, testing::Not(testing::Contains(testing::HasSubstr("protobuf expansion limit")))); +} + +TEST(FunctionTest, AotInliningLimitsUnclaimedCallsSharingClaimedFunction) { + auto model = CreateRecursiveFunctionExpansionModel(20, 15); + std::vector log_messages; + const auto status = InitializeFunctionExpansionModel( + std::move(model), log_messages, + {.claim_first_function_call = true, + .disable_aot_inlining = false, + .node_limit = 500, + .byte_limit = std::nullopt}); + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("node expansion limit")); + EXPECT_THAT(log_messages, testing::Contains(testing::HasSubstr("node expansion limit"))); +} + +TEST(FunctionTest, AotInliningLimitsFunctionExpansionByProtoBytes) { + auto model = CreateFunctionExpansionModel(2, 20); + auto* function = model.mutable_functions(0); + auto* payload_node = function->add_node(); + payload_node->set_op_type("Constant"); + payload_node->add_output("payload"); + auto* payload_attribute = payload_node->add_attribute(); + payload_attribute->set_name("value"); + payload_attribute->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_TENSOR); + auto* payload_tensor = payload_attribute->mutable_t(); + payload_tensor->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_UINT8); + payload_tensor->add_dims(256 * 1024); + payload_tensor->set_raw_data(std::string(256 * 1024, 'x')); + + std::vector log_messages; + const auto status = InitializeFunctionExpansionModel( + std::move(model), log_messages, + {.claim_first_function_call = false, + .disable_aot_inlining = false, + .node_limit = std::nullopt, + .byte_limit = 1024 * 1024}); + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("protobuf expansion limit")); + EXPECT_THAT(log_messages, testing::Contains(testing::HasSubstr("protobuf expansion limit"))); + EXPECT_THAT(log_messages, testing::Not(testing::Contains(testing::HasSubstr("node expansion limit")))); +} + +TEST(FunctionTest, FallbackInliningEnforcesExpansionLimitWhenAotIsDisabled) { + auto model = CreateRecursiveFunctionExpansionModel(20, 14); + std::vector log_messages; + const auto status = InitializeFunctionExpansionModel( + std::move(model), log_messages, + {.claim_first_function_call = false, + .disable_aot_inlining = true, + .node_limit = 500, + .byte_limit = std::nullopt}); + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("node expansion limit")); +} + +TEST(FunctionTest, DefaultExpansionLimitPreservesOrdinaryLargeFunctions) { + auto model = CreateFunctionExpansionModel(200, 24); + std::vector log_messages; + const auto status = InitializeFunctionExpansionModel(std::move(model), log_messages); + EXPECT_TRUE(status.IsOK()) << status.ErrorMessage(); +} + +TEST(FunctionTest, AotInliningIgnoresFunctionMetadataForProtoBytes) { + auto model = CreateFunctionExpansionModel(1, 11); + model.mutable_functions(0)->set_doc_string(std::string(1024 * 1024, 'x')); + + std::vector log_messages; + const auto status = InitializeFunctionExpansionModel(std::move(model), log_messages); + EXPECT_TRUE(status.IsOK()) << status.ErrorMessage(); + EXPECT_THAT(log_messages, testing::Not(testing::Contains(testing::HasSubstr("protobuf expansion limit")))); +} + // A recursive/cyclic chain of model-local functions can be rejected by either layer: // ONNX 1.22+ detects the cycle in its own model checker ("Cycle detected in model-local // function references"), which runs before ORT's equivalent check ("must not be recursive").