From ddaad866058d0792aaee83018d3b58a6ded958d4 Mon Sep 17 00:00:00 2001 From: Scott McKay Date: Tue, 3 Dec 2019 08:29:21 +1000 Subject: [PATCH] CUDA Loop (#2444) * Implement CUDA Loop operator. * Add control flow node implicit input handling to the memcpy transformer and allocation planner. --- .../core/framework/allocation_planner.cc | 28 +++- .../core/optimizer/transformer_memcpy.cc | 35 +++-- .../core/providers/cpu/controlflow/loop.cc | 137 +++++++++++------- .../core/providers/cpu/controlflow/loop.h | 13 +- .../core/providers/cuda/controlflow/loop.cc | 82 +++++++++++ .../core/providers/cuda/controlflow/loop.h | 20 +++ .../providers/cuda/cuda_execution_provider.cc | 4 + .../providers/cpu/controlflow/loop_test.cc | 60 ++++---- 8 files changed, 277 insertions(+), 102 deletions(-) create mode 100644 onnxruntime/core/providers/cuda/controlflow/loop.cc create mode 100644 onnxruntime/core/providers/cuda/controlflow/loop.h diff --git a/onnxruntime/core/framework/allocation_planner.cc b/onnxruntime/core/framework/allocation_planner.cc index 3c77ada710..2b287c89f2 100644 --- a/onnxruntime/core/framework/allocation_planner.cc +++ b/onnxruntime/core/framework/allocation_planner.cc @@ -281,7 +281,8 @@ class PlannerImpl { return elt_type->Size(); } - static bool SameSize(const TensorShapeProto& shape1, const onnxruntime::NodeArg& arg1, const TensorShapeProto& shape2, const onnxruntime::NodeArg& arg2) { + static bool SameSize(const TensorShapeProto& shape1, const onnxruntime::NodeArg& arg1, + const TensorShapeProto& shape2, const onnxruntime::NodeArg& arg2) { const auto& ptype1 = arg1.Type(); const auto& ptype2 = arg2.Type(); auto type1_size = GetElementSize(ptype1); @@ -392,8 +393,8 @@ class PlannerImpl { // This is determined by the opkernel bound to the node. const KernelCreateInfo* kernel_create_info = nullptr; ORT_RETURN_IF_ERROR(kernel_registry_.SearchKernelRegistry(*pnode, &kernel_create_info)); - auto p_kernelDef = kernel_create_info->kernel_def.get(); - if (nullptr == p_kernelDef) { + auto p_kernel_def = kernel_create_info->kernel_def.get(); + if (nullptr == p_kernel_def) { std::ostringstream errormsg; errormsg << "No suitable kernel definition found for op " << pnode->OpType(); if (pnode->Op() != nullptr) { @@ -402,14 +403,18 @@ class PlannerImpl { if (!pnode->Name().empty()) errormsg << " (node " << pnode->Name() << ")"; return Status(ONNXRUNTIME, FAIL, errormsg.str()); } + auto exec_provider = execution_providers_.Get(*pnode); if (exec_provider == nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Can not find the execution provider ", pnode->GetExecutionProviderType()); } + bool is_implicit_input = false; + // increment UseCount and add location information if applicable for the provided input def - auto process_input = [&graph_inputs, &exec_provider, &p_kernelDef, this](const NodeArg& input, size_t arg_idx) { + auto process_input = [&graph_inputs, &exec_provider, &p_kernel_def, &is_implicit_input, + this](const NodeArg& input, size_t arg_idx) { const auto& name = input.Name(); UseCount(name)++; @@ -422,14 +427,19 @@ class PlannerImpl { return value && value->Name() == name; }) != outer_scope_node_args_.cend()) { OrtValueIndex index = Index(name); - plan_.SetLocation(static_cast(index), - exec_provider->GetAllocator(0, p_kernelDef->InputMemoryType(arg_idx))->Info()); + + // implicit inputs do not have an entry in the kernel def so we use the default memory type. + // matching logic is used in TransformerMemcpyImpl::ProcessDefs + OrtMemType mem_type = is_implicit_input ? OrtMemTypeDefault : p_kernel_def->InputMemoryType(arg_idx); + plan_.SetLocation(static_cast(index), exec_provider->GetAllocator(0, mem_type)->Info()); } return Status::OK(); }; ORT_RETURN_IF_ERROR(Node::ForEachWithIndex(pnode->InputDefs(), process_input)); + + is_implicit_input = true; ORT_RETURN_IF_ERROR(Node::ForEachWithIndex(pnode->ImplicitInputDefs(), process_input)); auto outputs = pnode->OutputDefs(); @@ -440,12 +450,14 @@ class PlannerImpl { OrtValueIndex index = Index(node_output->Name()); ProcessDef(index, node_output); ++UseCount(index); - plan_.SetLocation(static_cast(index), exec_provider->GetAllocator(0, p_kernelDef->OutputMemoryType(i))->Info()); + plan_.SetLocation(static_cast(index), + exec_provider->GetAllocator(0, p_kernel_def->OutputMemoryType(i))->Info()); } + // if sync is needed, mark allocation plan as create_fence_if_async=true // note that the input arg may come from an execution provider (i.e. CPU) that does not support async, // in which case create_fence_if_async would be ignored when creating MLValue - if (p_kernelDef->ExecQueueId() != 0) { + if (p_kernel_def->ExecQueueId() != 0) { pnode->ForEachDef([this](const onnxruntime::NodeArg& arg, bool /*is_input*/) { OrtValueIndex index = Index(arg.Name()); AllocPlan(index).create_fence_if_async = true; diff --git a/onnxruntime/core/optimizer/transformer_memcpy.cc b/onnxruntime/core/optimizer/transformer_memcpy.cc index 7e9254d900..0e753f386a 100644 --- a/onnxruntime/core/optimizer/transformer_memcpy.cc +++ b/onnxruntime/core/optimizer/transformer_memcpy.cc @@ -175,25 +175,32 @@ void TransformerMemcpyImpl::ProcessDefs(onnxruntime::Node& node, const KernelReg const KernelCreateInfo* kci = nullptr; kernel_registries.SearchKernelRegistry(node, &kci); - auto status = onnxruntime::Node::ForEachWithIndex( - node.InputDefs(), - [this, &kci, &initializers_consumed](const onnxruntime::NodeArg& arg, size_t index) { + bool is_implicit_input = false; + auto process_inputs = + [this, &kci, &initializers_consumed, &is_implicit_input](const onnxruntime::NodeArg& arg, size_t index) { // check if this NodeArg is an initializer defined in current outer graph level - const auto* initializer_tensor_proto = - GetInitializer(graph_, arg.Name(), true); - if (initializer_tensor_proto != nullptr) + const auto* initializer_tensor_proto = GetInitializer(graph_, arg.Name(), true); + if (initializer_tensor_proto != nullptr) { initializers_consumed[arg.Name()] = initializer_tensor_proto; - if (kci && kci->kernel_def->IsInputOnCpu(index)) - non_provider_input_defs_.insert(&arg); - else - provider_input_defs_.insert(&arg); - return Status::OK(); - }); + } + // implicit inputs have no location info in the kernel def, so treat them as provider inputs. + // PlannerImpl::ComputeUseCounts has matching logic so the allocation plan does the same thing + if (!is_implicit_input && kci && kci->kernel_def->IsInputOnCpu(index)) { + non_provider_input_defs_.insert(&arg); + } else { + provider_input_defs_.insert(&arg); + } + + return Status::OK(); + }; + + auto status = onnxruntime::Node::ForEachWithIndex(node.InputDefs(), process_inputs); ORT_ENFORCE(status.IsOK(), status.ErrorMessage()); - // we don't need to handle implicit input here as provider_ is never kCpuExecutionProvider, all control flow - // nodes are CPU based, and only control flow nodes have implicit inputs. + is_implicit_input = true; + status = onnxruntime::Node::ForEachWithIndex(node.ImplicitInputDefs(), process_inputs); + ORT_ENFORCE(status.IsOK(), status.ErrorMessage()); auto& output_defs = node.MutableOutputDefs(); for (size_t i = 0; i < output_defs.size(); ++i) { diff --git a/onnxruntime/core/providers/cpu/controlflow/loop.cc b/onnxruntime/core/providers/cpu/controlflow/loop.cc index a25a4b2cb9..8da807bc9f 100644 --- a/onnxruntime/core/providers/cpu/controlflow/loop.cc +++ b/onnxruntime/core/providers/cpu/controlflow/loop.cc @@ -150,7 +150,8 @@ class LoopImpl { public: LoopImpl(OpKernelContextInternal& context, const SessionState& session_state, - const Loop::Info& info); + const Loop::Info& info, + const Loop::ConcatOutput& concat_output_func); // Initialize by validating all the inputs, and allocating the output tensors Status Initialize(); @@ -181,8 +182,40 @@ class LoopImpl { // collection of OrtValue outputs from each loop iteration for the loop outputs. // the order from the subgraph matches the order from the loop output std::vector> loop_output_tensors_; + + const Loop::ConcatOutput& concat_output_func_; }; +static Status ConcatenateCpuOutput(std::vector& per_iteration_output, + void* output, size_t output_size_in_bytes) { + const auto& first_output = per_iteration_output.front().Get(); + const auto& per_iteration_shape = first_output.Shape(); + size_t bytes_per_iteration = first_output.SizeInBytes(); + + // we can't easily use a C++ template for the tensor element type, + // so use a span for some protection but work in bytes + gsl::span output_span = gsl::make_span(static_cast(output), + output_size_in_bytes); + + for (size_t i = 0, num_iterations = per_iteration_output.size(); i < num_iterations; ++i) { + auto& ort_value = per_iteration_output[i]; + auto& iteration_data = ort_value.Get(); + + // sanity check + if (bytes_per_iteration != iteration_data.SizeInBytes()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Inconsistent shape in loop output for output. ", + " Expected:", per_iteration_shape, " Got:", iteration_data.Shape()); + } + + auto src = gsl::make_span(static_cast(iteration_data.DataRaw()), + bytes_per_iteration); + auto dst = output_span.subspan(i * bytes_per_iteration, bytes_per_iteration); + gsl::copy(src, dst); + } + + return Status::OK(); +} + Loop::Loop(const OpKernelInfo& info) : OpKernel(info) { // make sure the attribute was present even though we don't need it here. // The GraphProto is loaded as a Graph instance by main Graph::Resolve, @@ -191,6 +224,8 @@ Loop::Loop(const OpKernelInfo& info) : OpKernel(info) { ONNX_NAMESPACE::GraphProto proto; ORT_ENFORCE(info.GetAttr("body", &proto).IsOK()); ORT_IGNORE_RETURN_VALUE(proto); + + concat_output_func_ = ConcatenateCpuOutput; } // we need this to be in the .cc so 'unique_ptr info_' can be handled @@ -211,7 +246,7 @@ common::Status Loop::SetupSubgraphExecutionInfo(const SessionState& session_stat feed_names.reserve(info_->num_subgraph_inputs + info_->num_implicit_inputs); // iter_num and cond subgraph inputs - created by the LoopImpl::Initialize so the name doesn't matter - // and we'll skip them when we call FindDevicesForValues + // as we skip them when we call FindDevicesForValues, and default them to always being on CPU. feed_names.push_back(info_->subgraph_input_names[0]); feed_names.push_back(info_->subgraph_input_names[1]); @@ -226,10 +261,10 @@ common::Status Loop::SetupSubgraphExecutionInfo(const SessionState& session_stat feed_names.push_back(entry->Name()); } - // iter_num and cond are created on CPU via MakeScalarMLValue so skip those (they will correctly default to CPU) - // use the SessionState from the control flow node for this lookup. - std::vector feed_locations; + // iter_num and cond are created on CPU via MakeScalarMLValue so skip those (they will correctly default to CPU). + // use the SessionState from the control flow node to find the remaining input locations. size_t start_at = 2; + std::vector feed_locations; ORT_RETURN_IF_ERROR(controlflow::detail::FindDevicesForValues(session_state, feed_names, feed_locations, start_at)); // now update the feed names to use the subgraph input names for the loop carried vars so that we can determine @@ -244,9 +279,32 @@ common::Status Loop::SetupSubgraphExecutionInfo(const SessionState& session_stat subgraph_session_state.GetOrtValueNameIdxMap(), ffm)); ORT_RETURN_IF_ERROR(utils::InitializeFeedFetchCopyInfo(subgraph_session_state, *ffm)); - // we don't provide pre-allocated fetches for Loop subgraph execution. - // use nullptr for all the fetch locations to represent that - std::vector fetch_locations(info_->num_subgraph_outputs, nullptr); + // setup the locations where we want the subgraph output to end up on + std::vector fetch_locations; + fetch_locations.reserve(info_->num_subgraph_outputs); + + // 'cond' is first output and we need it to be on CPU so we can read the latest value + const auto& cpu_allocator_info = session_state.GetExecutionProviders() + .Get(onnxruntime::kCpuExecutionProvider) + ->GetAllocator(0, OrtMemTypeDefault) + ->Info(); + fetch_locations.push_back(&cpu_allocator_info); + + // Loop state variables need to be where we can feed them in to the next iteration, so set the fetch location + // to match the feed location. + for (int i = 0; i < info_->num_loop_carried_vars; ++i) { + // +2 for both to skip the iter_num and cond input values + const auto& alloc_info = utils::FindMemoryInfoForValue(session_state, loop_inputs[i + 2]->Name()); + fetch_locations.push_back(&alloc_info); + } + + // remaining outputs we want where the matching Loop output will be allocated + const auto& loop_outputs = node.OutputDefs(); + for (size_t i = info_->num_loop_carried_vars, end = loop_outputs.size(); i < end; ++i) { + const auto& alloc_info = utils::FindMemoryInfoForValue(session_state, loop_outputs[i]->Name()); + fetch_locations.push_back(&alloc_info); + } + utils::FinalizeFeedFetchCopyInfo(subgraph_session_state, *ffm, feed_locations, fetch_locations); feeds_fetches_manager_ = std::move(ffm); @@ -255,12 +313,12 @@ common::Status Loop::SetupSubgraphExecutionInfo(const SessionState& session_stat } Status Loop::Compute(OpKernelContext* ctx) const { - auto ctx_internal = static_cast(ctx); + auto* ctx_internal = static_cast(ctx); auto* session_state = ctx_internal->SubgraphSessionState("body"); ORT_ENFORCE(session_state, "Subgraph SessionState was not found for 'body' attribute."); ORT_ENFORCE(feeds_fetches_manager_, "CreateFeedsFetchesManager must be called prior to execution of graph."); - LoopImpl loop_impl{*ctx_internal, *session_state, *info_}; + LoopImpl loop_impl{*ctx_internal, *session_state, *info_, concat_output_func_}; auto status = loop_impl.Initialize(); ORT_RETURN_IF_ERROR(status); @@ -272,11 +330,13 @@ Status Loop::Compute(OpKernelContext* ctx) const { LoopImpl::LoopImpl(OpKernelContextInternal& context, const SessionState& session_state, - const Loop::Info& subgraph_info) + const Loop::Info& subgraph_info, + const Loop::ConcatOutput& concat_output_func) : context_(context), session_state_(session_state), info_(subgraph_info), - implicit_inputs_(context_.GetImplicitInputs()) { + implicit_inputs_(context_.GetImplicitInputs()), + concat_output_func_(concat_output_func) { auto* max_trip_count_tensor = context.Input(0); max_trip_count_ = max_trip_count_tensor ? *max_trip_count_tensor->Data() : INT64_MAX; @@ -285,7 +345,7 @@ LoopImpl::LoopImpl(OpKernelContextInternal& context, } template -static OrtValue MakeScalarMLValue(AllocatorPtr& allocator, T value, bool is_1d) { +static OrtValue MakeScalarMLValue(const AllocatorPtr& allocator, T value, bool is_1d) { auto* data_type = DataTypeImpl::GetType(); std::unique_ptr p_tensor = onnxruntime::make_unique(data_type, is_1d ? TensorShape({1}) : TensorShape({}), @@ -318,17 +378,17 @@ Status LoopImpl::Initialize() { } } - AllocatorPtr allocator; - status = context_.GetTempSpaceAllocator(&allocator); - ORT_RETURN_IF_ERROR(status); - auto& subgraph_inputs = info_.subgraph.GetInputs(); auto iter_num_rank = subgraph_inputs[0]->Shape()->dim_size(); auto condition_rank = subgraph_inputs[1]->Shape()->dim_size(); - iter_num_mlvalue_ = MakeScalarMLValue(allocator, 0, iter_num_rank); - condition_mlvalue_ = MakeScalarMLValue(allocator, condition_, condition_rank); + // these need to be on CPU + auto cpu_allocator = session_state_.GetExecutionProviders() + .Get(onnxruntime::kCpuExecutionProvider) + ->GetAllocator(0, OrtMemTypeDefault); + iter_num_mlvalue_ = MakeScalarMLValue(cpu_allocator, 0, iter_num_rank); + condition_mlvalue_ = MakeScalarMLValue(cpu_allocator, condition_, condition_rank); loop_output_tensors_.resize(info_.num_outputs - info_.num_loop_carried_vars); @@ -371,38 +431,19 @@ void LoopImpl::SaveOutputsAndUpdateFeeds(const std::vector& last_outpu Status LoopImpl::ConcatenateLoopOutput(std::vector& per_iteration_output, int output_index) { const auto& first_output = per_iteration_output.front().Get(); - size_t bytes_per_iteration = first_output.SizeInBytes(); - const auto& per_iteration_shape = first_output.Shape(); - const auto& per_iteration_dims = per_iteration_shape.GetDims(); + const auto& per_iteration_dims = first_output.Shape().GetDims(); - // prepend number of iterations to the dimensions - auto num_iterations = gsl::narrow_cast(per_iteration_output.size()); - std::vector dims{num_iterations}; + std::vector dims; + dims.reserve(1 + per_iteration_output.size()); + + // first dimension is number of iterations + dims.push_back(gsl::narrow_cast(per_iteration_output.size())); std::copy(per_iteration_dims.cbegin(), per_iteration_dims.cend(), std::back_inserter(dims)); - TensorShape output_shape{dims}; + TensorShape output_shape{dims}; Tensor* output = context_.Output(output_index, output_shape); - // we can't easily use a C++ template for the tensor element type, - // so use a span for some protection but work in bytes - gsl::span output_span = gsl::make_span(static_cast(output->MutableDataRaw()), - output->SizeInBytes()); - - for (int64_t i = 0; i < num_iterations; ++i) { - auto& ort_value = per_iteration_output[i]; - auto& iteration_data = ort_value.Get(); - - // sanity check - if (bytes_per_iteration != iteration_data.SizeInBytes()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Inconsistent shape in loop output for output ", output_index, - " Expected:", per_iteration_shape, " Got:", iteration_data.Shape()); - } - - auto num_bytes = iteration_data.SizeInBytes(); - auto src = gsl::make_span(static_cast(iteration_data.DataRaw()), num_bytes); - auto dst = output_span.subspan(i * bytes_per_iteration, bytes_per_iteration); - gsl::copy(src, dst); - } + ORT_RETURN_IF_ERROR(concat_output_func_(per_iteration_output, output->MutableDataRaw(), output->SizeInBytes())); return Status::OK(); } @@ -438,9 +479,7 @@ Status LoopImpl::Execute(const FeedsFetchesManager& ffm) { auto copy_tensor_from_mlvalue_to_output = [this](const OrtValue& input, int output_idx) { auto& data = input.Get(); Tensor* output = context_.Output(output_idx, data.Shape()); - auto src = gsl::make_span(static_cast(data.DataRaw()), data.SizeInBytes()); - auto dst = gsl::make_span(static_cast(output->MutableDataRaw()), output->SizeInBytes()); - gsl::copy(src, dst); + session_state_.GetDataTransferMgr().CopyTensor(input.Get(), *output); }; // copy to Loop output diff --git a/onnxruntime/core/providers/cpu/controlflow/loop.h b/onnxruntime/core/providers/cpu/controlflow/loop.h index 478051a9d1..dbdb2efbcd 100644 --- a/onnxruntime/core/providers/cpu/controlflow/loop.h +++ b/onnxruntime/core/providers/cpu/controlflow/loop.h @@ -12,7 +12,7 @@ namespace onnxruntime { -class Loop final : public OpKernel, public controlflow::IControlFlowKernel { +class Loop : public OpKernel, public controlflow::IControlFlowKernel { public: Loop(const OpKernelInfo& info); @@ -26,9 +26,20 @@ class Loop final : public OpKernel, public controlflow::IControlFlowKernel { struct Info; ~Loop(); + // function to concatenate the OrtValue instances from each Loop iteration into a single output buffer. + // @param per_iteration_output OrtValue instances from each iteration. Never empty. All should have the same shape. + // @param output Pre-allocated output buffer. On device specific to the ExecutionProvider running the Loop node. + using ConcatOutput = std::function& per_iteration_output, + void* output, size_t output_size_in_bytes)>; + + protected: + // derived class can provide implementation for handling concatenation of Loop output on a different device + void SetConcatOutputFunc(const ConcatOutput& concat_output_func) { concat_output_func_ = concat_output_func; } + private: // Info and FeedsFetchesManager re-used for each subgraph execution. std::unique_ptr info_; std::unique_ptr feeds_fetches_manager_; + ConcatOutput concat_output_func_; }; } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/controlflow/loop.cc b/onnxruntime/core/providers/cuda/controlflow/loop.cc new file mode 100644 index 0000000000..4793e9d1a1 --- /dev/null +++ b/onnxruntime/core/providers/cuda/controlflow/loop.cc @@ -0,0 +1,82 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/cuda/controlflow/loop.h" +#include "core/providers/cuda/cuda_common.h" + +using namespace ONNX_NAMESPACE; +using namespace onnxruntime::common; + +namespace onnxruntime { +namespace cuda { + +ONNX_OPERATOR_VERSIONED_KERNEL_EX(Loop, + kOnnxDomain, + 1, 10, + kCudaExecutionProvider, + KernelDefBuilder() + .InputMemoryType(0) // 'M' needs to be on CPU + .InputMemoryType(1) // 'cond' needs to be on CPU + .TypeConstraint("I", DataTypeImpl::GetTensorType()) + .TypeConstraint("B", DataTypeImpl::GetTensorType()) + .TypeConstraint("V", DataTypeImpl::AllFixedSizeTensorTypes()), + Loop); + +// zero variadic argument support was added in opset 11. using same implementation as for previous version +ONNX_OPERATOR_KERNEL_EX(Loop, + kOnnxDomain, + 11, + kCudaExecutionProvider, + KernelDefBuilder() + .InputMemoryType(0) // 'M' needs to be on CPU + .InputMemoryType(1) // 'cond' needs to be on CPU + .TypeConstraint("I", DataTypeImpl::GetTensorType()) + .TypeConstraint("B", DataTypeImpl::GetTensorType()) + .TypeConstraint("V", DataTypeImpl::AllFixedSizeTensorTypes()), + Loop); + +static Status ConcatenateGpuOutput(std::vector& per_iteration_output, + void* output, size_t output_size_in_bytes) { + const auto& first_output = per_iteration_output.front().Get(); + const auto& per_iteration_shape = first_output.Shape(); + size_t bytes_per_iteration = first_output.SizeInBytes(); + + void* cur_output = output; + for (size_t i = 0, num_iterations = per_iteration_output.size(); i < num_iterations; ++i) { + auto& ort_value = per_iteration_output[i]; + auto& iteration_data = ort_value.Get(); + + // sanity check + if (bytes_per_iteration != iteration_data.SizeInBytes()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Inconsistent shape in loop output for output. ", + " Expected:", per_iteration_shape, " Got:", iteration_data.Shape()); + } + + CUDA_RETURN_IF_ERROR(cudaMemcpy(cur_output, iteration_data.DataRaw(), bytes_per_iteration, + cudaMemcpyDeviceToDevice)); + + cur_output = static_cast((static_cast(cur_output) + bytes_per_iteration)); + } + + ORT_ENFORCE(static_cast(cur_output) - static_cast(output) == output_size_in_bytes, + "Concatenation did not fill output buffer as expected."); + + return Status::OK(); +} + +Loop::Loop(const OpKernelInfo& info) : onnxruntime::Loop(info) { + SetConcatOutputFunc(ConcatenateGpuOutput); +} + +Status Loop::Compute(OpKernelContext* ctx) const { + // call the base CPU version. + // we have this CUDA implementation so the inputs/outputs stay on GPU where possible. + // the logic to run the subgraph must be on CPU either way. + // technically we don't need this override of Compute, but it will be optimized out and it's easier to debug + // that this implementation is being called with it. + auto status = onnxruntime::Loop::Compute(ctx); + return status; +} + +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/controlflow/loop.h b/onnxruntime/core/providers/cuda/controlflow/loop.h new file mode 100644 index 0000000000..4f3ccb7cd6 --- /dev/null +++ b/onnxruntime/core/providers/cuda/controlflow/loop.h @@ -0,0 +1,20 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/common/common.h" +#include "core/providers/cpu/controlflow/loop.h" + +namespace onnxruntime { +namespace cuda { + +// Use the CPU implementation for the logic +class Loop final : public onnxruntime::Loop { + public: + Loop(const OpKernelInfo& info); + + Status Compute(OpKernelContext* ctx) const override; +}; +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index c0604f036c..a6535c7eeb 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -547,6 +547,7 @@ class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDom class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, 10, If); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 8, 8, Scan); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 9, 10, Scan); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, 10, Loop); // opset 10 class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, 10, float, AveragePool); @@ -589,6 +590,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, double, Gemm); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, MLFloat16, Gemm); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, If); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, Loop); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, NonMaxSuppression); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, Range); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, float, ReduceL1); @@ -1002,6 +1004,7 @@ static void RegisterCudaKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, // opset 10 BuildKernelCreateInfo, @@ -1045,6 +1048,7 @@ static void RegisterCudaKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/test/providers/cpu/controlflow/loop_test.cc b/onnxruntime/test/providers/cpu/controlflow/loop_test.cc index c0a8d913a8..e6813e2a07 100644 --- a/onnxruntime/test/providers/cpu/controlflow/loop_test.cc +++ b/onnxruntime/test/providers/cpu/controlflow/loop_test.cc @@ -49,34 +49,29 @@ class LoopOpTester : public OpTester { std::vector& graph_input_defs, std::vector& graph_output_defs, std::vector>& /*add_attribute_funcs*/) override { - // add outer_scope_0 node - { - TypeProto float_scalar; - float_scalar.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT); - auto mutable_dim = float_scalar.mutable_tensor_type()->mutable_shape()->add_dim(); - mutable_dim->set_dim_value(1); + // last entry in graph_input_defs is outer_scope_0 + // last entry in graph_output_defs is the output from casting outer_scope_0 + // all other inputs/outputs go to Loop + // We do this to validate the handling of graph inputs that are used as both implicit inputs in a Loop node + // and directly by another node. + TypeProto float_scalar; + float_scalar.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT); + auto mutable_dim = float_scalar.mutable_tensor_type()->mutable_shape()->add_dim(); + mutable_dim->set_dim_value(1); - { - auto& output_arg = graph.GetOrCreateNodeArg("outer_scope_0", &float_scalar); - auto& constant = graph.AddNode("outer_scope_constant", "Constant", "Constant in outer scope", {}, - {&output_arg}); - - TensorProto value_tensor; - value_tensor.add_dims(1); - value_tensor.add_float_data(kOuterNodeAddValue); - value_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - - constant.AddAttribute("value", value_tensor); - } - } + auto& cast_node = graph.AddNode("cast", "Cast", "Use graph input in main graph", + {graph_input_defs.back()}, {graph_output_defs.back()}); + cast_node.AddAttribute("to", int64_t(TensorProto_DataType_INT64)); // add Loop node - { - auto& loop_node = graph.AddNode("loop", "Loop", "Loop node", graph_input_defs, graph_output_defs); + std::vector loop_input_defs = graph_input_defs; + std::vector loop_output_defs = graph_output_defs; + loop_input_defs.pop_back(); + loop_output_defs.pop_back(); + auto& loop_node = graph.AddNode("loop", "Loop", "Loop node", loop_input_defs, loop_output_defs); - auto body = create_subgraph_(options_); - loop_node.AddAttribute("body", body); - } + auto body = create_subgraph_(options_); + loop_node.AddAttribute("body", body); } private: @@ -346,11 +341,14 @@ void RunTest(int64_t max_iterations, test.AddInput("loop_var_0_orig", {1}, {0.f}); test.AddInput("loop_var_1_orig", {1}, {0.f}); + test.AddInput("outer_scope_0", {1}, {kOuterNodeAddValue}); test.AddOutput("loop_var_0_final", {1}, {loop_var_0_final}); test.AddOutput("loop_var_1_final", loop_var_1_final_shape, loop_var_1_final); test.AddOutput("loop_out_0_final", loop_out_0_final_shape, loop_out_0_final); + test.AddOutput("outer_scope_0_out", {1}, {int64_t(kOuterNodeAddValue)}); + if (options.mixed_execution_providers) { // we want the CUDA provider to be first, and the CPU provider second. all except the Loop node should run on // CUDA given that, which creates the scenario where we need to copy to/from CPU to execute the Loop node correctly. @@ -547,8 +545,10 @@ TEST(Loop, InfiniteLoopTermination) { test.AddInput("M", {1}, {INT64_MAX}); test.AddInput("cond", {1}, {true}); test.AddInput("fake", {1}, {0.f}); + test.AddInput("outer_scope_0", {1}, {kOuterNodeAddValue}); test.AddOutput("loop_var_0_final", {1}, {0.f}); + test.AddOutput("outer_scope_0_out", {1}, {int64_t(kOuterNodeAddValue)}); OrtRunOptions session_run_options; session_run_options.run_tag = "Loop.InfiniteLoopTermination"; @@ -648,7 +648,7 @@ TEST(Loop, Opset11WithNoVariadicInputsAndOutputs) { TypeProto bool_scalar; bool_scalar.mutable_tensor_type()->set_elem_type(TensorProto_DataType_BOOL); bool_scalar.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); - + // graph output types // constant_out TypeProto float_scalar; @@ -670,7 +670,7 @@ TEST(Loop, Opset11WithNoVariadicInputsAndOutputs) { graph.AddNode("cond_in_identity", "Identity", "Forward cond_in to cond_out", inputs, outputs); } - // produce constant_out + // produce constant_out { outputs = {&constant_out}; @@ -683,9 +683,9 @@ TEST(Loop, Opset11WithNoVariadicInputsAndOutputs) { attr_proto.set_type(AttributeProto_AttributeType_TENSOR); auto* constant_attribute_tensor_proto = attr_proto.mutable_t(); - constant_attribute_tensor_proto->mutable_dims()->Clear(); // scalar - constant_attribute_tensor_proto->set_data_type(TensorProto_DataType_FLOAT); //float scalar - *constant_attribute_tensor_proto->mutable_float_data()->Add() = 1.0f; //float scalar with value 1.0f + constant_attribute_tensor_proto->mutable_dims()->Clear(); // scalar + constant_attribute_tensor_proto->set_data_type(TensorProto_DataType_FLOAT); //float scalar + *constant_attribute_tensor_proto->mutable_float_data()->Add() = 1.0f; //float scalar with value 1.0f constant_node.AddAttribute("value", attr_proto); } @@ -704,7 +704,7 @@ TEST(Loop, Opset11WithNoVariadicInputsAndOutputs) { test.AddAttribute("body", body); test.AddInput("M", {1}, {1}); test.AddInput("cond", {1}, {true}); - // This 'Loop' has no variadic inputs to test the spec of 'Loop' opset 11 which allows + // This 'Loop' has no variadic inputs to test the spec of 'Loop' opset 11 which allows // 'Loop' to be used without variadic inputs test.AddOutput("loop_scan_out", {1}, {1.0f});