mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Optimize sequence type usage on CUDA [2/n] (#8720)
This commit is contained in:
parent
9053e1522d
commit
17b0664e34
5 changed files with 60 additions and 33 deletions
|
|
@ -526,7 +526,7 @@ Do not modify directly.*
|
|||
|Identity|*in* input:**T**<br> *out* output:**T**<br><br>or<br><br>*in* input:**V**<br> *out* output:**V**|14+|**V** = seq(tensor(bfloat16)), seq(tensor(bool)), seq(tensor(double)), seq(tensor(float)), seq(tensor(float16)), seq(tensor(int16)), seq(tensor(int32)), seq(tensor(int64)), seq(tensor(int8)), seq(tensor(uint16)), seq(tensor(uint32)), seq(tensor(uint64)), seq(tensor(uint8)), tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
|
||||
|||13|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
|
||||
|||[1, 12]|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
|
||||
|If|*in* cond:**B**<br> *out* outputs:**V**|13+|**B** = tensor(bool)<br/> **V** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
|
||||
|If|*in* cond:**B**<br> *out* outputs:**V**|13+|**B** = tensor(bool)<br/> **V** = seq(tensor(bfloat16)), seq(tensor(bool)), seq(tensor(double)), seq(tensor(float)), seq(tensor(float16)), seq(tensor(int16)), seq(tensor(int32)), seq(tensor(int64)), seq(tensor(int8)), seq(tensor(string)), seq(tensor(uint16)), seq(tensor(uint32)), seq(tensor(uint64)), seq(tensor(uint8)), tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
|
||||
|||[11, 12]|**B** = tensor(bool)<br/> **V** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
|
||||
|||[1, 10]|**B** = tensor(bool)<br/> **V** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
|
||||
|ImageScaler|*in* input:**T**<br> *out* output:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
|
||||
|
|
@ -546,7 +546,7 @@ Do not modify directly.*
|
|||
|LogSoftmax|*in* input:**T**<br> *out* output:**T**|13+|**T** = tensor(double), tensor(float), tensor(float16)|
|
||||
|||[11, 12]|**T** = tensor(double), tensor(float), tensor(float16)|
|
||||
|||[1, 10]|**T** = tensor(double), tensor(float), tensor(float16)|
|
||||
|Loop|*in* M:**I**<br> *in* cond:**B**<br> *in* v_initial:**V**<br> *out* v_final_and_scan_outputs:**V**|13+|**B** = tensor(bool)<br/> **I** = tensor(int64)<br/> **V** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
|
||||
|Loop|*in* M:**I**<br> *in* cond:**B**<br> *in* v_initial:**V**<br> *out* v_final_and_scan_outputs:**V**|13+|**B** = tensor(bool)<br/> **I** = tensor(int64)<br/> **V** = seq(tensor(bfloat16)), seq(tensor(bool)), seq(tensor(double)), seq(tensor(float)), seq(tensor(float16)), seq(tensor(int16)), seq(tensor(int32)), seq(tensor(int64)), seq(tensor(int8)), seq(tensor(string)), seq(tensor(uint16)), seq(tensor(uint32)), seq(tensor(uint64)), seq(tensor(uint8)), tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
|
||||
|||[11, 12]|**B** = tensor(bool)<br/> **I** = tensor(int64)<br/> **V** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
|
||||
|||[1, 10]|**B** = tensor(bool)<br/> **I** = tensor(int64)<br/> **V** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
|
||||
|MatMul|*in* A:**T**<br> *in* B:**T**<br> *out* Y:**T**|13+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
|
||||
|
|
|
|||
|
|
@ -492,31 +492,46 @@ Status LoopImpl::Execute(const FeedsFetchesManager& ffm) {
|
|||
|
||||
// As the loop carried variables may change shape across iterations there's no way to avoid a copy
|
||||
// as we need the final shape.
|
||||
auto copy_tensor_from_mlvalue_to_output = [this](const OrtValue& input, int output_idx) {
|
||||
auto type = input.Type();
|
||||
if (type == DataTypeImpl::GetType<Tensor>()) {
|
||||
auto& data = input.Get<Tensor>();
|
||||
auto copy_mlvalue_to_output = [this](OrtValue& input, int output_idx, int64_t iter_num_value) {
|
||||
if (input.IsTensor()) {
|
||||
const auto& data = input.Get<Tensor>();
|
||||
Tensor* output = context_.Output(output_idx, data.Shape());
|
||||
// Safely use the IDataTransfer abstraction as we only allow using
|
||||
// Loop on CUDA if the copy stream is the same as the compute stream.
|
||||
// So there is no explicit sync required between the compute and copy streams
|
||||
// to avoid data races.
|
||||
session_state_.GetDataTransferMgr().CopyTensor(input.Get<Tensor>(), *output);
|
||||
} else if (type == DataTypeImpl::GetType<TensorSeq>()) {
|
||||
std::vector<Tensor> tensors;
|
||||
} else if (input.IsTensorSequence()) {
|
||||
if (iter_num_value != 0) {
|
||||
// We can move the subgraph outputs directly into the Loop's outputs.
|
||||
TensorSeq* output = context_.Output<TensorSeq>(output_idx);
|
||||
*output = std::move(*input.GetMutable<TensorSeq>());
|
||||
} else {
|
||||
// We can't move the Loop's inputs directly into the Loop's outputs
|
||||
// as operator inputs are read-only. Hence, we need to make a copy.
|
||||
std::vector<Tensor> tensors;
|
||||
|
||||
auto& data = input.Get<TensorSeq>();
|
||||
TensorSeq* output = context_.Output<TensorSeq>(output_idx);
|
||||
output->SetType(data.DataType());
|
||||
auto& data = input.Get<TensorSeq>();
|
||||
TensorSeq* output = context_.Output<TensorSeq>(output_idx);
|
||||
output->SetType(data.DataType());
|
||||
|
||||
AllocatorPtr alloc;
|
||||
auto status = context_.GetTempSpaceAllocator(&alloc);
|
||||
if (!status.IsOK()) {
|
||||
ORT_THROW("Unable to get an allocator");
|
||||
AllocatorPtr alloc;
|
||||
auto status = context_.GetTempSpaceAllocator(&alloc);
|
||||
if (!status.IsOK()) {
|
||||
ORT_THROW("Unable to get an allocator");
|
||||
}
|
||||
for (auto it = data.begin(), end = data.end(); it != end; ++it) {
|
||||
Tensor tmp(it->DataType(), onnxruntime::TensorShape(it->Shape()), alloc);
|
||||
// Safely use the IDataTransfer abstraction as we only allow using
|
||||
// Loop on CUDA if the copy stream is the same as the compute stream.
|
||||
// So there is no explicit sync required between the compute and copy streams
|
||||
// to avoid data races.
|
||||
session_state_.GetDataTransferMgr().CopyTensor(*it, tmp);
|
||||
tensors.push_back(std::move(tmp));
|
||||
}
|
||||
|
||||
output->SetElements(std::move(tensors));
|
||||
}
|
||||
for (auto it = data.begin(), end = data.end(); it != end; ++it) {
|
||||
Tensor tmp(it->DataType(), onnxruntime::TensorShape(it->Shape()), alloc);
|
||||
session_state_.GetDataTransferMgr().CopyTensor(*it, tmp);
|
||||
tensors.push_back(std::move(tmp));
|
||||
}
|
||||
|
||||
output->SetElements(std::move(tensors));
|
||||
}
|
||||
};
|
||||
|
||||
|
|
@ -524,7 +539,7 @@ Status LoopImpl::Execute(const FeedsFetchesManager& ffm) {
|
|||
if (iter_num_value != 0) {
|
||||
for (int i = 0; i < info_.num_loop_carried_vars; ++i) {
|
||||
// need to allocate Loop output and copy OrtValue from fetches
|
||||
copy_tensor_from_mlvalue_to_output(fetches[i + 1], i); // skip cond
|
||||
copy_mlvalue_to_output(fetches[i + 1], i, iter_num_value); // skip cond
|
||||
}
|
||||
|
||||
for (int i = info_.num_loop_carried_vars; i < info_.num_outputs; ++i) {
|
||||
|
|
@ -538,7 +553,7 @@ Status LoopImpl::Execute(const FeedsFetchesManager& ffm) {
|
|||
// no iterations.
|
||||
// copy input loop carried vars to output.
|
||||
for (int i = 0; i < info_.num_loop_carried_vars; ++i) {
|
||||
copy_tensor_from_mlvalue_to_output(feeds[i + 2], i); // skip iter# and cond
|
||||
copy_mlvalue_to_output(feeds[i + 2], i, iter_num_value); // skip iter# and cond
|
||||
}
|
||||
|
||||
// create empty outputs for loop outputs using the subgraph output shapes for the rank
|
||||
|
|
|
|||
|
|
@ -34,9 +34,7 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(If,
|
|||
.TypeConstraint("V", DataTypeImpl::AllFixedSizeTensorTypes()),
|
||||
If);
|
||||
|
||||
// sequence tensors were also supported in addition to existing support for tensors in opset-13,
|
||||
// but we do not support sequence tensors in the cuda If kernel because there are no ops that handle
|
||||
// sequence tensors on CUDA and supporting it for If doesn't add value while that is the case
|
||||
// opset-13 supports sequence type for If's subgraph outputs
|
||||
ONNX_OPERATOR_KERNEL_EX(If,
|
||||
kOnnxDomain,
|
||||
13,
|
||||
|
|
@ -44,7 +42,7 @@ ONNX_OPERATOR_KERNEL_EX(If,
|
|||
(*KernelDefBuilder::Create())
|
||||
.InputMemoryType(OrtMemTypeCPUInput, 0) // 'cond' needs to be on CPU
|
||||
.TypeConstraint("B", DataTypeImpl::GetTensorType<bool>())
|
||||
.TypeConstraint("V", DataTypeImpl::AllFixedSizeTensorTypes()),
|
||||
.TypeConstraint("V", DataTypeImpl::AllTensorAndSequenceTensorTypes()),
|
||||
If);
|
||||
|
||||
Status If::Compute(OpKernelContext* ctx) const {
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@
|
|||
#include "core/providers/cuda/controlflow/loop.h"
|
||||
#include "core/providers/cuda/cuda_common.h"
|
||||
#include "core/providers/cuda/cuda_fwd.h"
|
||||
#include "core/framework/ort_value.h"
|
||||
#include "core/providers/cuda/cuda_execution_provider.h"
|
||||
|
||||
using namespace ONNX_NAMESPACE;
|
||||
using namespace onnxruntime::common;
|
||||
|
|
@ -38,9 +38,7 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(Loop,
|
|||
.TypeConstraint("V", DataTypeImpl::AllFixedSizeTensorTypes()),
|
||||
Loop);
|
||||
|
||||
// sequence tensors were also supported in addition to existing support for tensors in opset-13,
|
||||
// but we do not support sequence tensors in the cuda Loop kernel because there are no ops that handle
|
||||
// sequence tensors on CUDA and supporting it for Loop doesn't add value while that is the case
|
||||
// opset-13 supports sequence type for loop carried dependencies
|
||||
ONNX_OPERATOR_KERNEL_EX(Loop,
|
||||
kOnnxDomain,
|
||||
13,
|
||||
|
|
@ -50,7 +48,7 @@ ONNX_OPERATOR_KERNEL_EX(Loop,
|
|||
.InputMemoryType(OrtMemTypeCPUInput, 1) // 'cond' needs to be on CPU
|
||||
.TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>())
|
||||
.TypeConstraint("B", DataTypeImpl::GetTensorType<bool>())
|
||||
.TypeConstraint("V", DataTypeImpl::AllFixedSizeTensorTypes()),
|
||||
.TypeConstraint("V", DataTypeImpl::AllTensorAndSequenceTensorTypes()),
|
||||
Loop);
|
||||
|
||||
static Status ConcatenateGpuOutput(void* stream, std::vector<OrtValue>& per_iteration_output,
|
||||
|
|
@ -83,6 +81,21 @@ static Status ConcatenateGpuOutput(void* stream, std::vector<OrtValue>& per_iter
|
|||
}
|
||||
|
||||
Loop::Loop(const OpKernelInfo& info) : onnxruntime::Loop(info) {
|
||||
// We use the IDataTransfer abstraction to perform copies in the Loop implementation.
|
||||
// By default, the GPUDataTransfer class is setup to use the same stream as the EP's compute stream
|
||||
// while performing copies to/from CUDA (do_copy_on_default_stream = true). This is good as we wouldn't
|
||||
// have to do any explicit syncs between the copy and compute streams.
|
||||
// However, there is a user-facing flag that allows users to use a dedicated stream just for copying.
|
||||
// To support using Loop for that case, we would have to do a sync between the copy stream and
|
||||
// the compute stream to avoid data races. At the very least, we need to expose an interface in IDataTransfer
|
||||
// to use a caller provided stream for Loop to provide for the GPUDataTransfer instance to use.
|
||||
// Currently, using a dedicated copy stream has larger negative implications (see comment in GPUDataTransfer's
|
||||
// constructor implementation), and so it is not in a usable state. When it becomes usable again,
|
||||
// we will re-visit this limitation in Loop.
|
||||
bool do_copy_on_default_stream = static_cast<const CUDAExecutionProvider*>(info.GetExecutionProvider())->DoCopyOnDefaultStream();
|
||||
ORT_ENFORCE(do_copy_on_default_stream,
|
||||
"Using Loop operator on CUDA while using a dedicated stream for copying "
|
||||
"(a stream that is different than the compute stream) is currently not supported");
|
||||
SetConcatOutputFunc(ConcatenateGpuOutput);
|
||||
SetComputeStream(static_cast<void*>(info.GetExecutionProvider()->GetComputeStream()));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -73,6 +73,7 @@ class CUDAExecutionProvider : public IExecutionProvider {
|
|||
int GetDeviceId() const override { return info_.device_id; }
|
||||
const cudaDeviceProp& GetDeviceProp() const { return device_prop_; };
|
||||
int GetCudnnConvAlgo() const { return info_.cudnn_conv_algo_search; }
|
||||
bool DoCopyOnDefaultStream() const { return info_.do_copy_in_default_stream; }
|
||||
|
||||
ProviderOptions GetProviderOptions() const override {
|
||||
return CUDAExecutionProviderInfo::ToProviderOptions(info_);
|
||||
|
|
|
|||
Loading…
Reference in a new issue