From b0019ac7fef58c40d731850c74403fe706e85c77 Mon Sep 17 00:00:00 2001 From: Ke Zhang Date: Thu, 9 Jan 2020 16:52:34 -0800 Subject: [PATCH] add interface to copy batch tensors. (#2807) * add interface to copy batch tensors. * onnxruntime --- include/onnxruntime/core/session/onnxruntime_c_api.h | 12 ++++++------ onnxruntime/core/framework/data_transfer.cc | 8 ++++++++ onnxruntime/core/framework/data_transfer.h | 1 + onnxruntime/core/session/onnxruntime_c_api.cc | 10 +--------- 4 files changed, 16 insertions(+), 15 deletions(-) diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h index c5a42d66cf..ed07c3bbb2 100644 --- a/include/onnxruntime/core/session/onnxruntime_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_c_api.h @@ -253,7 +253,7 @@ struct OrtApi { // TODO: document the path separator convention? '/' vs '\' // TODO: should specify the access characteristics of model_path. Is this read only during the - // execution of OrtCreateSession, or does the OrtSession retain a handle to the file/directory + // execution of CreateSession, or does the OrtSession retain a handle to the file/directory // and continue to access throughout the OrtSession lifetime? // What sort of access is needed to model_path : read or read/write? OrtStatus*(ORT_API_CALL* CreateSession)(_In_ const OrtEnv* env, _In_ const ORTCHAR_T* model_path, @@ -432,22 +432,22 @@ struct OrtApi { OrtStatus*(ORT_API_CALL* GetTensorMutableData)(_Inout_ OrtValue* value, _Outptr_ void** out)NO_EXCEPTION; /** - * \param value A tensor created from OrtCreateTensor... function. + * \param value A tensor created from CreateTensor... function. * \param s each A string array. Each string in this array must be null terminated. * \param s_len length of s */ OrtStatus*(ORT_API_CALL* FillStringTensor)(_Inout_ OrtValue* value, _In_ const char* const* s, size_t s_len)NO_EXCEPTION; /** - * \param value A tensor created from OrtCreateTensor... function. + * \param value A tensor created from CreateTensor... function. * \param len total data length, not including the trailing '\0' chars. */ OrtStatus*(ORT_API_CALL* GetStringTensorDataLength)(_In_ const OrtValue* value, _Out_ size_t* len)NO_EXCEPTION; /** * \param s string contents. Each string is NOT null-terminated. - * \param value A tensor created from OrtCreateTensor... function. - * \param s_len total data length, get it from OrtGetStringTensorDataLength + * \param value A tensor created from CreateTensor... function. + * \param s_len total data length, get it from GetStringTensorDataLength */ OrtStatus*(ORT_API_CALL* GetStringTensorContent)(_In_ const OrtValue* value, _Out_ void* s, size_t s_len, _Out_ size_t* offsets, size_t offsets_len)NO_EXCEPTION; @@ -470,7 +470,7 @@ struct OrtApi { OrtStatus*(ORT_API_CALL* SetTensorElementType)(_Inout_ OrtTensorTypeAndShapeInfo*, enum ONNXTensorElementDataType type)NO_EXCEPTION; /** - * \param info Created from OrtCreateTensorTypeAndShapeInfo() function + * \param info Created from CreateTensorTypeAndShapeInfo() function * \param dim_values An array with length of `dim_count`. Its elements can contain negative values. * \param dim_count length of dim_values */ diff --git a/onnxruntime/core/framework/data_transfer.cc b/onnxruntime/core/framework/data_transfer.cc index 465655f03d..80099c2e4e 100644 --- a/onnxruntime/core/framework/data_transfer.cc +++ b/onnxruntime/core/framework/data_transfer.cc @@ -9,6 +9,14 @@ common::Status IDataTransfer::CopyTensor(const Tensor& src, Tensor& dst) const { return CopyTensor(src, dst, 0); } +common::Status IDataTransfer::CopyTensors(const Tensor* src, Tensor* dst, int size) const { + ORT_ENFORCE(nullptr != src && nullptr != dst); + for (int i = 0; i < size; ++i) { + ORT_RETURN_IF_ERROR(CopyTensor(src[i], dst[i], 0)); + } + return Status::OK(); +} + bool CPUDataTransfer::CanCopy(const OrtDevice& src_device, const OrtDevice& dst_device) const { return src_device.Type() == OrtDevice::CPU && dst_device.Type() == OrtDevice::CPU; } diff --git a/onnxruntime/core/framework/data_transfer.h b/onnxruntime/core/framework/data_transfer.h index 1e707b196b..798525149b 100644 --- a/onnxruntime/core/framework/data_transfer.h +++ b/onnxruntime/core/framework/data_transfer.h @@ -17,6 +17,7 @@ class IDataTransfer { virtual common::Status CopyTensor(const Tensor& src, Tensor& dst) const; virtual common::Status CopyTensor(const Tensor& src, Tensor& dst, int exec_queue_id) const = 0; + virtual common::Status CopyTensors(const Tensor* src, Tensor* dst, int size) const; }; class CPUDataTransfer : public IDataTransfer { diff --git a/onnxruntime/core/session/onnxruntime_c_api.cc b/onnxruntime/core/session/onnxruntime_c_api.cc index 3a3dab335b..99ff0212ad 100644 --- a/onnxruntime/core/session/onnxruntime_c_api.cc +++ b/onnxruntime/core/session/onnxruntime_c_api.cc @@ -807,16 +807,8 @@ ORT_API_STATUS_IMPL(OrtApis::AllocatorGetInfo, _In_ const OrtAllocator* ptr, _Ou API_IMPL_END } -/////////////////////////////////////////////////////////////////////////// -// Code to handle non-tensor types -// OrtGetValueCount -// OrtGetVaue -// OrtCreateValue -/////////////////////////////////////////////////////////////////////////// const int NUM_MAP_INDICES = 2; -//////////////////// -// OrtGetValueCount template OrtStatus* OrtGetNumSequenceElements(const OrtValue* p_ml_value, size_t* out) { auto& data = p_ml_value->Get(); @@ -867,7 +859,7 @@ ORT_API_STATUS_IMPL(OrtApis::GetValueCount, const OrtValue* value, size_t* out) } /////////////////// -// OrtGetValue +// OrtGetValueImplSeqOfMap template static OrtStatus* OrtGetValueImplSeqOfMap(const OrtValue* p_ml_value, int index, OrtValue** out) { using TKey = typename T::value_type::key_type;