Optimize sequence type usage on CUDA [2/n] (#8720)

This commit is contained in:
Hariharan Seshadri 2021-08-24 10:40:28 -07:00 committed by GitHub
parent 9053e1522d
commit 17b0664e34
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 60 additions and 33 deletions

View file

@ -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)|

View file

@ -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

View file

@ -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 {

View file

@ -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()));
}

View file

@ -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_);