Add External Outputs Flag for YieldOp (#6789)

* add external outputs flag for YieldOp

* use kPreExisting

* add ut for mem_pattern

* fix ut after merge from master

include/onnxruntime/core/framework/kernel_def_builder.h

@@ -82,6 +82,8 @@ class KernelDef {
 
   bool AllocateInputsContiguously() const { return allocate_inputs_contiguously_; }
 
+  bool ExternalOutputs() const { return external_outputs_; }
+
   OrtMemType OutputMemoryType(size_t output_index) const {
     auto it = output_memory_type_args_.find(output_index);
     if (it == output_memory_type_args_.end())
@@ -148,6 +150,9 @@ class KernelDef {
   // Require input tensors to be allocated contiguously.
   bool allocate_inputs_contiguously_ = false;
 
+  // Whether the outputs are from external.
+  bool external_outputs_ = false;
+
   // The memory types of inputs/outputs of this kernel
   MemTypeMap input_memory_type_args_;
   MemTypeMap output_memory_type_args_;
@@ -258,6 +263,14 @@ class KernelDefBuilder {
     return *this;
   }
 
+  /**
+     Specify that this kernel's outputs are passed from external.
+  */
+  KernelDefBuilder& ExternalOutputs() {
+    kernel_def_->external_outputs_ = true;
+    return *this;
+  }
+
   /**
      Specify that this kernel requires an input arg
      in certain memory type (instead of the default, device memory).

onnxruntime/core/framework/allocation_planner.cc

@@ -600,6 +600,15 @@ class PlannerImpl {
     return Status::OK();
   }
 
+  bool ExternalOutputs(const Node& node) const {
+    const KernelCreateInfo& ci = GetKernelCreateInfo(kernel_create_info_map_, node.Index());
+    if (ci.kernel_def == nullptr) {
+      return false;
+    }
+
+    return ci.kernel_def->ExternalOutputs();
+  }
+
   // Should only be used after ProcessDef()
   Status ComputeReusePlan() {
     std::vector<SequentialExecutionPlan::NodeExecutionPlan>& execution_plan(plan_.execution_plan);
@@ -647,6 +656,8 @@ class PlannerImpl {
       const auto* pnode = graph_viewer_.GetNode(step.node_index);
       // node outputs.
       const auto& output_defs = pnode->OutputDefs();
+      // External outputs flag.
+      bool external_outputs = ExternalOutputs(*pnode);
       // output_arg_def_index is the index of ArgDefs in pnode's output list.
       // At the i-th iteration, we build the allocation plan for the i-th
       // NodeArg in pnode's output list. Allocation plan remains untouched for
@@ -702,6 +713,12 @@ 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;
+#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
+          AllocPlan(current).life_interval.second = execution_plan.size();
+#endif
         } else if (IsNonTensor(*node_output)) {
           // we do not try sharing-optimization for non-tensors
           AllocPlan(current).alloc_kind = AllocKind::kAllocate;

onnxruntime/test/framework/execution_frame_test.cc

@@ -254,6 +254,89 @@ TEST_F(ExecutionFrameTest, MemPatternTest) {
   ASSERT_EQ(p->GetBlock(4)->offset_, kAllocAlignment);
 }
 
+#if defined(ENABLE_TRAINING) || defined(ENABLE_TRAINING_OPS)
+TEST_F(ExecutionFrameTest, MemPatternWithExternalOutputsTest) {
+  auto cpu_xp = CreateCPUExecutionProvider();
+  auto xp_type = cpu_xp->Type();
+  std::unordered_map<std::string, int> domain_to_version;
+  domain_to_version[onnxruntime::kOnnxDomain] = 12;
+  domain_to_version[onnxruntime::kMSDomain] = 1;
+  onnxruntime::Model model("test", true, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(),
+                           domain_to_version, {}, DefaultLoggingManager().DefaultLogger());
+  onnxruntime::Graph& graph = model.MainGraph();
+  TypeProto tensor_float;
+  tensor_float.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT);
+  onnxruntime::NodeArg input_def("X", &tensor_float),
+      yield_out_def("T", &tensor_float),
+      gemm_out_def("Y", &tensor_float);
+
+  ONNX_NAMESPACE::AttributeProto required_grad;
+  const std::string attribute_name = "required_grad";
+  required_grad.set_name(attribute_name);
+  required_grad.set_type(ONNX_NAMESPACE::AttributeProto::INTS);
+  required_grad.add_ints(static_cast<int64_t>(0));
+  NodeAttributes attributes({{attribute_name, required_grad}});
+  graph.AddNode("node1", "YieldOp", "yield", ArgMap{&input_def}, ArgMap{&yield_out_def}, &attributes, kMSDomain)
+      .SetExecutionProviderType(xp_type);
+  // Add another node after YieldOp as YieldOp should not be graph output.
+  graph.AddNode("node2", "MatMul", "gemm1", ArgMap{&yield_out_def, &input_def}, ArgMap{&gemm_out_def})
+      .SetExecutionProviderType(xp_type);
+
+  ASSERT_STATUS_OK(graph.Resolve());
+
+  KernelRegistryManager kernel_registry_manager;
+
+  ExecutionProviders execution_providers;
+  execution_providers.Add(xp_type, std::move(cpu_xp));
+  ASSERT_STATUS_OK(kernel_registry_manager.RegisterKernels(execution_providers));
+
+  DataTransferManager dtm;
+  profiling::Profiler profiler;
+  SessionState state(graph, execution_providers, true, &tp_, nullptr, dtm,
+                     DefaultLoggingManager().DefaultLogger(), profiler);
+
+  ASSERT_STATUS_OK(state.FinalizeSessionState(ORT_TSTR(""), kernel_registry_manager));
+
+  const OrtValueNameIdxMap& mlvalue_name_idx_map(state.GetOrtValueNameIdxMap());
+
+  int x_idx = -1, t_idx = -1, y_idx = -1;
+  ASSERT_TRUE(mlvalue_name_idx_map.GetIdx("X", x_idx).IsOK());
+  ASSERT_TRUE(mlvalue_name_idx_map.GetIdx("T", t_idx).IsOK());
+  ASSERT_TRUE(mlvalue_name_idx_map.GetIdx("Y", y_idx).IsOK());
+
+  auto cpu_allocator = execution_providers.Get(xp_type)->GetAllocator(0, OrtMemTypeDefault);
+
+  OrtValue x_value, t_value;
+  CreateMLValue<float>(cpu_allocator,
+                       std::vector<int64_t>{2, 2},
+                       std::vector<float>(4, 2.0f), &x_value);
+  CreateMLValue<float>(cpu_allocator,
+                       std::vector<int64_t>{2, 2},
+                       std::vector<float>(4, 1.0f), &t_value);
+
+  vector<OrtValue> outputs;
+  ExecutionFrame frame({x_idx}, {x_value}, {y_idx}, outputs, {}, state);
+
+  ASSERT_FALSE(frame.GetMutableNodeInputOrOutputMLValue(t_idx)->IsTensor());
+  ASSERT_STATUS_OK(frame.SetOutputMLValue(t_idx, t_value));
+  ASSERT_TRUE(frame.GetMutableNodeInputOrOutputMLValue(t_idx)->IsTensor());
+
+  OrtValue& y_value = *frame.GetMutableNodeInputOrOutputMLValue(y_idx);
+  ASSERT_STATUS_OK(frame.AllocateMLValueTensorSelfOwnBuffer(y_value, y_idx,
+                                                            DataTypeImpl::GetType<float>(),
+                                                            cpu_allocator->Info(),
+                                                            TensorShape(std::vector<int64_t>{2, 2})));
+
+  MemoryPatternGroup pattern;
+  ASSERT_STATUS_OK(frame.GeneratePatterns(&pattern));
+
+  ASSERT_EQ(pattern.patterns.size(), pattern.locations.size());
+  ASSERT_EQ(pattern.patterns.size(), 1u);
+  auto p = pattern.GetPatterns(cpu_allocator->Info());
+  ASSERT_EQ(p->PeakSize(), 0u);  // Peak size is 0.
+}
+#endif
+
 TEST(ExecutionFrameTestWithoutSessionState, BadModelInvalidDimParamUsage) {
   // load model with 2 Scan ops that both incorrectly use shapes of { 'None', 'None' } for their outputs.
   // as 'None' is not a special value it's treated as a variable name, leading to a runtime error when we

orttraining/orttraining/training_ops/cpu/controlflow/yield.cc

@@ -15,7 +15,7 @@ ONNX_OPERATOR_KERNEL_EX(
     kCpuExecutionProvider,
     KernelDefBuilder()
         .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
-        .VariadicAlias(0, 0),  // TODO: this is a hack to avoid allocating output buffer
+        .ExternalOutputs(),
     YieldOp);
 
 Status YieldOp::Compute(OpKernelContext* ctx) const {

orttraining/orttraining/training_ops/cuda/controlflow/yield.cc

@@ -14,7 +14,7 @@ ONNX_OPERATOR_KERNEL_EX(
     kCudaExecutionProvider,
     KernelDefBuilder()
         .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
-        .VariadicAlias(0, 0),  // TODO: this is a hack to avoid allocating output buffer
+        .ExternalOutputs(),
     onnxruntime::contrib::YieldOp);
 
 }  // namespace cuda