From 3d79b1f06eff2c8c43ab6f549a49771307fb6e56 Mon Sep 17 00:00:00 2001 From: cao lei Date: Mon, 20 Feb 2023 09:47:57 -0800 Subject: [PATCH] Create new stream for data copy for IOBidning input scenario (#14719) ### Description Create new stream for data copy for IOBidning input scenario ### Motivation and Context Previously in bindInput(), a nullptr Stream is passed to copy data cross device. This caused the default stream is used thus hurt the performance. This PR is to fix https://github.com/microsoft/onnxruntime/issues/14484 --------- Co-authored-by: Lei Cao --- onnxruntime/core/framework/utils.cc | 133 ++++++++++++++++------------ 1 file changed, 74 insertions(+), 59 deletions(-) diff --git a/onnxruntime/core/framework/utils.cc b/onnxruntime/core/framework/utils.cc index f88d098454..6d8105329b 100644 --- a/onnxruntime/core/framework/utils.cc +++ b/onnxruntime/core/framework/utils.cc @@ -507,65 +507,6 @@ static common::Status CopyInputsAcrossDevices(const SessionState& session_state, return Status::OK(); } -// public method to do a single copy. used by external partners -common::Status CopyOneInputAcrossDevices(const SessionState& session_state, const std::string& input_name, - const OrtValue& orig_mlvalue, OrtValue& new_mlvalue) { - if (!orig_mlvalue.IsTensor() && !orig_mlvalue.IsSparseTensor()) { - new_mlvalue = orig_mlvalue; - return Status::OK(); - } - - MLValueCopyInfo copy_info; - // Sets copy_info.target_device. - ORT_RETURN_IF_ERROR(CalculateStaticCopyInfoForFeed(session_state, input_name, copy_info)); -#if !defined(DISABLE_SPARSE_TENSORS) - copy_info.source_device = (orig_mlvalue.IsTensor()) - ? orig_mlvalue.Get().Location().device - : orig_mlvalue.Get().Location().device; -#else - copy_info.source_device = orig_mlvalue.Get().Location().device; -#endif - - // copy_info.target_device is not set leaving to be equal to CPU. - return BatchOrCopyMLValue(session_state, copy_info, orig_mlvalue, new_mlvalue, nullptr); -} - -static common::Status CopyOutputsAcrossDevices(const SessionState& session_state, - gsl::span fetches, - std::vector& user_fetches, - gsl::span copy_info, - gsl::span fetch_streams) { - auto num_outputs = fetches.size(); - user_fetches.resize(num_outputs); - - std::vector batched_data_transfers; -#if !defined(DISABLE_SPARSE_TENSORS) - std::vector batched_sparse_data_transfers; -#endif - - for (size_t idx = 0; idx < num_outputs; ++idx) { -#if !defined(DISABLE_SPARSE_TENSORS) - ORT_RETURN_IF_ERROR(BatchOrCopyMLValue(session_state, copy_info[idx], fetches[idx], user_fetches[idx], fetch_streams[idx], - &batched_data_transfers, &batched_sparse_data_transfers)); -#else - ORT_RETURN_IF_ERROR(BatchOrCopyMLValue(session_state, copy_info[idx], fetches[idx], user_fetches[idx], fetch_streams[idx], - &batched_data_transfers)); -#endif - } - - if (!batched_data_transfers.empty()) { - ORT_RETURN_IF_ERROR(session_state.GetDataTransferMgr().CopyTensors(batched_data_transfers)); - } - -#if !defined(DISABLE_SPARSE_TENSORS) - if (!batched_sparse_data_transfers.empty()) { - ORT_RETURN_IF_ERROR(session_state.GetDataTransferMgr().CopySparseTensors(batched_sparse_data_transfers)); - } -#endif - - return Status::OK(); -} - #ifdef ORT_ENABLE_STREAM struct DeviceStreamCollectionHolder { DeviceStreamCollectionHolder( @@ -608,6 +549,80 @@ static void UpdateWithParentStream(DeviceStreamCollection& device_stream_collect } #endif +// public method to do a single copy. used by external partners +common::Status CopyOneInputAcrossDevices(const SessionState& session_state, const std::string& input_name, + const OrtValue& orig_mlvalue, OrtValue& new_mlvalue) { + if (!orig_mlvalue.IsTensor() && !orig_mlvalue.IsSparseTensor()) { + new_mlvalue = orig_mlvalue; + return Status::OK(); + } + + MLValueCopyInfo copy_info; + // Sets copy_info.target_device. + ORT_RETURN_IF_ERROR(CalculateStaticCopyInfoForFeed(session_state, input_name, copy_info)); +#if !defined(DISABLE_SPARSE_TENSORS) + copy_info.source_device = (orig_mlvalue.IsTensor()) + ? orig_mlvalue.Get().Location().device + : orig_mlvalue.Get().Location().device; +#else + copy_info.source_device = orig_mlvalue.Get().Location().device; +#endif + + Stream* device_stream = nullptr; +#ifdef ORT_ENABLE_STREAM + DeviceStreamCollectionHolder device_stream_collection_holder(session_state); + if (device_stream_collection_holder.p_ != nullptr) { + DeviceStreamCollection* device_stream_collection = device_stream_collection_holder.p_.get(); + gsl::span streams = device_stream_collection->GetStreams(); + for (Stream* stream : streams) { + if (stream && stream->GetDevice().Type() != OrtDevice::CPU) { + device_stream = stream; + break; + } + } + } +#endif + + // copy_info.target_device is not set leaving to be equal to CPU. + return BatchOrCopyMLValue(session_state, copy_info, orig_mlvalue, new_mlvalue, device_stream); +} + +static common::Status CopyOutputsAcrossDevices(const SessionState& session_state, + gsl::span fetches, + std::vector& user_fetches, + gsl::span copy_info, + gsl::span fetch_streams) { + auto num_outputs = fetches.size(); + user_fetches.resize(num_outputs); + + std::vector batched_data_transfers; +#if !defined(DISABLE_SPARSE_TENSORS) + std::vector batched_sparse_data_transfers; +#endif + + for (size_t idx = 0; idx < num_outputs; ++idx) { +#if !defined(DISABLE_SPARSE_TENSORS) + ORT_RETURN_IF_ERROR(BatchOrCopyMLValue(session_state, copy_info[idx], fetches[idx], user_fetches[idx], fetch_streams[idx], + &batched_data_transfers, &batched_sparse_data_transfers)); +#else + ORT_RETURN_IF_ERROR(BatchOrCopyMLValue(session_state, copy_info[idx], fetches[idx], user_fetches[idx], fetch_streams[idx], + &batched_data_transfers)); +#endif + } + + if (!batched_data_transfers.empty()) { + ORT_RETURN_IF_ERROR(session_state.GetDataTransferMgr().CopyTensors(batched_data_transfers)); + } + +#if !defined(DISABLE_SPARSE_TENSORS) + if (!batched_sparse_data_transfers.empty()) { + ORT_RETURN_IF_ERROR(session_state.GetDataTransferMgr().CopySparseTensors(batched_sparse_data_transfers)); + } +#endif + + return Status::OK(); +} + static common::Status ExecuteGraphImpl(const SessionState& session_state, const FeedsFetchesManager& feeds_fetches_manager,