diff --git a/onnxruntime/core/framework/execution_frame.cc b/onnxruntime/core/framework/execution_frame.cc index 59efae597ceb2..96a26dbc254fa 100644 --- a/onnxruntime/core/framework/execution_frame.cc +++ b/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().Shape() +#if !defined(DISABLE_SPARSE_TENSORS) + : p_ort_value->Get().DenseShape(); +#endif + LOGS_DEFAULT(INFO) << "existing_shape=" << existing_shape.ToString() << ", computed_shape=" << (shape ? shape->ToString() : ""); 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(); 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).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"); diff --git a/onnxruntime/test/framework/execution_frame_test.cc b/onnxruntime/test/framework/execution_frame_test.cc index bbfe22a2d4bc7..6f321e19971cd 100644 --- a/onnxruntime/test/framework/execution_frame_test.cc +++ b/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 input_data = {5.0f}; + OrtValue input; + Tensor::InitOrtValue(DataTypeImpl::GetType(), 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(), TensorShape({}), allocator, preallocated_output); + const void* preallocated_buffer = preallocated_output.Get().DataRaw(); + std::vector 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().Shape(), TensorShape({1})); + EXPECT_EQ(results[0].Get().DataRaw(), preallocated_buffer); + EXPECT_EQ(results[0].Get().DataAsSpan()[0], 5.0f); +} + } // namespace test } // namespace onnxruntime