From fff1ed9bfc28b6e1cbb1c29364b8872e2929a5df Mon Sep 17 00:00:00 2001 From: Zhang Lei Date: Thu, 12 Dec 2019 16:38:06 -0800 Subject: [PATCH] Optimize cuda scatter() on 2D compatible. (#2628) * Optimize cuda scatter() on 2D compatible. * Add some comments. --- .../providers/cuda/tensor/scatter_elements.cc | 33 ++- .../cuda/tensor/scatter_elements_impl.cu | 207 ++++++++++++++---- .../cuda/tensor/scatter_elements_impl.h | 10 +- 3 files changed, 185 insertions(+), 65 deletions(-) diff --git a/onnxruntime/core/providers/cuda/tensor/scatter_elements.cc b/onnxruntime/core/providers/cuda/tensor/scatter_elements.cc index 92466ec124..0032f148f8 100755 --- a/onnxruntime/core/providers/cuda/tensor/scatter_elements.cc +++ b/onnxruntime/core/providers/cuda/tensor/scatter_elements.cc @@ -41,37 +41,35 @@ ONNX_OPERATOR_KERNEL_EX( const T* update_data = updates_tensor->template Data(); \ if (utils::IsPrimitiveDataType(Tin_type)) { \ const int32_t* indices_data = indices_tensor->template Data(); \ - ScatterElementsImpl( \ + return ScatterElementsImpl( \ rank, \ reinterpret_cast::MappedType*>(input_data), \ input_data_size, \ - gpu_input_dims.GpuPtr(), \ - gpu_input_strides.GpuPtr(), \ + buffer_input_dims, \ + buffer_input_strides, \ indices_data, \ indices_size, \ - gpu_indices_dims.GpuPtr(), \ - fdm_indices_strides.GpuPtr(), \ + buffer_indices_dims, \ + fdm_indices_strides, \ reinterpret_cast::MappedType*>(update_data), \ axis, \ reinterpret_cast::MappedType*>(output_data)); \ - return Status::OK(); \ } \ if (utils::IsPrimitiveDataType(Tin_type)) { \ const int64_t* indices_data = indices_tensor->template Data(); \ - ScatterElementsImpl( \ + return ScatterElementsImpl( \ rank, \ reinterpret_cast::MappedType*>(input_data), \ input_data_size, \ - gpu_input_dims.GpuPtr(), \ - gpu_input_strides.GpuPtr(), \ + buffer_input_dims, \ + buffer_input_strides, \ indices_data, \ indices_size, \ - gpu_indices_dims.GpuPtr(), \ - fdm_indices_strides.GpuPtr(), \ + buffer_indices_dims, \ + fdm_indices_strides, \ reinterpret_cast::MappedType*>(update_data), \ axis, \ reinterpret_cast::MappedType*>(output_data)); \ - return Status::OK(); \ } \ } @@ -122,19 +120,14 @@ Status ScatterElements::ComputeInternal(OpKernelContext* context) const { int rank = (int)input_dims.size(); auto* output_tensor = context->Output(0, input_data_shape); - CudaAsyncBuffer gpu_input_dims(this, input_dims); + CudaAsyncBuffer buffer_input_dims(this, input_dims); TensorPitches input_strides(input_dims); - CudaAsyncBuffer gpu_input_strides(this, input_strides); + CudaAsyncBuffer buffer_input_strides(this, input_strides); - CudaAsyncBuffer gpu_indices_dims(this, indices_dims); + CudaAsyncBuffer buffer_indices_dims(this, indices_dims); CudaAsyncBuffer fdm_indices_strides(this, rank); ORT_ENFORCE(CalculateFdmStrides(fdm_indices_strides.CpuSpan(), indices_dims)); - ORT_RETURN_IF_ERROR(gpu_input_dims.CopyToGpu()); - ORT_RETURN_IF_ERROR(gpu_input_strides.CopyToGpu()); - ORT_RETURN_IF_ERROR(gpu_indices_dims.CopyToGpu()); - ORT_RETURN_IF_ERROR(fdm_indices_strides.CopyToGpu()); - MLDataType Tin_type = indices_tensor->DataType(); MLDataType T_type = data_tensor->DataType(); diff --git a/onnxruntime/core/providers/cuda/tensor/scatter_elements_impl.cu b/onnxruntime/core/providers/cuda/tensor/scatter_elements_impl.cu index 5220493c7d..9041ffc5c7 100755 --- a/onnxruntime/core/providers/cuda/tensor/scatter_elements_impl.cu +++ b/onnxruntime/core/providers/cuda/tensor/scatter_elements_impl.cu @@ -7,6 +7,33 @@ namespace onnxruntime { namespace cuda { +template +__global__ void _ScatterElementsKernel2D( + const int max_dim, // max dim on the scattered axis + const T* input_data, + const Tin* indices_data, + const int64_t indices_size, + const fast_divmod indices_stride_row, + const T* updates, + const int64_t output_row_size, + T* output_data) { + CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(indices_index, indices_size); + + int row, col, data_idx; + indices_stride_row.divmod(indices_index, row, col); + int dim = (int)(indices_data[indices_index]); + if (dim >= -max_dim && dim < max_dim) { + if (dim < 0) dim += max_dim; + if (OUTERAXIS) { + data_idx = dim * output_row_size + col; + } else { + data_idx = row * output_row_size + dim; + } + output_data[data_idx] = updates[indices_index]; + } + // else invalid index +} + template __global__ void _ScatterElementsKernel( const int rank, @@ -28,7 +55,7 @@ __global__ void _ScatterElementsKernel( if (i == axis) { dim = (int)(indices_data[indices_index]); if (dim < -input_dims[i] || dim >= input_dims[i]) { - return; // Invalid index + return; // Invalid index } if (dim < 0) dim += input_dims[i]; } @@ -37,61 +64,162 @@ __global__ void _ScatterElementsKernel( output_data[data_idx] = updates[indices_index]; } +// From the innermost axis (largest) check equality of dim value of input and indices. +// If same, merge it and continue. Otherwise, copy remaining. The scatter axis need +// to be keep. +static int CompactInputIndicesDims( + int rank, int axis, int64_t* input_dims, int64_t* indices_dims, + std::vector& eff_input_dims, + std::vector& eff_indices_dims) { + eff_input_dims.clear(); + eff_indices_dims.clear(); + + bool could_continue_merge = true; + if (axis < rank - 1) { + eff_input_dims.push_back(1); + eff_indices_dims.push_back(1); + int i = rank - 1; + for (; i > axis; --i) { + if (input_dims[i] == indices_dims[i]) { + eff_input_dims.back() *= input_dims[i]; + eff_indices_dims.back() *= indices_dims[i]; + } else { + could_continue_merge = false; + break; + } + } + if (eff_input_dims.back() == 1) { + eff_input_dims.pop_back(); + eff_indices_dims.pop_back(); + } + if (!could_continue_merge) { + for (; i > axis; --i) { + eff_input_dims.push_back(input_dims[i]); + eff_indices_dims.push_back(indices_dims[i]); + } + } + } + could_continue_merge = could_continue_merge && (input_dims[axis] == indices_dims[axis]); + eff_input_dims.push_back(input_dims[axis]); + eff_indices_dims.push_back(indices_dims[axis]); + int new_axis = (int)(eff_input_dims.size()); + if (axis > 0) { + if (could_continue_merge) { + eff_input_dims.push_back(1); + eff_indices_dims.push_back(1); + } + int i = axis - 1; + for (; i >= 0 && could_continue_merge; --i) { + if (input_dims[i] == indices_dims[i]) { + eff_input_dims.back() *= input_dims[i]; + eff_indices_dims.back() *= indices_dims[i]; + } else { + could_continue_merge = false; + break; + } + } + if (new_axis < (int)eff_indices_dims.size() && eff_input_dims.back() == 1) { + eff_input_dims.pop_back(); + eff_indices_dims.pop_back(); + } + if (!could_continue_merge) { + for (; i >= 0 && could_continue_merge; --i) { + eff_input_dims.push_back(input_dims[i]); + eff_indices_dims.push_back(indices_dims[i]); + } + } + } + new_axis = eff_input_dims.size() - new_axis; + std::reverse(eff_input_dims.begin(), eff_input_dims.end()); + std::reverse(eff_indices_dims.begin(), eff_indices_dims.end()); + return new_axis; +} + template -void ScatterElementsImpl( - const int rank, +Status ScatterElementsImpl2D( const T* input_data, - const int64_t input_size, - const int64_t* input_dims, - const int64_t* input_strides, + const std::vector& input_dims, const Tin* indices_data, const int64_t indices_size, - const int64_t* indices_dims, - const fast_divmod* indices_strides, + const std::vector& indices_dims, const T* updates, const int axis, T* output_data) { + int blocksPerGrid = gsl::narrow_cast(CeilDiv(indices_size, GridDim::maxThreadsPerBlock)); + fast_divmod indices_stride_row(indices_dims[1]); + if (axis == 0) { + _ScatterElementsKernel2D<<>>( + gsl::narrow_cast(input_dims[0]), input_data, + indices_data, indices_size, indices_stride_row, + updates, input_dims[1], output_data); + } else { + _ScatterElementsKernel2D<<>>( + gsl::narrow_cast(input_dims[1]), input_data, + indices_data, indices_size, indices_stride_row, + updates, input_dims[1], output_data); + } + return Status::OK(); +} +template +Status ScatterElementsImpl( + const int rank, + const T* input_data, + const int64_t input_size, + CudaKernel::CudaAsyncBuffer& buffer_input_dims, + CudaKernel::CudaAsyncBuffer& buffer_input_strides, + const Tin* indices_data, + const int64_t indices_size, + CudaKernel::CudaAsyncBuffer& buffer_indices_dims, + CudaKernel::CudaAsyncBuffer& fdm_indices_strides, + const T* updates, + const int axis, + T* output_data) { if (input_data != output_data) { - cudaMemcpyAsync(output_data, input_data, input_size * sizeof(T), cudaMemcpyDeviceToDevice, 0); + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(output_data, input_data, input_size * sizeof(T), cudaMemcpyDeviceToDevice, 0)); } if (indices_size > 0) { - int blocksPerGrid = (int)((indices_size + GridDim::maxThreadsPerBlock - 1) / GridDim::maxThreadsPerBlock); + std::vector eff_input_dims; + std::vector eff_indices_dims; + int new_axis = CompactInputIndicesDims( + rank, axis, buffer_input_dims.CpuPtr(), buffer_indices_dims.CpuPtr(), eff_input_dims, eff_indices_dims); + if (eff_input_dims.size() == 2) { + return ScatterElementsImpl2D( + input_data, eff_input_dims, indices_data, indices_size, eff_indices_dims, updates, new_axis, output_data); + } + + ORT_RETURN_IF_ERROR(buffer_input_dims.CopyToGpu()); + ORT_RETURN_IF_ERROR(buffer_input_strides.CopyToGpu()); + ORT_RETURN_IF_ERROR(buffer_indices_dims.CopyToGpu()); + ORT_RETURN_IF_ERROR(fdm_indices_strides.CopyToGpu()); + int blocksPerGrid = gsl::narrow_cast(CeilDiv(indices_size, GridDim::maxThreadsPerBlock)); _ScatterElementsKernel<<>>( - rank, input_data, input_dims, input_strides, - indices_data, indices_size, indices_dims, indices_strides, + rank, input_data, buffer_input_dims.GpuPtr(), buffer_input_strides.GpuPtr(), + indices_data, indices_size, buffer_indices_dims.GpuPtr(), fdm_indices_strides.GpuPtr(), updates, axis, output_data); } + return Status::OK(); } -#define SPECIALIZED_IMPL(T) \ - template void ScatterElementsImpl( \ - const int rank, \ - const T* input_data, \ - const int64_t input_size, \ - const int64_t* input_dims, \ - const int64_t* input_strides, \ - const int32_t* indices_data, \ - const int64_t indices_size, \ - const int64_t* indices_dims, \ - const fast_divmod* indices_strides, \ - const T* updates, \ - const int axis, \ - T* output_data); \ - template void ScatterElementsImpl( \ - const int rank, \ - const T* input_data, \ - const int64_t input_size, \ - const int64_t* input_dims, \ - const int64_t* input_strides, \ - const int64_t* indices_data, \ - const int64_t indices_size, \ - const int64_t* indices_dims, \ - const fast_divmod* indices_strides, \ - const T* updates, \ - const int axis, \ - T* output_data); \ +#define SPECIALIZED_TINDEX_IMPL(T, TIndex) \ + template Status ScatterElementsImpl( \ + const int rank, \ + const T* input_data, \ + const int64_t input_size, \ + CudaKernel::CudaAsyncBuffer& buffer_input_dims, \ + CudaKernel::CudaAsyncBuffer& buffer_input_strides, \ + const TIndex* indices_data, \ + const int64_t indices_size, \ + CudaKernel::CudaAsyncBuffer& buffer_indices_dims, \ + CudaKernel::CudaAsyncBuffer& indices_strides, \ + const T* updates, \ + const int axis, \ + T* output_data) + +#define SPECIALIZED_IMPL(T) \ + SPECIALIZED_TINDEX_IMPL(T, int32_t); \ + SPECIALIZED_TINDEX_IMPL(T, int64_t); SPECIALIZED_IMPL(int8_t) SPECIALIZED_IMPL(int16_t) @@ -108,4 +236,3 @@ SPECIALIZED_IMPL(bool) } // namespace cuda } // namespace onnxruntime - diff --git a/onnxruntime/core/providers/cuda/tensor/scatter_elements_impl.h b/onnxruntime/core/providers/cuda/tensor/scatter_elements_impl.h index 7b1a89d334..f4fb1968e0 100755 --- a/onnxruntime/core/providers/cuda/tensor/scatter_elements_impl.h +++ b/onnxruntime/core/providers/cuda/tensor/scatter_elements_impl.h @@ -10,16 +10,16 @@ namespace onnxruntime { namespace cuda { template -void ScatterElementsImpl( +Status ScatterElementsImpl( const int rank, const T* input_data, const int64_t input_size, - const int64_t* input_dims, - const int64_t* input_strides, + CudaKernel::CudaAsyncBuffer& buffer_input_dims, + CudaKernel::CudaAsyncBuffer& buffer_input_strides, const Tin* indices_data, const int64_t indices_size, - const int64_t* indices_dims, - const fast_divmod* indices_strides, + CudaKernel::CudaAsyncBuffer& buffer_indices_dims, + CudaKernel::CudaAsyncBuffer& indices_strides, const T* updates, const int axis, T* output_data);