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

Fix a issue for cuda reduce Max Min (#474)

Browse source 
* Fix a issue for Reduce Max, Min. Per cudnn document, only Max/Min ops requires the indices output, it will report error if requesting indices for the other reduction ops.
This commit is contained in:
Hector Li 2019-02-19 17:12:04 -08:00 committed by Changming Sun
parent 9bc6503463
commit 4408342f0b
2 changed files with 25 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -123,7 +123,11 @@ Status ReduceKernel<allow_multi_axes>::ComputeImpl(OpKernelContext* ctx, cudnnRe
ORT_RETURN_IF_ERROR(output_tensor.Set(output_dims_cudnn, cudnn_type_X));
size_t workspace_bytes = 0;
CUDNN_RETURN_IF_ERROR(cudnnGetReductionWorkspaceSize(CudnnHandle(), reduce_desc, input_tensor, output_tensor, &workspace_bytes));
auto workspace_cuda = GetScratchBuffer<void>(workspace_bytes);
auto workspace_cuda = GetScratchBuffer<CudaT>(workspace_bytes);
size_t indices_bytes = 0;
CUDNN_RETURN_IF_ERROR(cudnnGetReductionIndicesSize(CudnnHandle(), reduce_desc, input_tensor, output_tensor, &indices_bytes));
auto indices_cuda = GetScratchBuffer<uint32_t>(indices_bytes);
// need to allocate a separate buffer for ArgMin/ArgMax comparsion output
auto output_count = Y->Shape().Size();
@ -139,11 +143,14 @@ Status ReduceKernel<allow_multi_axes>::ComputeImpl(OpKernelContext* ctx, cudnnRe
tmp_div, tmp_div,
input_data, input_count);
} else if (log_sum_exp_) {
// Reduce max
// Reduce max -- Max/Min will output indices data
CudnnReduceDescriptor reduce_max_desc;
ORT_RETURN_IF_ERROR(reduce_max_desc.Set(CUDNN_REDUCE_TENSOR_MAX, cudnn_type_X, CUDNN_REDUCE_TENSOR_NO_INDICES));
size_t indices_bytes_max = 0;
CUDNN_RETURN_IF_ERROR(cudnnGetReductionIndicesSize(CudnnHandle(), reduce_max_desc, input_tensor, output_tensor, &indices_bytes_max));
auto indices_cuda_max = GetScratchBuffer<uint32_t>(indices_bytes);
CUDNN_RETURN_IF_ERROR(cudnnReduceTensor(
CudnnHandle(), reduce_max_desc, nullptr, 0, workspace_cuda.get(), workspace_bytes,
CudnnHandle(), reduce_max_desc, indices_cuda_max.get(), indices_bytes_max, workspace_cuda.get(), workspace_bytes,
&one, input_tensor, reinterpret_cast<const CudaT*>(X->template Data<T>()),
&zero, output_tensor, reinterpret_cast<CudaT*>(Y->template MutableData<T>())));
@ -169,7 +176,7 @@ Status ReduceKernel<allow_multi_axes>::ComputeImpl(OpKernelContext* ctx, cudnnRe
// ReduceSum
CUDNN_RETURN_IF_ERROR(cudnnReduceTensor(
CudnnHandle(), reduce_desc, nullptr, 0, workspace_cuda.get(), workspace_bytes,
CudnnHandle(), reduce_desc, indices_cuda.get(), indices_bytes, workspace_cuda.get(), workspace_bytes,
&one, input_tensor, exp_result,
&zero, output_tensor, reinterpret_cast<CudaT*>(log_sum_result)));
@ -188,16 +195,18 @@ Status ReduceKernel<allow_multi_axes>::ComputeImpl(OpKernelContext* ctx, cudnnRe
return Status::OK();
}
CUDNN_RETURN_IF_ERROR(cudnnReduceTensor(
CudnnHandle(), reduce_desc, nullptr, 0, workspace_cuda.get(), workspace_bytes,
&one, input_tensor, calculate_sqt_ ? input_data : reinterpret_cast<const CudaT*>(X->template Data<T>()),
&zero, output_tensor, reinterpret_cast<CudaT*>(Y->template MutableData<T>())));
} else {
size_t indices_bytes = 0;
CUDNN_RETURN_IF_ERROR(cudnnGetReductionIndicesSize(CudnnHandle(), reduce_desc, input_tensor, output_tensor, &indices_bytes));
auto indices_cuda = GetScratchBuffer<void>(indices_bytes);
if (calculate_sqt_) {
CUDNN_RETURN_IF_ERROR(cudnnReduceTensor(
CudnnHandle(), reduce_desc, indices_cuda.get(), indices_bytes, workspace_cuda.get(), workspace_bytes,
&one, input_tensor, input_data,
&zero, output_tensor, reinterpret_cast<CudaT*>(Y->template MutableData<T>())));
} else {
CUDNN_RETURN_IF_ERROR(cudnnReduceTensor(
CudnnHandle(), reduce_desc, indices_cuda.get(), indices_bytes, workspace_cuda.get(), workspace_bytes,
&one, input_tensor, reinterpret_cast<const CudaT*>(X->template Data<T>()),
&zero, output_tensor, reinterpret_cast<CudaT*>(Y->template MutableData<T>())));
}
} else { // For ArgMax & ArgMin ops, use the indicies as the output with int64 type
if (temp_X) {
auto temp_output = GetScratchBuffer<float>(output_count);
CUDNN_RETURN_IF_ERROR(cudnnReduceTensor(

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

@ -17,6 +17,8 @@ class ReduceKernel : public CudaKernel, public ReduceKernelBase<allow_multi_axes
calculate_sqt_(false),
log_sum_exp_(false) {}
// Only Max Min need to set ReduceTensorIndices CUDNN_REDUCE_TENSOR_FLATTENED_INDICES as per cudnn library manual
// Only Max Min will have indices output, need to set the indices to nullptr for other ops
template <typename T, cudnnReduceTensorIndices_t ReduceTensorIndices = CUDNN_REDUCE_TENSOR_NO_INDICES>
Status ComputeImpl(OpKernelContext* ctx, cudnnReduceTensorOp_t cudnnReduceOp) const;