From 4b9c93faab222a2d6608857a1fb9e4f45aad9d4c Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 13 Apr 2026 16:30:45 -0700 Subject: [PATCH 1/5] Limit output size of constant folding --- .../onnxruntime_session_options_config_keys.h | 12 ++ .../core/optimizer/constant_folding.cc | 132 ++++++++++++++- .../test/optimizer/graph_transform_test.cc | 154 ++++++++++++++++++ 3 files changed, 297 insertions(+), 1 deletion(-) 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 9941224258506..c15cdc3c93079 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h @@ -481,3 +481,15 @@ static const char* const kOrtSessionOptionsRecordEpGraphAssignmentInfo = "sessio // - "0": disable. (default) // - "1": enable. static const char* const kOrtSessionOptionEpEnableWeightlessEpContextNodes = "ep.enable_weightless_ep_context_nodes"; + +// Maximum total output size in bytes that the constant folding optimizer is allowed to produce per node. +// Prevents malicious models from causing excessive memory allocation during optimization. +// If the estimated or actual output size of a constant-foldable node exceeds this limit, the node will +// not be constant folded and will instead be executed at runtime. +// +// Option values: +// - A positive integer (as string): Maximum allowed output size in bytes per constant-folded node. +// Default is "1073741824" (1 GB). +// - "0": Disable the size limit (not recommended for untrusted models). +static const char* const kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes = + "optimization.constant_folding_max_output_size_in_bytes"; diff --git a/onnxruntime/core/optimizer/constant_folding.cc b/onnxruntime/core/optimizer/constant_folding.cc index cb6d65342bc54..60614034d0ff0 100644 --- a/onnxruntime/core/optimizer/constant_folding.cc +++ b/onnxruntime/core/optimizer/constant_folding.cc @@ -11,6 +11,9 @@ #include "core/optimizer/utils.h" #include "core/framework/op_kernel.h" #include "core/framework/tensorprotoutils.h" +#include "core/session/onnxruntime_session_options_config_keys.h" +#include "core/common/safeint.h" +#include "core/common/parse_string.h" using namespace onnxruntime::common; @@ -140,6 +143,70 @@ static Status ConstantFoldIfNode(Graph& graph, Node& if_node, const logging::Log return status; } +// Default maximum output size per constant-folded node: 1 GB. +// This prevents malicious models from causing excessive memory allocation during optimization. +static constexpr int64_t kDefaultConstantFoldingMaxOutputSizeInBytes = 1024 * 1024 * 1024; + +// Estimate the total output size in bytes for a node using shape inference results. +// Returns -1 if the output size cannot be estimated (e.g., unknown shapes or types). +static int64_t EstimateNodeOutputSizeInBytes(const Node& node) { + SafeInt total_size = 0; + + for (const auto* output_def : node.OutputDefs()) { + if (!output_def->Exists()) { + continue; + } + + const auto* type_proto = output_def->TypeAsProto(); + if (type_proto == nullptr || !utils::HasTensorType(*type_proto)) { + return -1; // Cannot estimate non-tensor or unknown types + } + + const auto* shape = output_def->Shape(); + if (shape == nullptr) { + return -1; // Unknown shape + } + + auto elem_type = static_cast( + type_proto->tensor_type().elem_type()); + size_t element_size = utils::GetElementSizeOfTensor(elem_type); + if (element_size == 0) { + return -1; // Unknown element type + } + + SafeInt num_elements = 1; + for (int i = 0; i < shape->dim_size(); ++i) { + const auto& dim = shape->dim(i); + if (!utils::HasDimValue(dim)) { + return -1; // Symbolic dimension + } + int64_t dim_value = dim.dim_value(); + if (dim_value < 0) { + return -1; // Invalid dimension + } + num_elements *= dim_value; + } + + total_size += num_elements * static_cast(element_size); + } + + return total_size; +} + +// Get the configured max output size from session options, or use the default. +static int64_t GetConstantFoldingMaxOutputSize(const ConfigOptions& config_options) { + std::string max_size_str = config_options.GetConfigOrDefault( + kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes, + std::to_string(kDefaultConstantFoldingMaxOutputSizeInBytes)); + + int64_t max_size = 0; + if (!TryParseStringWithClassicLocale(max_size_str, max_size) || max_size < 0) { + max_size = kDefaultConstantFoldingMaxOutputSizeInBytes; + } + + return max_size; +} + Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const { bool have_updated_nodes = false; GraphViewer graph_viewer(graph); @@ -151,6 +218,8 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level, }; #endif + const int64_t max_output_size = GetConstantFoldingMaxOutputSize(config_options_); + for (NodeIndex i : order) { auto* node = graph.GetNode(i); if (!node || !AllowConstantFolding(*node)) { @@ -233,6 +302,30 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level, } } + // Check if the estimated output size exceeds the configured limit. + // This prevents malicious models from causing excessive memory allocation during constant folding. + if (max_output_size > 0) { + int64_t estimated_size = -1; + try { + estimated_size = EstimateNodeOutputSizeInBytes(*node); + } catch (const std::exception&) { + // SafeInt overflow means the size is astronomically large - definitely skip + LOGS(logger, WARNING) << "Integer overflow while estimating output size of " + << node->OpType() << " node '" << node->Name() + << "'. Skipping constant folding for this node."; + continue; + } + + if (estimated_size > max_output_size) { + LOGS(logger, WARNING) << "Skipping constant folding for " << node->OpType() + << " node '" << node->Name() + << "' because estimated output size (" << estimated_size + << " bytes) exceeds the limit (" << max_output_size << " bytes)."; + continue; + } + // If estimated_size is -1, we couldn't estimate; we'll check actual size after execution. + } + #if !defined(DISABLE_SPARSE_TENSORS) // Create execution frame for executing constant nodes. OptimizerExecutionFrame::Info info({node}, constant_inputs, graph.ModelPath(), execution_provider_, @@ -312,7 +405,17 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level, #pragma warning(disable : 6387) #endif OpKernelContext op_kernel_context(&frame, kernel.get(), /*stream*/ nullptr, nullptr, logger); - ORT_RETURN_IF_ERROR(kernel->Compute(&op_kernel_context)); + + // Wrap Compute in try/catch so that overflows (e.g., SafeInt) or other failures in a + // single node don't abort the entire constant folding pass. + try { + ORT_RETURN_IF_ERROR(kernel->Compute(&op_kernel_context)); + } catch (const std::exception& ex) { + LOGS(logger, WARNING) << "Exception during constant folding of " << node->OpType() + << " node '" << node->Name() << "': " << ex.what() + << ". Skipping constant folding for this node."; + continue; + } #ifdef _WIN32 #pragma warning(pop) #endif @@ -320,6 +423,33 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level, std::vector fetches; ORT_RETURN_IF_ERROR(frame.GetOutputs(fetches)); + // Post-execution size check: verify actual output sizes don't exceed the limit. + // This catches cases where pre-execution shape inference couldn't determine the output size. + if (max_output_size > 0) { + SafeInt actual_total_size = 0; + bool size_exceeded = false; + try { + for (size_t fetch_idx = 0; fetch_idx < fetches.size(); ++fetch_idx) { + if (fetches[fetch_idx].IsTensor()) { + const auto& tensor = fetches[fetch_idx].Get(); + actual_total_size += tensor.SizeInBytes(); + } + } + size_exceeded = actual_total_size > max_output_size; + } catch (const std::exception&) { + // SafeInt overflow means total size is astronomically large + size_exceeded = true; + } + + if (size_exceeded) { + LOGS(logger, WARNING) << "Skipping constant folding for " << node->OpType() + << " node '" << node->Name() + << "' because actual output size exceeds the limit (" + << max_output_size << " bytes)."; + continue; + } + } + // Go over all output node args and substitute them with the newly computed tensors, which will be // added to the graph as initializers. ORT_ENFORCE(fetches.size() == fetch_to_output_idx.size()); diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index ce53703eabf18..52e5feda1c5e8 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -1473,6 +1473,160 @@ TEST_F(GraphTransformationTests, ConstantFoldingIfConstantInliningEdgesWithMiddl ASSERT_TRUE(dest_edges.find(2) != dest_edges.end()); } +// Test that constant folding respects the output size limit and skips nodes +// whose output would exceed it. This is a security measure against malicious +// models that could cause memory exhaustion during optimization. +TEST_F(GraphTransformationTests, ConstantFoldingOutputSizeLimit) { + // Build a model with an Expand node: scalar input [1.0] expanded by shape [1024, 1024]. + // Output = 1024*1024 * 4 bytes = 4 MB of float data. + // With a 1 MB limit, this should NOT be constant folded. + // With a 8 MB limit, this SHOULD be constant folded. + + auto build_model = [&](ModelTestBuilder& builder) { + auto* input_data = builder.MakeInitializer({1}, {1.0f}); + auto* shape_data = builder.MakeInitializer({2}, {1024, 1024}); + auto* output_arg = builder.MakeOutput(); + + builder.AddNode("Expand", {input_data, shape_data}, {output_arg}); + }; + + // Test 1: With a 1 MB limit, the Expand node should NOT be folded (output is ~4 MB). + { + auto pre_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count["Expand"] == 1); + return Status::OK(); + }; + + auto post_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + // Expand should remain because output is too large + TEST_RETURN_IF_NOT(op_to_count["Expand"] == 1); + return Status::OK(); + }; + + std::unique_ptr e = std::make_unique(CPUExecutionProviderInfo()); + ConfigOptions config_options; + // Set limit to 1 MB + ASSERT_STATUS_OK(config_options.AddConfigEntry( + kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes, "1048576")); + + ASSERT_STATUS_OK(TestGraphTransformer(build_model, 14, *logger_, + std::make_unique(*e.get(), false, config_options), + TransformerLevel::Level1, 1, + pre_graph_checker, post_graph_checker)); + } + + // Test 2: With an 8 MB limit, the Expand node SHOULD be folded (output is ~4 MB). + { + auto pre_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count["Expand"] == 1); + return Status::OK(); + }; + + auto post_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + // Expand should be folded since output is within limit + TEST_RETURN_IF_NOT(op_to_count["Expand"] == 0); + return Status::OK(); + }; + + std::unique_ptr e = std::make_unique(CPUExecutionProviderInfo()); + ConfigOptions config_options; + // Set limit to 8 MB + ASSERT_STATUS_OK(config_options.AddConfigEntry( + kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes, "8388608")); + + ASSERT_STATUS_OK(TestGraphTransformer(build_model, 14, *logger_, + std::make_unique(*e.get(), false, config_options), + TransformerLevel::Level1, 1, + pre_graph_checker, post_graph_checker)); + } +} + +// Test that the default constant folding limit prevents extremely large outputs. +TEST_F(GraphTransformationTests, ConstantFoldingDefaultLimitBlocksLargeExpand) { + // Build a model with a ConstantOfShape node producing a huge output. + // Shape = [16384, 16384] = 268M elements * 4 bytes = 1 GB. + // Default limit is 1 GB, so use a smaller limit to test the blocking behavior. + + auto build_model = [&](ModelTestBuilder& builder) { + auto* shape_data = builder.MakeInitializer({2}, {16384, 16384}); + auto* output_arg = builder.MakeOutput(); + + auto& node = builder.AddNode("ConstantOfShape", {shape_data}, {output_arg}); + // Default value is float 0.0 + ONNX_NAMESPACE::AttributeProto value_attr; + value_attr.set_name("value"); + value_attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_TENSOR); + auto* tensor = value_attr.mutable_t(); + tensor->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + tensor->add_dims(1); + tensor->add_float_data(0.0f); + node.AddAttributeProto(std::move(value_attr)); + }; + + auto pre_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count["ConstantOfShape"] == 1); + return Status::OK(); + }; + + auto post_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + // ConstantOfShape should remain because output is too large (1 GB > 512 MB limit) + TEST_RETURN_IF_NOT(op_to_count["ConstantOfShape"] == 1); + return Status::OK(); + }; + + std::unique_ptr e = std::make_unique(CPUExecutionProviderInfo()); + ConfigOptions config_options; + // Set limit to 512 MB so the 1 GB output is blocked + ASSERT_STATUS_OK(config_options.AddConfigEntry( + kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes, "536870912")); + + ASSERT_STATUS_OK(TestGraphTransformer(build_model, 14, *logger_, + std::make_unique(*e.get(), false, config_options), + TransformerLevel::Level1, 1, + pre_graph_checker, post_graph_checker)); +} + +// Test that small constant folding still works with the size limit. +TEST_F(GraphTransformationTests, ConstantFoldingSmallOutputAllowed) { + // Build a model with a small Expand: scalar -> [4, 4] = 16 * 4 = 64 bytes. + // This is well within even a small limit and should be folded. + + auto build_model = [&](ModelTestBuilder& builder) { + auto* input_data = builder.MakeInitializer({1}, {42.0f}); + auto* shape_data = builder.MakeInitializer({2}, {4, 4}); + auto* output_arg = builder.MakeOutput(); + + builder.AddNode("Expand", {input_data, shape_data}, {output_arg}); + }; + + auto pre_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count["Expand"] == 1); + return Status::OK(); + }; + + auto post_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + // Small Expand should be constant folded + TEST_RETURN_IF_NOT(op_to_count["Expand"] == 0); + return Status::OK(); + }; + + std::unique_ptr e = std::make_unique(CPUExecutionProviderInfo()); + const ConfigOptions empty_config_options; + + ASSERT_STATUS_OK(TestGraphTransformer(build_model, 14, *logger_, + std::make_unique(*e.get(), false, empty_config_options), + TransformerLevel::Level1, 1, + pre_graph_checker, post_graph_checker)); +} + // Check transformations in the case of a subgraph with constant inputs. TEST_F(GraphTransformationTests, SubgraphWithConstantInputs) { constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "constant-subgraph.onnx"; From 1ec4d997429539b49c788eda22449bbe7dbf7bb1 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 13 Apr 2026 16:59:59 -0700 Subject: [PATCH 2/5] Merge main --- onnxruntime/core/common/safeint.h | 45 ++++++++ onnxruntime/core/providers/cpu/rnn/rnn.cc | 58 ++++++---- onnxruntime/test/common/safeint_test.cc | 38 +++++++ .../test/providers/cpu/rnn/rnn_op_test.cc | 103 ++++++++++++++++++ .../templates/linux-wasm-ci.yml | 4 +- 5 files changed, 227 insertions(+), 21 deletions(-) create mode 100644 onnxruntime/test/common/safeint_test.cc diff --git a/onnxruntime/core/common/safeint.h b/onnxruntime/core/common/safeint.h index 6aba5871ac62e..7062e99f6ce3b 100644 --- a/onnxruntime/core/common/safeint.h +++ b/onnxruntime/core/common/safeint.h @@ -36,3 +36,48 @@ class SafeIntExceptionHandler { #if defined(__GNUC__) #pragma GCC diagnostic pop #endif + +#include + +namespace onnxruntime { + +template +using remove_cvref_t = std::remove_cv_t>; + +template +inline constexpr bool is_supported_integer_v = + std::is_integral_v> && !std::is_same_v, bool>; + +//------------------------------------------------------------------------------ +// Safe multiplication of two or more integer values into an explicit result type R. +// Throws OnnxRuntimeException on overflow. +//------------------------------------------------------------------------------ +template +[[nodiscard]] R SafeMul(T a, U b, Rest... rest) { + static_assert(is_supported_integer_v, + "SafeMul requires an integral result type (excluding bool)"); + static_assert(is_supported_integer_v && is_supported_integer_v, + "SafeMul requires integral operand types (excluding bool)"); + static_assert((is_supported_integer_v && ...), + "SafeMul requires integral operand types (excluding bool)"); + + // SafeMultiply(T, U, T&) requires the first argument and result to share + // the same type. Cast the first operand to R so the result is directly in R. + R cast_a{}; + if (!SafeCast(a, cast_a)) { + SafeIntDefaultExceptionHandler::SafeIntOnOverflow(); + } + + R result{}; + if (!SafeMultiply(cast_a, b, result)) { + SafeIntDefaultExceptionHandler::SafeIntOnOverflow(); + } + + if constexpr (sizeof...(rest) > 0) { + return SafeMul(result, rest...); + } else { + return result; + } +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/rnn/rnn.cc b/onnxruntime/core/providers/cpu/rnn/rnn.cc index 6865571eb7a13..7b27befeb3e66 100644 --- a/onnxruntime/core/providers/cpu/rnn/rnn.cc +++ b/onnxruntime/core/providers/cpu/rnn/rnn.cc @@ -3,6 +3,7 @@ #include "core/providers/cpu/rnn/rnn.h" +#include "core/common/narrow.h" #include "core/common/safeint.h" #include "core/framework/op_kernel_context_internal.h" #include "core/providers/cpu/rnn/rnn_activation_functors.h" @@ -84,15 +85,32 @@ void ApplyActivationToBatches(const Tensor* sequence_lens, const T* h_prev, T* Y template void Assign_Y_h(const T* Y_buffer_data, Tensor* Y_h, const Tensor* sequence_lens, int64_t num_directions, int direction, bool isReverse, int64_t batch_size, int64_t seq_length, int64_t hidden_size) { + if (seq_length == 0) { + // No sequence data was processed; zero out Y_h for this direction. + const size_t y_h_direction_size = SafeMul(batch_size, hidden_size); + const size_t Y_h_direction_offset = SafeMul(direction, y_h_direction_size); + math::Set(y_h_direction_size, T{0}, + Y_h->MutableData() + Y_h_direction_offset, &CPUMathUtil::Instance()); + return; + } + for (int batch = 0; batch < batch_size; batch++) { int64_t last_time_step = isReverse ? 0 : seq_length - 1; - if (nullptr != sequence_lens && !isReverse) + if (nullptr != sequence_lens && !isReverse) { last_time_step = sequence_lens->Data()[batch] - 1; + if (last_time_step < 0) { + // sequence_lens[batch] == 0: no data was processed for this batch; zero out Y_h. + int64_t Y_h_offset = direction * batch_size * hidden_size + batch * hidden_size; + math::Set(narrow(hidden_size), T{0}, + Y_h->MutableData() + Y_h_offset, &CPUMathUtil::Instance()); + continue; + } + } int64_t y_offset = last_time_step * num_directions * batch_size * hidden_size + direction * batch_size * hidden_size + batch * hidden_size; int64_t Y_h_offset = direction * batch_size * hidden_size + batch * hidden_size; - math::CopyVector(static_cast(hidden_size), Y_buffer_data + y_offset, + math::CopyVector(narrow(hidden_size), Y_buffer_data + y_offset, Y_h->MutableData() + Y_h_offset, &CPUMathUtil::Instance()); } @@ -109,7 +127,7 @@ void ClearMissingFrames(T* Y_buffer_data, const Tensor* sequence_lens, seq * num_directions * batch_size * hidden_size + direction * batch_size * hidden_size + batch * hidden_size; - math::Set(onnxruntime::narrow(hidden_size), 0, Y_buffer_data + offset, &CPUMathUtil::Instance()); + math::Set(narrow(hidden_size), 0, Y_buffer_data + offset, &CPUMathUtil::Instance()); } } } @@ -155,7 +173,7 @@ Status RNN::Compute(OpKernelContext* ctx) const { ORT_RETURN_IF_ERROR(ctx->GetTempSpaceAllocator(&alloc)); // X * W^t, each direction has shape of [seq_length, batch_size, hidden_size] - auto x_matmul_data = alloc->Alloc(SafeInt(sizeof(float)) * seq_length * batch_size * hidden_size_); + auto x_matmul_data = alloc->Alloc(SafeMul(sizeof(float), seq_length, batch_size, hidden_size_)); BufferUniquePtr x_matmul_buffer(x_matmul_data, BufferDeleter(alloc)); auto* x_matmul_w_buffer_data = static_cast(x_matmul_buffer.get()); @@ -165,7 +183,7 @@ Status RNN::Compute(OpKernelContext* ctx) const { if (Y != nullptr) Y_buffer_data = Y->MutableData(); else { - Y_data = alloc->Alloc(SafeInt(sizeof(float)) * seq_length * num_directions * batch_size * hidden_size_); + Y_data = alloc->Alloc(SafeMul(sizeof(float), seq_length, num_directions, batch_size, hidden_size_)); Y_matmul_buffer = BufferUniquePtr(Y_data, BufferDeleter(alloc)); Y_buffer_data = static_cast(Y_matmul_buffer.get()); } @@ -177,20 +195,20 @@ Status RNN::Compute(OpKernelContext* ctx) const { bool isReverse = direction_ == "reverse" || direction == 1; if (B != nullptr) { - EigenMatrixMapRowMajor(x_matmul_w_buffer_data, seq_length * SafeInt(batch_size), onnxruntime::narrow(hidden_size_)).rowwise() = - ConstEigenVectorMap(B->Data() + direction * 2 * hidden_size_, onnxruntime::narrow(hidden_size_)).transpose() + - ConstEigenVectorMap(B->Data() + direction * 2 * hidden_size_ + hidden_size_, onnxruntime::narrow(hidden_size_)).transpose(); + EigenMatrixMapRowMajor(x_matmul_w_buffer_data, SafeMul(seq_length, batch_size), narrow(hidden_size_)).rowwise() = + ConstEigenVectorMap(B->Data() + direction * 2 * hidden_size_, narrow(hidden_size_)).transpose() + + ConstEigenVectorMap(B->Data() + direction * 2 * hidden_size_ + hidden_size_, narrow(hidden_size_)).transpose(); } else { - math::Set(seq_length * batch_size * SafeInt(hidden_size_), 0, x_matmul_w_buffer_data, &CPUMathUtil::Instance()); + math::Set(SafeMul(seq_length, batch_size, hidden_size_), 0, x_matmul_w_buffer_data, &CPUMathUtil::Instance()); } // X * W[direction]^t + B math::Gemm( CblasNoTrans, CblasTrans, - static_cast(seq_length * batch_size), - static_cast(hidden_size_), - static_cast(input_size), + SafeMul(seq_length, batch_size), + narrow(hidden_size_), + narrow(input_size), 1, X.Data(), W.Data() + direction * hidden_size_ * input_size, @@ -202,7 +220,7 @@ Status RNN::Compute(OpKernelContext* ctx) const { int64_t time_step = isReverse ? (seq_length - t - 1) : t; int64_t Y_frame_offset = (time_step * num_directions + direction) * Y_frame_size; float* Y_buffer_data_current_frame = Y_buffer_data + Y_frame_offset; - auto y_frame_mat = EigenMatrixMapRowMajor(Y_buffer_data_current_frame, onnxruntime::narrow(batch_size), onnxruntime::narrow(hidden_size_)); + auto y_frame_mat = EigenMatrixMapRowMajor(Y_buffer_data_current_frame, narrow(batch_size), narrow(hidden_size_)); const float* h_prev = nullptr; if (t == 0) { @@ -224,9 +242,9 @@ Status RNN::Compute(OpKernelContext* ctx) const { math::Gemm( CblasNoTrans, CblasTrans, - static_cast(batch_size), - static_cast(hidden_size_), - static_cast(hidden_size_), + narrow(batch_size), + narrow(hidden_size_), + narrow(hidden_size_), 1, h_prev, R.Data() + direction * hidden_size_ * hidden_size_, @@ -234,11 +252,11 @@ Status RNN::Compute(OpKernelContext* ctx) const { Y_buffer_data_current_frame, tp, &mlas_backend_kernel_selector_config_); } else { - math::Set(batch_size * SafeInt(hidden_size_), 0, Y_buffer_data_current_frame, &CPUMathUtil::Instance()); + math::Set(SafeMul(batch_size, hidden_size_), 0, Y_buffer_data_current_frame, &CPUMathUtil::Instance()); } // X[time_step] * W^t + H_t_1 * R^t - y_frame_mat += EigenMatrixMapRowMajor(&x_matmul_w_buffer_data[time_step * Y_frame_size], onnxruntime::narrow(batch_size), onnxruntime::narrow(hidden_size_)); + y_frame_mat += EigenMatrixMapRowMajor(&x_matmul_w_buffer_data[time_step * Y_frame_size], narrow(batch_size), narrow(hidden_size_)); // apply activation ApplyActivationToBatches(sequence_lens, h_prev, Y_buffer_data_current_frame, @@ -258,10 +276,10 @@ Status RNN::Compute(OpKernelContext* ctx) const { } if (Y != nullptr) - DumpMatrix("Y", Y_buffer_data, (int)(seq_length * num_directions * batch_size), (int)hidden_size_); + DumpMatrix("Y", Y_buffer_data, SafeMul(seq_length, num_directions, batch_size), narrow(hidden_size_)); if (Y_h != nullptr) - DumpMatrix("Y_h", Y_h->Data(), (int)(num_directions * batch_size), (int)hidden_size_); + DumpMatrix("Y_h", Y_h->Data(), SafeMul(num_directions, batch_size), narrow(hidden_size_)); return Status::OK(); } diff --git a/onnxruntime/test/common/safeint_test.cc b/onnxruntime/test/common/safeint_test.cc new file mode 100644 index 0000000000000..ced9bd94975d2 --- /dev/null +++ b/onnxruntime/test/common/safeint_test.cc @@ -0,0 +1,38 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/common/safeint.h" + +#include +#include +#include + +#include "gtest/gtest.h" + +namespace onnxruntime::test { + +static_assert(is_supported_integer_v); +static_assert(is_supported_integer_v); +static_assert(!is_supported_integer_v); + +TEST(SafeIntTest, SafeMulMultipliesOperands) { + EXPECT_EQ(SafeMul(size_t{2}, 3U), size_t{6}); + EXPECT_EQ(SafeMul(-2, 3, 4), -24); +} + +TEST(SafeIntTest, SafeMulHandlesSameVariableOperands) { + const int value = 7; + EXPECT_EQ(SafeMul(value, value), 49); +} + +#ifndef ORT_NO_EXCEPTIONS +TEST(SafeIntTest, SafeMulThrowsOnInitialCastOverflow) { + EXPECT_THROW((void)SafeMul(-1, 2), OnnxRuntimeException); +} + +TEST(SafeIntTest, SafeMulThrowsOnMultiplyOverflow) { + EXPECT_THROW((void)SafeMul(std::numeric_limits::max(), 2), OnnxRuntimeException); +} +#endif + +} // namespace onnxruntime::test diff --git a/onnxruntime/test/providers/cpu/rnn/rnn_op_test.cc b/onnxruntime/test/providers/cpu/rnn/rnn_op_test.cc index 38734ab9f668f..382d1869a02f6 100644 --- a/onnxruntime/test/providers/cpu/rnn/rnn_op_test.cc +++ b/onnxruntime/test/providers/cpu/rnn/rnn_op_test.cc @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include + #include "core/providers/cpu/rnn/rnn.h" #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" @@ -883,5 +885,106 @@ TEST(RNNTest, RNN_with_invalid_activation_load_failure) { {kCudaExecutionProvider, kTensorrtExecutionProvider}); } +// Test that seq_length == 0 produces zero-filled Y and Y_h without crashing. +TEST(RNNTest, RNN_seq_length_zero) { + auto cpu = DefaultCpuExecutionProvider(); + if (!cpu) GTEST_SKIP() << "CPU EP not available in this build."; + + OpTester test("RNN"); + int64_t num_directions = 1, input_size = 2, hidden_size = 3, batch_size = 2, seq_length = 0; + + test.AddAttribute("activations", vector(num_directions, "Tanh")); + test.AddAttribute("direction", "forward"); + test.AddAttribute("hidden_size", hidden_size); + + std::vector X_dims = {seq_length, batch_size, input_size}; + std::vector X_data{}; + test.AddInput("X", X_dims, X_data); + + std::vector W_dims = {num_directions, hidden_size, input_size}; + std::vector W_data({-0.1f, 0.2f, 1.f, -2.f, -1.f, 3.f}); + test.AddInput("W", W_dims, W_data); + + std::vector R_dims = {num_directions, hidden_size, hidden_size}; + std::vector R_data(hidden_size * hidden_size, 0.f); + test.AddInput("R", R_dims, R_data); + + // Y: shape [0, 1, 2, 3] -> empty + std::vector Y_dims = {seq_length, num_directions, batch_size, hidden_size}; + std::vector Y_data{}; + test.AddOutput("Y", Y_dims, Y_data); + + // Y_h: shape [1, 2, 3] -> all zeros + std::vector Y_h_dims{num_directions, batch_size, hidden_size}; + std::vector Y_h_data(num_directions * batch_size * hidden_size, 0.f); + test.AddOutput("Y_h", Y_h_dims, Y_h_data); + test.ConfigEp(std::move(cpu)).RunWithConfig(); +} + +// Test that per-batch sequence_lens containing 0 produces zero-filled Y_h for those batches. +TEST(RNNTest, RNN_forward_sequence_lens_with_zero) { + auto cpu = DefaultCpuExecutionProvider(); + if (!cpu) GTEST_SKIP() << "CPU EP not available in this build."; + + OpTester test("RNN"); + int64_t num_directions = 1, input_size = 2, hidden_size = 3, batch_size = 2, seq_length = 2; + + test.AddAttribute("activations", vector(num_directions, "Tanh")); + test.AddAttribute("direction", "forward"); + test.AddAttribute("hidden_size", hidden_size); + + // X shape: [seq_length=2, batch_size=2, input_size=2] + std::vector X_dims = {seq_length, batch_size, input_size}; + std::vector X_data({0.1f, 0.2f, + 0.3f, 0.4f, + 0.5f, 0.6f, + 0.7f, 0.8f}); + test.AddInput("X", X_dims, X_data); + + std::vector W_dims = {num_directions, hidden_size, input_size}; + std::vector W_data({-0.1f, 0.2f, 1.f, -2.f, -1.f, 3.f}); + test.AddInput("W", W_dims, W_data); + + std::vector R_dims = {num_directions, hidden_size, hidden_size}; + std::vector R_data(hidden_size * hidden_size, 0.f); + test.AddInput("R", R_dims, R_data); + + std::vector B_dims = {num_directions, 2 * hidden_size}; + std::vector B_data(2 * hidden_size, 0.f); + test.AddInput("B", B_dims, B_data); + + // batch 0 has sequence_lens=2, batch 1 has sequence_lens=0 + std::vector sequence_lens_dims{batch_size}; + std::vector sequence_lens_data{2, 0}; + test.AddInput("sequence_lens", sequence_lens_dims, sequence_lens_data); + + std::vector initial_h_dims = {num_directions, batch_size, hidden_size}; + std::vector initial_h_data(num_directions * batch_size * hidden_size, 0.f); + test.AddInput("initial_h", initial_h_dims, initial_h_data); + + // Y output is optional; skip it to keep test simple. + test.AddOptionalOutputEdge(); + + // Y_h: shape [1, 2, 3] + // batch 0 gets the result of forward pass at last time step (seq_length-1=1). + // batch 1 has sequence_lens=0 so Y_h should be zero. + // + // For batch 0: + // time_step 0: X=[0.1, 0.2], Y = tanh(X * W^T) = tanh([-0.1*0.1+0.2*0.2, 1*0.1-2*0.2, -1*0.1+3*0.2]) + // = tanh([0.03, -0.3, 0.5]) + // time_step 1: X=[0.5, 0.6], Y = tanh(X * W^T + H_prev * R^T) + // R is zero, so Y = tanh([-0.1*0.5+0.2*0.6, 1*0.5-2*0.6, -1*0.5+3*0.6]) + // = tanh([0.07, -0.7, 1.3]) + float y_h_batch0_f0 = std::tanh(0.07f); + float y_h_batch0_f1 = std::tanh(-0.7f); + float y_h_batch0_f2 = std::tanh(1.3f); + + std::vector Y_h_dims{num_directions, batch_size, hidden_size}; + std::vector Y_h_data{y_h_batch0_f0, y_h_batch0_f1, y_h_batch0_f2, + 0.f, 0.f, 0.f}; + test.AddOutput("Y_h", Y_h_dims, Y_h_data); + test.ConfigEp(std::move(cpu)).RunWithConfig(); +} + } // namespace test } // namespace onnxruntime diff --git a/tools/ci_build/github/azure-pipelines/templates/linux-wasm-ci.yml b/tools/ci_build/github/azure-pipelines/templates/linux-wasm-ci.yml index 8eafeddf28f14..849a54345a33b 100644 --- a/tools/ci_build/github/azure-pipelines/templates/linux-wasm-ci.yml +++ b/tools/ci_build/github/azure-pipelines/templates/linux-wasm-ci.yml @@ -84,10 +84,12 @@ jobs: git submodule update --init --recursive workingDirectory: '$(Build.SourcesDirectory)' displayName: 'Checkout submodules' - - template: templates/setup-feeds-and-python-steps.yml + + - template: setup-feeds-and-python-steps.yml parameters: versionSpec: "3.12" architecture: $(buildArch) + - task: NodeTool@0 inputs: versionSpec: '22.x' From aeff70f00b65c6eded3a15dc93510eaa46a77f44 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 13 Apr 2026 17:00:10 -0700 Subject: [PATCH 3/5] Add safe int in Expand and a const folding test --- .../core/providers/cpu/tensor/expand.cc | 14 +++---- .../test/optimizer/graph_transform_test.cc | 40 +++++++++++++++++++ 2 files changed, 47 insertions(+), 7 deletions(-) diff --git a/onnxruntime/core/providers/cpu/tensor/expand.cc b/onnxruntime/core/providers/cpu/tensor/expand.cc index b0c636281bc7a..6d299282f3e60 100644 --- a/onnxruntime/core/providers/cpu/tensor/expand.cc +++ b/onnxruntime/core/providers/cpu/tensor/expand.cc @@ -94,8 +94,8 @@ Status Expand::Compute(OpKernelContext* context) const { auto input_dim = input_dims_iter > -1 ? input_dims[input_dims_iter] : 1; auto output_dim = output_dims[output_dims_iter]; - input_count *= input_dim; - output_count *= output_dim; + input_count = SafeInt(input_count) * input_dim; + output_count = SafeInt(output_count) * output_dim; if (0 == input_count || 0 == output_count) { return Status::OK(); @@ -106,26 +106,26 @@ Status Expand::Compute(OpKernelContext* context) const { input_dim_group[onnxruntime::narrow(dim_group_start)] = input_count; output_dim_group[onnxruntime::narrow(dim_group_start)] = output_count; expand_dim_size[onnxruntime::narrow(dim_group_start)] = output_count / input_count / last_dim_size; - last_dim_size *= expand_dim_size[onnxruntime::narrow(dim_group_start)]; + last_dim_size = SafeInt(last_dim_size) * expand_dim_size[onnxruntime::narrow(dim_group_start)]; } } auto distribute_count = input_dim_group[onnxruntime::narrow(dim_group_start)] / input_dim_group[SafeInt(max_dims_size) - 1]; std::vector output_offsets(onnxruntime::narrow(distribute_count), 0); int64_t copy_len = input_dim_group[SafeInt(max_dims_size) - 1]; - auto copy_byte = copy_len * sizeof(T); + size_t copy_byte = SafeInt(copy_len) * sizeof(T); auto distribute_fn = [&](ptrdiff_t i_start, ptrdiff_t i_end) { for (auto i = i_start; i < i_end; i++) { - auto input_offset = i * copy_len; + int64_t input_offset = SafeInt(i) * copy_len; int64_t output_offset = 0; for (auto j = dim_group_start + 1, remains = input_offset; j < max_dims_size; ++j) { auto current_count = remains / input_dim_group[onnxruntime::narrow(j)]; - output_offset += current_count * output_dim_group[onnxruntime::narrow(j)]; + output_offset = SafeInt(output_offset) + SafeInt(current_count) * output_dim_group[onnxruntime::narrow(j)]; remains = remains % input_dim_group[onnxruntime::narrow(j)]; } // for j - memcpy(output_data + output_offset, input_data + input_offset, onnxruntime::narrow(copy_byte)); + memcpy(output_data + output_offset, input_data + input_offset, copy_byte); output_offsets[onnxruntime::narrow(i)] = output_offset; } // for i }; // distribute_fn diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 52e5feda1c5e8..75b9d738d6540 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -1627,6 +1627,46 @@ TEST_F(GraphTransformationTests, ConstantFoldingSmallOutputAllowed) { pre_graph_checker, post_graph_checker)); } +// Test that constant folding gracefully handles an Expand node whose output shape +// dimensions would cause integer overflow. This simulates the attack vector where +// a malicious model embeds constant initializers with extreme shape values, causing +// kernel Compute() to execute during graph optimization. The SafeInt-protected +// arithmetic in Expand::Compute (or TensorShape overflow) should be caught by the +// try/catch around Compute, and the node should NOT be constant folded. +TEST_F(GraphTransformationTests, ConstantFoldingExpandOverflowDimsSkipped) { + constexpr int64_t kLargeDim = int64_t(1) << 32; // 4294967296 + + auto build_model = [&](ModelTestBuilder& builder) { + auto* input_data = builder.MakeInitializer({1}, {1.0f}); + auto* shape_data = builder.MakeInitializer({2}, {kLargeDim, kLargeDim}); + auto* output_arg = builder.MakeOutput(); + + builder.AddNode("Expand", {input_data, shape_data}, {output_arg}); + }; + + auto pre_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count["Expand"] == 1); + return Status::OK(); + }; + + auto post_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + // Expand should remain because the overflow prevents constant folding. + TEST_RETURN_IF_NOT(op_to_count["Expand"] == 1); + return Status::OK(); + }; + + std::unique_ptr e = + std::make_unique(CPUExecutionProviderInfo()); + const ConfigOptions empty_config_options; + + ASSERT_STATUS_OK(TestGraphTransformer(build_model, 14, *logger_, + std::make_unique(*e.get(), false, empty_config_options), + TransformerLevel::Level1, 1, + pre_graph_checker, post_graph_checker)); +} + // Check transformations in the case of a subgraph with constant inputs. TEST_F(GraphTransformationTests, SubgraphWithConstantInputs) { constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "constant-subgraph.onnx"; From 331709c8336a893e4327d19f38b3a3da07e2c97c Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 15 Apr 2026 15:06:07 -0700 Subject: [PATCH 4/5] fix: check IsAllocated() before accessing tensor in constant folding size check The post-execution size check was calling IsTensor() and then Get().SizeInBytes() on OrtValue entries from fetches. For optional outputs that are not produced, IsTensor() returns true (the type is set) but the data pointer is null, causing a SEGFAULT when SizeInBytes() dereferences the null tensor's dtype_ member. Add IsAllocated() check to skip unallocated (optional/None) outputs. --- onnxruntime/core/optimizer/constant_folding.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/core/optimizer/constant_folding.cc b/onnxruntime/core/optimizer/constant_folding.cc index 60614034d0ff0..8a7cb07697e0e 100644 --- a/onnxruntime/core/optimizer/constant_folding.cc +++ b/onnxruntime/core/optimizer/constant_folding.cc @@ -430,7 +430,7 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level, bool size_exceeded = false; try { for (size_t fetch_idx = 0; fetch_idx < fetches.size(); ++fetch_idx) { - if (fetches[fetch_idx].IsTensor()) { + if (fetches[fetch_idx].IsAllocated() && fetches[fetch_idx].IsTensor()) { const auto& tensor = fetches[fetch_idx].Get(); actual_total_size += tensor.SizeInBytes(); } From de3824dec21b1437b3028c080e25fb475907d3e9 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 16 Apr 2026 16:32:20 -0700 Subject: [PATCH 5/5] fix: address constant folding review feedback --- .../onnxruntime_session_options_config_keys.h | 24 +-- .../core/optimizer/constant_folding.cc | 167 ++++++++++++++---- .../test/optimizer/graph_transform_test.cc | 7 +- 3 files changed, 148 insertions(+), 50 deletions(-) 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 c15cdc3c93079..4572f4b68d108 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h @@ -111,6 +111,18 @@ static const char* const kOrtSessionOptionsMemoryOptimizerProbeConfig = "optimiz // Default is an empty string which means no optimizers are disabled. static const char* const kOrtSessionOptionsDisableSpecifiedOptimizers = "optimization.disable_specified_optimizers"; +// Maximum total output size in bytes that the constant folding optimizer is allowed to produce per node. +// Prevents malicious models from causing excessive memory allocation during optimization. +// If the estimated or actual output size of a constant-foldable node exceeds this limit, the node will +// not be constant folded and will instead be executed at runtime. +// +// Option values: +// - A positive integer (as string): Maximum allowed output size in bytes per constant-folded node. +// Default is "1073741824" (1 GB). +// - "0": Disable the size limit (not recommended for untrusted models). +static const char* const kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes = + "optimization.constant_folding_max_output_size_in_bytes"; + // It controls whether to run graph optimizations in loop or not. // // "0": disable. Graph Optimization Loop is disabled. @@ -481,15 +493,3 @@ static const char* const kOrtSessionOptionsRecordEpGraphAssignmentInfo = "sessio // - "0": disable. (default) // - "1": enable. static const char* const kOrtSessionOptionEpEnableWeightlessEpContextNodes = "ep.enable_weightless_ep_context_nodes"; - -// Maximum total output size in bytes that the constant folding optimizer is allowed to produce per node. -// Prevents malicious models from causing excessive memory allocation during optimization. -// If the estimated or actual output size of a constant-foldable node exceeds this limit, the node will -// not be constant folded and will instead be executed at runtime. -// -// Option values: -// - A positive integer (as string): Maximum allowed output size in bytes per constant-folded node. -// Default is "1073741824" (1 GB). -// - "0": Disable the size limit (not recommended for untrusted models). -static const char* const kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes = - "optimization.constant_folding_max_output_size_in_bytes"; diff --git a/onnxruntime/core/optimizer/constant_folding.cc b/onnxruntime/core/optimizer/constant_folding.cc index 8a7cb07697e0e..1b8bb57d6a74e 100644 --- a/onnxruntime/core/optimizer/constant_folding.cc +++ b/onnxruntime/core/optimizer/constant_folding.cc @@ -2,13 +2,13 @@ // Licensed under the MIT License. #include +#include #include "core/optimizer/constant_folding.h" #include "core/optimizer/initializer.h" #include "core/optimizer/utils.h" #include "core/graph/graph_utils.h" #include "core/optimizer/optimizer_execution_frame.h" -#include "core/optimizer/utils.h" #include "core/framework/op_kernel.h" #include "core/framework/tensorprotoutils.h" #include "core/session/onnxruntime_session_options_config_keys.h" @@ -147,47 +147,131 @@ static Status ConstantFoldIfNode(Graph& graph, Node& if_node, const logging::Log // This prevents malicious models from causing excessive memory allocation during optimization. static constexpr int64_t kDefaultConstantFoldingMaxOutputSizeInBytes = 1024 * 1024 * 1024; -// Estimate the total output size in bytes for a node using shape inference results. -// Returns -1 if the output size cannot be estimated (e.g., unknown shapes or types). -static int64_t EstimateNodeOutputSizeInBytes(const Node& node) { - SafeInt total_size = 0; +static size_t GetElementSizeForConstantFolding(ONNX_NAMESPACE::TensorProto_DataType elem_type) { + const size_t element_size = utils::GetElementSizeOfTensor(elem_type); + if (element_size != 0) { + return element_size; + } - for (const auto* output_def : node.OutputDefs()) { + // String tensors allocate storage for std::string slots even though the payload size is variable. + return elem_type == ONNX_NAMESPACE::TensorProto_DataType_STRING ? sizeof(std::string) : 0; +} + +static int64_t EstimateTensorElementCount(const ONNX_NAMESPACE::TensorShapeProto& shape) { + SafeInt num_elements = 1; + for (int i = 0; i < shape.dim_size(); ++i) { + const auto& dim = shape.dim(i); + if (!utils::HasDimValue(dim)) { + return -1; // Symbolic dimension + } + int64_t dim_value = dim.dim_value(); + if (dim_value < 0) { + return -1; // Invalid dimension + } + num_elements *= dim_value; + } + + return num_elements; +} + +static int64_t EstimateTensorSizeInBytes(const NodeArg& node_arg) { + const auto* type_proto = node_arg.TypeAsProto(); + if (type_proto == nullptr || !utils::HasTensorType(*type_proto)) { + return -1; // Cannot estimate non-tensor or unknown types + } + + const auto* shape = node_arg.Shape(); + if (shape == nullptr) { + return -1; // Unknown shape + } + + auto elem_type = static_cast( + type_proto->tensor_type().elem_type()); + size_t element_size = GetElementSizeForConstantFolding(elem_type); + if (element_size == 0) { + return -1; // Unknown element type + } + + int64_t num_elements = EstimateTensorElementCount(*shape); + if (num_elements < 0) { + return -1; + } + + return SafeInt(num_elements) * static_cast(element_size); +} + +static int64_t EstimateUniqueOutputSizeInBytes(const Node& node) { + const auto& input_defs = node.InputDefs(); + if (input_defs.empty() || input_defs[0] == nullptr) { + return -1; + } + + const int64_t input_num_elements = input_defs[0]->Shape() != nullptr + ? EstimateTensorElementCount(*input_defs[0]->Shape()) + : -1; + if (input_num_elements < 0) { + return -1; + } + + const auto* input_type_proto = input_defs[0]->TypeAsProto(); + if (input_type_proto == nullptr || !utils::HasTensorType(*input_type_proto)) { + return -1; + } + + auto input_elem_type = static_cast( + input_type_proto->tensor_type().elem_type()); + const size_t input_element_size = GetElementSizeForConstantFolding(input_elem_type); + if (input_element_size == 0) { + return -1; + } + + SafeInt total_size = 0; + const auto& output_defs = node.OutputDefs(); + for (size_t output_idx = 0; output_idx < output_defs.size(); ++output_idx) { + const auto* output_def = output_defs[output_idx]; if (!output_def->Exists()) { continue; } - const auto* type_proto = output_def->TypeAsProto(); - if (type_proto == nullptr || !utils::HasTensorType(*type_proto)) { - return -1; // Cannot estimate non-tensor or unknown types - } + const size_t element_size = output_idx == 0 ? input_element_size : sizeof(int64_t); + total_size += SafeInt(input_num_elements) * static_cast(element_size); + } - const auto* shape = output_def->Shape(); - if (shape == nullptr) { - return -1; // Unknown shape - } + return total_size; +} + +static int64_t EstimateIdentityOutputSizeInBytes(const Node& node) { + const auto& input_defs = node.InputDefs(); + if (input_defs.empty() || input_defs[0] == nullptr) { + return -1; + } + + return EstimateTensorSizeInBytes(*input_defs[0]); +} + +// Estimate the total output size in bytes for a node using shape inference results. +// Returns -1 if the output size cannot be estimated (e.g., unknown shapes or types). +static int64_t EstimateNodeOutputSizeInBytes(const Node& node) { + if (node.OpType() == "Identity" && node.Domain().empty()) { + return EstimateIdentityOutputSizeInBytes(node); + } + + if (node.OpType() == "Unique" && node.Domain().empty()) { + return EstimateUniqueOutputSizeInBytes(node); + } - auto elem_type = static_cast( - type_proto->tensor_type().elem_type()); - size_t element_size = utils::GetElementSizeOfTensor(elem_type); - if (element_size == 0) { - return -1; // Unknown element type + SafeInt total_size = 0; + for (const auto* output_def : node.OutputDefs()) { + if (!output_def->Exists()) { + continue; } - SafeInt num_elements = 1; - for (int i = 0; i < shape->dim_size(); ++i) { - const auto& dim = shape->dim(i); - if (!utils::HasDimValue(dim)) { - return -1; // Symbolic dimension - } - int64_t dim_value = dim.dim_value(); - if (dim_value < 0) { - return -1; // Invalid dimension - } - num_elements *= dim_value; + const int64_t output_size = EstimateTensorSizeInBytes(*output_def); + if (output_size < 0) { + return -1; } - total_size += num_elements * static_cast(element_size); + total_size += output_size; } return total_size; @@ -323,7 +407,12 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level, << " bytes) exceeds the limit (" << max_output_size << " bytes)."; continue; } - // If estimated_size is -1, we couldn't estimate; we'll check actual size after execution. + if (estimated_size < 0) { + LOGS(logger, INFO) << "Skipping constant folding for " << node->OpType() + << " node '" << node->Name() + << "' because output size could not be estimated before execution."; + continue; + } } #if !defined(DISABLE_SPARSE_TENSORS) @@ -406,16 +495,24 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level, #endif OpKernelContext op_kernel_context(&frame, kernel.get(), /*stream*/ nullptr, nullptr, logger); - // Wrap Compute in try/catch so that overflows (e.g., SafeInt) or other failures in a - // single node don't abort the entire constant folding pass. + // Skip the current node if Compute fails so one bad constant-fold candidate does not abort + // the entire constant folding pass. + Status compute_status = Status::OK(); try { - ORT_RETURN_IF_ERROR(kernel->Compute(&op_kernel_context)); + compute_status = kernel->Compute(&op_kernel_context); } catch (const std::exception& ex) { LOGS(logger, WARNING) << "Exception during constant folding of " << node->OpType() << " node '" << node->Name() << "': " << ex.what() << ". Skipping constant folding for this node."; continue; } + + if (!compute_status.IsOK()) { + LOGS(logger, WARNING) << "Failure during constant folding of " << node->OpType() + << " node '" << node->Name() << "': " << compute_status.ErrorMessage() + << ". Skipping constant folding for this node."; + continue; + } #ifdef _WIN32 #pragma warning(pop) #endif diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 75b9d738d6540..cf99e31d445ce 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -1545,11 +1545,12 @@ TEST_F(GraphTransformationTests, ConstantFoldingOutputSizeLimit) { } } -// Test that the default constant folding limit prevents extremely large outputs. -TEST_F(GraphTransformationTests, ConstantFoldingDefaultLimitBlocksLargeExpand) { +// Test that an explicitly configured constant folding output-size limit blocks +// folding a very large ConstantOfShape output. +TEST_F(GraphTransformationTests, ConstantFoldingConfiguredLimitBlocksLargeConstantOfShape) { // Build a model with a ConstantOfShape node producing a huge output. // Shape = [16384, 16384] = 268M elements * 4 bytes = 1 GB. - // Default limit is 1 GB, so use a smaller limit to test the blocking behavior. + // Use an explicit 512 MB limit so the 1 GB output is not folded. auto build_model = [&](ModelTestBuilder& builder) { auto* shape_data = builder.MakeInitializer({2}, {16384, 16384});