From 4238ce341a2389a2ced50430afbcf748610b22a5 Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Tue, 2 Mar 2021 11:38:18 +0800 Subject: [PATCH] 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 --- .../core/framework/kernel_def_builder.h | 13 +++ .../core/framework/allocation_planner.cc | 17 ++++ .../test/framework/execution_frame_test.cc | 83 +++++++++++++++++++ .../training_ops/cpu/controlflow/yield.cc | 2 +- .../training_ops/cuda/controlflow/yield.cc | 2 +- 5 files changed, 115 insertions(+), 2 deletions(-) diff --git a/include/onnxruntime/core/framework/kernel_def_builder.h b/include/onnxruntime/core/framework/kernel_def_builder.h index 965d553eb9..9a220feec6 100644 --- a/include/onnxruntime/core/framework/kernel_def_builder.h +++ b/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). diff --git a/onnxruntime/core/framework/allocation_planner.cc b/onnxruntime/core/framework/allocation_planner.cc index 067bee1aa5..f7c3fbc76a 100644 --- a/onnxruntime/core/framework/allocation_planner.cc +++ b/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& 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; diff --git a/onnxruntime/test/framework/execution_frame_test.cc b/onnxruntime/test/framework/execution_frame_test.cc index e6308b0e72..40ba188d7d 100644 --- a/onnxruntime/test/framework/execution_frame_test.cc +++ b/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 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(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(cpu_allocator, + std::vector{2, 2}, + std::vector(4, 2.0f), &x_value); + CreateMLValue(cpu_allocator, + std::vector{2, 2}, + std::vector(4, 1.0f), &t_value); + + vector 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(), + cpu_allocator->Info(), + TensorShape(std::vector{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 diff --git a/orttraining/orttraining/training_ops/cpu/controlflow/yield.cc b/orttraining/orttraining/training_ops/cpu/controlflow/yield.cc index cdbd90ceab..bffea2395b 100644 --- a/orttraining/orttraining/training_ops/cpu/controlflow/yield.cc +++ b/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 { diff --git a/orttraining/orttraining/training_ops/cuda/controlflow/yield.cc b/orttraining/orttraining/training_ops/cuda/controlflow/yield.cc index b2ddbc7c0c..32756afa8a 100644 --- a/orttraining/orttraining/training_ops/cuda/controlflow/yield.cc +++ b/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