Add New AllocKind for YieldOp Outputs, Run YieldOp with InferenceSession in UT (#7125)

* new allockind, add ut

* change macro

* fix win build

* rename alloc kind

* fix mem leak

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);

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

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.

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

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;

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);
 

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<std::string> output_names;
+    output_names.push_back("Y");
+    std::vector<OrtValue> 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<IOBinding> 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<OrtValue> user_outputs;
+    int64_t run_id;
+    ASSERT_STATUS_OK(training_agent.RunForward(run_options, *io_binding, user_outputs, run_id));
+    const std::vector<float> yield_input_expected{2.0f, 2.0f, 2.0f, 2.0f};
+    EXPECT_THAT(user_outputs[0].Get<Tensor>().DataAsSpan<float>(),
+                ::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<float> output_expected{4.0f, 4.0f, 4.0f, 4.0f};
+    EXPECT_THAT(io_binding->GetOutputs()[0].Get<Tensor>().DataAsSpan<float>(),
+                ::testing::ContainerEq(gsl::make_span(output_expected)));
+  }
 }
 #endif
 

orttraining/orttraining/training_ops/cpu/controlflow/ort_tasks.h

@@ -19,8 +19,8 @@ typedef std::pair<bool, std::vector<OrtValue>> 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);