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

Speed up GatherElements CPU Op (#11349)

Browse source 
* Specialize element copy by element size
This commit is contained in:
Ryan Hill 2022-04-27 17:57:56 -07:00 committed by GitHub
parent d64769c38e
commit 9e7753c109
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 73 additions and 146 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

215
onnxruntime/core/providers/cpu/tensor/gather_elements.cc
View file

@ -25,8 +25,6 @@ ONNX_CPU_OPERATOR_KERNEL(
DataTypeImpl::GetTensorType<int64_t>()}),
GatherElements);
static constexpr int kParallelizationThreshold = 10 * 1000;
// Some helpers needed for GatherElements op -
// The following method computes the offset in the flattened array
@ -35,15 +33,15 @@ static constexpr int kParallelizationThreshold = 10 * 1000;
// 'indices' value
// This prevents the need to compute this offset for every element within the same 'inner_dimension' chunk
// as this value just differs by 1 for the chunk elements and we can have this cached and re-use as needed
static inline int64_t compute_base_offset(const std::vector<int64_t>& shape, const TensorPitches& pitches, int64_t skip_axis) {
static inline size_t compute_base_offset(const TensorShapeVector& shape, const TensorPitches& pitches, int64_t skip_axis) {
// in this context, rank can never be < 1, so saving checking overhead
auto loop_size = static_cast<int64_t>(shape.size()) - 1;
int64_t base_offset = 0;
size_t base_offset = 0;
for (int64_t i = 0; i < loop_size; ++i) {
if (i != skip_axis)
base_offset += (shape[i] * pitches[i]);
base_offset += shape[i] * pitches[i];
}
return base_offset;
@ -65,7 +63,7 @@ static int64_t calculate_num_inner_dim(const TensorShape& dims) {
// Example 2: current_dims = [0, 0, x] tensor_dims = [1, 2, 2], then current_dims = [0, 1, x]
// current_dims = [0, 1, x] tensor_dims = [1, 2, 2], then current_dims = [0, 0, x]
static inline void increment_over_inner_dim(std::vector<int64_t>& current_dims, const TensorShape& tensor_dims) {
static inline void increment_over_inner_dim(TensorShapeVector& current_dims, const TensorShape& tensor_dims) {
// in this context, rank can never be < 1, so saving checking overhead
int64_t rank = static_cast<int64_t>(current_dims.size());
@ -88,17 +86,21 @@ static inline void increment_over_inner_dim(std::vector<int64_t>& current_dims,
}
}
template <typename Tin>
static inline int64_t GetNegativeIndexAdjustedValue(const Tin* indices_data, Tin index, int64_t axis,
const TensorShape& input_shape) {
int64_t retval = -1;
if (indices_data[index] < 0) {
retval = static_cast<int64_t>(indices_data[index] + input_shape[axis]);
} else {
retval = static_cast<int64_t>(indices_data[index]);
}
return retval;
}
#if defined(_MSC_VER)
#define FORCEINLINE __forceinline
#else
#define FORCEINLINE __attribute__((always_inline)) inline
#endif
template <typename T>
FORCEINLINE int64_t GetIndex(size_t i, const T* indices, int64_t axis_size) {
int64_t index = indices[i];
if (index < 0) // Handle negative indices
index += axis_size;
if (std::make_unsigned_t<T>(index) >= std::make_unsigned_t<T>(axis_size))
ORT_THROW("GatherElements op: Value in indices must be within bounds [-", axis_size, " , ", axis_size - 1, "]. Actual value is ", indices[i]);
return index;
};
#ifdef __GNUC__
#pragma GCC diagnostic push
@ -106,135 +108,68 @@ static inline int64_t GetNegativeIndexAdjustedValue(const Tin* indices_data, Tin
#pragma GCC diagnostic ignored "-Wclass-memaccess"
#endif
#endif
template <bool is_string, typename T, typename Tin>
static void core_impl(const Tensor* input_tensor, const Tensor* indices_tensor,
Tensor* output_tensor, int64_t axis, concurrency::ThreadPool* ttp) {
// get pointer to input data
// optimizer will remove the redundant if/else block based on 'is_string' template parameter
const T* input_data = nullptr;
if (is_string) {
input_data = input_tensor->Data<T>();
} else {
input_data = reinterpret_cast<const T*>(input_tensor->DataRaw());
}
// get pointer to output data
// optimizer will remove the redundant if/else block based on 'is_string' template parameter
T* output_data = nullptr;
if (is_string) {
output_data = output_tensor->MutableData<T>();
} else {
output_data = reinterpret_cast<T*>(output_tensor->MutableDataRaw());
}
template <typename Tin>
static void core_impl(const Tensor* input_tensor, const Tensor* indices_tensor, Tensor* output_tensor, int64_t axis) {
// Get input & output pointers
const int8_t* input_data = reinterpret_cast<const int8_t*>(input_tensor->DataRaw());
int8_t* output_data = reinterpret_cast<int8_t*>(output_tensor->MutableDataRaw());
size_t element_size = input_tensor->DataType()->Size();
bool is_string = input_tensor->IsDataTypeString();
const int64_t input_rank = static_cast<int64_t>(input_tensor->Shape().NumDimensions());
const TensorPitches input_shape_pitches(*input_tensor);
const auto& input_shape = input_tensor->Shape();
const TensorShape& indices_shape = indices_tensor->Shape();
const Tin* indices_data = indices_tensor->Data<Tin>();
size_t num_inner_dim = calculate_num_inner_dim(indices_shape);
size_t inner_dim_size = indices_shape[input_rank - 1];
const Tin* indices = indices_tensor->Data<Tin>();
// validate indices
auto num_elements = indices_tensor->Shape().Size();
int64_t lower_index_limit = -input_shape[axis];
int64_t upper_index_limit = input_shape[axis] - 1;
TensorShapeVector process_dims(input_rank, 0);
const TensorPitches input_shape_pitches(*input_tensor);
int64_t axis_pitch = input_shape_pitches[axis];
int64_t axis_size = input_tensor->Shape()[axis];
auto validation_fn = [indices_data, lower_index_limit, upper_index_limit](ptrdiff_t i) {
auto indices_val = indices_data[i];
if (indices_val < lower_index_limit || indices_val > upper_index_limit)
ORT_THROW("GatherElements op: Value in indices must be within bounds [",
lower_index_limit, " , ", upper_index_limit, "]. Actual value is ", indices_val);
auto DoAxis = [](auto* output, auto* input, auto* indices, size_t inner_dim_size, int64_t axis_size, int64_t axis_pitch) {
for (size_t i = 0; i < inner_dim_size; i++)
output[i] = input[GetIndex(i, indices, axis_size) * axis_pitch + i];
};
for (int64_t i = 0; i < num_elements; ++i) { // TODO: parallelize this? didn't give any benefit in my tests
validation_fn(i);
}
int64_t num_inner_dim = calculate_num_inner_dim(indices_shape);
int64_t inner_dim_size = indices_shape[input_rank - 1];
bool processing_inner_dim = (axis == input_rank - 1) ? true : false;
// Special case required for innermost axis, no axis_pitch multiply needed or adding i
auto DoInnermostAxis = [](auto* output, auto* input, auto* indices, size_t inner_dim_size, int64_t axis_size, int64_t /*axis_pitch*/) {
for (size_t i = 0; i < inner_dim_size; i++)
output[i] = input[GetIndex(i, indices, axis_size)];
};
int64_t base_offset = 0;
Tin indices_counter = 0;
size_t element_size = input_tensor->DataType()->Size();
std::vector<int64_t> process_dims(input_rank, 0);
int64_t output_counter = 0;
auto conditional_batch_call = [ttp, inner_dim_size](std::function<void(ptrdiff_t)> f) {
if (inner_dim_size < kParallelizationThreshold) { // TODO: tune this, arbitrary threshold
for (int64_t i = 0; i < inner_dim_size; ++i) {
f(i);
}
} else {
concurrency::ThreadPool::TryBatchParallelFor(ttp, inner_dim_size, f, 0);
auto MainLoop = [&](auto* output, auto* input, auto LoopFn) {
while (num_inner_dim-- != 0) {
LoopFn(output, input + compute_base_offset(process_dims, input_shape_pitches, axis), indices, inner_dim_size, axis_size, axis_pitch);
output += inner_dim_size;
indices += static_cast<Tin>(inner_dim_size);
increment_over_inner_dim(process_dims, indices_shape);
}
};
if (!processing_inner_dim) {
while (num_inner_dim-- != 0) {
base_offset = compute_base_offset(process_dims, input_shape_pitches, axis);
// Iterate over the elements based on the element size (or if it's a string). For everything but strings
// we do a binary copy, so handling the 1, 2, 4, and 8 byte sizes covers all cases
auto SwitchOnSizes = [&](auto LoopFn) {
if (is_string)
MainLoop(reinterpret_cast<std::string*>(output_data), reinterpret_cast<const std::string*>(input_data), LoopFn);
else if (element_size == sizeof(uint32_t))
MainLoop(reinterpret_cast<uint32_t*>(output_data), reinterpret_cast<const uint32_t*>(input_data), LoopFn);
else if (element_size == sizeof(uint16_t))
MainLoop(reinterpret_cast<uint16_t*>(output_data), reinterpret_cast<const uint16_t*>(input_data), LoopFn);
else if (element_size == sizeof(uint8_t))
MainLoop(reinterpret_cast<uint8_t*>(output_data), reinterpret_cast<const uint8_t*>(input_data), LoopFn);
else if (element_size == sizeof(uint64_t))
MainLoop(reinterpret_cast<uint64_t*>(output_data), reinterpret_cast<const uint64_t*>(input_data), LoopFn);
else
ORT_THROW("GatherElements op: Unsupported tensor type, size:", element_size);
};
// process 1 chunk of 'inner dimension' length
// optimizer will remove the redundant if/else block based on 'is_string' template parameter
if (is_string) {
auto fn = [input_data, output_data, base_offset, input_shape_pitches,
indices_data, indices_counter, axis, input_shape, output_counter](ptrdiff_t i) {
output_data[i + output_counter] =
input_data[base_offset +
(GetNegativeIndexAdjustedValue<Tin>(indices_data, static_cast<Tin>(i) + indices_counter, axis, input_shape) *
input_shape_pitches[axis]) +
i];
};
conditional_batch_call(fn);
output_counter += inner_dim_size;
} else {
auto fn = [input_data, output_data, base_offset, input_shape_pitches, element_size,
indices_data, indices_counter, axis, input_shape](ptrdiff_t i) {
memcpy(output_data + (i * element_size),
input_data + (base_offset + (GetNegativeIndexAdjustedValue<Tin>(indices_data, static_cast<Tin>(i) + indices_counter, axis, input_shape) * input_shape_pitches[axis]) + i) * element_size,
element_size);
};
conditional_batch_call(fn);
output_data += inner_dim_size * element_size;
}
indices_counter += static_cast<Tin>(inner_dim_size);
increment_over_inner_dim(process_dims, indices_shape);
}
}
// we special-case inner dim as we can weed-out some unnecessary computations in element offset calculations
else {
while (num_inner_dim-- != 0) {
base_offset = compute_base_offset(process_dims, input_shape_pitches, axis);
// process 1 chunk of 'inner dimension' length
if (is_string) {
auto fn = [input_data, output_data, base_offset,
indices_data, indices_counter, axis, input_shape, output_counter](ptrdiff_t i) {
// for innermost axis, input_shape_pitches[axis] = 1 (so no need to multiply)
output_data[i + output_counter] =
input_data[base_offset +
GetNegativeIndexAdjustedValue<Tin>(indices_data, static_cast<Tin>(i) + indices_counter, axis, input_shape)];
};
conditional_batch_call(fn);
output_counter += inner_dim_size;
} else {
// for innermost axis, input_shape_pitches[axis] = 1 (so no need to multiply)
auto fn = [input_data, output_data, base_offset, element_size,
indices_data, indices_counter, axis, input_shape](ptrdiff_t i) {
memcpy(output_data + (i * element_size),
input_data + (base_offset +
GetNegativeIndexAdjustedValue<Tin>(indices_data, static_cast<Tin>(i) + indices_counter, axis, input_shape)) *
element_size,
element_size);
};
conditional_batch_call(fn);
output_data += inner_dim_size * element_size;
}
indices_counter += static_cast<Tin>(inner_dim_size);
increment_over_inner_dim(process_dims, indices_shape);
}
}
if (axis == input_rank - 1)
SwitchOnSizes(DoInnermostAxis);
else
SwitchOnSizes(DoAxis);
}
#ifdef __GNUC__
#pragma GCC diagnostic pop
@ -298,18 +233,10 @@ Status GatherElements::Compute(OpKernelContext* context) const {
if (indices_shape.Size() == 0)
return Status::OK();
auto* ttp = context->GetOperatorThreadPool();
if (input_tensor->IsDataTypeString()) {
if (indices_tensor->IsDataType<int32_t>())
core_impl<true, std::string, int32_t>(input_tensor, indices_tensor, output_tensor, axis, ttp);
else
core_impl<true, std::string, int64_t>(input_tensor, indices_tensor, output_tensor, axis, ttp);
} else {
if (indices_tensor->IsDataType<int32_t>())
core_impl<false, int8_t, int32_t>(input_tensor, indices_tensor, output_tensor, axis, ttp);
else
core_impl<false, int8_t, int64_t>(input_tensor, indices_tensor, output_tensor, axis, ttp);
}
if (indices_tensor->IsDataType<int32_t>())
core_impl<int32_t>(input_tensor, indices_tensor, output_tensor, axis);
else
core_impl<int64_t>(input_tensor, indices_tensor, output_tensor, axis);
return Status::OK();
}

4
onnxruntime/test/providers/cpu/tensor/gather_elements_op_test.cc
View file

@ -89,7 +89,7 @@ void RunTypedTest() {
2, 2});
test5.AddOutput<T>("output", {2, 2},
{1, 1,
4, 4});
3, 3});
// skip nuphar, which will not throw error message but will ensure no out-of-bound access
// skip cuda as the cuda kernel won't throw the error message
// skip openvino which will not throw error message but will ensure no out-of-bound access
@ -246,7 +246,7 @@ void RunTypedTest<std::string>() {
-3, -3});
test4.AddOutput<std::string>("output", {2, 2},
{"a", "a",
"d", "d"});
"c", "c"});
// skip nuphar, which will not throw error message but will ensure no out-of-bound access
// skip Openvino, which will not throw error message but will ensure no out-of-bound access
test4.Run(OpTester::ExpectResult::kExpectFailure,