[MLAS] AVX-512 16-wide Erf kernel and NCHWc reorder transpose for MobileClip-S0 model

cmake/onnxruntime_mlas.cmake

@@ -248,6 +248,7 @@ function(setup_mlas_source_for_windows)
       ${MLAS_SRC_DIR}/intrinsics/avx512/quantize_avx512f.cpp
       ${MLAS_SRC_DIR}/intrinsics/avx512/sconv_nchw_depthwise_multiplier_greater_than_1_avx512f.cpp
       ${MLAS_SRC_DIR}/linear_attention_kernel_avx512f.cpp
+      ${MLAS_SRC_DIR}/intrinsics/avx512/reorder_avx512f.cpp
     )
 
     set_source_files_properties(${mlas_platform_srcs_avx512} PROPERTIES COMPILE_FLAGS "/arch:AVX512")
@@ -932,6 +933,7 @@ else()
           ${MLAS_SRC_DIR}/intrinsics/avx512/quantize_avx512f.cpp
           ${MLAS_SRC_DIR}/intrinsics/avx512/sconv_nchw_depthwise_multiplier_greater_than_1_avx512f.cpp
           ${MLAS_SRC_DIR}/linear_attention_kernel_avx512f.cpp
+          ${MLAS_SRC_DIR}/intrinsics/avx512/reorder_avx512f.cpp
         )
         set_source_files_properties(${mlas_platform_srcs_avx512f} PROPERTIES COMPILE_FLAGS "-mavx512f")
 

onnxruntime/core/mlas/lib/intrinsics/avx512/gelu_avx512f.cpp

@@ -205,6 +205,34 @@ MlasGeluErfKernelAvx512FExactImpl(
     }
 }
 
+void
+MlasErfKernelAvx512FImpl(
+    const float* Input,
+    float* Output,
+    size_t N
+    )
+{
+    const GeluAvx512BroadcastConstants Constants;
+    while (N >= 16) {
+        const __m512 X = _mm512_loadu_ps(Input);
+        const __m512 Result = MlasGeluErfAvx512(X, Constants);
+
+        _mm512_storeu_ps(Output, Result);
+
+        Input += 16;
+        Output += 16;
+        N -= 16;
+    }
+
+    if (N > 0) {
+        const __mmask16 TailMask = __mmask16((1u << static_cast<unsigned>(N)) - 1u);
+        const __m512 X = _mm512_maskz_loadu_ps(TailMask, Input);
+        const __m512 Result = MlasGeluErfAvx512(X, Constants);
+
+        _mm512_mask_storeu_ps(Output, TailMask, Result);
+    }
+}
+
 }  // namespace
 
 void
