fix output shape mismatch for pre allocated buffer

onnxruntime/core/framework/execution_frame.cc

@@ -17,6 +17,7 @@
 #include "core/framework/session_state.h"
 #include "core/framework/TensorSeq.h"
 #include "core/framework/utils.h"
+#include "core/common/logging/logging.h"
 #if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
 #include "core/framework/memory_info.h"
 #endif
@@ -196,16 +197,30 @@ Status IExecutionFrame::GetOrCreateNodeOutputMLValue(const int output_index, int
 #else
                                                 : *shape;  // unreachable, but satisfies compiler
 #endif
-        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
-                               "The output OrtValue provided for output '",
-                               node.OutputDefs()[output_index]->Name(),
-                               "' of node '", node.Name(),
-                               "' (", node.OpType(), ") has shape ", existing_shape,
-                               " but the computed output shape for this run is ", *shape,
-                               ". When calling Run() with pre-allocated output OrtValues on a model "
-                               "with dynamic output shapes, either supply unallocated output OrtValues "
-                               "or ensure the pre-allocated shapes match the expected output shapes "
-                               "for each run.");
+        LOGS_DEFAULT(VERBOSE) << "Output shape mismatch for pre-allocated fetch buffer.";
+        auto& tensor = *p_ort_value->GetMutable<Tensor>();
+        // Keep caller-provided buffers for dynamic outputs when only the runtime
+        // dimensions changed and the total element count is unchanged.
+        LOGS_DEFAULT(VERBOSE) << "Trying in-place reshape for caller-provided output buffer.";
+        const bool old_shape_is_singleton_vector = shape->NumDimensions() == 1 && shape->AsShapeVector()[0] == 1;  // {1}
+        const bool new_shape_is_scalar = tensor.Shape().NumDimensions() == 0;                                      // {}
+        const bool can_reshape_in_place = (old_shape_is_singleton_vector && new_shape_is_scalar);
+        if (can_reshape_in_place) {
+          LOGS_DEFAULT(VERBOSE) << "Reusing caller buffer by reshape. old_num_elements="
+                                << tensor.Shape().Size() << " new_num_elements=" << shape->Size();
+          tensor.Reshape(*shape);
+        } else {
+          return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                                 "The output OrtValue provided for output '",
+                                 node.OutputDefs()[output_index]->Name(),
+                                 "' of node '", node.Name(),
+                                 "' (", node.OpType(), ") has shape ", existing_shape,
+                                 " but the computed output shape for this run is ", *shape,
+                                 ". When calling Run() with pre-allocated output OrtValues on a model "
+                                 "with dynamic output shapes, either supply unallocated output OrtValues "
+                                 "or ensure the pre-allocated shapes match the expected output shapes "
+                                 "for each run.");
+        }
       }
     } else {
       // shape is nullptr for traditional ML output values

onnxruntime/test/framework/execution_frame_test.cc

@@ -743,5 +743,56 @@ TEST(ExecutionFrameTestInit, SparseInitializerAsOutput) {
 }
 #endif  // !defined(DISABLE_SPARSE_TENSORS)
 
+TEST(ExecutionFrameTestInit, FetchReusesPreallocatedScalarOutputForSingleElementVector) {
+  SessionOptions so;
+  so.enable_mem_pattern = true;
+
+  InferenceSession session(so, GetEnvironment());
+
+  onnxruntime::Model model("scalar_to_vector1_output_test", false, ModelMetaData(), PathString(),
+                           IOnnxRuntimeOpSchemaRegistryList(),
+                           {{kOnnxDomain, 12}}, {}, DefaultLoggingManager().DefaultLogger());
+  auto& graph = model.MainGraph();
+
+  TypeProto float_tensor;
+  float_tensor.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT);
+
+  auto& input_arg = graph.GetOrCreateNodeArg("X", &float_tensor);
+  auto& output_arg = graph.GetOrCreateNodeArg("Y", &float_tensor);
+  graph.AddNode("identity", "Identity", "identity", {&input_arg}, {&output_arg});
+  graph.SetInputs({&input_arg});
+  graph.SetOutputs({&output_arg});
+  ASSERT_STATUS_OK(graph.Resolve());
+
+  std::string serialized;
+  ASSERT_TRUE(model.ToProto().SerializeToString(&serialized));
+  std::istringstream model_stream(serialized);
+  ASSERT_STATUS_OK(session.Load(model_stream));
+  ASSERT_STATUS_OK(session.Initialize());
+
+  auto allocator = test::AllocatorManager::Instance().GetAllocator(CPU);
+
+  std::array<float, 1> input_data = {5.0f};
+  OrtValue input;
+  Tensor::InitOrtValue(DataTypeImpl::GetType<float>(), TensorShape({1}), input_data.data(), allocator->Info(), input);
+
+  // Pre-allocate scalar output ({}). Runtime output for this run is {1}.
+  OrtValue preallocated_output;
+  Tensor::InitOrtValue(DataTypeImpl::GetType<float>(), TensorShape({}), allocator, preallocated_output);
+  const void* preallocated_buffer = preallocated_output.Get<Tensor>().DataRaw();
+  std::vector<OrtValue> results = {preallocated_output};
+
+  RunOptions ro;
+  ASSERT_STATUS_OK(session.Run(ro,
+                               AsSpan({std::string("X")}), AsSpan({input}),
+                               AsSpan({std::string("Y")}), &results, nullptr));
+
+  ASSERT_EQ(results.size(), 1u);
+  ASSERT_TRUE(results[0].IsTensor());
+  EXPECT_EQ(results[0].Get<Tensor>().Shape(), TensorShape({1}));
+  EXPECT_EQ(results[0].Get<Tensor>().DataRaw(), preallocated_buffer);
+  EXPECT_EQ(results[0].Get<Tensor>().DataAsSpan<float>()[0], 5.0f);
+}
+
 }  // namespace test
 }  // namespace onnxruntime