diff --git a/include/onnxruntime/core/framework/ml_value.h b/include/onnxruntime/core/framework/ml_value.h index 96d6ee88a1..1217235df5 100644 --- a/include/onnxruntime/core/framework/ml_value.h +++ b/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 /** diff --git a/onnxruntime/core/framework/TensorSeq.h b/onnxruntime/core/framework/TensorSeq.h index a17a0866c8..8f69f6f972 100644 --- a/onnxruntime/core/framework/TensorSeq.h +++ b/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. diff --git a/onnxruntime/core/framework/provider_bridge_ort.cc b/onnxruntime/core/framework/provider_bridge_ort.cc index 155b8c8b6e..8dc938a055 100644 --- a/onnxruntime/core/framework/provider_bridge_ort.cc +++ b/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(); } + MLDataType DataTypeImpl__GetType_TensorSeq () override { return DataTypeImpl::GetType(); } MLDataType DataTypeImpl__GetType_bool() override { return DataTypeImpl::GetType(); } MLDataType DataTypeImpl__GetType_int8() override { return DataTypeImpl::GetType(); } MLDataType DataTypeImpl__GetType_uint8() override { return DataTypeImpl::GetType(); } @@ -651,8 +653,11 @@ struct ProviderHostImpl : ProviderHost { // OpKernelContext (wrapped) const Tensor* OpKernelContext__Input_Tensor(const OpKernelContext* p, int index) override { return p->Input(index); } + const TensorSeq* OpKernelContext__Input_TensorSeq(const OpKernelContext* p, int index) override { return p->Input(index); } const Tensor& OpKernelContext__RequiredInput_Tensor(const OpKernelContext* p, int index) override { return p->RequiredInput(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(index); } + TensorSeq* OpKernelContext__Output_TensorSeq(OpKernelContext* p, int index) override { return p->Output(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); }; diff --git a/onnxruntime/core/framework/utils.cc b/onnxruntime/core/framework/utils.cc index ec3c651adf..648a3be1c8 100644 --- a/onnxruntime/core/framework/utils.cc +++ b/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 p_tensor = std::make_unique(fetched_tensor.DataType(), - fetched_tensor.Shape(), - allocator); - auto ml_tensor = DataTypeImpl::GetType(); - output_mlvalue.Init(p_tensor.release(), - ml_tensor, - ml_tensor->GetDeleteFunc()); + const Tensor& source_tensor = source_mlvalue.Get(); + std::unique_ptr target_tensor = std::make_unique(source_tensor.DataType(), + source_tensor.Shape(), + allocator); + auto ml_tensor = DataTypeImpl::GetType(); + target_mlvalue.Init(target_tensor.release(), ml_tensor, ml_tensor->GetDeleteFunc()); + } else if (source_mlvalue.IsTensorSequence()) { + const TensorSeq& source_tensor_seq = source_mlvalue.Get(); + auto target_tensor_seq = std::make_unique(source_tensor_seq.DataType()); + std::vector 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(); + 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(); 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(); - - 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& target_tensor = *target_mlvalue.GetMutable(); + 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(); + const TensorSeq& target_tensor_seq = target_mlvalue.Get(); + 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(*target_iter), 0}); + } else { + ORT_RETURN_IF_ERROR(session_state.GetDataTransferMgr().CopyTensor(*source_iter, const_cast(*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(); } diff --git a/onnxruntime/core/providers/cuda/tensor/identity_op.cc b/onnxruntime/core/providers/cuda/tensor/identity_op.cc index 9281e3a42b..563bb45643 100644 --- a/onnxruntime/core/providers/cuda/tensor/identity_op.cc +++ b/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); } // namespace cuda diff --git a/onnxruntime/core/providers/cuda/tensor/identity_op.h b/onnxruntime/core/providers/cuda/tensor/identity_op.h index 4bac78d042..35f695c9d9 100644 --- a/onnxruntime/core/providers/cuda/tensor/identity_op.h +++ b/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(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(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(0); + if (nullptr == X) { + return Status(common::ONNXRUNTIME, common::FAIL, + "IdentityOp cuda: input tensor is missing."); + } + TensorSeq* Y = context->Output(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 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(); } }; diff --git a/onnxruntime/core/providers/shared_library/provider_api.h b/onnxruntime/core/providers/shared_library/provider_api.h index aa3304718f..a2333972c2 100644 --- a/onnxruntime/core/providers/shared_library/provider_api.h +++ b/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; diff --git a/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc b/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc index 058211379d..014e6a0dc0 100644 --- a/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc +++ b/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() { return Provider_GetHost()->DataTypeImpl__GetType_Tensor(); } template <> +MLDataType DataTypeImpl::GetType() { return Provider_GetHost()->DataTypeImpl__GetType_TensorSeq(); } +template <> MLDataType DataTypeImpl::GetType() { return Provider_GetHost()->DataTypeImpl__GetType_bool(); } template <> MLDataType DataTypeImpl::GetType() { return Provider_GetHost()->DataTypeImpl__GetType_int8(); } diff --git a/onnxruntime/core/providers/shared_library/provider_interfaces.h b/onnxruntime/core/providers/shared_library/provider_interfaces.h index 7261b5f1ea..d680a616c7 100644 --- a/onnxruntime/core/providers/shared_library/provider_interfaces.h +++ b/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 T* Output(int index); @@ -1466,11 +1479,21 @@ inline const Tensor* OpKernelContext::Input(int index) const { return g_host->OpKernelContext__Input_Tensor(this, index); } +template <> +inline const TensorSeq* OpKernelContext::Input(int index) const { + return g_host->OpKernelContext__Input_TensorSeq(this, index); +} + template <> inline Tensor* OpKernelContext::Output(int index) { return g_host->OpKernelContext__Output_Tensor(this, index); } +template <> +inline TensorSeq* OpKernelContext::Output(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() const { return g_host->Tensor__D template <> inline const MLFloat16* Tensor::Data() 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 Tensor::DataAsSpan() const { return g_host->Tensor__DataAsSpan_int64(this); }