From 41ddcd30a1fa9233b4d5314f0e94c7abbc147d93 Mon Sep 17 00:00:00 2001 From: Chen Fu <1316708+chenfucn@users.noreply.github.com> Date: Tue, 28 Mar 2023 08:22:55 -0700 Subject: [PATCH] Fp16 NHWC Max and Average Pooling (#15181) ### Description Max and average pooling operators for fp16, NHWC ### Motivation and Context Continue on the steps for fp16 inference support --- cmake/onnxruntime_mlas.cmake | 20 + .../contrib_ops/cpu/cpu_contrib_kernels.cc | 33 ++ .../contrib_ops/cpu/fp16/nhwc_pool_fp16.cc | 178 +++++++++ onnxruntime/core/mlas/inc/mlas.h | 41 ++ onnxruntime/core/mlas/lib/fp16_common.h | 14 + onnxruntime/core/mlas/lib/pooling_fp16.cpp | 355 ++++++++++++++++++ .../test/contrib_ops/nhwc_pool_in_op_test.cc | 314 ++++++++++++++++ 7 files changed, 955 insertions(+) create mode 100644 onnxruntime/contrib_ops/cpu/fp16/nhwc_pool_fp16.cc create mode 100644 onnxruntime/core/mlas/lib/pooling_fp16.cpp create mode 100644 onnxruntime/test/contrib_ops/nhwc_pool_in_op_test.cc diff --git a/cmake/onnxruntime_mlas.cmake b/cmake/onnxruntime_mlas.cmake index df3dac6d1a..42f4b8960e 100644 --- a/cmake/onnxruntime_mlas.cmake +++ b/cmake/onnxruntime_mlas.cmake @@ -25,6 +25,11 @@ if(CMAKE_CXX_COMPILER_ID MATCHES "MSVC") endif() +# +# All hardware agnostic source files here +# hardware specific files would cause trouble in +# multi-target build +# onnxruntime_add_static_library(onnxruntime_mlas ${MLAS_SRC_DIR}/platform.cpp ${MLAS_SRC_DIR}/threading.cpp @@ -62,6 +67,16 @@ set(ONNXRUNTIME_MLAS_LIBS onnxruntime_mlas) #TODO: set MASM flags properly function(setup_mlas_source_for_windows) + + # + # Sources common for all platforms. + # + target_sources(onnxruntime_mlas PRIVATE + ${MLAS_SRC_DIR}/activate_fp16.cpp + ${MLAS_SRC_DIR}/dwconv.cpp + ${MLAS_SRC_DIR}/pooling_fp16.cpp + ) + #The onnxruntime_target_platform variable was added by Windows AI team in onnxruntime_common.cmake #Don't use it for other platforms. if((onnxruntime_target_platform STREQUAL "ARM64") OR (onnxruntime_target_platform STREQUAL "ARM64EC")) @@ -332,6 +347,7 @@ else() ${MLAS_SRC_DIR}/activate_fp16.cpp ${MLAS_SRC_DIR}/dwconv.cpp ${MLAS_SRC_DIR}/halfgemm_kernel_neon.cpp + ${MLAS_SRC_DIR}/pooling_fp16.cpp ${MLAS_SRC_DIR}/qgemm_kernel_neon.cpp ${MLAS_SRC_DIR}/qgemm_kernel_udot.cpp ${MLAS_SRC_DIR}/qgemm_kernel_sdot.cpp @@ -339,6 +355,7 @@ else() set_source_files_properties(${MLAS_SRC_DIR}/aarch64/HalfGemmKernelNeon.S PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ") set_source_files_properties(${MLAS_SRC_DIR}/activate_fp16.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ") set_source_files_properties(${MLAS_SRC_DIR}/dwconv.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ") + set_source_files_properties(${MLAS_SRC_DIR}/pooling_fp16.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ") if(ONNXRUNTIME_MLAS_MULTI_ARCH) onnxruntime_add_static_library(onnxruntime_mlas_arm64 ${mlas_platform_srcs}) @@ -517,7 +534,10 @@ else() set_source_files_properties(${mlas_platform_srcs_avx512core} PROPERTIES COMPILE_FLAGS "-mavx512bw -mavx512dq -mavx512vl") set(mlas_platform_srcs + ${MLAS_SRC_DIR}/activate_fp16.cpp + ${MLAS_SRC_DIR}/dwconv.cpp ${MLAS_SRC_DIR}/dgemm.cpp + ${MLAS_SRC_DIR}/pooling_fp16.cpp ${MLAS_SRC_DIR}/qgemm_kernel_avx2.cpp ${mlas_platform_srcs_sse2} ${mlas_platform_srcs_avx} diff --git a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc b/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc index de4e699b52..ced26e99e4 100644 --- a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc @@ -80,6 +80,12 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, uint8_t, QGemm); // ******** End: Quantization ******************* // +#ifdef MLAS_F16VEC_INTRINSICS_SUPPORTED +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSInternalNHWCDomain, 11, MLFloat16, MaxPool); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSInternalNHWCDomain, 11, MLFloat16, AveragePool); +#endif + + // This section includes all op kernel declarations for former experimental ops which have now been removed from onnx. // To maintain backward compatibility these are added as contrib ops. // Note: the domain for all contrib ops should be MSDomain. However since these ops started out as onnx domain ops @@ -144,6 +150,27 @@ Status RegisterNchwcKernels(KernelRegistry& kernel_registry) { return Status::OK(); } + +#ifdef MLAS_F16VEC_INTRINSICS_SUPPORTED +Status RegisterFp16Kernels(KernelRegistry& kernel_registry) { + static const BuildKernelCreateInfoFn function_table[] = { + BuildKernelCreateInfo, + BuildKernelCreateInfo, + }; + + for (auto& function_table_entry : function_table) { + KernelCreateInfo info = function_table_entry(); + if (info.kernel_def != nullptr) { // filter disabled entries where type is void + ORT_RETURN_IF_ERROR(kernel_registry.Register(std::move(info))); + } + } + + return Status::OK(); +} +#endif + + + Status RegisterQuantizationKernels(KernelRegistry& kernel_registry) { static const BuildKernelCreateInfoFn function_table[] = { BuildKernelCreateInfo, // default entry to avoid the list become empty after ops-reducing @@ -270,6 +297,12 @@ Status RegisterCpuContribKernels(KernelRegistry& kernel_registry) { ORT_RETURN_IF_ERROR(RegisterQuantizationKernels(kernel_registry)); +#ifdef MLAS_F16VEC_INTRINSICS_SUPPORTED + if (MlasFp16AccelerationSupported()) { + ORT_RETURN_IF_ERROR(RegisterFp16Kernels(kernel_registry)); + } +#endif + return Status::OK(); } diff --git a/onnxruntime/contrib_ops/cpu/fp16/nhwc_pool_fp16.cc b/onnxruntime/contrib_ops/cpu/fp16/nhwc_pool_fp16.cc new file mode 100644 index 0000000000..a1cc1fa307 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/fp16/nhwc_pool_fp16.cc @@ -0,0 +1,178 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/common/common.h" +#include "core/framework/op_kernel.h" +#include "core/providers/cpu/nn/pool_attributes.h" +#include "core/common/safeint.h" +#include "core/util/math.h" +#include "core/mlas/inc/mlas.h" + +namespace onnxruntime { +namespace contrib { + +#ifdef MLAS_F16VEC_INTRINSICS_SUPPORTED + +/** + * @brief Pooling operator for type FP16, layout NHWC. + * Only max pool and average pool supported. + * + * Single threadded operation for now. + * + * TODO!! implemente thread partition similar with + * fp16 conv operator +*/ +class NhwcPoolFp16 : public OpKernel { + public: + explicit NhwcPoolFp16(const OpKernelInfo& info, const std::string& pooltype, bool is_max_pool) + : OpKernel(info), pool_attrs_(info, pooltype, info.node().SinceVersion()), is_max_pool_(is_max_pool) {} + + Status Compute(OpKernelContext* context) const override; + + protected: + PoolAttributes pool_attrs_; + bool is_max_pool_; // either max pool or average pool +}; + +Status NhwcPoolFp16::Compute(OpKernelContext* context) const { + const auto* X = context->Input(0); + const TensorShape& input_shape = X->Shape(); + + const size_t input_rank = input_shape.NumDimensions(); + ORT_RETURN_IF_NOT(input_rank >= 3, "Input dimension cannot be less than 3."); + + const int64_t N = input_shape[0]; + const int64_t C = input_shape[input_rank - 1]; + + ORT_ENFORCE(input_shape.Size() > 0 || N == 0, "Invalid input shape. Only N can be zero. Got:", input_shape); + + const size_t spatial_dims = input_rank - 2; + + // Compute the output size and effective padding for this pooling operation. + TensorShapeVector output_dims({N}); + TensorShapeVector pads = pool_attrs_.pads; + int64_t kernel_size = 1; + int64_t input_image_size = 1; + int64_t output_image_size = 1; + for (size_t dim = 0; dim < spatial_dims; ++dim) { + int64_t kernel = pool_attrs_.kernel_shape[dim]; + int64_t input_dim = input_shape[dim + 1]; + + kernel_size *= kernel; + input_image_size *= input_dim; + + int64_t output_dim = 0; + pool_attrs_.ComputeSizePadDilations(input_dim, + pool_attrs_.strides[dim], + kernel, + &pads.at(dim), + &pads.at(spatial_dims + dim), + pool_attrs_.dilations[dim], + &output_dim); + output_dims.push_back(output_dim); + + output_image_size *= output_dim; + } + output_dims.push_back(C); + + Tensor* Y = context->Output(0, output_dims); + + constexpr int64_t output_batch_count = 512; + + // Allocate indirection buffer pointers and prepare a padding vector for the + // im2col transform. + AllocatorPtr alloc; + ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&alloc)); + int64_t col_buffer_batch_count = std::min(output_image_size, output_batch_count); + auto* col_data = alloc->Alloc(SafeInt(sizeof(const MLFloat16*)) * kernel_size * col_buffer_batch_count); + BufferUniquePtr col_buffer(col_data, BufferDeleter(std::move(alloc))); + + const bool need_padding = !is_max_pool_ && pool_attrs_.count_include_pad; + std::vector padding_data; + if (need_padding) { + padding_data.resize(static_cast(C), MLFloat16()); + } + + const auto* Xdata = X->Data(); + auto* Ydata = Y->MutableData(); + + for (int64_t image_id = 0; image_id < N; ++image_id) { + for (int64_t output_start = 0; output_start < output_image_size;) { + int64_t output_count = std::min(output_image_size - output_start, output_batch_count); + math::Im2col()( + Xdata, + C, + input_shape.GetDims().data() + 1, + output_dims.data() + 1, + pool_attrs_.kernel_shape.data(), + pool_attrs_.strides.data(), + pool_attrs_.dilations.data(), + pads.data(), + static_cast(spatial_dims), + output_start, + output_count, + static_cast(col_buffer.get()), + need_padding ? padding_data.data() : nullptr); + if (is_max_pool_) { + MlasNhwcMaxPool( + static_cast(col_buffer.get()), + Ydata, + static_cast(C), + static_cast(output_count), + static_cast(kernel_size)); + } else { + MlasNhwcAvgPool( + static_cast(col_buffer.get()), + Ydata, + static_cast(C), + static_cast(output_count), + static_cast(kernel_size)); + } + + Ydata += output_count * C; + output_start += output_count; + } + + Xdata += input_image_size * C; + } + + return Status::OK(); +} + +class NhwcMaxPoolFp16 : public NhwcPoolFp16 { + public: + explicit NhwcMaxPoolFp16(const OpKernelInfo& info) + : NhwcPoolFp16(info, "MaxPool", true /*maxpool*/) {} +}; + +class NhwcAvgPoolFp16 : public NhwcPoolFp16 { + public: + explicit NhwcAvgPoolFp16(const OpKernelInfo& info) + : NhwcPoolFp16(info, "AveragePool", false /*not maxpool*/) {} +}; + +ONNX_OPERATOR_TYPED_KERNEL_EX( + MaxPool, + kMSInternalNHWCDomain, + 11, + MLFloat16, + kCpuExecutionProvider, + KernelDefBuilder() + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + NhwcMaxPoolFp16); + +ONNX_OPERATOR_TYPED_KERNEL_EX( + AveragePool, + kMSInternalNHWCDomain, + 11, + MLFloat16, + kCpuExecutionProvider, + KernelDefBuilder() + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + NhwcAvgPoolFp16); + + +#endif + +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/core/mlas/inc/mlas.h b/onnxruntime/core/mlas/inc/mlas.h index f3853a211f..f4cc6b1005 100644 --- a/onnxruntime/core/mlas/inc/mlas.h +++ b/onnxruntime/core/mlas/inc/mlas.h @@ -1641,3 +1641,44 @@ MlasTranspose( reinterpret_cast(Output), M, N); } + +#ifdef MLAS_F16VEC_INTRINSICS_SUPPORTED +/** + * @brief Max Pooling for fp16 NHWC + * @param Input Indirect buffer to activations + * @param Output Address of the result tensor + * @param Channels C in NHWC + * @param OutputCount Number of output pixels + * @param KernelSize Size of the kernel + * @return +*/ +void +MLASCALL +MlasNhwcMaxPool( + const MLAS_FP16* const* Input, + MLAS_FP16* Output, + size_t Channels, + size_t OutputCount, + size_t KernelSize + ); + +/** + * @brief Avg Pooling for fp16 nhwc + * @param Input Indirect buffer to activations + * @param Output Address of the output data + * @param Channels C in NHWC + * @param OutputCount Number of output pixels + * @param KernelSize size of the kernel + * @return +*/ +void +MLASCALL +MlasNhwcAvgPool( + const MLAS_FP16* const* Input, + MLAS_FP16* Output, + size_t Channels, + size_t OutputCount, + size_t KernelSize + ); + +#endif diff --git a/onnxruntime/core/mlas/lib/fp16_common.h b/onnxruntime/core/mlas/lib/fp16_common.h index 8b00190757..1fcab870af 100644 --- a/onnxruntime/core/mlas/lib/fp16_common.h +++ b/onnxruntime/core/mlas/lib/fp16_common.h @@ -148,6 +148,20 @@ MlasMultiplyFloat16x4(MLAS_FLOAT16X4 Vector1, MLAS_FLOAT16X4 Vector2) return vmul_f16(Vector1, Vector2); } +MLAS_FORCEINLINE +MLAS_FLOAT16X8 +MlasDivFloat16x8(MLAS_FLOAT16X8 Vector1, MLAS_FLOAT16X8 Vector2) +{ + return vdivq_f16(Vector1, Vector2); +} + +MLAS_FORCEINLINE +MLAS_FLOAT16X4 +MlasDivFloat16x4(MLAS_FLOAT16X4 Vector1, MLAS_FLOAT16X4 Vector2) +{ + return vdiv_f16(Vector1, Vector2); +} + MLAS_FORCEINLINE MLAS_FLOAT16X8 MlasMultiplyAddFloat16x8(MLAS_FLOAT16X8 Vector1, MLAS_FLOAT16X8 Vector2, MLAS_FLOAT16X8 Vector3) diff --git a/onnxruntime/core/mlas/lib/pooling_fp16.cpp b/onnxruntime/core/mlas/lib/pooling_fp16.cpp new file mode 100644 index 0000000000..98e84736cb --- /dev/null +++ b/onnxruntime/core/mlas/lib/pooling_fp16.cpp @@ -0,0 +1,355 @@ +/*++ + +Copyright (c) Microsoft Corporation. All rights reserved. + +Licensed under the MIT License. + +Module Name: + pooling_fp16.cpp + +Abstract: + This module implements the pooling operation for fp16 + tensors in NHWC format. +--*/ + +#include "mlasi.h" + +#include "fp16_common.h" + +#ifdef MLAS_F16VEC_INTRINSICS_SUPPORTED + + +template +typename AggregationType::CtxType +PoolCreateContext(size_t KernelSize); + +template +MLAS_FLOAT16X8 +PoolInit16x8(); + +template +MLAS_FLOAT16X4 +PoolInit16x4(); + +template +MLAS_FLOAT16X8 +PoolAggregate16x8(MLAS_FLOAT16X8 agg, MLAS_FLOAT16X8 element); + +template +MLAS_FLOAT16X4 +PoolAggregate16x4(MLAS_FLOAT16X4 agg, MLAS_FLOAT16X4 element); + +template +MLAS_FLOAT16X8 +PoolSummary16x8(MLAS_FLOAT16X8 agg, typename AggregationType::CtxType context); + +template +MLAS_FLOAT16X4 +PoolSummary16x4(MLAS_FLOAT16X4 agg, typename AggregationType::CtxType context); + + +struct MaxPoolAggregation { + typedef size_t CtxType; // useless type to satisfy compilers +}; + +template <> +MLAS_FORCEINLINE +size_t +PoolCreateContext(size_t KernelSize) +{ + MLAS_UNREFERENCED_PARAMETER(KernelSize); + return 0; +} + +template<> +MLAS_FORCEINLINE +MLAS_FLOAT16X8 +PoolInit16x8() +{ + // lowest fp16 -65504.0f + return MlasBroadcastFloat16x8(0xfbff); +} + +template <> +MLAS_FORCEINLINE +MLAS_FLOAT16X4 +PoolInit16x4() +{ + // lowest fp16 -65504.0f + return MlasBroadcastFloat16x4(0xfbff); +} + +template<> +MLAS_FORCEINLINE +MLAS_FLOAT16X8 +PoolAggregate16x8(MLAS_FLOAT16X8 agg, MLAS_FLOAT16X8 element) +{ + return MlasMaximumFloat16x8(agg, element); +} + +template<> +MLAS_FORCEINLINE +MLAS_FLOAT16X4 +PoolAggregate16x4(MLAS_FLOAT16X4 agg, MLAS_FLOAT16X4 element) +{ + return MlasMaximumFloat16x4(agg, element); +} + +template<> +MLAS_FORCEINLINE +MLAS_FLOAT16X8 +PoolSummary16x8(MLAS_FLOAT16X8 agg, size_t size) +{ + MLAS_UNREFERENCED_PARAMETER(size); + return agg; +} + +template<> +MLAS_FORCEINLINE +MLAS_FLOAT16X4 +PoolSummary16x4(MLAS_FLOAT16X4 agg, size_t size) +{ + MLAS_UNREFERENCED_PARAMETER(size); + return agg; +} + +struct AveragePoolAggregation { + typedef MLAS_FLOAT16X8 CtxType; +}; + +template<> +MLAS_FLOAT16X8 +PoolCreateContext(size_t KernelSize) +{ + return MlasBroadcastFloat16x8(MLAS_Float2Half(float(KernelSize))); +} + + +template <> +MLAS_FORCEINLINE MLAS_FLOAT16X8 +PoolInit16x8() +{ + return MlasZeroFloat16x8(); +} + +template <> +MLAS_FORCEINLINE MLAS_FLOAT16X4 +PoolInit16x4() +{ + return MlasZeroFloat16x4(); +} + + +template <> +MLAS_FORCEINLINE MLAS_FLOAT16X8 +PoolAggregate16x8(MLAS_FLOAT16X8 agg, MLAS_FLOAT16X8 element) +{ + return MlasAddFloat16x8(agg, element); +} + +template <> +MLAS_FORCEINLINE MLAS_FLOAT16X4 +PoolAggregate16x4(MLAS_FLOAT16X4 agg, MLAS_FLOAT16X4 element) +{ + return MlasAddFloat16x4(agg, element); +} + +template <> +MLAS_FORCEINLINE MLAS_FLOAT16X8 +PoolSummary16x8(MLAS_FLOAT16X8 agg, MLAS_FLOAT16X8 context) +{ + return MlasDivFloat16x8(agg, context); +} + +template <> +MLAS_FORCEINLINE MLAS_FLOAT16X4 +PoolSummary16x4(MLAS_FLOAT16X4 agg, MLAS_FLOAT16X8 context) +{ + return MlasDivFloat16x4(agg, MlasToLowHalfFloat16x4(context)); +} + + +template +MLAS_FORCEINLINE +void +MlasPoolFp16HWC( + const _mlas_fp16_* const* Input, + _mlas_fp16_* Output, + size_t Channels, + size_t OutputCount, + size_t KernelSize + ) +{ + while (OutputCount > 0) { + size_t ChannelOffset = 0; + size_t c = Channels; + + while (c >= 32) { + MLAS_FLOAT16X8 MaximumVector0 = PoolInit16x8(); + MLAS_FLOAT16X8 MaximumVector1 = MaximumVector0; + MLAS_FLOAT16X8 MaximumVector2 = MaximumVector0; + MLAS_FLOAT16X8 MaximumVector3 = MaximumVector0; + size_t cnt = 0; + + for (size_t k = 0; k < KernelSize; k++) { + if (Input[k] == nullptr) { + continue; + } + MLAS_FLOAT16X8 InputVector0 = MlasLoadFloat16x8(&Input[k][ChannelOffset]); + MLAS_FLOAT16X8 InputVector1 = MlasLoadFloat16x8(&Input[k][ChannelOffset + 8]); + MLAS_FLOAT16X8 InputVector2 = MlasLoadFloat16x8(&Input[k][ChannelOffset + 16]); + MLAS_FLOAT16X8 InputVector3 = MlasLoadFloat16x8(&Input[k][ChannelOffset + 24]); + + MaximumVector0 = PoolAggregate16x8(MaximumVector0, InputVector0); + MaximumVector1 = PoolAggregate16x8(MaximumVector1, InputVector1); + MaximumVector2 = PoolAggregate16x8(MaximumVector2, InputVector2); + MaximumVector3 = PoolAggregate16x8(MaximumVector3, InputVector3); + cnt++; + } + typename AggregationType::CtxType context = PoolCreateContext(cnt); + MaximumVector0 = PoolSummary16x8(MaximumVector0, context); + MaximumVector1 = PoolSummary16x8(MaximumVector1, context); + MaximumVector2 = PoolSummary16x8(MaximumVector2, context); + MaximumVector3 = PoolSummary16x8(MaximumVector3, context); + + MlasStoreFloat16x8(&Output[0], MaximumVector0); + MlasStoreFloat16x8(&Output[8], MaximumVector1); + MlasStoreFloat16x8(&Output[16], MaximumVector2); + MlasStoreFloat16x8(&Output[24], MaximumVector3); + + Output += 32; + ChannelOffset += 32; + c -= 32; + } + + if (c >= 16) { + MLAS_FLOAT16X8 MaximumVector0 = PoolInit16x8(); + MLAS_FLOAT16X8 MaximumVector1 = MaximumVector0; + size_t cnt = 0; + + for (size_t k = 0; k < KernelSize; k++) { + if (Input[k] == nullptr) { + continue; + } + MLAS_FLOAT16X8 InputVector0 = MlasLoadFloat16x8(&Input[k][ChannelOffset]); + MLAS_FLOAT16X8 InputVector1 = MlasLoadFloat16x8(&Input[k][ChannelOffset + 8]); + + MaximumVector0 = PoolAggregate16x8(MaximumVector0, InputVector0); + MaximumVector1 = PoolAggregate16x8(MaximumVector1, InputVector1); + cnt++; + } + typename AggregationType::CtxType context = PoolCreateContext(cnt); + MaximumVector0 = PoolSummary16x8(MaximumVector0, context); + MaximumVector1 = PoolSummary16x8(MaximumVector1, context); + + MlasStoreFloat16x8(&Output[0], MaximumVector0); + MlasStoreFloat16x8(&Output[8], MaximumVector1); + + Output += 16; + ChannelOffset += 16; + c -= 16; + } + + if (c >= 8) { + MLAS_FLOAT16X8 MaximumVector0 = PoolInit16x8(); + size_t cnt = 0; + + for (size_t k = 0; k < KernelSize; k++) { + if (Input[k] == nullptr) { + continue; + } + MLAS_FLOAT16X8 InputVector0 = MlasLoadFloat16x8(&Input[k][ChannelOffset]); + MaximumVector0 = PoolAggregate16x8(MaximumVector0, InputVector0); + cnt++; + } + typename AggregationType::CtxType context = PoolCreateContext(cnt); + MaximumVector0 = PoolSummary16x8(MaximumVector0, context); + + MlasStoreFloat16x8(&Output[0], MaximumVector0); + + Output += 8; + ChannelOffset += 8; + c -= 8; + } + + if (c >= 4) { + MLAS_FLOAT16X4 MaximumVector0 = PoolInit16x4(); + size_t cnt = 0; + + for (size_t k = 0; k < KernelSize; k++) { + if (Input[k] == nullptr) { + continue; + } + MLAS_FLOAT16X4 InputVector0 = MlasLoadFloat16x4(&Input[k][ChannelOffset]); + MaximumVector0 = PoolAggregate16x4(MaximumVector0, InputVector0); + cnt++; + } + typename AggregationType::CtxType context = PoolCreateContext(cnt); + MaximumVector0 = PoolSummary16x4(MaximumVector0, context); + + MlasStoreFloat16x4(&Output[0], MaximumVector0); + + Output += 4; + ChannelOffset += 4; + c -= 4; + } + + if (c > 0) { + // possible over read by 7 bytes + MLAS_FLOAT16X4 MaximumVector0 = PoolInit16x4(); + size_t cnt = 0; + + for (size_t k = 0; k < KernelSize; k++) { + if (Input[k] == nullptr) { + continue; + } + MLAS_FLOAT16X4 InputVector0 = MlasLoadFloat16x4(&Input[k][ChannelOffset]); + MaximumVector0 = PoolAggregate16x4(MaximumVector0, InputVector0); + cnt++; + } + typename AggregationType::CtxType context = PoolCreateContext(cnt); + MaximumVector0 = PoolSummary16x4(MaximumVector0, context); + + MlasStorePartialFloat16x4(&Output[0], MaximumVector0, c); + Output += c; + } + + Input += KernelSize; + OutputCount -= 1; + } +} + + +void +MLASCALL +MlasNhwcMaxPool( + const MLAS_FP16* const* Input, + MLAS_FP16* Output, + size_t Channels, + size_t OutputCount, + size_t KernelSize + ) +{ + const _mlas_fp16_* const* input_ptr = reinterpret_cast(Input); + _mlas_fp16_* output_ptr = reinterpret_cast<_mlas_fp16_*>(Output); + MlasPoolFp16HWC(input_ptr, output_ptr, Channels, OutputCount, KernelSize); +} + + +void +MLASCALL +MlasNhwcAvgPool( + const MLAS_FP16* const* Input, + MLAS_FP16* Output, + size_t Channels, + size_t OutputCount, + size_t KernelSize + ) +{ + const _mlas_fp16_* const* input_ptr = reinterpret_cast(Input); + _mlas_fp16_* output_ptr = reinterpret_cast<_mlas_fp16_*>(Output); + MlasPoolFp16HWC(input_ptr, output_ptr, Channels, OutputCount, KernelSize); +} + + +#endif // MLAS_F16VEC_INTRINSICS_SUPPORTED diff --git a/onnxruntime/test/contrib_ops/nhwc_pool_in_op_test.cc b/onnxruntime/test/contrib_ops/nhwc_pool_in_op_test.cc new file mode 100644 index 0000000000..a2a875cdb7 --- /dev/null +++ b/onnxruntime/test/contrib_ops/nhwc_pool_in_op_test.cc @@ -0,0 +1,314 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// +// Test module for NWHC fp16 internal operators +// + +#include +#include + +#include "core/util/math.h" +#include "core/mlas/inc/mlas.h" +#include "gtest/gtest.h" +#include "test/providers/provider_test_utils.h" + +namespace onnxruntime { +namespace test { +#ifdef MLAS_F16VEC_INTRINSICS_SUPPORTED + +class NhwcFp16PoolOpTester { + private: + bool is_max_pool_; // max or average pool + std::vector X_data_; + std::vector X_shape_; + std::vector kernel_shape_; + std::vector pads_; + std::vector strides_; + std::vector dilations_; + + static size_t ShapeSize(const std::vector& shape) { + return static_cast(std::accumulate(shape.cbegin(), shape.cend(), 1LL, std::multiplies())); + } + + static bool NextPosition(int64_t N, const int64_t* shape, int64_t* dims) { + // Loop over spatial axes in reverse order to choose an index, like counting. + bool incremented = false; + for (int64_t d_i = N - 1; d_i >= 0; --d_i) { + int64_t d_max = shape[d_i]; + ORT_ENFORCE(dims[d_i] < d_max); + if (dims[d_i] == d_max - 1) { + dims[d_i] = 0; + } else { // dims[d_i] < d_max - 1 + ++dims[d_i]; + incremented = true; + break; + } + } + return incremented; + } + + void ComputeExpectedOutput(std::vector& Y_data, std::vector& Y_shape) { + ORT_ENFORCE(X_shape_.size() >= 2 && X_shape_.size() == kernel_shape_.size() + 2); + + const size_t kernel_rank = kernel_shape_.size(); + + const int64_t batch_count = X_shape_[0]; + const int64_t channels = X_shape_[X_shape_.size() - 1]; + + std::vector pads(pads_); + if (pads.empty()) { + pads.resize(kernel_rank * 2, 0); + } + std::vector dilations(dilations_); + if (dilations.empty()) { + dilations.resize(kernel_rank, 1); + } + std::vector strides(strides_); + if (strides.empty()) { + strides.resize(kernel_rank, 1); + } + + const int64_t* input_shape = X_shape_.data() + 1; + + // Compute the expected shape of the output. + Y_shape.reserve(kernel_rank + 2); + Y_shape.push_back(batch_count); + for (size_t n = 0; n < kernel_rank; n++) { + Y_shape.push_back(((input_shape[n] + pads[n] + pads[kernel_rank + n]) - + (dilations[n] * (kernel_shape_[n] - 1) + 1)) / + strides[n] + + 1); + } + Y_shape.push_back(channels); + Y_data.resize(ShapeSize(Y_shape)); + + const int64_t* output_shape = Y_shape.data() + 1; + + const int64_t input_image_size = std::accumulate( + input_shape, input_shape + kernel_rank, 1LL, std::multiplies()); + + const MLFloat16* Xdata = X_data_.data(); + MLFloat16* Ydata = Y_data.data(); + + for (int64_t batch = 0; batch < batch_count; batch++) { + std::vector d_output(kernel_rank, 0); + std::vector d_kernel(kernel_rank, 0); + do { + std::vector accs(channels, is_max_pool_ ? std::numeric_limits::lowest() : 0.f); + size_t cnt = 0; + do { + int64_t input_offset = 0; + bool is_padding = false; + for (size_t axis = 0; axis < kernel_rank; ++axis) { + int64_t input_dim = d_kernel[axis] * dilations[axis] + d_output[axis] * strides[axis] - pads[axis]; + is_padding |= !math::is_a_ge_zero_and_a_lt_b(input_dim, input_shape[axis]); + input_offset *= input_shape[axis]; + input_offset += input_dim; + } + if (!is_padding) { + const MLFloat16* data_ptr = Xdata + input_offset * channels; + cnt++; + for (int64_t c = 0; c < channels; c++) { + if (is_max_pool_) { + accs[c] = std::max(accs[c], data_ptr[c].ToFloat()); + } else { + accs[c] += data_ptr[c].ToFloat(); + } + } + } + } while (NextPosition(kernel_rank, kernel_shape_.data(), d_kernel.data())); + for (int64_t c = 0; c < channels; c++) { + if (!is_max_pool_) { + accs[c] /= cnt; + } + Ydata[c] = MLFloat16(accs[c]); + } + Ydata += channels; + } while (NextPosition(kernel_rank, output_shape, d_output.data())); + Xdata += channels * input_image_size; + } + } + + public: + NhwcFp16PoolOpTester(bool is_max_pool) : is_max_pool_(is_max_pool) { + } + + void GenerateRandomInput(const std::vector& shape) { + constexpr float MinimumFillValue = -23.0f; + static size_t offset = 7; + size_t shape_size = ShapeSize(shape); + X_data_.resize(shape_size); + + for (size_t n = 0; n < shape_size; n++) { + offset = (offset + 31) % 47; + X_data_[n] = MLFloat16((MinimumFillValue + offset) / 16.0f); + } + X_shape_ = shape; + } + + void SetKernelShape(const std::vector& kernel_shape) { + kernel_shape_ = kernel_shape; + } + + void SetPads(const std::vector& pads) { + pads_ = pads; + } + + void SetStrides(const std::vector& strides) { + strides_ = strides; + } + + void SetDilations(const std::vector& dilations) { + dilations_ = dilations; + } + + void Run() { + std::vector Y_data; + std::vector Y_shape; + ComputeExpectedOutput(Y_data, Y_shape); + + OpTester test(is_max_pool_ ? "MaxPool" : "AveragePool", 11, onnxruntime::kMSInternalNHWCDomain); + test.AddInput("x", X_shape_, X_data_); + test.AddOutput("y", Y_shape, Y_data); + test.AddAttribute("kernel_shape", kernel_shape_); + if (!pads_.empty()) { + test.AddAttribute("pads", pads_); + } + if (!strides_.empty()) { + test.AddAttribute("strides", strides_); + } + if (!dilations_.empty()) { + test.AddAttribute("dilations", dilations_); + } + test.Run(OpTester::ExpectResult::kExpectSuccess, ""); + } +}; + + +TEST(NhwcFp16PoolOpTest, MaxPool1D) { + for (int64_t channels = 1; channels < 94; channels++) { + NhwcFp16PoolOpTester test(true); + test.GenerateRandomInput({1, 23, channels}); + test.SetKernelShape({5}); + test.SetPads({2, 2}); + test.Run(); + } +} + +TEST(NhwcFp16PoolOpTest, MaxPool2D) { + for (int64_t channels = 1; channels < 94; channels++) { + NhwcFp16PoolOpTester test(true); + test.GenerateRandomInput({1, 15, 19, channels}); + test.SetKernelShape({3, 5}); + test.SetPads({1, 1, 1, 1}); + test.Run(); + } +} + +TEST(NhwcFp16PoolOpTest, MaxPool3D) { + for (int64_t channels = 1; channels < 94; channels++) { + NhwcFp16PoolOpTester test(true); + test.GenerateRandomInput({1, 9, 13, 15, channels}); + test.SetKernelShape({2, 4, 6}); + test.SetPads({0, 0, 0, 1, 1, 1}); + test.Run(); + } +} + +TEST(NhwcFp16PoolOpTest, MaxPoolStrides) { + NhwcFp16PoolOpTester test(true); + test.GenerateRandomInput({4, 23, 19, 32}); + test.SetKernelShape({3, 3}); + test.SetStrides({2, 2}); + test.Run(); +} + +TEST(NhwcFp16PoolOpTest, MaxPoolDilations) { + NhwcFp16PoolOpTester test(true); + test.GenerateRandomInput({4, 23, 19, 32}); + test.SetKernelShape({3, 3}); + test.SetDilations({2, 2}); + test.Run(); +} + +TEST(NhwcFp16PoolOpTest, AvgPool1D) { + for (int64_t channels = 1; channels < 94; channels++) { + NhwcFp16PoolOpTester test(false); + test.GenerateRandomInput({1, 23, channels}); + test.SetKernelShape({5}); + test.SetPads({2, 2}); + test.Run(); + } +} + +TEST(NhwcFp16PoolOpTest, AvgPool2D) { + for (int64_t channels = 1; channels < 94; channels++) { + NhwcFp16PoolOpTester test(false); + test.GenerateRandomInput({1, 15, 19, channels}); + test.SetKernelShape({3, 5}); + test.SetPads({1, 1, 1, 1}); + test.Run(); + } +} + +TEST(NhwcFp16PoolOpTest, AvgPool3D) { + for (int64_t channels = 1; channels < 94; channels++) { + NhwcFp16PoolOpTester test(false); + test.GenerateRandomInput({1, 9, 13, 15, channels}); + test.SetKernelShape({2, 4, 6}); + test.SetPads({0, 0, 0, 1, 1, 1}); + test.Run(); + } +} + +TEST(NhwcFp16PoolOpTest, AvgPoolStrides) { + NhwcFp16PoolOpTester test(false); + test.GenerateRandomInput({4, 23, 19, 32}); + test.SetKernelShape({3, 3}); + test.SetStrides({2, 2}); + test.Run(); +} + +/****** + * AveragePool op does not support dilations until version 19 +TEST(NhwcFp16PoolOpTest, AvgPoolDilations) { + NhwcFp16PoolOpTester test(false); + test.GenerateRandomInput({4, 23, 19, 32}); + test.SetKernelShape({3, 3}); + test.SetDilations({2, 2}); + test.Run(); +} +*/ + +TEST(NhwcFp16PoolOpTest, AvgPoolIncludePadPixel) { + OpTester test("AveragePool", 11, onnxruntime::kMSInternalNHWCDomain); + + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{1, 1}); + test.AddAttribute("pads", std::vector{1, 1, 1, 1}); + test.AddAttribute("kernel_shape", std::vector{2, 2}); + test.AddAttribute("count_include_pad", (int64_t)1); + std::vector x_vals = { + MLFloat16(0.3337f), MLFloat16(0.8794f), MLFloat16(0.3375f), + MLFloat16(0.6666f), MLFloat16(0.4426f), MLFloat16(0.6474f), + MLFloat16(0.7675f), MLFloat16(0.8823f), MLFloat16(0.8852f)}; + + std::vector x_dims = {1, 3, 3, 1}; + std::vector expected_dims = {1, 4, 4, 1}; + std::vector expected_vals = { + MLFloat16(0.0834f), MLFloat16(0.3033f), MLFloat16(0.3042f), MLFloat16(0.0844f), + MLFloat16(0.2501f), MLFloat16(0.5806f), MLFloat16(0.5767f), MLFloat16(0.2462f), + MLFloat16(0.3585f), MLFloat16(0.6897f), MLFloat16(0.7144f), MLFloat16(0.3832f), + MLFloat16(0.1919f), MLFloat16(0.4124f), MLFloat16(0.4419f), MLFloat16(0.2213f)}; + + test.AddInput("X", x_dims, x_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); +} + + + +#endif +} // namespace test +} // namespace onnxruntime