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

Improve cuda expand() opeator's performance. (#2624)

Browse source
This commit is contained in:
Zhang Lei 2019-12-12 16:36:06 -08:00 committed by GitHub
parent ac08b58867
commit 8631b70c73
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 209 additions and 113 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -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<int64_t>& x_dims, vector<int64_t>& y_dims) {
vector<int64_t> x_reverse;
vector<int64_t> y_reverse;
int xi = gsl::narrow_cast<int>(x_dims.size()) - 1;
for (int yi = gsl::narrow_cast<int>(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<int>(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<Tensor>(0);
const auto& input1 = *ctx->Input<Tensor>(1);
const auto& input_data_tensor = *ctx->Input<Tensor>(0);
const auto& input_shape_tensor = *ctx->Input<Tensor>(1);
// new shape to be expanded to
const auto* p_shape = input1.template Data<int64_t>();
std::vector<int64_t> output_dims{p_shape, p_shape + input1.Shape().Size()};
const auto* p_shape = input_shape_tensor.template Data<int64_t>();
std::vector<int64_t> 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<int>(output_dims.size());
CudaAsyncBuffer<fast_divmod> fdm_output_strides(this, rank);
ORT_ENFORCE(CalculateFdmStrides(fdm_output_strides.CpuSpan(), output_dims));
CudaAsyncBuffer<int64_t> 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<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();
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<int>(input_shape[i]));
out_span[i] = fast_divmod(static_cast<int>(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<int>(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<int>(output_shape.Size()),
gsl::narrow_cast<int>(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,

186
onnxruntime/core/providers/cuda/tensor/expand_impl.cu
View file

@ -8,84 +8,144 @@
namespace onnxruntime {
namespace cuda {
template <typename T, int NumThreadsPerBlock, int NumElementsPerThread>
__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 <typename T>
__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<int>(CeilDiv(count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
CUDA_LONG N = static_cast<CUDA_LONG>(count);
_FillFromDataPtrKernel<T, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>
<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(output_data, input_data, N);
}
template <typename T>
__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 <typename T>
__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<TYPE*>(output_data), \
reinterpret_cast<const TYPE*>(input_data), \
static_cast<int64_t>(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<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>( \
N, reinterpret_cast<const TYPE*>(input_data), reinterpret_cast<TYPE*>(output_data), \
fdm_output_stride0, input_view_stride0, input_view_stride1); \
break
int blocksPerGrid = gsl::narrow_cast<int>(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<float>(N) / GridDim::maxThreadsPerBlock));
CudaKernel::CudaAsyncBuffer<fast_divmod>& fdm_output_strides,
CudaKernel::CudaAsyncBuffer<int64_t>& input_view_strides) {
const int rank = static_cast<int>(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<int>(input_view_strides.CpuSpan()[0]),
static_cast<int>(input_view_strides.CpuSpan()[1]));
}
int blocksPerGrid = gsl::narrow_cast<int>(CeilDiv(N_output, GridDim::maxThreadsPerBlock));
fdm_output_strides.CopyToGpu();
input_view_strides.CopyToGpu();
#define EXPAND_ON(TYPE) \
case sizeof(TYPE): \
ExpandKernel<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>( \
rank, N_output, reinterpret_cast<const TYPE*>(input_data), reinterpret_cast<TYPE*>(output_data), \
fdm_output_strides.GpuPtr(), input_view_strides.GpuPtr()); \
break
switch (element_size) {
case sizeof(uint8_t):
ExpandKernel<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
rank, N, N_input,
reinterpret_cast<const ToCudaType<uint8_t>::MappedType*>(input_data),
reinterpret_cast<ToCudaType<uint8_t>::MappedType*>(output_data),
fdm_input_dims, fdm_output_dims, fdm_output_subdim_size);
break;
case sizeof(uint16_t):
ExpandKernel<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
rank, N, N_input,
reinterpret_cast<const ToCudaType<uint16_t>::MappedType*>(input_data),
reinterpret_cast<ToCudaType<uint16_t>::MappedType*>(output_data),
fdm_input_dims, fdm_output_dims, fdm_output_subdim_size);
break;
case sizeof(uint32_t):
ExpandKernel<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
rank, N, N_input,
reinterpret_cast<const ToCudaType<uint32_t>::MappedType*>(input_data),
reinterpret_cast<ToCudaType<uint32_t>::MappedType*>(output_data),
fdm_input_dims, fdm_output_dims, fdm_output_subdim_size);
break;
case sizeof(uint64_t):
ExpandKernel<<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
rank, N, N_input,
reinterpret_cast<const ToCudaType<uint64_t>::MappedType*>(input_data),
reinterpret_cast<ToCudaType<uint64_t>::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");
}

12
onnxruntime/core/providers/cuda/tensor/expand_impl.h
View file

@ -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<fast_divmod>& fdm_output_strides,
CudaKernel::CudaAsyncBuffer<int64_t>& input_view_strides);
} // namespace cuda
} // namespace onnxruntime