Add validation of shape when re-using a buffer in ExecutionFrame (#1356)

* Check for empty string as dim_param in allocation planner.
* Validate shape is compatible at runtime when re-using Tensor.

onnxruntime/core/framework/allocation_planner.cc

@@ -51,6 +51,7 @@ std::ostream& operator<<(std::ostream& out, std::pair<const SequentialExecutionP
   std::unordered_map<int, std::string> index_to_name;
 
   out << "Allocation Plan:\n";
+  out << "(ort_value_idx) output_name : <allocation plan>\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;

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<Tensor>();
+
+  // 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

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<int64_t> dims_X = {10, 6};
+  std::vector<float> values_X;
+  values_X.reserve(60);
+  for (int i = 0; i < 60; ++i) {
+    values_X.push_back(float(i));
+  }
+
+  OrtValue ml_value;
+  CreateMLValue<float>(TestCPUExecutionProvider()->GetAllocator(0, OrtMemTypeDefault), dims_X, values_X, &ml_value);
+  NameMLValMap feeds;
+  feeds.insert(std::make_pair("X", ml_value));
+
+  // prepare outputs
+  std::vector<std::string> output_names;
+  output_names.push_back("Y");
+  std::vector<OrtValue> 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