diff --git a/onnxruntime/core/framework/execution_frame.cc b/onnxruntime/core/framework/execution_frame.cc index 59efae597ceb2..15001355c8366 100644 --- a/onnxruntime/core/framework/execution_frame.cc +++ b/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(); + // 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 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