From 0f46b0864626349a4582ae65cd57e25eaef85977 Mon Sep 17 00:00:00 2001 From: Zhang Lei Date: Tue, 27 Jul 2021 20:56:59 -0700 Subject: [PATCH] improve the qlinear avg pool perf (#8514) *) use context buffer allocator, remove init cost of vector *) using lookup table to dequantize large input *) fall back to global average pool if it is --- .../cpu/qlinear_global_average_pool.cc | 4 +- .../cpu/qlinear_global_average_pool.h | 13 ++++ .../contrib_ops/cpu/qlinear_lookup_table.cc | 16 +++-- .../contrib_ops/cpu/qlinear_lookup_table.h | 3 +- onnxruntime/contrib_ops/cpu/qlinear_pool.cc | 61 ++++++++++++++----- .../test/contrib_ops/qlinear_pool_test.cc | 42 +++++++++++++ 6 files changed, 114 insertions(+), 25 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/qlinear_global_average_pool.cc b/onnxruntime/contrib_ops/cpu/qlinear_global_average_pool.cc index 780c44b36c..0a75348e3f 100644 --- a/onnxruntime/contrib_ops/cpu/qlinear_global_average_pool.cc +++ b/onnxruntime/contrib_ops/cpu/qlinear_global_average_pool.cc @@ -14,7 +14,7 @@ using onnxruntime::concurrency::ThreadPool; namespace onnxruntime { namespace contrib { -Status ComputeAveragePool( +Status ComputeQLinearGlobalAvgPool( const uint8_t* x, float x_scale, uint8_t x_zero_point, @@ -112,7 +112,7 @@ Status QLinearGlobalAveragePool::Compute(OpKernelContext* context) const { auto dtype = X.GetElementType(); switch (dtype) { case ONNX_NAMESPACE::TensorProto_DataType_UINT8: - return ComputeAveragePool(X.Data(), x_scale, *(tensor_x_zero_point->Data()), + return ComputeQLinearGlobalAvgPool(X.Data(), x_scale, *(tensor_x_zero_point->Data()), Y.MutableData(), y_scale, *(tensor_y_zero_point->Data()), N, C, image_size, channels_last_, tp); default: diff --git a/onnxruntime/contrib_ops/cpu/qlinear_global_average_pool.h b/onnxruntime/contrib_ops/cpu/qlinear_global_average_pool.h index ae5647780b..f46ed1c0e4 100644 --- a/onnxruntime/contrib_ops/cpu/qlinear_global_average_pool.h +++ b/onnxruntime/contrib_ops/cpu/qlinear_global_average_pool.h @@ -21,5 +21,18 @@ class QLinearGlobalAveragePool final : public OpKernel { bool channels_last_; }; +Status ComputeQLinearGlobalAvgPool( + const uint8_t* x, + float x_scale, + uint8_t x_zero_point, + uint8_t* y, + float y_scale, + uint8_t y_zero_point, + int64_t N, + int64_t C, + int64_t image_size, + bool channels_last, + concurrency::ThreadPool* tp); + } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/qlinear_lookup_table.cc b/onnxruntime/contrib_ops/cpu/qlinear_lookup_table.cc index 6de919dc54..44ea18398c 100644 --- a/onnxruntime/contrib_ops/cpu/qlinear_lookup_table.cc +++ b/onnxruntime/contrib_ops/cpu/qlinear_lookup_table.cc @@ -9,17 +9,18 @@ namespace onnxruntime { namespace contrib { -void QLinearLookupTableTransform(const uint8_t* x, const uint8_t* table, uint8_t* y, size_t n) { +template +void QLinearLookupTableTransform(const uint8_t* x, const TOutput* table, TOutput* y, size_t n) { for (; n >= 4; n -= 4) { const size_t x_value0 = x[0]; const size_t x_value1 = x[1]; const size_t x_value2 = x[2]; const size_t x_value3 = x[3]; x += 4; - const uint8_t table_value0 = table[x_value0]; - const uint8_t table_value1 = table[x_value1]; - const uint8_t table_value2 = table[x_value2]; - const uint8_t table_value3 = table[x_value3]; + const TOutput table_value0 = table[x_value0]; + const TOutput table_value1 = table[x_value1]; + const TOutput table_value2 = table[x_value2]; + const TOutput table_value3 = table[x_value3]; y[0] = table_value0; y[1] = table_value1; @@ -29,11 +30,14 @@ void QLinearLookupTableTransform(const uint8_t* x, const uint8_t* table, uint8_t } for (; n != 0; --n) { const size_t x_value0 = *x++; - const uint8_t table_value0 = table[x_value0]; + const TOutput table_value0 = table[x_value0]; *y++ = table_value0; } } +template void QLinearLookupTableTransform(const uint8_t* x, const uint8_t* table, uint8_t* y, size_t n); +template void QLinearLookupTableTransform(const uint8_t* x, const float* table, float* y, size_t n); + template void QlinearBuildLookupTable(uint8_t* table, const Tensor* tensor_x_scale, diff --git a/onnxruntime/contrib_ops/cpu/qlinear_lookup_table.h b/onnxruntime/contrib_ops/cpu/qlinear_lookup_table.h index 2df4d47226..f2ad74578e 100644 --- a/onnxruntime/contrib_ops/cpu/qlinear_lookup_table.h +++ b/onnxruntime/contrib_ops/cpu/qlinear_lookup_table.h @@ -34,7 +34,8 @@ void QlinearBuildLookupTable(uint8_t* table, const Tensor* tensor_y_zero_point, const LookupTableScalarTransformer& value_transformer); -void QLinearLookupTableTransform(const uint8_t* x, const uint8_t* table, uint8_t* y, size_t n); +template +void QLinearLookupTableTransform(const uint8_t* x, const TOutput* table, TOutput* y, size_t n); } // namespace contrib diff --git a/onnxruntime/contrib_ops/cpu/qlinear_pool.cc b/onnxruntime/contrib_ops/cpu/qlinear_pool.cc index 0d39e85797..6edd75d8a3 100644 --- a/onnxruntime/contrib_ops/cpu/qlinear_pool.cc +++ b/onnxruntime/contrib_ops/cpu/qlinear_pool.cc @@ -3,6 +3,9 @@ #include "qlinear_pool.h" +#include "contrib_ops/cpu/qlinear_lookup_table.h" +#include "contrib_ops/cpu/qlinear_global_average_pool.h" +#include "core/common/safeint.h" #include "core/util/math_cpuonly.h" #include "core/providers/common.h" #include "core/platform/threadpool.h" @@ -11,6 +14,8 @@ #include +#include + namespace onnxruntime { using concurrency::ThreadPool; @@ -485,6 +490,24 @@ struct QLinearPoolNhwc3DTask final { } }; +template +void dequantize_array(int64_t N, const T8Bits* input, float scale, T8Bits zero_point, float* output, ThreadPool* tp) { + if (N > 512) { + float dequantize_lookup[256]; + for (int i = 0; i < 256; ++i) { + T8Bits x = static_cast(i); + dequantize_lookup[i] = dequantize_value(x, scale, zero_point); + } + ThreadPool::TryParallelFor(tp, (ptrdiff_t)N, 1.0f, [input, output, &dequantize_lookup](ptrdiff_t first, ptrdiff_t last) { + QLinearLookupTableTransform((const uint8_t*)(input + first), dequantize_lookup, output + first, (size_t)(last - first)); + }); + } else { + for (int64_t i = 0; i < N; ++i) { + *output++ = dequantize_value(input[i], scale, zero_point); + } + } +} + Status QLinearAveragePool::Compute(OpKernelContext* context) const { const auto tensor_x_scale = context->Input(1); const auto tensor_x_zero_point = context->Input(2); @@ -543,30 +566,36 @@ Status QLinearAveragePool::Compute(OpKernelContext* context) const { Tensor* Y = context->Output(0, output_dims); const auto* X_data = X->Data(); auto* Y_data = Y->MutableData(); - ThreadPool* tp = context->GetOperatorThreadPool(); - std::vector x_data_fp32; + + // Check for special case which could fall back to global average pool + bool fallback_to_global = std::equal(x_shape.GetDims().begin() + 2, x_shape.GetDims().end(), kernel_shape.begin()) && + std::all_of(pads.begin(), pads.end(), [](int64_t dim) { return dim == 0LL; }); + if (fallback_to_global) { + return ComputeQLinearGlobalAvgPool(X_data, x_scale, x_zero_point, Y_data, y_scale, y_zero_point, + batch_count, channels, kernel_size, channels_last_, tp); + } + + AllocatorPtr allocator; + ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&allocator)); + float* x_data_fp32 = nullptr; + BufferUniquePtr x_data_fp32_guard; if (kernel_shape.size() <= 3) { - x_data_fp32.resize(x_shape.Size()); - ThreadPool::TryParallelFor(tp, x_shape.Size(), 1.0f, [=, &x_data_fp32](ptrdiff_t first, ptrdiff_t last) { - const auto* x8 = X_data + first; - float* x32 = x_data_fp32.data() + first; - for (ptrdiff_t i = 0, sz = last - first; i < sz; ++i) { - *x32++ = dequantize_value(x8[i], x_scale, x_zero_point); - } - }); + x_data_fp32 = (float*)allocator->Alloc(SafeInt(x_shape.Size()) * sizeof(float)); + x_data_fp32_guard = BufferUniquePtr(x_data_fp32, BufferDeleter(allocator)); + dequantize_array(x_shape.Size(), X_data, x_scale, x_zero_point, x_data_fp32, tp); } switch (kernel_shape.size()) { case 1: { if (channels_last_) { QLinearPoolNhwc1DTask avg_pool_task_1d = { - x_data_fp32.data(), Y_data, y_scale, y_zero_point, channels, + x_data_fp32, Y_data, y_scale, y_zero_point, channels, pooled_height, strides[0], height, kernel_shape, pads, pool_context_, pool_attrs_}; ThreadPool::TryParallelFor(tp, y_image_size * batch_count, avg_pool_task_1d.Cost(), avg_pool_task_1d); } else { QLinearPool1DTask avg_pool_task_1d = { - x_data_fp32.data(), Y_data, y_scale, y_zero_point, x_image_size, y_image_size, + x_data_fp32, Y_data, y_scale, y_zero_point, x_image_size, y_image_size, pooled_height, strides[0], height, kernel_shape, pads, pool_context_, pool_attrs_}; ThreadPool::TryParallelFor(tp, total_channels, avg_pool_task_1d.Cost(), avg_pool_task_1d); } @@ -576,13 +605,13 @@ Status QLinearAveragePool::Compute(OpKernelContext* context) const { case 2: { if (channels_last_) { QLinearPoolNhwc2DTask avg_pool_task_2d = { - x_data_fp32.data(), Y_data, y_scale, y_zero_point, x_image_size, y_image_size, kernel_size, channels, + x_data_fp32, Y_data, y_scale, y_zero_point, x_image_size, y_image_size, kernel_size, channels, pooled_height, pooled_width, strides[0], strides[1], height, width, kernel_shape, pads, pool_context_, pool_attrs_}; ThreadPool::TryParallelFor(tp, y_image_size * batch_count, avg_pool_task_2d.Cost(), avg_pool_task_2d); } else { QLinearPool2DTask avg_pool_task_2d = { - x_data_fp32.data(), Y_data, y_scale, y_zero_point, x_image_size, y_image_size, + x_data_fp32, Y_data, y_scale, y_zero_point, x_image_size, y_image_size, pooled_height, pooled_width, strides[0], strides[1], height, width, kernel_shape, pads, pool_context_, pool_attrs_}; ThreadPool::TryParallelFor(tp, total_channels, avg_pool_task_2d.Cost(), avg_pool_task_2d); } @@ -592,14 +621,14 @@ Status QLinearAveragePool::Compute(OpKernelContext* context) const { case 3: { if (channels_last_) { QLinearPoolNhwc3DTask avg_pool_task_3d = { - x_data_fp32.data(), Y_data, y_scale, y_zero_point, x_image_size, y_image_size, kernel_size, channels, + x_data_fp32, Y_data, y_scale, y_zero_point, x_image_size, y_image_size, kernel_size, channels, pooled_height, pooled_width, pooled_depth, strides[0], strides[1], strides[2], height, width, depth, kernel_shape, pads, pool_context_, pool_attrs_}; ThreadPool::TryParallelFor(tp, y_image_size * batch_count, avg_pool_task_3d.Cost(), avg_pool_task_3d); } else { QLinearPool3DTask avg_pool_task_3d = { - x_data_fp32.data(), Y_data, y_scale, y_zero_point, x_image_size, y_image_size, + x_data_fp32, Y_data, y_scale, y_zero_point, x_image_size, y_image_size, pooled_height, pooled_width, pooled_depth, strides[0], strides[1], strides[2], height, width, depth, kernel_shape, pads, pool_context_, pool_attrs_}; ThreadPool::TryParallelFor(tp, total_channels, avg_pool_task_3d.Cost(), avg_pool_task_3d); diff --git a/onnxruntime/test/contrib_ops/qlinear_pool_test.cc b/onnxruntime/test/contrib_ops/qlinear_pool_test.cc index 5edd0a2cd3..85b8c998e7 100644 --- a/onnxruntime/test/contrib_ops/qlinear_pool_test.cc +++ b/onnxruntime/test/contrib_ops/qlinear_pool_test.cc @@ -420,5 +420,47 @@ TEST(QLinearPoolTest, AveragePool3D_IncludePadPixel_nhwc) { 1); // count_include_pad } + +TEST(QLinearPoolTest, AveragePool2D_BigImage) { + RunQLinearAveragePoolNchwU8( + {1, 1, 32, 64}, // x shape + {1, 1, 32, 64}, // expected y shape + {3, 3}, // kernel shape + {1, 1}, // strides + {1, 1, 1, 1}, // pads + 1); // count_include_pad +} + +TEST(QLinearPoolTest, AveragePool2D_BigImage_nhwc) { + RunQLinearAveragePoolNhwcU8( + {1, 1, 32, 64}, // x shape + {1, 1, 32, 64}, // expected y shape + {3, 3}, // kernel shape + {1, 1}, // strides + {1, 1, 1, 1}, // pads + 1); // count_include_pad +} + +TEST(QLinearPoolTest, AveragePool2D_Global) { + RunQLinearAveragePoolNchwU8( + {1, 2, 32, 16}, // x shape + {1, 2, 1, 1}, // expected y shape + {32, 16}, // kernel shape + {1, 1}, // strides + {0, 0, 0, 0}, // pads + 1); // count_include_pad +} + +TEST(QLinearPoolTest, AveragePool2D_Global_nhwc) { + RunQLinearAveragePoolNhwcU8( + {1, 2, 32, 16}, // x shape + {1, 2, 1, 1}, // expected y shape + {32, 16}, // kernel shape + {1, 1}, // strides + {0, 0, 0, 0}, // pads + 1); // count_include_pad +} + + } // namespace test } // namespace onnxruntime