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

Fix the issue that it run into alloc failed on multiple cuda device. We have some place hard code the allocator always use device 0. (#1815)

Browse source 
Fix the issue that it run into alloc failed on multiple cuda device. We have some place hard code the allocator always use device 0.
This commit is contained in:
Hector Li 2019-09-12 15:31:59 -07:00 committed by GitHub
parent 18f7377269
commit a0ba25f98f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
19 changed files with 84 additions and 88 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
onnxruntime/contrib_ops/cuda/activation/activations.cc
View file

@ -27,7 +27,7 @@ namespace cuda {
Status x<T>::ComputeInternal(OpKernelContext* context) const { \
UnaryElementwisePreparation p; \
UnaryElementwise::Prepare(context, &p); \
CudaAsyncBuffer<Ctx##x> func_ctx(this, 0, MakeFuncCtx(), 1); \
CudaAsyncBuffer<Ctx##x> func_ctx(this, MakeFuncCtx(), 1); \
if (!std::is_same<CtxNull, Ctx##x>::value) ORT_RETURN_IF_ERROR(func_ctx.CopyToGpu()); \
Impl_##x<typename ToCudaType<T>::MappedType>( \
reinterpret_cast<const typename ToCudaType<T>::MappedType*>(p.input_tensor->template Data<T>()), \

2
onnxruntime/core/providers/cuda/activation/activations.cc
View file

@ -23,7 +23,7 @@ namespace cuda {
Status x<T>::ComputeInternal(OpKernelContext* context) const { \
UnaryElementwisePreparation p; \
UnaryElementwise::Prepare(context, &p); \
CudaAsyncBuffer<Ctx##x> func_ctx(this, 0, MakeFuncCtx(), 1); \
CudaAsyncBuffer<Ctx##x> func_ctx(this, MakeFuncCtx(), 1); \
if (!std::is_same<CtxNull, Ctx##x>::value) ORT_RETURN_IF_ERROR(func_ctx.CopyToGpu()); \
Impl_##x<typename ToCudaType<T>::MappedType>( \
reinterpret_cast<const typename ToCudaType<T>::MappedType*>(p.input_tensor->template Data<T>()), \

18
onnxruntime/core/providers/cuda/cuda_common.h
View file

@ -45,8 +45,8 @@ class CudaKernel : public OpKernel {
virtual Status ComputeInternal(OpKernelContext* p_op_kernel_context) const = 0;
template <typename T>
inline IAllocatorUniquePtr<T> AllocateBufferOnCPUPinned(int id, size_t count_or_bytes) const {
AllocatorPtr allocator = provider_->GetAllocator(id, OrtMemTypeCPU);
inline IAllocatorUniquePtr<T> AllocateBufferOnCPUPinned(size_t count_or_bytes) const {
AllocatorPtr allocator = provider_->GetAllocator(CPU_ALLOCATOR_DEVICE_ID, OrtMemTypeCPU);
if (!allocator)
return nullptr;
return IAllocator::MakeUniquePtr<T>(allocator, count_or_bytes);
@ -68,24 +68,24 @@ class CudaKernel : public OpKernel {
public:
CudaAsyncBuffer(const CudaKernel* op_kernel) : gpu_copy_(nullptr), count_(0), op_kernel_(op_kernel) {}
CudaAsyncBuffer(const CudaKernel* op_kernel, int device_id, size_t count) : CudaAsyncBuffer(op_kernel) {
AllocCpuPtr(device_id, count);
CudaAsyncBuffer(const CudaKernel* op_kernel, size_t count) : CudaAsyncBuffer(op_kernel) {
AllocCpuPtr(count);
}
CudaAsyncBuffer(const CudaKernel* op_kernel, int device_id, const T& value, size_t count)
: CudaAsyncBuffer(op_kernel, device_id, count) {
CudaAsyncBuffer(const CudaKernel* op_kernel, const T& value, size_t count)
: CudaAsyncBuffer(op_kernel, count) {
T* p = CpuPtr();
for (size_t i = 0; i != count; ++i) {
*p++ = value;
}
}
CudaAsyncBuffer(const CudaKernel* op_kernel, int device_id, const std::vector<T>& vec) : CudaAsyncBuffer(op_kernel, device_id, vec.size()) {
CudaAsyncBuffer(const CudaKernel* op_kernel, const std::vector<T>& vec) : CudaAsyncBuffer(op_kernel, vec.size()) {
memcpy(CpuPtr(), vec.data(), vec.size() * sizeof(T));
}
void AllocCpuPtr(int id, size_t count) {
cpu_pinned_copy_ = op_kernel_->AllocateBufferOnCPUPinned<T>(id, count);
void AllocCpuPtr( size_t count) {
cpu_pinned_copy_ = op_kernel_->AllocateBufferOnCPUPinned<T>(count);
if (cpu_pinned_copy_ == nullptr)
throw std::runtime_error("alloc failed");
count_ = count;

15
onnxruntime/core/providers/cuda/cuda_execution_provider.cc
View file

@ -67,21 +67,24 @@ CUDAExecutionProvider::CUDAExecutionProvider(const CUDAExecutionProviderInfo& in
CUDA_CALL_THROW(cudaSetDevice(device_id_));
DeviceAllocatorRegistrationInfo default_memory_info(
{OrtMemTypeDefault, [](int id) { return std::make_unique<CUDAAllocator>(id, CUDA); }, std::numeric_limits<size_t>::max()});
{OrtMemTypeDefault, [](int device_id) { return std::make_unique<CUDAAllocator>(device_id, CUDA); }, std::numeric_limits<size_t>::max()});
InsertAllocator(CreateAllocator(default_memory_info, device_id_));
DeviceAllocatorRegistrationInfo pinned_memory_info(
{OrtMemTypeCPUOutput, [](int) { return std::make_unique<CUDAPinnedAllocator>(0, CUDA_PINNED); }, std::numeric_limits<size_t>::max()});
InsertAllocator(CreateAllocator(pinned_memory_info, device_id_));
{OrtMemTypeCPUOutput, [](int device_id) { return std::make_unique<CUDAPinnedAllocator>(device_id, CUDA_PINNED); }, std::numeric_limits<size_t>::max()});
InsertAllocator(CreateAllocator(pinned_memory_info, CPU_ALLOCATOR_DEVICE_ID));
// TODO: this is actually used for the cuda kernels which explicitly ask for inputs from CPU.
// This will be refactored/removed when allocator and execution provider are decoupled.
DeviceAllocatorRegistrationInfo cpu_memory_info({OrtMemTypeCPUInput, [](int) { return std::make_unique<CPUAllocator>(std::make_unique<OrtMemoryInfo>("CUDA_CPU", OrtAllocatorType::OrtDeviceAllocator, OrtDevice(), 0, OrtMemTypeCPUInput)); }, std::numeric_limits<size_t>::max()});
InsertAllocator(CreateAllocator(cpu_memory_info));
DeviceAllocatorRegistrationInfo cpu_memory_info({
OrtMemTypeCPUInput,
[](int device_id) { return std::make_unique<CPUAllocator>(std::make_unique<OrtMemoryInfo>("CUDA_CPU", OrtAllocatorType::OrtDeviceAllocator, OrtDevice(), device_id, OrtMemTypeCPUInput)); },
std::numeric_limits<size_t>::max()});
InsertAllocator(CreateAllocator(cpu_memory_info, CPU_ALLOCATOR_DEVICE_ID));
}
CUDAExecutionProvider::~CUDAExecutionProvider() {
auto cpu_alloc = GetAllocator(0, OrtMemTypeCPU);
auto cpu_alloc = GetAllocator(CPU_ALLOCATOR_DEVICE_ID, OrtMemTypeCPU);
std::lock_guard<OrtMutex> lock(deferred_release_cpu_ptr_mutex_);
auto it = deferred_release_cpu_ptr_.begin();
while (it != deferred_release_cpu_ptr_.end()) {

6
onnxruntime/core/providers/cuda/cuda_execution_provider.h
View file

@ -13,6 +13,8 @@
namespace onnxruntime {
const int CPU_ALLOCATOR_DEVICE_ID = 0;
// Information needed to construct CUDA execution providers.
struct CUDAExecutionProviderInfo {
int device_id{0};
@ -24,7 +26,7 @@ class CUDAExecutionProvider : public IExecutionProvider {
explicit CUDAExecutionProvider(const CUDAExecutionProviderInfo& info);
virtual ~CUDAExecutionProvider();
AllocatorPtr GetAllocator(int id, OrtMemType mem_type = OrtMemTypeDefault) const override;
AllocatorPtr GetAllocator(int id, OrtMemType mem_type) const override;
Status Sync() const override;
@ -52,7 +54,7 @@ class CUDAExecutionProvider : public IExecutionProvider {
if (count_or_bytes == 0)
return nullptr;
return IAllocator::MakeUniquePtr<T>(GetAllocator(OrtMemTypeDefault), count_or_bytes);
return IAllocator::MakeUniquePtr<T>(GetAllocator(device_id_, OrtMemTypeDefault), count_or_bytes);
}
virtual std::shared_ptr<KernelRegistry> GetKernelRegistry() const override;

36
onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc
View file

@ -10,7 +10,7 @@ namespace onnxruntime {
namespace cuda {
template <>
Status BinaryElementwise<ShouldNotBroadcast>::Prepare(OpKernelContext* context, int /*device_id*/, BinaryElementwisePreparation* p) const {
Status BinaryElementwise<ShouldNotBroadcast>::Prepare(OpKernelContext* context, BinaryElementwisePreparation* p) const {
p->lhs_tensor = context->Input<Tensor>(0);
p->rhs_tensor = context->Input<Tensor>(1);
if (!(p->lhs_tensor->Shape() == p->rhs_tensor->Shape()))
@ -47,7 +47,7 @@ Status BinaryElementwise<ShouldNotBroadcast>::Prepare(OpKernelContext* context,
}
Status BinaryElementwiseBroadcastPrepare(
int device_id, const Tensor* lhs_tensor,
const Tensor* lhs_tensor,
const Tensor* rhs_tensor,
Tensor* output_tensor,
BinaryElementwisePreparation* p,
@ -61,13 +61,13 @@ Status BinaryElementwiseBroadcastPrepare(
p->output_tensor = output_tensor;
const auto& output_shape = output_tensor->Shape();
ORT_RETURN_IF_ERROR(p->BinaryElementwiseBroadcastPrepareHelper(device_id, lhs_shape, rhs_shape, output_shape));
ORT_RETURN_IF_ERROR(p->BinaryElementwiseBroadcastPrepareHelper(lhs_shape, rhs_shape, output_shape));
return Status::OK();
}
template <>
Status BinaryElementwise<ShouldBroadcast>::Prepare(OpKernelContext* context, int device_id, BinaryElementwisePreparation* p) const {
Status BinaryElementwise<ShouldBroadcast>::Prepare(OpKernelContext* context, BinaryElementwisePreparation* p) const {
auto lhs_tensor = context->Input<Tensor>(0);
auto rhs_tensor = context->Input<Tensor>(1);
const auto& lhs_shape = lhs_tensor->Shape();
@ -77,7 +77,7 @@ Status BinaryElementwise<ShouldBroadcast>::Prepare(OpKernelContext* context, int
ORT_RETURN_IF_ERROR(ComputeOutputShape(Node().Name(), lhs_shape, rhs_shape, output_shape));
auto output_tensor = context->Output(0, output_shape);
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(device_id, lhs_tensor, rhs_tensor, output_tensor, p));
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(lhs_tensor, rhs_tensor, output_tensor, p));
return Status::OK();
}
@ -92,14 +92,14 @@ Status BinaryElementwise<ShouldBroadcast>::Prepare(OpKernelContext* context, int
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
x<T>);
#define BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED(x, ver, T) \
#define BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED(x, ver, T) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
x, \
kOnnxDomain, \
ver, \
T, \
kCudaExecutionProvider, \
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<T>()) \
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<T>()) \
.TypeConstraint("T1", DataTypeImpl::GetTensorType<bool>()), \
x<T>);
@ -118,7 +118,7 @@ Status BinaryElementwise<ShouldBroadcast>::Prepare(OpKernelContext* context, int
template <> \
Status x<T>::ComputeInternal(OpKernelContext* context) const { \
BinaryElementwisePreparation prepare(this); \
Prepare(context, 0, &prepare); \
Prepare(context, &prepare); \
ORT_RETURN_IF_ERROR(prepare.CopyToGpu()); \
Impl_##x<typename ToCudaType<T>::MappedType>( \
prepare.output_rank_or_simple_broadcast, \
@ -138,7 +138,7 @@ Status BinaryElementwise<ShouldBroadcast>::Prepare(OpKernelContext* context, int
BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED(name, ver, T) \
BINARY_ELEMENTWISE_COMPUTE(name, T)
#define BINARY_LOGICALOP_TYPED(name, ver, T) \
#define BINARY_LOGICALOP_TYPED(name, ver, T) \
BINARY_ELEMENTWISE_LOGICALOP_REGISTER_KERNEL_TYPED(name, ver, T) \
BINARY_ELEMENTWISE_COMPUTE(name, T)
@ -248,7 +248,7 @@ Status Sum<T>::ComputeInternal(OpKernelContext* context) const {
BinaryElementwisePreparation prepare(this);
if (input_count == 2) {
// special case for 2 tensors to avoid memset zero
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(0, context->Input<Tensor>(0), context->Input<Tensor>(1), output_tensor, &prepare));
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(context->Input<Tensor>(0), context->Input<Tensor>(1), output_tensor, &prepare));
Impl_Add<CudaT>(
prepare.output_rank_or_simple_broadcast,
prepare.lhs_padded_strides.GpuPtr(),
@ -264,7 +264,7 @@ Status Sum<T>::ComputeInternal(OpKernelContext* context) const {
// for more than 2 inputs, we need to accumulate into output tensor, as the shape from input0 + input1 might be different from output shape
CUDA_RETURN_IF_ERROR(cudaMemset(output_tensor->MutableDataRaw(), 0, output_shape.Size() * sizeof(CudaT)));
for (int index = 0; index < input_count; index++) {
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(0, output_tensor, context->Input<Tensor>(index), output_tensor, &prepare));
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(output_tensor, context->Input<Tensor>(index), output_tensor, &prepare));
Impl_Add<CudaT>(
prepare.output_rank_or_simple_broadcast,
prepare.lhs_padded_strides.GpuPtr(),
@ -308,7 +308,7 @@ Status Max<T>::ComputeInternal(OpKernelContext* context) const {
// More than 2 inputs, set output to 0, add input0 to output, so that input0 can be broadcast with output shape correctly
CUDA_RETURN_IF_ERROR(cudaMemset(output_tensor->MutableDataRaw(), 0, output_shape.Size() * sizeof(CudaT)));
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(0, output_tensor, context->Input<Tensor>(0), output_tensor, &prepare));
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(output_tensor, context->Input<Tensor>(0), output_tensor, &prepare));
Impl_Add<CudaT>(
prepare.output_rank_or_simple_broadcast,
prepare.lhs_padded_strides.GpuPtr(),
@ -321,7 +321,7 @@ Status Max<T>::ComputeInternal(OpKernelContext* context) const {
reinterpret_cast<CudaT*>(prepare.output_tensor->template MutableData<T>()),
prepare.output_tensor->Shape().Size());
for (int index = 1; index < input_count; index++) {
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(0, output_tensor, context->Input<Tensor>(index), output_tensor, &prepare));
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(output_tensor, context->Input<Tensor>(index), output_tensor, &prepare));
Impl_Max<CudaT>(
prepare.output_rank_or_simple_broadcast,
prepare.lhs_padded_strides.GpuPtr(),
@ -364,7 +364,7 @@ Status Min<T>::ComputeInternal(OpKernelContext* context) const {
// More than 2 inputs, set output to 0, add input0 to output, so that input0 can be broadcast with output shape correctly
CUDA_RETURN_IF_ERROR(cudaMemset(output_tensor->MutableDataRaw(), 0, output_shape.Size() * sizeof(CudaT)));
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(0, output_tensor, context->Input<Tensor>(0), output_tensor, &prepare));
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(output_tensor, context->Input<Tensor>(0), output_tensor, &prepare));
Impl_Add<CudaT>(
prepare.output_rank_or_simple_broadcast,
prepare.lhs_padded_strides.GpuPtr(),
@ -377,7 +377,7 @@ Status Min<T>::ComputeInternal(OpKernelContext* context) const {
reinterpret_cast<CudaT*>(prepare.output_tensor->template MutableData<T>()),
prepare.output_tensor->Shape().Size());
for (int index = 1; index < input_count; index++) {
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(0, output_tensor, context->Input<Tensor>(index), output_tensor, &prepare));
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(output_tensor, context->Input<Tensor>(index), output_tensor, &prepare));
Impl_Min<CudaT>(
prepare.output_rank_or_simple_broadcast,
prepare.lhs_padded_strides.GpuPtr(),
@ -410,7 +410,7 @@ Status Greater<T>::ComputeInternal(OpKernelContext* context) const {
Tensor* output_tensor = context->Output(0, output_shape);
BinaryElementwisePreparation prepare(this);
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(0, input0, input1, output_tensor, &prepare));
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(input0, input1, output_tensor, &prepare));
IAllocatorUniquePtr<T> output_buffer = GetScratchBuffer<T>(output_size);
Impl_Greater<CudaT>(
@ -446,7 +446,7 @@ Status Equal<T>::ComputeInternal(OpKernelContext* context) const {
Tensor* output_tensor = context->Output(0, output_shape);
BinaryElementwisePreparation prepare(this);
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(0, input0, input1, output_tensor, &prepare));
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(input0, input1, output_tensor, &prepare));
IAllocatorUniquePtr<T> output_buffer = GetScratchBuffer<T>(output_size);
Impl_Equal<CudaT>(
@ -484,7 +484,7 @@ Status Less<T>::ComputeInternal(OpKernelContext* context) const {
Tensor* output_tensor = context->Output(0, output_shape);
BinaryElementwisePreparation prepare(this);
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(0, input0, input1, output_tensor, &prepare));
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(input0, input1, output_tensor, &prepare));
IAllocatorUniquePtr<T> output_buffer = GetScratchBuffer<T>(output_size);
Impl_Less<CudaT>(

10
onnxruntime/core/providers/cuda/math/binary_elementwise_ops.h
View file

@ -34,7 +34,7 @@ struct BinaryElementwisePreparation {
return Status::OK();
}
Status BinaryElementwiseBroadcastPrepareHelper(int device_id, const TensorShape& lhs_shape,
Status BinaryElementwiseBroadcastPrepareHelper(const TensorShape& lhs_shape,
const TensorShape& rhs_shape,
const TensorShape& output_shape) {
size_t lhs_rank = lhs_shape.NumDimensions();
@ -82,20 +82,20 @@ struct BinaryElementwisePreparation {
if (lhs_shape != output_shape) {
// compute strides with 1 more dim than out_rank, and use strides[0] == strides[1]
// to decide if dim0 needs broadcast
lhs_padded_strides.AllocCpuPtr(device_id, out_rank + 1);
lhs_padded_strides.AllocCpuPtr(out_rank + 1);
ORT_RETURN_IF_NOT(TensorPitches::Calculate(lhs_padded_strides.CpuSpan(), lhs_shape.GetDims()));
if (lhs_shape[0] > 1 && lhs_rank == out_rank)
lhs_padded_strides.CpuPtr()[0] = 0;
}
if (rhs_shape != output_shape) {
rhs_padded_strides.AllocCpuPtr(device_id, out_rank + 1);
rhs_padded_strides.AllocCpuPtr(out_rank + 1);
ORT_RETURN_IF_NOT(TensorPitches::Calculate(rhs_padded_strides.CpuSpan(), rhs_shape.GetDims()));
if (rhs_shape[0] > 1 && rhs_rank == out_rank)
rhs_padded_strides.CpuPtr()[0] = 0;
}
fdm_output_strides.AllocCpuPtr(device_id, out_rank);
fdm_output_strides.AllocCpuPtr(out_rank);
ORT_RETURN_IF_NOT(CalculateFdmStrides(fdm_output_strides.CpuSpan(), output_shape.GetDims()));
return Status::OK();
}
@ -117,7 +117,7 @@ class BinaryElementwise : public CudaKernel {
Status ComputeInternal(OpKernelContext*) const override {
return Status(common::ONNXRUNTIME, common::FAIL); // should not reach here
}
Status Prepare(OpKernelContext* context, int device_id, BinaryElementwisePreparation* p) const;
Status Prepare(OpKernelContext* context, BinaryElementwisePreparation* p) const;
};
template <typename T>

7
onnxruntime/core/providers/cuda/math/matmul.cc
View file

@ -67,10 +67,9 @@ Status MatMul<T>::ComputeInternal(OpKernelContext* ctx) const {
static_cast<int>(helper.N())));
return Status::OK();
}
int device_id = GetDeviceId();
CudaAsyncBuffer<const CudaT*> left_arrays(this, device_id, helper.LeftOffsets().size());
CudaAsyncBuffer<const CudaT*> right_arrays(this, device_id, helper.RightOffsets().size());
CudaAsyncBuffer<CudaT*> output_arrays(this, device_id, helper.OutputOffsets().size());
CudaAsyncBuffer<const CudaT*> left_arrays(this, helper.LeftOffsets().size());
CudaAsyncBuffer<const CudaT*> right_arrays(this, helper.RightOffsets().size());
CudaAsyncBuffer<CudaT*> output_arrays(this, helper.OutputOffsets().size());
MatMulComputeHelper::OffsetToArrays(reinterpret_cast<const CudaT*>(left_X->template Data<T>()), helper.LeftOffsets(), left_arrays.CpuSpan());
MatMulComputeHelper::OffsetToArrays(reinterpret_cast<const CudaT*>(right_X->template Data<T>()), helper.RightOffsets(), right_arrays.CpuSpan());
MatMulComputeHelper::OffsetToArrays(reinterpret_cast<CudaT*>(Y->template MutableData<T>()), helper.OutputOffsets(), output_arrays.CpuSpan());

2
onnxruntime/core/providers/cuda/reduction/reduction_ops.cc
View file

@ -159,7 +159,7 @@ Status ReduceKernel<allow_multi_axes>::ComputeImpl(OpKernelContext* ctx, cudnnRe
auto exp_result = GetScratchBuffer<T>(input_count).get();
auto log_sum_result = GetScratchBuffer<T>(output_count).get();
BinaryElementwisePreparation prepare(this);
prepare.BinaryElementwiseBroadcastPrepareHelper(0, input_shape, output_shape, input_shape);
prepare.BinaryElementwiseBroadcastPrepareHelper(input_shape, output_shape, input_shape);
prepare.CopyToGpu();
Impl_Sub<CudaT>(prepare.output_rank_or_simple_broadcast,
prepare.lhs_padded_strides.GpuPtr(),

4
onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc
View file

@ -348,7 +348,7 @@ Status CudnnRnnBase<T>::ComputeInternal(OpKernelContext* ctx) const {
}
if ((CUDNN_RNN_RELU == rnn_mode_ || CUDNN_RNN_TANH == rnn_mode_) && sequence_lens_data != nullptr && y_h_data != nullptr && y_data != nullptr) {
CudaAsyncBuffer<int32_t> sequence_lens_buffer(this, GetDeviceId(), batch_size);
CudaAsyncBuffer<int32_t> sequence_lens_buffer(this, batch_size);
memcpy(sequence_lens_buffer.CpuPtr(), sequence_lens_data, batch_size * sizeof(int32_t));
sequence_lens_buffer.CopyToGpu();
RnnMaskImpl(gsl::narrow_cast<int32_t>(num_directions_),
@ -371,7 +371,7 @@ void CudnnRnnBase<T>::SetZeroSequences(const int64_t zero_seq_index_cache_size,
T* y_h_data,
T* y_c_data) const {
typedef typename ToCudaType<T>::MappedType CudaT;
CudaAsyncBuffer<int32_t> zero_seq_index_cache_async_buffer(this, GetDeviceId(), zero_seq_index_cache_size);
CudaAsyncBuffer<int32_t> zero_seq_index_cache_async_buffer(this, zero_seq_index_cache_size);
memcpy(zero_seq_index_cache_async_buffer.CpuPtr(), zero_seq_index_cache.data(), zero_seq_index_cache_size * sizeof(int32_t));
zero_seq_index_cache_async_buffer.CopyToGpu();
MaskZeroSequences(gsl::narrow_cast<int32_t>(hidden_size_),

9
onnxruntime/core/providers/cuda/tensor/concat.cc
View file

@ -25,10 +25,9 @@ Status Concat::ComputeInternal(OpKernelContext* ctx) const {
if (p.output_num_elements == 0)
return Status::OK();
int device_id = GetDeviceId();
std::vector<int64_t> concat_sizes(input_count);
CudaAsyncBuffer<const void*> input_ptr(this, device_id, input_count);
CudaAsyncBuffer<const void*> input_ptr(this, input_count);
gsl::span<const void*> input_ptr_cpuspan = input_ptr.CpuSpan();
std::vector<int64_t> axis_dimension_input_output_mapping(p.output_tensor->Shape()[p.axis]);
int index = 0;
@ -45,9 +44,9 @@ Status Concat::ComputeInternal(OpKernelContext* ctx) const {
concat_sizes_range[i] += concat_sizes_range[i - 1];
}
CudaAsyncBuffer<int64_t> concat_sizes_gpu(this, device_id, concat_sizes);
CudaAsyncBuffer<int64_t> axis_dimension_input_output_mapping_gpu(this, device_id, axis_dimension_input_output_mapping);
CudaAsyncBuffer<int64_t> concat_sizes_range_gpu(this, device_id, concat_sizes_range);
CudaAsyncBuffer<int64_t> concat_sizes_gpu(this, concat_sizes);
CudaAsyncBuffer<int64_t> axis_dimension_input_output_mapping_gpu(this, axis_dimension_input_output_mapping);
CudaAsyncBuffer<int64_t> concat_sizes_range_gpu(this, concat_sizes_range);
concat_sizes_gpu.CopyToGpu();
axis_dimension_input_output_mapping_gpu.CopyToGpu();
concat_sizes_range_gpu.CopyToGpu();

7
onnxruntime/core/providers/cuda/tensor/expand.cc
View file

@ -11,7 +11,6 @@ namespace cuda {
Status Expand::ComputeInternal(OpKernelContext* ctx) const {
const auto& input0 = *ctx->Input<Tensor>(0);
const auto& input1 = *ctx->Input<Tensor>(1);
int device_id = GetDeviceId();
// new shape to be expanded to
const auto* p_shape = input1.template Data<int64_t>();
@ -34,9 +33,9 @@ Status Expand::ComputeInternal(OpKernelContext* ctx) const {
}
// create fast_divmod using dimension values
CudaAsyncBuffer<fast_divmod> fdm_input_dims(this, device_id, rank);
CudaAsyncBuffer<fast_divmod> fdm_output_dims(this, device_id, rank);
CudaAsyncBuffer<fast_divmod> fdm_output_subdim_size(this, device_id, rank);
CudaAsyncBuffer<fast_divmod> fdm_input_dims(this, rank);
CudaAsyncBuffer<fast_divmod> fdm_output_dims(this, rank);
CudaAsyncBuffer<fast_divmod> fdm_output_subdim_size(this, rank);
{
auto in_span = fdm_input_dims.CpuSpan();
auto out_span = fdm_output_dims.CpuSpan();

2
onnxruntime/core/providers/cuda/tensor/gather.cc
View file

@ -62,7 +62,7 @@ Status Gather::ComputeInternal(OpKernelContext* context) const {
// Put the output_block_size and block_size into div_strides
// for divmod calling in _GatherKernel to calculate the input index
CudaAsyncBuffer<fast_divmod> div_strides(this, 0, 2);
CudaAsyncBuffer<fast_divmod> div_strides(this, 2);
gsl::span<fast_divmod> div_strides_span = div_strides.CpuSpan();
div_strides_span[0] = fast_divmod(gsl::narrow_cast<int>(output_block_size));
div_strides_span[1] = fast_divmod(gsl::narrow_cast<int>(block_size));

11
onnxruntime/core/providers/cuda/tensor/pad.cc
View file

@ -24,12 +24,11 @@ Status Pad<T>::ComputeInternal(OpKernelContext* ctx) const {
const auto& input_tensor = *ctx->Input<Tensor>(0);
auto const& input_shape = input_tensor.Shape();
auto dimension_count = input_shape.NumDimensions();
int device_id = GetDeviceId();
CudaAsyncBuffer<int64_t> input_dims(this, device_id, input_shape.GetDims());
CudaAsyncBuffer<int64_t> input_strides(this, device_id, dimension_count);
CudaAsyncBuffer<int64_t> lower_pads(this, device_id, dimension_count);
CudaAsyncBuffer<int64_t> upper_pads(this, device_id, dimension_count);
CudaAsyncBuffer<fast_divmod> fdm_output_strides(this, device_id, dimension_count);
CudaAsyncBuffer<int64_t> input_dims(this, input_shape.GetDims());
CudaAsyncBuffer<int64_t> input_strides(this, dimension_count);
CudaAsyncBuffer<int64_t> lower_pads(this, dimension_count);
CudaAsyncBuffer<int64_t> upper_pads(this, dimension_count);
CudaAsyncBuffer<fast_divmod> fdm_output_strides(this, dimension_count);
TensorPitches::Calculate(input_strides.CpuSpan(), input_shape.GetDims());
std::vector<int64_t> output_dims(input_shape.GetDims());

9
onnxruntime/core/providers/cuda/tensor/slice.cc
View file

@ -68,28 +68,27 @@ Status Slice<Tind, dynamic>::ComputeInternal(OpKernelContext* ctx) const {
if (output_size == 0) {
return Status::OK();
}
int device_id = GetDeviceId();
CudaAsyncBuffer<int64_t> starts_buffer(this, device_id, dimension_count);
CudaAsyncBuffer<int64_t> starts_buffer(this, dimension_count);
gsl::span<int64_t> starts_buffer_span = starts_buffer.CpuSpan();
for (int i = 0; i < dimension_count; ++i) {
starts_buffer_span[i] = starts[i];
}
starts_buffer.CopyToGpu();
CudaAsyncBuffer<int64_t> steps_buffer(this, device_id, dimension_count);
CudaAsyncBuffer<int64_t> steps_buffer(this, dimension_count);
gsl::span<int64_t> steps_buffer_span = steps_buffer.CpuSpan();
for (int i = 0; i < dimension_count; ++i) {
steps_buffer_span[i] = steps[i];
}
steps_buffer.CopyToGpu();
CudaAsyncBuffer<int64_t> input_strides(this, device_id, dimension_count);
CudaAsyncBuffer<int64_t> input_strides(this, dimension_count);
ORT_ENFORCE(TensorPitches::Calculate(input_strides.CpuSpan(), input_dimensions));
input_strides.CopyToGpu();
TensorPitches output_pitches(output_dims);
CudaAsyncBuffer<fast_divmod> div_strides(this, device_id, dimension_count);
CudaAsyncBuffer<fast_divmod> div_strides(this, dimension_count);
gsl::span<fast_divmod> div_strides_span = div_strides.CpuSpan();
for (int i = 0; i < dimension_count; ++i) {
div_strides_span[i] = fast_divmod(gsl::narrow_cast<int>(output_pitches[i]));

9
onnxruntime/core/providers/cuda/tensor/split.cc
View file

@ -41,8 +41,7 @@ Status Split::ComputeInternal(OpKernelContext* ctx) const {
auto& input_dims = input_shape.GetDims();
std::vector<int64_t> output_dimensions{input_dims};
int device_id = GetDeviceId();
CudaAsyncBuffer<void*> output_ptr(this, device_id, num_outputs);
CudaAsyncBuffer<void*> output_ptr(this, num_outputs);
gsl::span<void*> output_ptr_span = output_ptr.CpuSpan();
std::vector<int64_t> axis_dimension_input_output_mapping(input_dims[axis]);
int index = 0;
@ -60,17 +59,17 @@ Status Split::ComputeInternal(OpKernelContext* ctx) const {
}
output_ptr.CopyToGpu();
CudaAsyncBuffer<int64_t> split_sizes_gpu(this, device_id, split_sizes);
CudaAsyncBuffer<int64_t> split_sizes_gpu(this, split_sizes);
split_sizes_gpu.CopyToGpu();
std::vector<int64_t> split_sizes_range(split_sizes);
for (int i = 1; i < split_sizes_range.size(); ++i) {
split_sizes_range[i] += split_sizes_range[i - 1];
}
CudaAsyncBuffer<int64_t> split_sizes_range_gpu(this, device_id, split_sizes_range);
CudaAsyncBuffer<int64_t> split_sizes_range_gpu(this, split_sizes_range);
split_sizes_range_gpu.CopyToGpu();
CudaAsyncBuffer<int64_t> axis_dimension_input_output_mapping_gpu(this, device_id, axis_dimension_input_output_mapping);
CudaAsyncBuffer<int64_t> axis_dimension_input_output_mapping_gpu(this, axis_dimension_input_output_mapping);
axis_dimension_input_output_mapping_gpu.CopyToGpu();
size_t element_size = input_tensor->DataType()->Size();

7
onnxruntime/core/providers/cuda/tensor/tile.cc
View file

@ -43,10 +43,9 @@ Status Tile<T>::ComputeInternal(OpKernelContext* ctx) const {
T* output_data = output_tensor.template MutableData<T>();
const T* input_data = input_tensor.template Data<T>();
int device_id = GetDeviceId();
CudaAsyncBuffer<int64_t> input_strides(this, device_id, rank);
CudaAsyncBuffer<fast_divmod> fdm_input_shape(this, device_id, rank);
CudaAsyncBuffer<fast_divmod> fdm_output_strides(this, device_id, rank);
CudaAsyncBuffer<int64_t> input_strides(this, rank);
CudaAsyncBuffer<fast_divmod> fdm_input_shape(this, rank);
CudaAsyncBuffer<fast_divmod> fdm_output_strides(this, rank);
ORT_ENFORCE(TensorPitches::Calculate(input_strides.CpuSpan(), input_shape));
ORT_ENFORCE(CalculateFdmStrides(fdm_output_strides.CpuSpan(), output_dims));

7
onnxruntime/core/providers/cuda/tensor/transpose.cc
View file

@ -92,10 +92,9 @@ Status Transpose<T>::ComputeInternal(OpKernelContext* ctx) const {
return Status::OK();
}
int device_id = GetDeviceId();
CudaAsyncBuffer<int64_t> input_strides(this, device_id, rank);
CudaAsyncBuffer<size_t> perm(this, device_id, *p_perm);
CudaAsyncBuffer<fast_divmod> fdm_output_strides(this, device_id, rank);
CudaAsyncBuffer<int64_t> input_strides(this, rank);
CudaAsyncBuffer<size_t> perm(this, *p_perm);
CudaAsyncBuffer<fast_divmod> fdm_output_strides(this, rank);
ORT_ENFORCE(TensorPitches::Calculate(input_strides.CpuSpan(), input_dims));
ORT_ENFORCE(CalculateFdmStrides(fdm_output_strides.CpuSpan(), output_dims));

9
onnxruntime/core/providers/cuda/tensor/upsample.cc
View file

@ -62,13 +62,12 @@ Status Upsample<T>::BaseCompute(OpKernelContext* context, const std::vector<floa
typedef typename ToCudaType<T>::MappedType CudaT;
// kernel
int device_id = GetDeviceId();
TensorPitches input_pitches(X_dims);
CudaAsyncBuffer<int64_t> input_strides(this, device_id, rank);
CudaAsyncBuffer<int64_t> input_strides(this, rank);
gsl::span<int64_t> input_stride_span = input_strides.CpuSpan();
TensorPitches output_pitches(Y_dims);
CudaAsyncBuffer<fast_divmod> output_div_pitches(this, device_id, rank);
CudaAsyncBuffer<fast_divmod> output_div_pitches(this, rank);
gsl::span<fast_divmod> div_strides_span = output_div_pitches.CpuSpan();
for (int i = 0; i < rank; ++i) {
@ -81,7 +80,7 @@ Status Upsample<T>::BaseCompute(OpKernelContext* context, const std::vector<floa
size_t output_count = Y->Shape().Size();
if (is_resize) {
CudaAsyncBuffer<float> scales_vals(this, device_id, scales);
CudaAsyncBuffer<float> scales_vals(this, scales);
scales_vals.CopyToGpu();
ResizeImpl(mode_,
rank,
@ -93,7 +92,7 @@ Status Upsample<T>::BaseCompute(OpKernelContext* context, const std::vector<floa
reinterpret_cast<CudaT*>(Y->template MutableData<T>()),
output_count);
} else {
CudaAsyncBuffer<fast_divmod> scales_div(this, device_id, rank);
CudaAsyncBuffer<fast_divmod> scales_div(this, rank);
gsl::span<fast_divmod> scales_div_span = scales_div.CpuSpan();
for (int i = 0; i < rank; ++i) {