diff --git a/include/onnxruntime/core/framework/alloc_kind.h b/include/onnxruntime/core/framework/alloc_kind.h index 4534d08470..c7a953a44b 100644 --- a/include/onnxruntime/core/framework/alloc_kind.h +++ b/include/onnxruntime/core/framework/alloc_kind.h @@ -28,7 +28,8 @@ enum class AllocKind { kPreExisting = 2, kAllocateStatically = 3, kAllocateOutput = 4, - kShare = 5 + kShare = 5, + kAllocatedExternally = 6 }; std::ostream& operator<<(std::ostream& out, AllocKind alloc_kind); diff --git a/onnxruntime/core/framework/allocation_planner.cc b/onnxruntime/core/framework/allocation_planner.cc index 8c72823f88..0c308faf72 100644 --- a/onnxruntime/core/framework/allocation_planner.cc +++ b/onnxruntime/core/framework/allocation_planner.cc @@ -40,6 +40,9 @@ std::ostream& operator<<(std::ostream& out, AllocKind alloc_kind) { case AllocKind::kShare: out << "Share"; break; + case AllocKind::kAllocatedExternally: + out << "AllocatedExternally"; + break; case AllocKind::kNotSet: out << "NotSet"; break; @@ -717,7 +720,7 @@ class PlannerImpl { } } else if (external_outputs) { ORT_ENFORCE(!IsNonTensor(*node_output), "Only tensors are supported for external outputs for now."); - AllocPlan(current).alloc_kind = AllocKind::kPreExisting; + AllocPlan(current).alloc_kind = AllocKind::kAllocatedExternally; #if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE) AllocPlan(current).life_interval.second = execution_plan.size(); #endif diff --git a/onnxruntime/core/framework/execution_frame.cc b/onnxruntime/core/framework/execution_frame.cc index fc7241ea2b..f3e9118ec9 100644 --- a/onnxruntime/core/framework/execution_frame.cc +++ b/onnxruntime/core/framework/execution_frame.cc @@ -51,6 +51,11 @@ Status IExecutionFrame::SetOutputMLValue(int index, const OrtValue& ort_value) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "invalid index ", ort_value_idx); } + if (!IsAllocatedExternally(ort_value_idx)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "SetOutputMLValue() is not allowed for OrtValue index ", ort_value_idx, + " as its allocation kind is not kAllocatedExternally."); + } + all_values_[ort_value_idx] = ort_value; return Status::OK(); } @@ -632,6 +637,11 @@ const AllocPlanPerValue& ExecutionFrame::GetAllocationPlan(int ort_value_idx) { return alloc_plan[ort_value_idx]; } +bool ExecutionFrame::IsAllocatedExternally(int ort_value_idx) { + const auto& allocation_plan = GetAllocationPlan(ort_value_idx); + return allocation_plan.alloc_kind == AllocKind::kAllocatedExternally; +} + void ExecutionFrame::TraceAllocate(int ort_value_idx, size_t size) { if (planner_) { // don't trace the output tensors. diff --git a/onnxruntime/core/framework/execution_frame.h b/onnxruntime/core/framework/execution_frame.h index 39473f0ea3..01f5e5680e 100644 --- a/onnxruntime/core/framework/execution_frame.h +++ b/onnxruntime/core/framework/execution_frame.h @@ -98,6 +98,10 @@ class IExecutionFrame { virtual Status CopyTensor(const Tensor& src, Tensor& dest) const = 0; + virtual bool IsAllocatedExternally(int /*ort_value_idx*/) { + return false; + } + const NodeIndexInfo& node_index_info_; // All the intermediate values for the entire graph. @@ -185,6 +189,8 @@ class ExecutionFrame final : public IExecutionFrame { const AllocPlanPerValue& GetAllocationPlan(int ort_value_idx); + bool IsAllocatedExternally(int ort_value_idx) override; + const SessionState& session_state_; // map of index to custom allocator diff --git a/onnxruntime/core/framework/memory_info.cc b/onnxruntime/core/framework/memory_info.cc index 293acf17f3..3b72dbf929 100644 --- a/onnxruntime/core/framework/memory_info.cc +++ b/onnxruntime/core/framework/memory_info.cc @@ -34,6 +34,7 @@ void MemoryInfo::GenerateTensorMap(const SequentialExecutionPlan* execution_plan //If the tensor is using memory outside of the scope, do not store it if (execution_plan->allocation_plan[mem_info.reused_buffer].alloc_kind == AllocKind::kPreExisting) continue; if (execution_plan->allocation_plan[mem_info.reused_buffer].alloc_kind == AllocKind::kAllocateOutput) continue; + if (execution_plan->allocation_plan[mem_info.reused_buffer].alloc_kind == AllocKind::kAllocatedExternally) continue; mem_info.inplace_reuse = (execution_plan->allocation_plan[value_idx].inplace_reuse != -1 && execution_plan->allocation_plan[value_idx].inplace_reuse != value_idx); mem_info.alloc_kind = execution_plan->allocation_plan[value_idx].alloc_kind; mem_info.location = execution_plan->allocation_plan[value_idx].location; diff --git a/onnxruntime/test/framework/allocation_planner_test.cc b/onnxruntime/test/framework/allocation_planner_test.cc index 8b35ca101a..de82e1b54d 100644 --- a/onnxruntime/test/framework/allocation_planner_test.cc +++ b/onnxruntime/test/framework/allocation_planner_test.cc @@ -428,7 +428,7 @@ TEST_F(PlannerTest, ExternalOutputsTest) { // check allocation kind: CheckAllocKind(X1, AllocKind::kPreExisting); - CheckAllocKind(X2, AllocKind::kPreExisting); + CheckAllocKind(X2, AllocKind::kAllocatedExternally); CheckAllocKind(X3, AllocKind::kAllocate); CheckAllocKind(X4, AllocKind::kAllocateOutput); diff --git a/onnxruntime/test/framework/execution_frame_test.cc b/onnxruntime/test/framework/execution_frame_test.cc index 52147f18e2..43ba139a11 100644 --- a/onnxruntime/test/framework/execution_frame_test.cc +++ b/onnxruntime/test/framework/execution_frame_test.cc @@ -16,6 +16,11 @@ #include "gtest/gtest.h" #include "gmock/gmock.h" +#ifdef ENABLE_TRAINING +#include "core/session/IOBinding.h" +#include "orttraining/core/agent/training_agent.h" +#endif + using namespace ONNX_NAMESPACE; using namespace std; @@ -254,7 +259,7 @@ TEST_F(ExecutionFrameTest, MemPatternTest) { ASSERT_EQ(p->GetBlock(4)->offset_, kAllocAlignment); } -#if defined(ENABLE_TRAINING) || defined(ENABLE_TRAINING_OPS) +#ifdef ENABLE_TRAINING TEST_F(ExecutionFrameTest, MemPatternWithExternalOutputsTest) { auto cpu_xp = CreateCPUExecutionProvider(); auto xp_type = cpu_xp->Type(); @@ -327,6 +332,52 @@ TEST_F(ExecutionFrameTest, MemPatternWithExternalOutputsTest) { ASSERT_EQ(pattern.patterns.size(), 1u); auto p = pattern.GetPatterns(cpu_allocator->Info()); ASSERT_EQ(p->PeakSize(), 0u); // Peak size is 0. + + SessionOptions so; + so.session_logid = "MemPatternWithExternalOutputsTest"; + InferenceSession session_obj{so, GetEnvironment()}; + std::stringstream buffer; + model.ToProto().SerializeToOstream(&buffer); + ASSERT_STATUS_OK(session_obj.Load(buffer)); + ASSERT_STATUS_OK(session_obj.Initialize()); + + { + // Run with original InferenceSession::Run, it should fail due to the YieldOp. + NameMLValMap feeds; + feeds.insert(std::make_pair("X", x_value)); + + // prepare outputs + std::vector output_names; + output_names.push_back("Y"); + std::vector fetches; + + RunOptions run_options; + auto st = session_obj.Run(run_options, feeds, output_names, &fetches); + EXPECT_FALSE(st.IsOK()); + EXPECT_THAT(st.ErrorMessage(), testing::HasSubstr("Non-zero status code returned while running YieldOp node.")); + } + + { + // Run with new RunForward/RunBackward. + training::TrainingAgent training_agent(&session_obj); + unique_ptr io_binding; + ASSERT_STATUS_OK(session_obj.NewIOBinding(&io_binding)); + io_binding->BindInput("X", x_value); + OrtValue output; + io_binding->BindOutput("Y", output); + RunOptions run_options; + std::vector user_outputs; + int64_t run_id; + ASSERT_STATUS_OK(training_agent.RunForward(run_options, *io_binding, user_outputs, run_id)); + const std::vector yield_input_expected{2.0f, 2.0f, 2.0f, 2.0f}; + EXPECT_THAT(user_outputs[0].Get().DataAsSpan(), + ::testing::ContainerEq(gsl::make_span(yield_input_expected))); + ASSERT_STATUS_OK(training_agent.RunBackward(run_id, {t_value})); + // The output is MatMul(x_value, t_value); + const std::vector output_expected{4.0f, 4.0f, 4.0f, 4.0f}; + EXPECT_THAT(io_binding->GetOutputs()[0].Get().DataAsSpan(), + ::testing::ContainerEq(gsl::make_span(output_expected))); + } } #endif diff --git a/orttraining/orttraining/training_ops/cpu/controlflow/ort_tasks.h b/orttraining/orttraining/training_ops/cpu/controlflow/ort_tasks.h index bb78e0a935..56993a97d5 100644 --- a/orttraining/orttraining/training_ops/cpu/controlflow/ort_tasks.h +++ b/orttraining/orttraining/training_ops/cpu/controlflow/ort_tasks.h @@ -19,8 +19,8 @@ typedef std::pair> BackwardReturnType; class OrtTasks final { public: static OrtTasks& GetInstance() { - static OrtTasks* instance_ = new OrtTasks; - return *instance_; + static OrtTasks instance_; + return instance_; } void CreateBackgroundTask(int64_t run_id);