diff --git a/onnxruntime/core/framework/allocation_planner.cc b/onnxruntime/core/framework/allocation_planner.cc index 9ffcc9b555..552d702b80 100644 --- a/onnxruntime/core/framework/allocation_planner.cc +++ b/onnxruntime/core/framework/allocation_planner.cc @@ -51,6 +51,7 @@ std::ostream& operator<<(std::ostream& out, std::pair index_to_name; out << "Allocation Plan:\n"; + out << "(ort_value_idx) output_name : \n"; auto plan_size = plan.allocation_plan.size(); for (auto& name_index : session_state.GetOrtValueNameIdxMap()) { @@ -256,8 +257,11 @@ class PlannerImpl { const auto& val2 = shape2.dim(i); if (val1.has_dim_value() && val2.has_dim_value() && (val1.dim_value() == val2.dim_value())) continue; // same known dimension - if (val1.has_dim_param() && val2.has_dim_param() && (val1.dim_param() == val2.dim_param())) - continue; // same unknown dimension + if (val1.has_dim_param() && val2.has_dim_param()) { + const auto& val1_param = val1.dim_param(); + if (val1_param == val2.dim_param() && !val1_param.empty()) + continue; // same unknown dimension + } return false; } return true; diff --git a/onnxruntime/core/framework/execution_frame.cc b/onnxruntime/core/framework/execution_frame.cc index b0831b465e..feae6cddb8 100644 --- a/onnxruntime/core/framework/execution_frame.cc +++ b/onnxruntime/core/framework/execution_frame.cc @@ -316,6 +316,16 @@ Status ExecutionFrame::AllocateMLValueTensorPreAllocateBuffer(OrtValue& ort_valu OrtValue& ort_value_reuse = GetMutableMLValue(ort_value_index_reuse); auto* reuse_tensor = ort_value_reuse.GetMutable(); + + // check number of elements matches. shape may not be an exact match (e.g. Reshape op) + if (reuse_tensor->Shape().Size() != shape.Size()) { + // could be an allocation planner bug (less likely) or the model incorrectly uses something like 'None' + // as a dim_param in multiple places making the planner think those shapes are equal. + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "Shape mismatch attempting to re-use buffer. ", reuse_tensor->Shape(), " != ", shape, + ". Validate usage of dim_param in shapes in the model."); + } + void* reuse_buffer = reuse_tensor->MutableDataRaw(); // create fence on reused ort_value if needed diff --git a/onnxruntime/test/framework/execution_frame_test.cc b/onnxruntime/test/framework/execution_frame_test.cc index a77ce02805..48e2ce8c3b 100644 --- a/onnxruntime/test/framework/execution_frame_test.cc +++ b/onnxruntime/test/framework/execution_frame_test.cc @@ -6,8 +6,12 @@ #include "core/framework/session_state.h" #include "core/graph/model.h" #include "core/providers/cpu/cpu_execution_provider.h" +#include "core/session/inference_session.h" #include "test_utils.h" +#include "test/test_environment.h" + #include "gtest/gtest.h" +#include "gmock/gmock.h" using namespace ONNX_NAMESPACE; using namespace std; @@ -259,5 +263,43 @@ TEST(ExecutionFrameTest, MemPatternTest) { EXPECT_EQ(p->GetBlock(3)->offset_, 0); EXPECT_EQ(p->GetBlock(4)->offset_, 64); } + +TEST(ExecutionFrameTest, BadModelInvalidDimParamUsage) { + // load model with 2 Scan ops that both incorrectly use shapes of { 'None', 'None' } for their outputs. + // as 'None' is not a special value it's treated as a variable name, leading to a runtime error when we + // attempt to re-use the output from the first Scan node for the second. validate we detect this and error out. + SessionOptions so; + so.session_logid = "BadModelInvalidDimParamUsage"; + + InferenceSession session_object{so, &DefaultLoggingManager()}; + Status st; + ASSERT_TRUE((st = session_object.Load("testdata/invalid_dim_param_value_repetition.onnx")).IsOK()) << st; + ASSERT_TRUE((st = session_object.Initialize()).IsOK()) << st; + + std::vector dims_X = {10, 6}; + std::vector values_X; + values_X.reserve(60); + for (int i = 0; i < 60; ++i) { + values_X.push_back(float(i)); + } + + OrtValue ml_value; + CreateMLValue(TestCPUExecutionProvider()->GetAllocator(0, OrtMemTypeDefault), dims_X, values_X, &ml_value); + NameMLValMap feeds; + feeds.insert(std::make_pair("X", ml_value)); + + // prepare outputs + std::vector output_names; + output_names.push_back("Y"); + std::vector fetches; + + // Now run + RunOptions run_options; + st = session_object.Run(run_options, feeds, output_names, &fetches); + + EXPECT_FALSE(st.IsOK()) << st; + EXPECT_THAT(st.ErrorMessage(), testing::HasSubstr("Shape mismatch attempting to re-use buffer.")); +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/testdata/invalid_dim_param_value_repetition.onnx b/onnxruntime/test/testdata/invalid_dim_param_value_repetition.onnx new file mode 100644 index 0000000000..f7d0624635 Binary files /dev/null and b/onnxruntime/test/testdata/invalid_dim_param_value_repetition.onnx differ