@@ -217,3 +245,14 @@ MlasGeluErfKernelAvx512F(
 {
     MlasGeluErfKernelAvx512FExactImpl(Input, Output, N);
 }
+
+void
+MLASCALL
+MlasErfKernelAvx512F(
+    const float* Input,
+    float* Output,
+    size_t N
+    )
+{
+    MlasErfKernelAvx512FImpl(Input, Output, N);
+}

onnxruntime/core/mlas/lib/intrinsics/avx512/reorder_avx512f.cpp

@@ -0,0 +1,144 @@
+/*++
+
+Copyright (c) Microsoft Corporation. All rights reserved.
+
+Licensed under the MIT License.
+
+Module Name:
+
+    reorder_avx512f.cpp
+
+Abstract:
+
+    This module implements AVX-512 accelerated helpers for the NCHW<->NCHWc
+    reorder routines when the NCHWc block size is 16. The baseline reorder path
+    transposes 4x4 tiles with SSE2, requiring four transposes per 16-wide block;
+    these routines transpose full 16x16 tiles in one pass.
+
+    The transposes are pure data-movement: for a tile,
+    Output[p*DstStride + c] = Input[c*SrcStride + p].
+
+--*/
+
+#include <immintrin.h>
+
+#include "mlasi.h"
+
+namespace {
+
+//
+// Transpose a 16x16 float tile.
+//
+//  Input  row c (0..15): 16 contiguous floats at Input + c*SrcStride.
+//  Output row p (0..15): 16 contiguous floats at Output + p*DstStride, with
+//                        Output[p*DstStride + c] = Input[c*SrcStride + p].
+//
+MLAS_FORCEINLINE
+void
+MlasReorderTranspose16x16Avx512F(
+    const float* Input,
+    float* Output,
+    size_t SrcStride,
+    size_t DstStride
+    )
+{
+    __m512 r[16];
+    for (int i = 0; i < 16; i++) {
+        r[i] = _mm512_loadu_ps(Input + i * SrcStride);
+    }
+
+    __m512 t[16];
+    for (int i = 0; i < 8; i++) {
+        t[2 * i + 0] = _mm512_unpacklo_ps(r[2 * i], r[2 * i + 1]);
+        t[2 * i + 1] = _mm512_unpackhi_ps(r[2 * i], r[2 * i + 1]);
+    }
+
+    __m512 u[16];
+    for (int i = 0; i < 4; i++) {
+        u[4 * i + 0] = _mm512_castpd_ps(_mm512_unpacklo_pd(_mm512_castps_pd(t[4 * i + 0]), _mm512_castps_pd(t[4 * i + 2])));
+        u[4 * i + 1] = _mm512_castpd_ps(_mm512_unpackhi_pd(_mm512_castps_pd(t[4 * i + 0]), _mm512_castps_pd(t[4 * i + 2])));
+        u[4 * i + 2] = _mm512_castpd_ps(_mm512_unpacklo_pd(_mm512_castps_pd(t[4 * i + 1]), _mm512_castps_pd(t[4 * i + 3])));
+        u[4 * i + 3] = _mm512_castpd_ps(_mm512_unpackhi_pd(_mm512_castps_pd(t[4 * i + 1]), _mm512_castps_pd(t[4 * i + 3])));
+    }
+
+    __m512 v[16];
+    for (int i = 0; i < 4; i++) {
+        v[i + 0]  = _mm512_shuffle_f32x4(u[i], u[i + 4], 0x88);
+        v[i + 4]  = _mm512_shuffle_f32x4(u[i], u[i + 4], 0xDD);
+        v[i + 8]  = _mm512_shuffle_f32x4(u[i + 8], u[i + 12], 0x88);
+        v[i + 12] = _mm512_shuffle_f32x4(u[i + 8], u[i + 12], 0xDD);
+    }
+
+    for (int i = 0; i < 8; i++) {
+        _mm512_storeu_ps(Output + (i + 0) * DstStride, _mm512_shuffle_f32x4(v[i], v[i + 8], 0x88));
+        _mm512_storeu_ps(Output + (i + 8) * DstStride, _mm512_shuffle_f32x4(v[i], v[i + 8], 0xDD));
+    }
+}
+
+}  // namespace
+
+//
+// Reorder one full 16-channel NCHW block into NCHWc (block size 16).
+//
+//  Source S: 16 channels, each InputSize contiguous spatial floats (stride InputSize).
+//  Dest   D: InputSize spatial rows, each 16 contiguous channel floats (stride 16).
+//  D[p*16 + c] = S[c*InputSize + p].
+//
+// Equivalent to the InputChannelsThisIteration == BlockSize == 16 case of the
+// scalar MlasReorderInputNchw inner loops.
+//
+void
+MLASCALL
+MlasReorderInputNchwBlock16Avx512F(
+    const float* S,
+    float* D,
+    size_t InputSize
+    )
+{
+    size_t p = 0;
+    for (; p + 16 <= InputSize; p += 16) {
+        MlasReorderTranspose16x16Avx512F(S + p, D + p * 16, InputSize, 16);
+    }
+    for (; p < InputSize; p++) {
+        float* d = D + p * 16;
+        const float* s = S + p;
+        for (int c = 0; c < 16; c++) {
+            d[c] = s[c * InputSize];
+        }
+    }
+}
+
+//
+// Reorder one full 16-channel NCHWc block into NCHW (block size 16), for a
+// contiguous run of OutputSize spatial positions.
+//
+//  Source S: OutputSize spatial rows, each 16 contiguous channel floats (stride 16).
+//  Dest   D: 16 channels, each OutputSize contiguous spatial floats (stride OutputSize).
+//  D[c*OutputSize + p] = S[p*16 + c].
+//
+// Equivalent to the OutputChannelsThisIteration == BlockSize == 16 case of the
+// scalar MlasReorderOutputNchwThreaded inner loops.
+//
+void
+MLASCALL
+MlasReorderOutputNchwBlock16Avx512F(
+    const float* S,
+    float* D,
+    size_t OutputSize
+    )
+{
+    size_t p = 0;
+    for (; p + 16 <= OutputSize; p += 16) {
+        // Source rows are spatial positions (stride 16 channels); dest rows are
+        // channels (stride OutputSize). Transpose with SrcStride=16,
+        // DstStride=OutputSize maps Output[c*OutputSize + p] = Input[p*16 + c].
+        MlasReorderTranspose16x16Avx512F(S + p * 16, D + p, 16, OutputSize);
+    }
+    for (; p < OutputSize; p++) {
+        const float* s = S + p * 16;
+        float* d = D + p;
+        for (int c = 0; c < 16; c++) {
+            d[c * OutputSize] = s[c];
+        }
+    }
+}

onnxruntime/core/mlas/lib/mlasi.h

@@ -1360,9 +1360,32 @@ extern "C" {
     MLAS_QUANTIZE_LINEAR_S8_KERNEL MlasQuantizeLinearS8KernelAvx512F;
     MLAS_QUANTIZE_LINEAR_U8_KERNEL MlasQuantizeLinearU8KernelAvx512F;
     MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasGeluErfKernelAvx512F;
+    MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasErfKernelAvx512F;
     MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasSiluKernelAvx512F;
 #endif
 
+#if defined(MLAS_TARGET_AMD64)
+//
+// AVX-512 accelerated NCHW<->NCHWc reorder block helpers (block size 16).
+// Declared unconditionally for AMD64; only invoked when NchwcBlockSize == 16.
+//
+void
+MLASCALL
+MlasReorderInputNchwBlock16Avx512F(
+    const float* S,
+    float* D,
+    size_t InputSize
+    );
+
+void
+MLASCALL
+MlasReorderOutputNchwBlock16Avx512F(
+    const float* S,
+    float* D,
+    size_t OutputSize
+    );
+#endif
+
     MLAS_REDUCE_MAXIMUM_FLOAT_KERNEL MlasReduceMaximumF32Kernel;
     MLAS_REDUCE_MINIMUM_MAXIMUM_FLOAT_KERNEL MlasReduceMinimumMaximumF32Kernel;
 #if defined(MLAS_TARGET_RISCV64) && defined(MLAS_USE_RVV)

onnxruntime/core/mlas/lib/platform.cpp

@@ -555,6 +555,7 @@ Return Value:
 
                 if (((Cpuid7[1] & 0x10000) != 0) && ((xcr0 & 0xE0) == 0xE0)) {
                     this->GeluErfKernelRoutine = MlasGeluErfKernelAvx512F;
+                    this->ErfKernelRoutine = MlasErfKernelAvx512F;
                     this->SiluKernelRoutine = MlasSiluKernelAvx512F;
                     this->GemmFloatKernel = MlasGemmFloatKernelAvx512F;
                     this->GemmDoubleKernel = MlasGemmDoubleKernelAvx512F;

onnxruntime/core/mlas/lib/reorder.cpp

@@ -260,6 +260,19 @@ Return Value:
         const size_t InputChannelsThisIteration = std::min(i, BlockSize);
         i -= InputChannelsThisIteration;
 
+#if defined(MLAS_TARGET_AMD64)
+        //
+        // Fast path: on AVX-512 (BlockSize == 16) a full 16-channel block is a
+        // 16-wide transpose. Bit-exact with the SSE2 4x4 path below.
+        //
+        if (BlockSize == 16 && InputChannelsThisIteration == 16) {
+            MlasReorderInputNchwBlock16Avx512F(S, D, InputSize);
+            S += BlockSize * InputSize;
+            D += BlockSize * InputSize;
+            continue;
+        }
+#endif
+
         const float* s = S;
         float* d = D;
         size_t InputSizeRemaining = InputSize;
@@ -502,6 +515,19 @@ Return Value:
             BlockSize : OutputChannels - BlockSize * LastTaskInBatchIndex;
         const size_t AlignedOutputChannelsThisIteration = OutputChannelsThisIteration & (~3);
 
+#if defined(MLAS_TARGET_AMD64)
+        //
+        // Fast path: on AVX-512 (BlockSize == 16) a full 16-channel block is a
+        // 16-wide transpose. Bit-exact with the SSE2 4x4 path below.
+        //
+        if (BlockSize == 16 && OutputChannelsThisIteration == 16) {
+            MlasReorderOutputNchwBlock16Avx512F(S, D, OutputSize);
+            S += BlockSize * OutputSize;
+            D += OutputChannelsThisIteration * OutputSize;
+            continue;
+        }
+#endif
+
         const float* s = S;
         float* d = D;
         size_t OutputSizeRemaining = OutputSize;