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

Optimize cuda scatter() on 2D compatible. (#2628)

Browse source 
* Optimize cuda scatter() on 2D compatible.

* Add some comments.
This commit is contained in:
Zhang Lei 2019-12-12 16:38:06 -08:00 committed by GitHub
parent e31be23c42
commit fff1ed9bfc
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 185 additions and 65 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
onnxruntime/core/providers/cuda/tensor/scatter_elements.cc
View file

@ -41,37 +41,35 @@ ONNX_OPERATOR_KERNEL_EX(
const T* update_data = updates_tensor->template Data<T>(); \
if (utils::IsPrimitiveDataType<int32_t>(Tin_type)) { \
const int32_t* indices_data = indices_tensor->template Data<int32_t>(); \
ScatterElementsImpl( \
return ScatterElementsImpl( \
rank, \
reinterpret_cast<const ToCudaType<T>::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<const ToCudaType<T>::MappedType*>(update_data), \
axis, \
reinterpret_cast<ToCudaType<T>::MappedType*>(output_data)); \
return Status::OK(); \
} \
if (utils::IsPrimitiveDataType<int64_t>(Tin_type)) { \
const int64_t* indices_data = indices_tensor->template Data<int64_t>(); \
ScatterElementsImpl( \
return ScatterElementsImpl( \
rank, \
reinterpret_cast<const ToCudaType<T>::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<const ToCudaType<T>::MappedType*>(update_data), \
axis, \
reinterpret_cast<ToCudaType<T>::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<int64_t> gpu_input_dims(this, input_dims);
CudaAsyncBuffer<int64_t> buffer_input_dims(this, input_dims);
TensorPitches input_strides(input_dims);
CudaAsyncBuffer<int64_t> gpu_input_strides(this, input_strides);
CudaAsyncBuffer<int64_t> buffer_input_strides(this, input_strides);
CudaAsyncBuffer<int64_t> gpu_indices_dims(this, indices_dims);
CudaAsyncBuffer<int64_t> buffer_indices_dims(this, indices_dims);
CudaAsyncBuffer<fast_divmod> 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();

207
onnxruntime/core/providers/cuda/tensor/scatter_elements_impl.cu
View file

@ -7,6 +7,33 @@
namespace onnxruntime {
namespace cuda {
template <typename T, typename Tin, bool OUTERAXIS>
__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 <typename T, typename Tin>
__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<int64_t>& eff_input_dims,
std::vector<int64_t>& 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 <typename T, typename Tin>
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<int64_t>& input_dims,
const Tin* indices_data,
const int64_t indices_size,
const int64_t* indices_dims,
const fast_divmod* indices_strides,
const std::vector<int64_t>& indices_dims,
const T* updates,
const int axis,
T* output_data) {
int blocksPerGrid = gsl::narrow_cast<int>(CeilDiv(indices_size, GridDim::maxThreadsPerBlock));
fast_divmod indices_stride_row(indices_dims[1]);
if (axis == 0) {
_ScatterElementsKernel2D<T, Tin, true><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
gsl::narrow_cast<int>(input_dims[0]), input_data,
indices_data, indices_size, indices_stride_row,
updates, input_dims[1], output_data);
} else {
_ScatterElementsKernel2D<T, Tin, false><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
gsl::narrow_cast<int>(input_dims[1]), input_data,
indices_data, indices_size, indices_stride_row,
updates, input_dims[1], output_data);
}
return Status::OK();
}
template <typename T, typename Tin>
Status ScatterElementsImpl(
const int rank,
const T* input_data,
const int64_t input_size,
CudaKernel::CudaAsyncBuffer<int64_t>& buffer_input_dims,
CudaKernel::CudaAsyncBuffer<int64_t>& buffer_input_strides,
const Tin* indices_data,
const int64_t indices_size,
CudaKernel::CudaAsyncBuffer<int64_t>& buffer_indices_dims,
CudaKernel::CudaAsyncBuffer<fast_divmod>& 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<int64_t> eff_input_dims;
std::vector<int64_t> 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<int>(CeilDiv(indices_size, GridDim::maxThreadsPerBlock));
_ScatterElementsKernel<T, Tin><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
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<T, int32_t>( \
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<T, int64_t>( \
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<T, TIndex>( \
const int rank, \
const T* input_data, \
const int64_t input_size, \
CudaKernel::CudaAsyncBuffer<int64_t>& buffer_input_dims, \
CudaKernel::CudaAsyncBuffer<int64_t>& buffer_input_strides, \
const TIndex* indices_data, \
const int64_t indices_size, \
CudaKernel::CudaAsyncBuffer<int64_t>& buffer_indices_dims, \
CudaKernel::CudaAsyncBuffer<fast_divmod>& 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

10
onnxruntime/core/providers/cuda/tensor/scatter_elements_impl.h
View file

@ -10,16 +10,16 @@ namespace onnxruntime {
namespace cuda {
template <typename T, typename Tin>
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<int64_t>& buffer_input_dims,
CudaKernel::CudaAsyncBuffer<int64_t>& buffer_input_strides,
const Tin* indices_data,
const int64_t indices_size,
const int64_t* indices_dims,
const fast_divmod* indices_strides,
CudaKernel::CudaAsyncBuffer<int64_t>& buffer_indices_dims,
CudaKernel::CudaAsyncBuffer<fast_divmod>& indices_strides,
const T* updates,
const int axis,
T* output_data);