CUDA Loop (#2444)

* Implement CUDA Loop operator.

* Add control flow node implicit input handling to the memcpy transformer and allocation planner.

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<size_t>(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<size_t>(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<size_t>(index), exec_provider->GetAllocator(0, p_kernelDef->OutputMemoryType(i))->Info());
+        plan_.SetLocation(static_cast<size_t>(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;

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

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<std::vector<OrtValue>> loop_output_tensors_;
+
+  const Loop::ConcatOutput& concat_output_func_;
 };
 
+static Status ConcatenateCpuOutput(std::vector<OrtValue>& per_iteration_output,
+                                   void* output, size_t output_size_in_bytes) {
+  const auto& first_output = per_iteration_output.front().Get<Tensor>();
+  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<gsl::byte> output_span = gsl::make_span<gsl::byte>(static_cast<gsl::byte*>(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<Tensor>();
+
+    // 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<const gsl::byte>(static_cast<const gsl::byte*>(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<ONNX_NAMESPACE::GraphProto>("body", &proto).IsOK());
   ORT_IGNORE_RETURN_VALUE(proto);
+
+  concat_output_func_ = ConcatenateCpuOutput;
 }
 
 // we need this to be in the .cc so 'unique_ptr<Info> 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<OrtDevice> 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<OrtDevice> 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<const OrtMemoryInfo*> fetch_locations(info_->num_subgraph_outputs, nullptr);
+  // setup the locations where we want the subgraph output to end up on
+  std::vector<const OrtMemoryInfo*> 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<OpKernelContextInternal*>(ctx);
+  auto* ctx_internal = static_cast<OpKernelContextInternal*>(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<Tensor>(0);
   max_trip_count_ = max_trip_count_tensor ? *max_trip_count_tensor->Data<int64_t>() : INT64_MAX;
 
@@ -285,7 +345,7 @@ LoopImpl::LoopImpl(OpKernelContextInternal& context,
 }
 
 template <typename T>
-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<T>();
   std::unique_ptr<Tensor> p_tensor = onnxruntime::make_unique<Tensor>(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<int64_t>(allocator, 0, iter_num_rank);
-  condition_mlvalue_ = MakeScalarMLValue<bool>(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<int64_t>(cpu_allocator, 0, iter_num_rank);
+  condition_mlvalue_ = MakeScalarMLValue<bool>(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<OrtValue>& last_outpu
 
 Status LoopImpl::ConcatenateLoopOutput(std::vector<OrtValue>& per_iteration_output, int output_index) {
   const auto& first_output = per_iteration_output.front().Get<Tensor>();
-  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<int64_t>(per_iteration_output.size());
-  std::vector<int64_t> dims{num_iterations};
+  std::vector<int64_t> dims;
+  dims.reserve(1 + per_iteration_output.size());
+
+  // first dimension is number of iterations
+  dims.push_back(gsl::narrow_cast<int64_t>(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<gsl::byte> output_span = gsl::make_span<gsl::byte>(static_cast<gsl::byte*>(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<Tensor>();
-
-    // 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<const gsl::byte>(static_cast<const gsl::byte*>(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>();
     Tensor* output = context_.Output(output_idx, data.Shape());
-    auto src = gsl::make_span<const gsl::byte>(static_cast<const gsl::byte*>(data.DataRaw()), data.SizeInBytes());
-    auto dst = gsl::make_span<gsl::byte>(static_cast<gsl::byte*>(output->MutableDataRaw()), output->SizeInBytes());
-    gsl::copy(src, dst);
+    session_state_.GetDataTransferMgr().CopyTensor(input.Get<Tensor>(), *output);
   };
 
   // copy to Loop output

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<Status(std::vector<OrtValue>& 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> info_;
   std::unique_ptr<FeedsFetchesManager> feeds_fetches_manager_;
+  ConcatOutput concat_output_func_;
 };
 }  // namespace onnxruntime

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<OrtMemTypeCPUInput>(0)  // 'M' needs to be on CPU
+                                      .InputMemoryType<OrtMemTypeCPUInput>(1)  // 'cond' needs to be on CPU
+                                      .TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>())
+                                      .TypeConstraint("B", DataTypeImpl::GetTensorType<bool>())
+                                      .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<OrtMemTypeCPUInput>(0)  // 'M' needs to be on CPU
+                            .InputMemoryType<OrtMemTypeCPUInput>(1)  // 'cond' needs to be on CPU
+                            .TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>())
+                            .TypeConstraint("B", DataTypeImpl::GetTensorType<bool>())
+                            .TypeConstraint("V", DataTypeImpl::AllFixedSizeTensorTypes()),
+                        Loop);
+
+static Status ConcatenateGpuOutput(std::vector<OrtValue>& per_iteration_output,
+                                   void* output, size_t output_size_in_bytes) {
+  const auto& first_output = per_iteration_output.front().Get<Tensor>();
+  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<Tensor>();
+
+    // 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<void*>((static_cast<gsl::byte*>(cur_output) + bytes_per_iteration));
+  }
+
+  ORT_ENFORCE(static_cast<gsl::byte*>(cur_output) - static_cast<gsl::byte*>(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

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

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<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, 9, TopK)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 8, 8, Scan)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 9, 10, Scan)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, 10, Loop)>,
 
       // opset 10
       BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, 10, float, AveragePool)>,
@@ -1045,6 +1048,7 @@ static void RegisterCudaKernels(KernelRegistry& kernel_registry) {
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, float, Gemm)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, double, Gemm)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, If)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, Loop)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, NonMaxSuppression)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, Range)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, float, ReduceL1)>,

onnxruntime/test/providers/cpu/controlflow/loop_test.cc

@@ -49,34 +49,29 @@ class LoopOpTester : public OpTester {
                 std::vector<onnxruntime::NodeArg*>& graph_input_defs,
                 std::vector<onnxruntime::NodeArg*>& graph_output_defs,
                 std::vector<std::function<void(onnxruntime::Node& node)>>& /*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<onnxruntime::NodeArg*> loop_input_defs = graph_input_defs;
+    std::vector<onnxruntime::NodeArg*> 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<float>("loop_var_0_orig", {1}, {0.f});
   test.AddInput<float>("loop_var_1_orig", {1}, {0.f});
+  test.AddInput<float>("outer_scope_0", {1}, {kOuterNodeAddValue});
 
   test.AddOutput<float>("loop_var_0_final", {1}, {loop_var_0_final});
   test.AddOutput<float>("loop_var_1_final", loop_var_1_final_shape, loop_var_1_final);
   test.AddOutput<float>("loop_out_0_final", loop_out_0_final_shape, loop_out_0_final);
 
+  test.AddOutput<int64_t>("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<int64_t>("M", {1}, {INT64_MAX});
   test.AddInput<bool>("cond", {1}, {true});
   test.AddInput<float>("fake", {1}, {0.f});
+  test.AddInput<float>("outer_scope_0", {1}, {kOuterNodeAddValue});
 
   test.AddOutput<float>("loop_var_0_final", {1}, {0.f});
+  test.AddOutput<int64_t>("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<GraphProto>("body", body);
   test.AddInput<int64_t>("M", {1}, {1});
   test.AddInput<bool>("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<float>("loop_scan_out", {1}, {1.0f});