fix output shape mismatch for pre allocated buffer

onnxruntime/core/framework/execution_frame.cc

@@ -177,9 +177,24 @@ Status IExecutionFrame::GetOrCreateNodeOutputMLValue(const int output_index, int
       bool shape_matched = true;
 
       if (p_ort_value->IsTensor()) {
+        const TensorShape& existing_shape = p_ort_value->IsTensor()
+                                                ? p_ort_value->Get<Tensor>().Shape()
+#if !defined(DISABLE_SPARSE_TENSORS)
+                                                : p_ort_value->Get<SparseTensor>().DenseShape();
+#endif
+        LOGS_DEFAULT(INFO) << "existing_shape=" << existing_shape.ToString() << ", computed_shape=" << (shape ? shape->ToString() : "<null>");
         ORT_RETURN_IF_NOT(shape != nullptr, "shape must not be null for tensor output that is already allocated");
         const Tensor& tensor = p_ort_value->Get<Tensor>();
         shape_matched = (tensor.Shape() == *shape);
+        LOGS_DEFAULT(INFO) << "existing_shape size=" << existing_shape.Size() << ", computed_shape size =" << shape->Size() << std::endl;
+        // Compare number of elements
+        if (existing_shape.Size() == shape->Size()) {
+          // Reuse buffer, update shape in-place
+          const_cast<Tensor&>(tensor).Reshape(*shape);
+          shape_matched = true;
+        } else {
+          shape_matched = false;
+        }
       } else if (p_ort_value->IsSparseTensor()) {
 #if !defined(DISABLE_SPARSE_TENSORS)
         ORT_RETURN_IF_NOT(shape != nullptr, "shape must not be null for sparse tensor output that is already allocated");

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