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

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

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<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSInternalNHWCDomain, 11, MLFloat16, MaxPool)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSInternalNHWCDomain, 11, MLFloat16, AveragePool)>,
+  };
+
+  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<void>,  // 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();
 }
 

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<Tensor>(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<size_t>(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<MLFloat16> padding_data;
+  if (need_padding) {
+    padding_data.resize(static_cast<size_t>(C), MLFloat16());
+  }
+
+  const auto* Xdata = X->Data<MLFloat16>();
+  auto* Ydata = Y->MutableData<MLFloat16>();
+
+  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<MLFloat16, StorageOrder::NHWC>()(
+          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<ptrdiff_t>(spatial_dims),
+          output_start,
+          output_count,
+          static_cast<MLFloat16 const**>(col_buffer.get()),
+          need_padding ? padding_data.data() : nullptr);
+      if (is_max_pool_) {
+        MlasNhwcMaxPool(
+            static_cast<MLFloat16 const**>(col_buffer.get()),
+            Ydata,
+            static_cast<size_t>(C),
+            static_cast<size_t>(output_count),
+            static_cast<size_t>(kernel_size));
+      } else {
+        MlasNhwcAvgPool(
+            static_cast<MLFloat16 const**>(col_buffer.get()),
+            Ydata,
+            static_cast<size_t>(C),
+            static_cast<size_t>(output_count),
+            static_cast<size_t>(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<MLFloat16>()),
+    NhwcMaxPoolFp16);
+
+ONNX_OPERATOR_TYPED_KERNEL_EX(
+    AveragePool,
+    kMSInternalNHWCDomain,
+    11,
+    MLFloat16,
+    kCpuExecutionProvider,
+    KernelDefBuilder()
+        .TypeConstraint("T", DataTypeImpl::GetTensorType<MLFloat16>()),
+    NhwcAvgPoolFp16);
+
+
+#endif
+
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/core/mlas/inc/mlas.h

@@ -1641,3 +1641,44 @@ MlasTranspose(
         reinterpret_cast<uint16_t*>(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

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)

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>
+typename AggregationType::CtxType
+PoolCreateContext(size_t KernelSize);
+
+template <typename AggregationType>
+MLAS_FLOAT16X8
+PoolInit16x8();
+
+template <typename AggregationType>
+MLAS_FLOAT16X4
+PoolInit16x4();
+
+template <typename AggregationType>
+MLAS_FLOAT16X8
+PoolAggregate16x8(MLAS_FLOAT16X8 agg, MLAS_FLOAT16X8 element);
+
+template <typename AggregationType>
+MLAS_FLOAT16X4
+PoolAggregate16x4(MLAS_FLOAT16X4 agg, MLAS_FLOAT16X4 element);
+
+template <typename AggregationType>
+MLAS_FLOAT16X8
+PoolSummary16x8(MLAS_FLOAT16X8 agg, typename AggregationType::CtxType context);
+
+template <typename AggregationType>
+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<MaxPoolAggregation>(size_t KernelSize)
+{
+    MLAS_UNREFERENCED_PARAMETER(KernelSize);
+    return 0;
+}
+
+template<>
+MLAS_FORCEINLINE
+MLAS_FLOAT16X8
+PoolInit16x8<MaxPoolAggregation>()
+{
+    // lowest fp16 -65504.0f
+    return MlasBroadcastFloat16x8(0xfbff);
+}
+
+template <>
+MLAS_FORCEINLINE
+MLAS_FLOAT16X4
+PoolInit16x4<MaxPoolAggregation>()
+{
+    // lowest fp16 -65504.0f
+    return MlasBroadcastFloat16x4(0xfbff);
+}
+
+template<>
+MLAS_FORCEINLINE 
+MLAS_FLOAT16X8
+PoolAggregate16x8<MaxPoolAggregation>(MLAS_FLOAT16X8 agg, MLAS_FLOAT16X8 element)
+{
+    return MlasMaximumFloat16x8(agg, element);
+}
+
+template<>
+MLAS_FORCEINLINE 
+MLAS_FLOAT16X4
+PoolAggregate16x4<MaxPoolAggregation>(MLAS_FLOAT16X4 agg, MLAS_FLOAT16X4 element)
+{
+    return MlasMaximumFloat16x4(agg, element);
+}
+
+template<>
+MLAS_FORCEINLINE 
+MLAS_FLOAT16X8
+PoolSummary16x8<MaxPoolAggregation>(MLAS_FLOAT16X8 agg, size_t size)
+{
+    MLAS_UNREFERENCED_PARAMETER(size);
+    return agg;
+}
+
+template<>
+MLAS_FORCEINLINE 
+MLAS_FLOAT16X4
+PoolSummary16x4<MaxPoolAggregation>(MLAS_FLOAT16X4 agg, size_t size)
+{
+    MLAS_UNREFERENCED_PARAMETER(size);
+    return agg;
+}
+
+struct AveragePoolAggregation {
+    typedef MLAS_FLOAT16X8 CtxType;
+};
+
+template<>
+MLAS_FLOAT16X8
+PoolCreateContext<AveragePoolAggregation>(size_t KernelSize)
+{
+    return MlasBroadcastFloat16x8(MLAS_Float2Half(float(KernelSize)));
+}
+
+
+template <>
+MLAS_FORCEINLINE MLAS_FLOAT16X8
+PoolInit16x8<AveragePoolAggregation>()
+{
+    return MlasZeroFloat16x8();
+}
+
+template <>
+MLAS_FORCEINLINE MLAS_FLOAT16X4
+PoolInit16x4<AveragePoolAggregation>()
+{
+    return MlasZeroFloat16x4();
+}
+
+
+template <>
+MLAS_FORCEINLINE MLAS_FLOAT16X8
+PoolAggregate16x8<AveragePoolAggregation>(MLAS_FLOAT16X8 agg, MLAS_FLOAT16X8 element)
+{
+    return MlasAddFloat16x8(agg, element);
+}
+
+template <>
+MLAS_FORCEINLINE MLAS_FLOAT16X4
+PoolAggregate16x4<AveragePoolAggregation>(MLAS_FLOAT16X4 agg, MLAS_FLOAT16X4 element)
+{
+    return MlasAddFloat16x4(agg, element);
+}
+
+template <>
+MLAS_FORCEINLINE MLAS_FLOAT16X8
+PoolSummary16x8<AveragePoolAggregation>(MLAS_FLOAT16X8 agg, MLAS_FLOAT16X8 context)
+{
+    return MlasDivFloat16x8(agg, context);
+}
+
+template <>
+MLAS_FORCEINLINE MLAS_FLOAT16X4
+PoolSummary16x4<AveragePoolAggregation>(MLAS_FLOAT16X4 agg, MLAS_FLOAT16X8 context)
+{
+    return MlasDivFloat16x4(agg, MlasToLowHalfFloat16x4(context));
+}
+
+
+template<typename AggregationType>
+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<AggregationType>();
+            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<AggregationType>(MaximumVector0, InputVector0);
+                MaximumVector1 = PoolAggregate16x8<AggregationType>(MaximumVector1, InputVector1);
+                MaximumVector2 = PoolAggregate16x8<AggregationType>(MaximumVector2, InputVector2);
+                MaximumVector3 = PoolAggregate16x8<AggregationType>(MaximumVector3, InputVector3);
+                cnt++;
+            }
+            typename AggregationType::CtxType context = PoolCreateContext<AggregationType>(cnt);
+            MaximumVector0 = PoolSummary16x8<AggregationType>(MaximumVector0, context);
+            MaximumVector1 = PoolSummary16x8<AggregationType>(MaximumVector1, context);
+            MaximumVector2 = PoolSummary16x8<AggregationType>(MaximumVector2, context);
+            MaximumVector3 = PoolSummary16x8<AggregationType>(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<AggregationType>();
+            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<AggregationType>(MaximumVector0, InputVector0);
+                MaximumVector1 = PoolAggregate16x8<AggregationType>(MaximumVector1, InputVector1);
+                cnt++;
+            }
+            typename AggregationType::CtxType context = PoolCreateContext<AggregationType>(cnt);
+            MaximumVector0 = PoolSummary16x8<AggregationType>(MaximumVector0, context);
+            MaximumVector1 = PoolSummary16x8<AggregationType>(MaximumVector1, context);
+
+            MlasStoreFloat16x8(&Output[0], MaximumVector0);
+            MlasStoreFloat16x8(&Output[8], MaximumVector1);
+
+            Output += 16;
+            ChannelOffset += 16;
+            c -= 16;
+        }
+
+        if (c >= 8) {
+            MLAS_FLOAT16X8 MaximumVector0 = PoolInit16x8<AggregationType>();
+            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<AggregationType>(MaximumVector0, InputVector0);
+                cnt++;
+            }
+            typename AggregationType::CtxType context = PoolCreateContext<AggregationType>(cnt);
+            MaximumVector0 = PoolSummary16x8<AggregationType>(MaximumVector0, context);
+
+            MlasStoreFloat16x8(&Output[0], MaximumVector0);
+
+            Output += 8;
+            ChannelOffset += 8;
+            c -= 8;
+        }
+
+        if (c >= 4) {
+            MLAS_FLOAT16X4 MaximumVector0 = PoolInit16x4<AggregationType>();
+            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<AggregationType>(MaximumVector0, InputVector0);
+                cnt++;
+            }
+            typename AggregationType::CtxType context = PoolCreateContext<AggregationType>(cnt);
+            MaximumVector0 = PoolSummary16x4<AggregationType>(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<AggregationType>();
+            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<AggregationType>(MaximumVector0, InputVector0);
+                cnt++;
+            }
+            typename AggregationType::CtxType context = PoolCreateContext<AggregationType>(cnt);
+            MaximumVector0 = PoolSummary16x4<AggregationType>(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<const _mlas_fp16_* const*>(Input);
+    _mlas_fp16_* output_ptr = reinterpret_cast<_mlas_fp16_*>(Output);
+    MlasPoolFp16HWC<MaxPoolAggregation>(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<const _mlas_fp16_* const*>(Input);
+    _mlas_fp16_* output_ptr = reinterpret_cast<_mlas_fp16_*>(Output);
+    MlasPoolFp16HWC<AveragePoolAggregation>(input_ptr, output_ptr, Channels, OutputCount, KernelSize);
+}
+
+
+#endif  // MLAS_F16VEC_INTRINSICS_SUPPORTED

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 <algorithm>
+#include <random>
+
+#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<MLFloat16> X_data_;
+  std::vector<int64_t>   X_shape_;
+  std::vector<int64_t>   kernel_shape_;
+  std::vector<int64_t>   pads_;
+  std::vector<int64_t>   strides_;
+  std::vector<int64_t>   dilations_;
+
+  static size_t ShapeSize(const std::vector<int64_t>& shape) {
+    return static_cast<size_t>(std::accumulate(shape.cbegin(), shape.cend(), 1LL, std::multiplies<int64_t>()));
+  }
+
+  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<MLFloat16>& Y_data, std::vector<int64_t>& 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<int64_t> pads(pads_);
+    if (pads.empty()) {
+      pads.resize(kernel_rank * 2, 0);
+    }
+    std::vector<int64_t> dilations(dilations_);
+    if (dilations.empty()) {
+      dilations.resize(kernel_rank, 1);
+    }
+    std::vector<int64_t> 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<int64_t>());
+
+    const MLFloat16* Xdata = X_data_.data();
+    MLFloat16* Ydata = Y_data.data();
+
+    for (int64_t batch = 0; batch < batch_count; batch++) {
+      std::vector<int64_t> d_output(kernel_rank, 0);
+      std::vector<int64_t> d_kernel(kernel_rank, 0);
+      do {
+        std::vector<float> accs(channels, is_max_pool_ ? std::numeric_limits<float>::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<int64_t>& 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<int64_t>& kernel_shape) {
+    kernel_shape_ = kernel_shape;
+  }
+
+  void SetPads(const std::vector<int64_t>& pads) {
+    pads_ = pads;
+  }
+
+  void SetStrides(const std::vector<int64_t>& strides) {
+    strides_ = strides;
+  }
+
+  void SetDilations(const std::vector<int64_t>& dilations) {
+    dilations_ = dilations;
+  }
+
+  void Run() {
+    std::vector<MLFloat16> Y_data;
+    std::vector<int64_t> Y_shape;
+    ComputeExpectedOutput(Y_data, Y_shape);
+
+    OpTester test(is_max_pool_ ? "MaxPool" : "AveragePool", 11, onnxruntime::kMSInternalNHWCDomain);
+    test.AddInput<MLFloat16>("x", X_shape_, X_data_);
+    test.AddOutput<MLFloat16>("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<int64_t>{1, 1});
+  test.AddAttribute("pads", std::vector<int64_t>{1, 1, 1, 1});
+  test.AddAttribute("kernel_shape", std::vector<int64_t>{2, 2});
+  test.AddAttribute("count_include_pad", (int64_t)1);
+  std::vector<MLFloat16> 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<int64_t> x_dims = {1, 3, 3, 1};
+  std::vector<int64_t> expected_dims = {1, 4, 4, 1};
+  std::vector<MLFloat16> 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<MLFloat16>("X", x_dims, x_vals);
+  test.AddOutput<MLFloat16>("Y", expected_dims, expected_vals);
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
+}
+
+
+
+#endif
+}  // namespace test
+}  // namespace onnxruntime