Add sequence support for identity on GPU (#7810)

* Add sequence supprot for identity on GPU

* implement TensorSeq in provider interface

* fix definition err

* Add new interface to TensorSeq

* fix comments

* fix comments

* fix mac warning

* move TensorSeq forward declaration

* add TensorSeq header

* remove declaration

* fix minor format

* fix minor format

* define TensorSeq as struct

Co-authored-by: RandySheriffH <rashuai@microsoft.com>

include/onnxruntime/core/framework/ml_value.h

@@ -10,11 +10,11 @@
 #include "core/framework/allocator.h"
 #include "core/framework/data_types.h"
 #include "core/framework/tensor.h"
+#include "core/framework/TensorSeq.h"
 #endif
 
 namespace onnxruntime {
 class SparseTensor;
-class TensorSeq;
 }  // namespace onnxruntime
 
 /**

onnxruntime/core/framework/TensorSeq.h

@@ -62,6 +62,12 @@ class TensorSeq {
     return tensors_[i];
   }
 
+  void Add(Tensor&& tensor) {
+    ORT_ENFORCE(IsSameDataType(tensor),
+                "TensorSeq: tensor to be added has a different data type.");
+    tensors_.push_back(std::move(tensor));
+  }
+
  private:
   // A sequence must be associated with only one data type and all tensors in the seq must be of that type
   // One other alternative of storing the data type of a seq is to templatize the TensorSeq class.

onnxruntime/core/framework/provider_bridge_ort.cc

@@ -19,6 +19,7 @@
 #include "core/session/ort_apis.h"
 #include "core/util/math.h"
 #include "core/framework/tensorprotoutils.h"
+#include "core/framework/TensorSeq.h"
 
 #include "core/framework/fallback_cpu_capability.h"
 #include "core/framework/random_generator.h"
@@ -485,6 +486,7 @@ struct ProviderHostImpl : ProviderHost {
 
   // DataTypeImpl (wrapped)
   MLDataType DataTypeImpl__GetType_Tensor() override { return DataTypeImpl::GetType<Tensor>(); }
+  MLDataType DataTypeImpl__GetType_TensorSeq () override { return DataTypeImpl::GetType<TensorSeq>(); }
   MLDataType DataTypeImpl__GetType_bool() override { return DataTypeImpl::GetType<bool>(); }
   MLDataType DataTypeImpl__GetType_int8() override { return DataTypeImpl::GetType<int8_t>(); }
   MLDataType DataTypeImpl__GetType_uint8() override { return DataTypeImpl::GetType<uint8_t>(); }
@@ -651,8 +653,11 @@ struct ProviderHostImpl : ProviderHost {
 
   // OpKernelContext (wrapped)
   const Tensor* OpKernelContext__Input_Tensor(const OpKernelContext* p, int index) override { return p->Input<Tensor>(index); }
+  const TensorSeq* OpKernelContext__Input_TensorSeq(const OpKernelContext* p, int index) override { return p->Input<TensorSeq>(index); }
   const Tensor& OpKernelContext__RequiredInput_Tensor(const OpKernelContext* p, int index) override { return p->RequiredInput<Tensor>(index); }
+  MLDataType OpKernelContext__InputType(const OpKernelContext* p, int index) override { return p->InputType(index); }
   Tensor* OpKernelContext__Output_Tensor(OpKernelContext* p, int index) override { return p->Output<Tensor>(index); }
+  TensorSeq* OpKernelContext__Output_TensorSeq(OpKernelContext* p, int index) override { return p->Output<TensorSeq>(index); }
   Tensor* OpKernelContext__Output(OpKernelContext* p, int index, const TensorShape& shape) override { return p->Output(index, shape); }
   Tensor& OpKernelContext__RequiredOutput(OpKernelContext* p, int index, const TensorShape& shape) override { return p->RequiredOutput(index, shape); }
   int OpKernelContext__InputCount(const OpKernelContext* p) override { return p->InputCount(); }
@@ -749,6 +754,13 @@ struct ProviderHostImpl : ProviderHost {
   int32_t Tensor__GetElementType(const Tensor* p) override { return p->GetElementType(); }
   MLDataType Tensor__DataType(const Tensor* p) override { return p->DataType(); }
 
+  // TensorSeq(wrapped)
+  MLDataType TensorSeq__DataType(const TensorSeq* p) noexcept override { return p->DataType(); }
+  void TensorSeq__SetType(TensorSeq* p, MLDataType data_type) override { p->SetType(data_type); }
+  size_t TensorSeq__Size(const TensorSeq* p) noexcept override { return p->Size(); }
+  const Tensor& TensorSeq__Get(const TensorSeq* p, size_t i) override { return p->Get(i); }
+  void TensorSeq__Add(TensorSeq* p, Tensor&& tensor) override { p->Add(std::move(tensor)); }
+
   // AllocatorManager (direct)
   void AllocatorManager__InsertAllocator(AllocatorManager* p, AllocatorPtr allocator) override { p->AllocatorManager::InsertAllocator(allocator); }
   AllocatorPtr AllocatorManager__GetAllocator(const AllocatorManager* p, int id, OrtMemType mem_type) override { return p->AllocatorManager::GetAllocator(id, mem_type); };

onnxruntime/core/framework/utils.cc

@@ -18,6 +18,7 @@
 #include "core/framework/sequential_executor.h"
 #include "core/framework/tensorprotoutils.h"
 #include "core/mlas/inc/mlas.h"
+#include "core/framework/TensorSeq.h"
 #ifdef ENABLE_TRAINING
 #include "core/framework/orttraining_partial_executor.h"
 #endif
@@ -111,19 +112,33 @@ bool ProviderIsCpuBased(const std::string& provider_type) {
 }
 
 static common::Status AllocateHelper(const AllocatorPtr& allocator,
-                                     const Tensor& fetched_tensor, OrtValue& output_mlvalue) {
+                                     const OrtValue& source_mlvalue,
+                                     OrtValue& target_mlvalue) {
   if (!allocator) {
-    return Status(common::ONNXRUNTIME, common::FAIL, "invalid allocator");
+    return Status(common::ONNXRUNTIME, common::FAIL, "invalid allocator.");
   }
+  if (source_mlvalue.IsTensor()) {
 
-  std::unique_ptr<Tensor> p_tensor = std::make_unique<Tensor>(fetched_tensor.DataType(),
-                                                                      fetched_tensor.Shape(),
-                                                                      allocator);
-  auto ml_tensor = DataTypeImpl::GetType<Tensor>();
-  output_mlvalue.Init(p_tensor.release(),
-                      ml_tensor,
-                      ml_tensor->GetDeleteFunc());
+    const Tensor& source_tensor = source_mlvalue.Get<Tensor>();
+    std::unique_ptr<Tensor> target_tensor = std::make_unique<Tensor>(source_tensor.DataType(),
+                                                                     source_tensor.Shape(),
+                                                                     allocator);
+    auto ml_tensor = DataTypeImpl::GetType<Tensor>();
+    target_mlvalue.Init(target_tensor.release(), ml_tensor, ml_tensor->GetDeleteFunc());
 
+  } else if (source_mlvalue.IsTensorSequence()) {
+    const TensorSeq& source_tensor_seq = source_mlvalue.Get<TensorSeq>();
+    auto target_tensor_seq = std::make_unique<TensorSeq>(source_tensor_seq.DataType());
+    std::vector<Tensor> tensors;
+    for (auto iter = source_tensor_seq.begin(); iter != source_tensor_seq.end(); ++iter) {
+      tensors.emplace_back(iter->DataType(), onnxruntime::TensorShape(iter->Shape()), allocator);
+    }
+    target_tensor_seq->SetElements(std::move(tensors)); 
+    auto ml_tensor_seq = DataTypeImpl::GetType<TensorSeq>();
+    target_mlvalue.Init(target_tensor_seq.release(), ml_tensor_seq, ml_tensor_seq->GetDeleteFunc());
+  } else {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Unsupported OrtValue type.");
+  }
   return Status::OK();
 }
 
@@ -159,22 +174,40 @@ static Status BatchOrCopyMLValue(const SessionState& session_state,
     return Status::OK();
   }
 
-  auto& source_tensor = source_mlvalue.Get<Tensor>();
   if (!target_mlvalue.IsAllocated()) {
     auto allocator = session_state.GetAllocator(copy_info.target_device);
     ORT_ENFORCE(allocator != nullptr, "Failed to find allocator for device ", copy_info.target_device.ToString());
-
-    ORT_RETURN_IF_ERROR(utils::AllocateHelper(allocator, source_tensor, target_mlvalue));
+    ORT_RETURN_IF_ERROR(utils::AllocateHelper(allocator, source_mlvalue, target_mlvalue));
   }
 
-  Tensor* p_output_tensor = target_mlvalue.GetMutable<Tensor>();
-
-  if (copy_pairs != nullptr) {
-    copy_pairs->push_back({source_tensor, *p_output_tensor, 0});
+  if (source_mlvalue.IsTensor()) {
+    const Tensor& source_tensor = source_mlvalue.Get<Tensor>();
+    Tensor& target_tensor = *target_mlvalue.GetMutable<Tensor>();
+    if (copy_pairs != nullptr) {
+      copy_pairs->push_back({source_tensor, target_tensor, 0});
+    } else {
+      ORT_RETURN_IF_ERROR(session_state.GetDataTransferMgr().CopyTensor(source_tensor, target_tensor));
+    }
+  } else if (source_mlvalue.IsTensorSequence()) {
+    const TensorSeq& source_tensor_seq = source_mlvalue.Get<TensorSeq>();
+    const TensorSeq& target_tensor_seq = target_mlvalue.Get<TensorSeq>();
+    ORT_ENFORCE(source_tensor_seq.Size() == target_tensor_seq.Size(),
+      "source and target tensor sequence have different number of elements.");
+    auto source_iter = source_tensor_seq.begin();
+    auto target_iter = target_tensor_seq.begin();
+    while (source_iter != source_tensor_seq.end() &&
+           target_iter != target_tensor_seq.end()) {
+      if (copy_pairs != nullptr) {
+        copy_pairs->push_back({*source_iter, const_cast<Tensor&>(*target_iter), 0});
+      } else {
+        ORT_RETURN_IF_ERROR(session_state.GetDataTransferMgr().CopyTensor(*source_iter, const_cast<Tensor&>(*target_iter)));
+      }
+      ++source_iter;
+      ++target_iter;
+    }//while
   } else {
-    ORT_RETURN_IF_ERROR(session_state.GetDataTransferMgr().CopyTensor(source_tensor, *p_output_tensor));
+    return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Unsupported OrtValue type to copy between device.");
   }
-
   return Status::OK();
 }
 

onnxruntime/core/providers/cuda/tensor/identity_op.cc

@@ -57,7 +57,7 @@ ONNX_OPERATOR_KERNEL_EX(
     14,
     kCudaExecutionProvider,
     (*KernelDefBuilder::Create())
-        .TypeConstraint("V", DataTypeImpl::AllFixedSizeTensorTypes())
+        .TypeConstraint("V", DataTypeImpl::AllTensorAndSequenceTensorTypes())
         .Alias(0, 0),
     IdentityOp<false>);
 }  // namespace cuda

onnxruntime/core/providers/cuda/tensor/identity_op.h

@@ -15,34 +15,75 @@ class IdentityOp final : public CudaKernel {
   }
 
   Status ComputeInternal(OpKernelContext* context) const override {
-    const Tensor* X = context->Input<Tensor>(0);
-    if (X == nullptr) return Status(common::ONNXRUNTIME, common::FAIL, "input count mismatch");
-    const TensorShape& shape = X->Shape();
-    Tensor* Y = context->Output(0, shape);
-    auto X_type = X->DataType();
-
-    const void* source = X->DataRaw(X_type);
-    void* target = Y->MutableDataRaw(X_type);
-    //If source and target pointers are not equal, we need to copy the data.
-    if (target != source) {
-      CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(target, source, X->Shape().Size() * X->DataType()->Size(), cudaMemcpyDeviceToDevice, Stream()));
-    }
-
-    if (is_dropout) {
-      Tensor* mask = context->Output(1, shape);
-      // a 'nullptr' returned would make it an unused optional output
-      if (mask != nullptr) {
-        // Opset 7 differs with Opset 10 in that the type of the 'mask'
-        // output is tied with the type of the input in Opset 7 whereas
-        // the type of 'mask' in Opset 10 is 'bool' always
-        // so we have a common solution
-        void* mask_data = mask->MutableDataRaw();
-        // In 'test'/'inference' mode, there are no input values dropped out
-        // so fill the buffer with 0/false
-        CUDA_RETURN_IF_ERROR(cudaMemsetAsync(mask_data, 0, mask->SizeInBytes(), Stream()));
+    auto X_ml_type = context->InputType(0);
+    if (X_ml_type->IsTensorType()) {
+      const Tensor* X = context->Input<Tensor>(0);
+      if (nullptr == X) {
+        return Status(common::ONNXRUNTIME, common::FAIL,
+                      "IdentityOp cuda: input count mismatch.");
       }
-    }
+      const TensorShape& shape = X->Shape();
+      Tensor* Y = context->Output(0, shape);
+      if (nullptr == Y) {
+        return Status(common::ONNXRUNTIME, common::FAIL,
+                      "IdentityOp cuda: failed to allocate output tensor.");
+      }
+      auto X_type = X->DataType();
 
+      const void* source = X->DataRaw(X_type);
+      void* target = Y->MutableDataRaw(X_type);
+      //If source and target pointers are not equal, we need to copy the data.
+      if (target != source) {
+        CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(target, source, X->Shape().Size() * X->DataType()->Size(), cudaMemcpyDeviceToDevice, Stream()));
+      }
+
+      if (is_dropout) {
+        Tensor* mask = context->Output(1, shape);
+        // a 'nullptr' returned would make it an unused optional output
+        if (mask != nullptr) {
+          // Opset 7 differs with Opset 10 in that the type of the 'mask'
+          // output is tied with the type of the input in Opset 7 whereas
+          // the type of 'mask' in Opset 10 is 'bool' always
+          // so we have a common solution
+          void* mask_data = mask->MutableDataRaw();
+          // In 'test'/'inference' mode, there are no input values dropped out
+          // so fill the buffer with 0/false
+          CUDA_RETURN_IF_ERROR(cudaMemsetAsync(mask_data, 0, mask->SizeInBytes(), Stream()));
+        }
+      }
+    } else if (X_ml_type->IsTensorSequenceType()) {
+      const TensorSeq* X = context->Input<TensorSeq>(0);
+      if (nullptr == X) {
+        return Status(common::ONNXRUNTIME, common::FAIL,
+                      "IdentityOp cuda: input tensor is missing.");
+      }
+      TensorSeq* Y = context->Output<TensorSeq>(0);
+      if (nullptr == Y) {
+        return Status(common::ONNXRUNTIME, common::FAIL,
+                      "IdentityOp cuda: failed to allocate output tensor sequence.");
+      }
+      auto X_type = X->DataType();
+      Y->SetType(X_type);
+      AllocatorPtr alloc;
+      auto status = context->GetTempSpaceAllocator(&alloc);
+      if (!status.IsOK()) {
+        return Status(common::ONNXRUNTIME, common::FAIL,
+                      "IdentityOp cuda: unable to get an allocator.");
+      }
+      auto X_size = X->Size();
+      for (size_t i = 0; i < X_size; ++i) {
+        const Tensor& source_tensor = X->Get(i);
+        std::unique_ptr<Tensor> target_tensor = Tensor::Create(X_type, source_tensor.Shape(), alloc);
+        CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(target_tensor->MutableDataRaw(),
+                                             source_tensor.DataRaw(),
+                                             source_tensor.SizeInBytes(),
+                                             cudaMemcpyDeviceToDevice, Stream()));
+        Y->Add(std::move(*target_tensor));
+      }
+    } else {
+      return Status(common::ONNXRUNTIME, common::FAIL,
+                    "IdentityOp cuda: unsupported input type.");
+    }
     return Status::OK();
   }
 };

onnxruntime/core/providers/shared_library/provider_api.h

@@ -169,6 +169,7 @@ struct OpKernelContext;
 struct OpKernelInfo;
 struct PrimitiveDataTypeBase;
 struct Tensor;
+struct TensorSeq;
 
 class UnsqueezeBase;
 class SliceBase;

onnxruntime/core/providers/shared_library/provider_bridge_provider.cc

@@ -98,6 +98,8 @@ AllocatorPtr AllocatorManager::GetAllocator(int id, OrtMemType mem_type) const {
 template <>
 MLDataType DataTypeImpl::GetType<Tensor>() { return Provider_GetHost()->DataTypeImpl__GetType_Tensor(); }
 template <>
+MLDataType DataTypeImpl::GetType<TensorSeq>() { return Provider_GetHost()->DataTypeImpl__GetType_TensorSeq(); }
+template <>
 MLDataType DataTypeImpl::GetType<bool>() { return Provider_GetHost()->DataTypeImpl__GetType_bool(); }
 template <>
 MLDataType DataTypeImpl::GetType<int8_t>() { return Provider_GetHost()->DataTypeImpl__GetType_int8(); }

onnxruntime/core/providers/shared_library/provider_interfaces.h

@@ -409,6 +409,7 @@ struct ProviderHost {
 
   // DataTypeImpl
   virtual MLDataType DataTypeImpl__GetType_Tensor() = 0;
+  virtual MLDataType DataTypeImpl__GetType_TensorSeq() = 0;
   virtual MLDataType DataTypeImpl__GetType_bool() = 0;
   virtual MLDataType DataTypeImpl__GetType_int8() = 0;
   virtual MLDataType DataTypeImpl__GetType_uint8() = 0;
@@ -562,10 +563,13 @@ struct ProviderHost {
 
   // OpKernelContext
   virtual const Tensor* OpKernelContext__Input_Tensor(const OpKernelContext* p, int index) = 0;
+  virtual const TensorSeq* OpKernelContext__Input_TensorSeq(const OpKernelContext* p, int index) = 0;
   virtual const Tensor& OpKernelContext__RequiredInput_Tensor(const OpKernelContext* p, int index) = 0;
   virtual Tensor* OpKernelContext__Output_Tensor(OpKernelContext* p, int index) = 0;
+  virtual TensorSeq* OpKernelContext__Output_TensorSeq(OpKernelContext* p, int index) = 0;
   virtual Tensor* OpKernelContext__Output(OpKernelContext* p, int index, const TensorShape& shape) = 0;
   virtual Tensor& OpKernelContext__RequiredOutput(OpKernelContext* p, int index, const TensorShape& shape) = 0;
+  virtual MLDataType OpKernelContext__InputType(const OpKernelContext* p, int index) = 0; 
   virtual int OpKernelContext__InputCount(const OpKernelContext* p) = 0;
   virtual int OpKernelContext__OutputCount(const OpKernelContext* p) = 0;
   virtual Status OpKernelContext__GetTempSpaceAllocator(const OpKernelContext* p, AllocatorPtr* output) = 0;
@@ -660,6 +664,13 @@ struct ProviderHost {
   virtual int32_t Tensor__GetElementType(const Tensor* p) = 0;
   virtual MLDataType Tensor__DataType(const Tensor* p) = 0;
 
+  // TensorSeq
+  virtual MLDataType TensorSeq__DataType(const TensorSeq* p) noexcept = 0;
+  virtual void TensorSeq__SetType(TensorSeq* p, MLDataType data_type) = 0;
+  virtual size_t TensorSeq__Size(const TensorSeq* p) noexcept = 0;
+  virtual const Tensor& TensorSeq__Get(const TensorSeq* p, size_t i) = 0;
+  virtual void TensorSeq__Add(TensorSeq* p, Tensor&& tensor) = 0;
+
   // AllocatorManager
   virtual void AllocatorManager__InsertAllocator(AllocatorManager* p, AllocatorPtr allocator) = 0;
   virtual AllocatorPtr AllocatorManager__GetAllocator(const AllocatorManager* p, int id, OrtMemType mem_type) = 0;
@@ -1445,6 +1456,8 @@ struct OpKernelContext final {
   const T* Input(int index) const;
   int InputCount() const { return g_host->OpKernelContext__InputCount(this); }
 
+  MLDataType InputType(int index) const { return g_host->OpKernelContext__InputType(this, index); }
+
   template <typename T>
   T* Output(int index);
 
@@ -1466,11 +1479,21 @@ inline const Tensor* OpKernelContext::Input<Tensor>(int index) const {
   return g_host->OpKernelContext__Input_Tensor(this, index);
 }
 
+template <>
+inline const TensorSeq* OpKernelContext::Input<TensorSeq>(int index) const {
+  return g_host->OpKernelContext__Input_TensorSeq(this, index);
+}
+
 template <>
 inline Tensor* OpKernelContext::Output<Tensor>(int index) {
   return g_host->OpKernelContext__Output_Tensor(this, index);
 }
 
+template <>
+inline TensorSeq* OpKernelContext::Output<TensorSeq>(int index) {
+  return g_host->OpKernelContext__Output_TensorSeq(this, index);
+}
+
 template <>
 inline const Tensor& OpKernelContext::RequiredInput(int index) const {
   return g_host->OpKernelContext__RequiredInput_Tensor(this, index);
@@ -1663,6 +1686,15 @@ inline const BFloat16* Tensor::Data<BFloat16>() const { return g_host->Tensor__D
 template <>
 inline const MLFloat16* Tensor::Data<MLFloat16>() const { return g_host->Tensor__Data_MLFloat16(this); }
 
+//TensorSeq
+struct TensorSeq final {
+  MLDataType DataType() const noexcept { return g_host->TensorSeq__DataType(this); }
+  void SetType(MLDataType elem_type) { g_host->TensorSeq__SetType(this, elem_type); }
+  size_t Size() const noexcept { return g_host->TensorSeq__Size(this); }
+  const Tensor& Get(size_t i) const { return g_host->TensorSeq__Get(this, i); }
+  void Add(Tensor&& tensor) { g_host->TensorSeq__Add(this, std::move(tensor)); }
+};
+
 template <>
 inline gsl::span<const int64_t> Tensor::DataAsSpan() const { return g_host->Tensor__DataAsSpan_int64(this); }