From 8631b70c7355534138273303d380eec2ca0890f2 Mon Sep 17 00:00:00 2001 From: Zhang Lei Date: Thu, 12 Dec 2019 16:36:06 -0800 Subject: [PATCH] Improve cuda expand() opeator's performance. (#2624) --- .../core/providers/cuda/tensor/expand.cc | 124 +++++++----- .../core/providers/cuda/tensor/expand_impl.cu | 186 ++++++++++++------ .../core/providers/cuda/tensor/expand_impl.h | 12 +- 3 files changed, 209 insertions(+), 113 deletions(-) diff --git a/onnxruntime/core/providers/cuda/tensor/expand.cc b/onnxruntime/core/providers/cuda/tensor/expand.cc index 72ae71cec4..347c36529b 100644 --- a/onnxruntime/core/providers/cuda/tensor/expand.cc +++ b/onnxruntime/core/providers/cuda/tensor/expand.cc @@ -5,69 +5,105 @@ #include "expand_impl.h" #include "core/providers/cpu/tensor/utils.h" +using std::vector; + namespace onnxruntime { namespace cuda { +// Logically expanded y could just be a view of x. +static void CalcEffectiveDims(vector& x_dims, vector& y_dims) { + vector x_reverse; + vector y_reverse; + + int xi = gsl::narrow_cast(x_dims.size()) - 1; + for (int yi = gsl::narrow_cast(y_dims.size()) - 1; yi >= 0; --yi, --xi) { + int64_t xdim = (xi >= 0) ? x_dims[xi] : 1; + int64_t ydim = y_dims[yi]; + if (xdim == ydim || xdim == 1) { + x_reverse.push_back(xdim); + y_reverse.push_back(ydim); + } + else { // xdim < ydim && xdim > 1, split + ydim /= xdim; + x_reverse.push_back(xdim); + y_reverse.push_back(xdim); + x_reverse.push_back(1); + y_reverse.push_back(ydim); + } + } + + x_dims.clear(); + y_dims.clear(); + x_dims.push_back(1); + y_dims.push_back(1); + // compact the dims, remove (x=1, y=1), merge (x=1, y1*y2...) + for (int i = gsl::narrow_cast(y_reverse.size()) - 1; i >= 0; --i) { + if (x_reverse[i] == 1) { + if (y_reverse[i] == 1) { + continue; + } + if (x_dims.back() == 1) { + y_dims.back() *= y_reverse[i]; + } + else { + x_dims.push_back(1); + y_dims.push_back(y_reverse[i]); + } + } + else { // x_reverse[i] == y_reverse[i] + if (x_dims.back() == y_dims.back()) { + x_dims.back() *= x_reverse[i]; + y_dims.back() *= y_reverse[i]; + } + else { + x_dims.push_back(x_reverse[i]); + y_dims.push_back(y_reverse[i]); + } + } + } +} + Status Expand::ComputeInternal(OpKernelContext* ctx) const { - const auto& input0 = *ctx->Input(0); - const auto& input1 = *ctx->Input(1); + const auto& input_data_tensor = *ctx->Input(0); + const auto& input_shape_tensor = *ctx->Input(1); // new shape to be expanded to - const auto* p_shape = input1.template Data(); - std::vector output_dims{p_shape, p_shape + input1.Shape().Size()}; + const auto* p_shape = input_shape_tensor.template Data(); + std::vector output_dims{p_shape, p_shape + input_shape_tensor.Shape().Size()}; TensorShape output_shape(output_dims); - ORT_RETURN_IF_ERROR(ComputeOutputShape(Node().Name(), input0.Shape(), output_dims, output_shape)); - auto rank = output_shape.NumDimensions(); + ORT_RETURN_IF_ERROR(ComputeOutputShape(Node().Name(), input_data_tensor.Shape(), output_dims, output_shape)); auto& output_tensor = *ctx->Output(0, output_shape); - if (0 == output_shape.Size()) { return Status::OK(); } - auto input_shape = input0.Shape().GetDims(); + output_dims = output_shape.GetDims(); + auto input_dims = input_data_tensor.Shape().GetDims(); - // pad input_dims with 1 to make ranks match - for (size_t i = 0; i < rank - input_shape.size(); i++) { - input_shape.insert(input_shape.begin(), 1); + CalcEffectiveDims(input_dims, output_dims); + int rank = gsl::narrow_cast(output_dims.size()); + + CudaAsyncBuffer fdm_output_strides(this, rank); + ORT_ENFORCE(CalculateFdmStrides(fdm_output_strides.CpuSpan(), output_dims)); + + CudaAsyncBuffer input_view_strides(this, rank); + TensorPitches::Calculate(input_view_strides.CpuSpan(), input_dims); + for (int i = 0; i < rank; ++i) { + if (input_dims[i] == 1) input_view_strides.CpuSpan()[i] = 0; } - // create fast_divmod using dimension values - CudaAsyncBuffer fdm_input_dims(this, rank); - CudaAsyncBuffer fdm_output_dims(this, rank); - CudaAsyncBuffer fdm_output_subdim_size(this, rank); - { - auto in_span = fdm_input_dims.CpuSpan(); - auto out_span = fdm_output_dims.CpuSpan(); - auto sdm_span = fdm_output_subdim_size.CpuSpan(); - auto subdim_size = output_shape.Size(); - for (size_t i = 0; i < rank; i++) { - in_span[i] = fast_divmod(static_cast(input_shape[i])); - out_span[i] = fast_divmod(static_cast(output_shape[i])); - // output_shape[i] won't be 0 here, it's covered in (0 == output_shape.Size()) - // a null output will be returned for that case - subdim_size /= output_shape[i]; - sdm_span[i] = static_cast(subdim_size); - } - } - ORT_RETURN_IF_ERROR(fdm_input_dims.CopyToGpu()); - ORT_RETURN_IF_ERROR(fdm_output_dims.CopyToGpu()); - ORT_RETURN_IF_ERROR(fdm_output_subdim_size.CopyToGpu()); - - ExpandImpl( - input0.DataType()->Size(), - output_shape.NumDimensions(), - output_shape.Size(), - input0.Shape().Size(), - input0.DataRaw(), + return ExpandImpl( + input_data_tensor.DataType()->Size(), + gsl::narrow_cast(output_shape.Size()), + gsl::narrow_cast(input_data_tensor.Shape().Size()), + input_data_tensor.DataRaw(), output_tensor.MutableDataRaw(), - fdm_input_dims.GpuPtr(), - fdm_output_dims.GpuPtr(), - fdm_output_subdim_size.GpuPtr()); - - return Status::OK(); + fdm_output_strides, + input_view_strides); } + ONNX_OPERATOR_KERNEL_EX( Expand, kOnnxDomain, diff --git a/onnxruntime/core/providers/cuda/tensor/expand_impl.cu b/onnxruntime/core/providers/cuda/tensor/expand_impl.cu index bbad602ec0..38e71052a5 100644 --- a/onnxruntime/core/providers/cuda/tensor/expand_impl.cu +++ b/onnxruntime/core/providers/cuda/tensor/expand_impl.cu @@ -8,84 +8,144 @@ namespace onnxruntime { namespace cuda { +template +__global__ void _FillFromDataPtrKernel(T* output_data, const T* input_data, CUDA_LONG N) { + CUDA_LONG id = NumElementsPerThread * blockDim.x * blockIdx.x + threadIdx.x; + T val = *input_data; +#pragma unroll + for (int i = 0; i < NumElementsPerThread; i++) { + if (id < N) { + output_data[id] = val; + id += NumThreadsPerBlock; + } + } +} + template -__global__ void ExpandKernel( - const size_t rank, - const size_t N, - const size_t N_input, +void FillFromDataPtr(T* output_data, const T* input_data, int64_t count) { + int blocksPerGrid = gsl::narrow_cast(CeilDiv(count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread)); + CUDA_LONG N = static_cast(count); + _FillFromDataPtrKernel + <<>>(output_data, input_data, N); +} + +template +__global__ void ExpandKernel2D( + const int N, const T* input_data, T* output_data, - const fast_divmod* fdm_input_dims, - const fast_divmod* fdm_output_dims, - const fast_divmod* fdm_output_subdim_size) { + const fast_divmod fdm_output_stride0, + const int input_view_stride0, + const int input_view_stride1) { CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N); - // initialize - auto output_index = id; - auto input_index = 0; - auto input_subdim_size = N_input; - auto out_coord = output_index; - // use striding when tensor is larger than grid - int stride = blockDim.x * gridDim.x; + int dim0, dim1; + fdm_output_stride0.divmod(id, dim0, dim1); + output_data[id] = input_data[dim0 * input_view_stride0 + dim1 * input_view_stride1]; +} - // translate indices to coordinates. copy expanded dims from source - while (output_index < N) { - for (int64_t i = 0; i < rank; i++) { - input_subdim_size = fdm_input_dims[i].div(input_subdim_size); - auto new_out_coord = fdm_output_subdim_size[i].div(out_coord); - auto in_coord = (new_out_coord > (fdm_input_dims[i].d_ - 1)) ? fdm_input_dims[i].d_ - 1 : new_out_coord; - input_index += input_subdim_size * in_coord; - out_coord -= new_out_coord * fdm_output_subdim_size[i].d_; - } - output_data[output_index] = input_data[input_index]; - output_index += stride; - out_coord = output_index; - input_subdim_size = N_input; - input_index = 0; +template +__global__ void ExpandKernel( + const int rank, + const int N, + const T* input_data, + T* output_data, + const fast_divmod* fdm_output_strides, + const int64_t* input_view_strides) { + CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N); + + int dim, r = id, input_index = 0; + for (int i = 0; i < rank; ++i) { + fdm_output_strides[i].divmod(r, dim, r); + input_index += dim * input_view_strides[i]; } + output_data[id] = input_data[input_index]; +} + +Status ExpandByFill(const size_t element_size, const int N, const void* input_data, void* output_data) { +#define EXPAND_FILL_ON(TYPE) \ + case sizeof(TYPE): \ + FillFromDataPtr(reinterpret_cast(output_data), \ + reinterpret_cast(input_data), \ + static_cast(N)); \ + break + + switch (element_size) { + EXPAND_FILL_ON(int8_t); + EXPAND_FILL_ON(int16_t); + EXPAND_FILL_ON(int32_t); + EXPAND_FILL_ON(int64_t); + default: + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Type not supported for Expand operator"); + } + return Status::OK(); +} + +Status Expand2D( + const size_t element_size, + const int N, + const void* input_data, + void* output_data, + const fast_divmod fdm_output_stride0, + const int input_view_stride0, + const int input_view_stride1) { +#define EXPAND2D_ON(TYPE) \ + case sizeof(TYPE): \ + ExpandKernel2D<<>>( \ + N, reinterpret_cast(input_data), reinterpret_cast(output_data), \ + fdm_output_stride0, input_view_stride0, input_view_stride1); \ + break + + int blocksPerGrid = gsl::narrow_cast(CeilDiv(N, GridDim::maxThreadsPerBlock)); + switch (element_size) { + EXPAND2D_ON(int8_t); + EXPAND2D_ON(int16_t); + EXPAND2D_ON(int32_t); + EXPAND2D_ON(int64_t); + default: + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Type not supported for Expand operator"); + } + return Status::OK(); } Status ExpandImpl( const size_t element_size, - const size_t rank, - const size_t N, - const size_t N_input, + const int N_output, + const int N_input, const void* input_data, void* output_data, - const fast_divmod* fdm_input_dims, - const fast_divmod* fdm_output_dims, - const fast_divmod* fdm_output_subdim_size) { - int blocksPerGrid = (int)(ceil(static_cast(N) / GridDim::maxThreadsPerBlock)); + CudaKernel::CudaAsyncBuffer& fdm_output_strides, + CudaKernel::CudaAsyncBuffer& input_view_strides) { + const int rank = static_cast(fdm_output_strides.count()); + if (rank == 1) { + if (N_input == N_output) { + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(output_data, input_data, N_output * element_size, cudaMemcpyDeviceToDevice)); + } else { // N_input == 1 + return ExpandByFill(element_size, N_output, input_data, output_data); + } + } else if (rank == 2) { + return Expand2D(element_size, N_output, input_data, output_data, + fdm_output_strides.CpuSpan()[0], + static_cast(input_view_strides.CpuSpan()[0]), + static_cast(input_view_strides.CpuSpan()[1])); + } + + int blocksPerGrid = gsl::narrow_cast(CeilDiv(N_output, GridDim::maxThreadsPerBlock)); + fdm_output_strides.CopyToGpu(); + input_view_strides.CopyToGpu(); + +#define EXPAND_ON(TYPE) \ + case sizeof(TYPE): \ + ExpandKernel<<>>( \ + rank, N_output, reinterpret_cast(input_data), reinterpret_cast(output_data), \ + fdm_output_strides.GpuPtr(), input_view_strides.GpuPtr()); \ + break switch (element_size) { - case sizeof(uint8_t): - ExpandKernel<<>>( - rank, N, N_input, - reinterpret_cast::MappedType*>(input_data), - reinterpret_cast::MappedType*>(output_data), - fdm_input_dims, fdm_output_dims, fdm_output_subdim_size); - break; - case sizeof(uint16_t): - ExpandKernel<<>>( - rank, N, N_input, - reinterpret_cast::MappedType*>(input_data), - reinterpret_cast::MappedType*>(output_data), - fdm_input_dims, fdm_output_dims, fdm_output_subdim_size); - break; - case sizeof(uint32_t): - ExpandKernel<<>>( - rank, N, N_input, - reinterpret_cast::MappedType*>(input_data), - reinterpret_cast::MappedType*>(output_data), - fdm_input_dims, fdm_output_dims, fdm_output_subdim_size); - break; - case sizeof(uint64_t): - ExpandKernel<<>>( - rank, N, N_input, - reinterpret_cast::MappedType*>(input_data), - reinterpret_cast::MappedType*>(output_data), - fdm_input_dims, fdm_output_dims, fdm_output_subdim_size); - break; + EXPAND_ON(uint8_t); + EXPAND_ON(uint16_t); + EXPAND_ON(uint32_t); + EXPAND_ON(uint64_t); default: return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Type not supported for Expand operator"); } diff --git a/onnxruntime/core/providers/cuda/tensor/expand_impl.h b/onnxruntime/core/providers/cuda/tensor/expand_impl.h index c127d72477..543dedaa4f 100644 --- a/onnxruntime/core/providers/cuda/tensor/expand_impl.h +++ b/onnxruntime/core/providers/cuda/tensor/expand_impl.h @@ -6,20 +6,20 @@ #include "core/providers/cuda/shared_inc/cuda_utils.h" #include "core/framework/data_types.h" #include "core/common/common.h" +#include "core/providers/cuda/cuda_common.h" namespace onnxruntime { namespace cuda { Status ExpandImpl( const size_t element_size, - const size_t shape_rank, - const size_t N, - const size_t N_input, + const int N_output, + const int N_input, const void* input_data, void* output_data, - const fast_divmod* fdm_input_dims, - const fast_divmod* fdm_output_dims, - const fast_divmod* fdm_output_subdim_size); + CudaKernel::CudaAsyncBuffer& fdm_output_strides, + CudaKernel::CudaAsyncBuffer& input_view_strides); + } // namespace cuda } // namespace onnxruntime