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

improve the qlinear avg pool perf (#8514)

Browse source 
*) 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
This commit is contained in:
Zhang Lei 2021-07-27 20:56:59 -07:00 committed by GitHub
parent 56441dcd88
commit 0f46b08646
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 114 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
onnxruntime/contrib_ops/cpu/qlinear_global_average_pool.cc
View file

@ -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<uint8_t>(), x_scale, *(tensor_x_zero_point->Data<uint8_t>()),
return ComputeQLinearGlobalAvgPool(X.Data<uint8_t>(), x_scale, *(tensor_x_zero_point->Data<uint8_t>()),
Y.MutableData<uint8_t>(), y_scale, *(tensor_y_zero_point->Data<uint8_t>()),
N, C, image_size, channels_last_, tp);
default:

13
onnxruntime/contrib_ops/cpu/qlinear_global_average_pool.h
View file

@ -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

16
onnxruntime/contrib_ops/cpu/qlinear_lookup_table.cc
View file

@ -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 <typename TOutput>
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 <typename T>
void QlinearBuildLookupTable(uint8_t* table,
const Tensor* tensor_x_scale,

3
onnxruntime/contrib_ops/cpu/qlinear_lookup_table.h
View file

@ -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 <typename TOutput>
void QLinearLookupTableTransform(const uint8_t* x, const TOutput* table, TOutput* y, size_t n);
} // namespace contrib

61
onnxruntime/contrib_ops/cpu/qlinear_pool.cc
View file

@ -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 <functional>
#include <iostream>
namespace onnxruntime {
using concurrency::ThreadPool;
@ -485,6 +490,24 @@ struct QLinearPoolNhwc3DTask final {
}
};
template <typename T8Bits>
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<T8Bits>(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<Tensor>(1);
const auto tensor_x_zero_point = context->Input<Tensor>(2);
@ -543,30 +566,36 @@ Status QLinearAveragePool::Compute(OpKernelContext* context) const {
Tensor* Y = context->Output(0, output_dims);
const auto* X_data = X->Data<uint8_t>();
auto* Y_data = Y->MutableData<uint8_t>();
ThreadPool* tp = context->GetOperatorThreadPool();
std::vector<float> 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<size_t>(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<uint8_t, onnxruntime::AveragePool> 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<uint8_t, onnxruntime::AveragePool> 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<uint8_t, onnxruntime::AveragePool> 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<uint8_t, onnxruntime::AveragePool> 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<uint8_t, onnxruntime::AveragePool> 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<uint8_t, onnxruntime::AveragePool> 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);

42
onnxruntime/test/contrib_ops/qlinear_pool_test.cc
View file

@ -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