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

Range CUDA inputs should be in CPU (#4871)

Browse source 
Co-authored-by: Sherlock Huang <bahuang@OrtTrainingDev3.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
This commit is contained in:
Sherlock 2020-08-20 15:52:02 -07:00 committed by GitHub
parent c5342b5417
commit a0271f619a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 13 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
onnxruntime/core/providers/cuda/generator/range.cc
View file

@ -18,11 +18,15 @@ ONNX_OPERATOR_KERNEL_EX(
kOnnxDomain,
11,
kCudaExecutionProvider,
KernelDefBuilder().TypeConstraint("T", {DataTypeImpl::GetTensorType<float>(),
DataTypeImpl::GetTensorType<double>(),
DataTypeImpl::GetTensorType<int16_t>(),
DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
KernelDefBuilder()
.InputMemoryType<OrtMemTypeCPUInput>(0) // start
.InputMemoryType<OrtMemTypeCPUInput>(1) // limit
.InputMemoryType<OrtMemTypeCPUInput>(2) // delta
.TypeConstraint("T", {DataTypeImpl::GetTensorType<float>(),
DataTypeImpl::GetTensorType<double>(),
DataTypeImpl::GetTensorType<int16_t>(),
DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
Range);
template <typename T>
@ -47,17 +51,13 @@ static Status ComputeRange(OpKernelContext* ctx) {
delta_tensor_ptr->Shape());
}
// Start, Limit and Delta are stored in GPU. So we need copy it to CPU to read.
// It is better to store these tensors in pinned memory or CPU for better performance.
T start;
CUDA_RETURN_IF_ERROR(cudaMemcpy(&start, start_tensor.template Data<T>(), sizeof(T), cudaMemcpyDeviceToHost));
T limit;
CUDA_RETURN_IF_ERROR(cudaMemcpy(&limit, limit_tensor.template Data<T>(), sizeof(T), cudaMemcpyDeviceToHost));
// Start, Limit and Delta are stored in CPU.
T start = *(start_tensor.template Data<T>());
T limit = *(limit_tensor.template Data<T>());
T delta = T(1);
if (delta_tensor_ptr != nullptr) {
CUDA_RETURN_IF_ERROR(cudaMemcpy(&delta, delta_tensor_ptr->template Data<T>(), sizeof(T), cudaMemcpyDeviceToHost));
delta = *(delta_tensor_ptr->template Data<T>());
}
if (delta == T(0)) {