From 9e7753c109d4eb8e7efe5d98088dd9e6b7e12432 Mon Sep 17 00:00:00 2001 From: Ryan Hill <38674843+RyanUnderhill@users.noreply.github.com> Date: Wed, 27 Apr 2022 17:57:56 -0700 Subject: [PATCH] Speed up GatherElements CPU Op (#11349) * Specialize element copy by element size --- .../providers/cpu/tensor/gather_elements.cc | 215 ++++++------------ .../cpu/tensor/gather_elements_op_test.cc | 4 +- 2 files changed, 73 insertions(+), 146 deletions(-) diff --git a/onnxruntime/core/providers/cpu/tensor/gather_elements.cc b/onnxruntime/core/providers/cpu/tensor/gather_elements.cc index b09148440d..0d2742f766 100644 --- a/onnxruntime/core/providers/cpu/tensor/gather_elements.cc +++ b/onnxruntime/core/providers/cpu/tensor/gather_elements.cc @@ -25,8 +25,6 @@ ONNX_CPU_OPERATOR_KERNEL( DataTypeImpl::GetTensorType()}), 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& 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(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& 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(current_dims.size()); @@ -88,17 +86,21 @@ static inline void increment_over_inner_dim(std::vector& current_dims, } } -template -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(indices_data[index] + input_shape[axis]); - } else { - retval = static_cast(indices_data[index]); - } - return retval; -} +#if defined(_MSC_VER) +#define FORCEINLINE __forceinline +#else +#define FORCEINLINE __attribute__((always_inline)) inline +#endif + +template +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(index) >= std::make_unsigned_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 -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(); - } else { - input_data = reinterpret_cast(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(); - } else { - output_data = reinterpret_cast(output_tensor->MutableDataRaw()); - } +template +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(input_tensor->DataRaw()); + int8_t* output_data = reinterpret_cast(output_tensor->MutableDataRaw()); + size_t element_size = input_tensor->DataType()->Size(); + bool is_string = input_tensor->IsDataTypeString(); const int64_t input_rank = static_cast(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(); + 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(); - // 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 process_dims(input_rank, 0); - int64_t output_counter = 0; - - auto conditional_batch_call = [ttp, inner_dim_size](std::function 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(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(output_data), reinterpret_cast(input_data), LoopFn); + else if (element_size == sizeof(uint32_t)) + MainLoop(reinterpret_cast(output_data), reinterpret_cast(input_data), LoopFn); + else if (element_size == sizeof(uint16_t)) + MainLoop(reinterpret_cast(output_data), reinterpret_cast(input_data), LoopFn); + else if (element_size == sizeof(uint8_t)) + MainLoop(reinterpret_cast(output_data), reinterpret_cast(input_data), LoopFn); + else if (element_size == sizeof(uint64_t)) + MainLoop(reinterpret_cast(output_data), reinterpret_cast(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(indices_data, static_cast(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(indices_data, static_cast(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(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(indices_data, static_cast(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(indices_data, static_cast(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(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()) - core_impl(input_tensor, indices_tensor, output_tensor, axis, ttp); - else - core_impl(input_tensor, indices_tensor, output_tensor, axis, ttp); - } else { - if (indices_tensor->IsDataType()) - core_impl(input_tensor, indices_tensor, output_tensor, axis, ttp); - else - core_impl(input_tensor, indices_tensor, output_tensor, axis, ttp); - } + if (indices_tensor->IsDataType()) + core_impl(input_tensor, indices_tensor, output_tensor, axis); + else + core_impl(input_tensor, indices_tensor, output_tensor, axis); return Status::OK(); } diff --git a/onnxruntime/test/providers/cpu/tensor/gather_elements_op_test.cc b/onnxruntime/test/providers/cpu/tensor/gather_elements_op_test.cc index 783d2eca1a..32bd512ae3 100644 --- a/onnxruntime/test/providers/cpu/tensor/gather_elements_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/gather_elements_op_test.cc @@ -89,7 +89,7 @@ void RunTypedTest() { 2, 2}); test5.AddOutput("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() { -3, -3}); test4.AddOutput("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,