saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-23 19:32:23 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Create new stream for data copy for IOBidning input scenario (#14719)

Browse source 
### 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 <leca@microsoft.com>
This commit is contained in:
cao lei 2023-02-20 09:47:57 -08:00 committed by GitHub
parent 1ea360148f
commit 3d79b1f06e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 74 additions and 59 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

133
onnxruntime/core/framework/utils.cc
View file

@ -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<Tensor>().Location().device
: orig_mlvalue.Get<SparseTensor>().Location().device;
#else
copy_info.source_device = orig_mlvalue.Get<Tensor>().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<const OrtValue> fetches,
std::vector<OrtValue>& user_fetches,
gsl::span<const MLValueCopyInfo> copy_info,
gsl::span<Stream* const> fetch_streams) {
auto num_outputs = fetches.size();
user_fetches.resize(num_outputs);
std::vector<IDataTransfer::SrcDstPair> batched_data_transfers;
#if !defined(DISABLE_SPARSE_TENSORS)
std::vector<IDataTransfer::SparseSrcDstPair> 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<Tensor>().Location().device
: orig_mlvalue.Get<SparseTensor>().Location().device;
#else
copy_info.source_device = orig_mlvalue.Get<Tensor>().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<Stream*> 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<const OrtValue> fetches,
std::vector<OrtValue>& user_fetches,
gsl::span<const MLValueCopyInfo> copy_info,
gsl::span<Stream* const> fetch_streams) {
auto num_outputs = fetches.size();
user_fetches.resize(num_outputs);
std::vector<IDataTransfer::SrcDstPair> batched_data_transfers;
#if !defined(DISABLE_SPARSE_TENSORS)
std::vector<IDataTransfer::SparseSrcDstPair> 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,