Optimize FP16 LayerNorm and RMSNorm on x86

cmake/onnxruntime_mlas.cmake

@@ -847,7 +847,7 @@ else()
         set(mlas_platform_srcs_avx2
           ${MLAS_SRC_DIR}/layernorm_kernel_avx2.cpp
         )
-        set_source_files_properties(${mlas_platform_srcs_avx2} PROPERTIES COMPILE_FLAGS "-mavx2 -mfma")
+        set_source_files_properties(${mlas_platform_srcs_avx2} PROPERTIES COMPILE_FLAGS "-mavx2 -mfma -mf16c")
 
         set(mlas_platform_srcs
           ${mlas_platform_srcs_sse2}

onnxruntime/core/mlas/inc/mlas.h

@@ -1742,6 +1742,25 @@ MlasLayerNormF32(
     bool Simplified
 );
 
+/**
+ * @brief Compute LayerNorm or RMSNorm (simplified) for one row of IEEE FP16 data.
+ *        Scale and bias are supplied as float32 to match the CPU LayerNorm
+ *        prepacking path. Returns false when no optimized kernel is available.
+ */
+bool
+MLASCALL
+MlasLayerNormF16(
+    const uint16_t* Input,
+    const float* Scale,
+    const float* Bias,
+    uint16_t* Output,
+    float* MeanOut,
+    float* InvStdDevOut,
+    size_t NormSize,
+    float Epsilon,
+    bool Simplified
+);
+
 /**
  * @brief Supply matrices data information to half precision gemm functions
  */

onnxruntime/core/mlas/lib/layernorm.cpp

@@ -59,3 +59,26 @@ bool
     kernel(Input, Scale, Bias, Output, MeanOut, InvStdDevOut, NormSize, Epsilon, Simplified);
     return true;
 }
+
+bool
+    MLASCALL
+    MlasLayerNormF16(
+        const uint16_t* Input,
+        const float* Scale,
+        const float* Bias,
+        uint16_t* Output,
+        float* MeanOut,
+        float* InvStdDevOut,
+        size_t NormSize,
+        float Epsilon,
+        bool Simplified
+    )
+{
+    auto kernel = GetMlasPlatform().LayerNormF16Kernel;
+    if (kernel == nullptr || NormSize < 16) {
+        return false;
+    }
+
+    kernel(Input, Scale, Bias, Output, MeanOut, InvStdDevOut, NormSize, Epsilon, Simplified);
+    return true;
+}

onnxruntime/core/mlas/lib/layernorm_kernel_avx2.cpp

@@ -33,6 +33,7 @@ Module Name:
 --*/
 
 #include "mlasi.h"
+#include "mlas_float16.h"
 
 #if defined(MLAS_TARGET_AMD64) || defined(MLAS_TARGET_IX86)
 
@@ -247,4 +248,168 @@ MlasLayerNormKernelAvx2(
     }
 }
 
+void MLASCALL
+MlasLayerNormKernelF16Avx2(
+    const uint16_t* Input,
+    const float* Scale,
+    const float* Bias,
+    uint16_t* Output,
+    float* MeanOut,
+    float* InvStdDevOut,
+    size_t NormSize,
+    float Epsilon,
+    bool Simplified)
+{
+    assert(!Simplified || Bias == nullptr);
+
+    const size_t n = NormSize;
+    float mean_val = 0.0f;
+    float inv_denom;
+
+    if (Simplified) {
+        __m256 vsumsq = _mm256_setzero_ps();
+        size_t i = 0;
+        double sum = 0.0;
+        if (MeanOut != nullptr) {
+            __m256d vsum = _mm256_setzero_pd();
+            for (; i + 8 <= n; i += 8) {
+                const __m256 vx = _mm256_cvtph_ps(
+                    _mm_loadu_si128(reinterpret_cast<const __m128i*>(Input + i)));
+                vsumsq = _mm256_fmadd_ps(vx, vx, vsumsq);
+                vsum = _mm256_add_pd(vsum, _mm256_cvtps_pd(_mm256_castps256_ps128(vx)));
+                vsum = _mm256_add_pd(vsum, _mm256_cvtps_pd(_mm256_extractf128_ps(vx, 1)));
+            }
+
+            __m128d rd = _mm_add_pd(
+                _mm256_castpd256_pd128(vsum),
+                _mm256_extractf128_pd(vsum, 1));
+            rd = _mm_add_sd(rd, _mm_unpackhi_pd(rd, rd));
+            sum = _mm_cvtsd_f64(rd);
+        } else {
+            for (; i + 8 <= n; i += 8) {
+                const __m256 vx = _mm256_cvtph_ps(
+                    _mm_loadu_si128(reinterpret_cast<const __m128i*>(Input + i)));
+                vsumsq = _mm256_fmadd_ps(vx, vx, vsumsq);
+            }
+        }
+
+        __m128 r = _mm_add_ps(
+            _mm256_castps256_ps128(vsumsq),
+            _mm256_extractf128_ps(vsumsq, 1));
+        r = _mm_add_ps(r, _mm_movehl_ps(r, r));
+        r = _mm_add_ss(r, _mm_movehdup_ps(r));
+        float sumsq = _mm_cvtss_f32(r);
+
+        for (; i < n; ++i) {
+            const float value = MLAS_Half2Float(Input[i]);
+            sumsq += value * value;
+            if (MeanOut != nullptr) {
+                sum += static_cast<double>(value);
+            }
+        }
+
+        if (MeanOut != nullptr) {
+            mean_val = static_cast<float>(sum / static_cast<double>(n));
+        }
+        inv_denom = 1.0f / sqrtf(sumsq / static_cast<float>(n) + Epsilon);
+    } else {
+        __m256d vsum = _mm256_setzero_pd();
+        size_t i = 0;
+        for (; i + 8 <= n; i += 8) {
+            const __m256 vx = _mm256_cvtph_ps(
+                _mm_loadu_si128(reinterpret_cast<const __m128i*>(Input + i)));
+            vsum = _mm256_add_pd(vsum, _mm256_cvtps_pd(_mm256_castps256_ps128(vx)));
+            vsum = _mm256_add_pd(vsum, _mm256_cvtps_pd(_mm256_extractf128_ps(vx, 1)));
+        }
+
+        __m128d rd = _mm_add_pd(
+            _mm256_castpd256_pd128(vsum),
+            _mm256_extractf128_pd(vsum, 1));
+        rd = _mm_add_sd(rd, _mm_unpackhi_pd(rd, rd));
+        double sum = _mm_cvtsd_f64(rd);
+
+        for (; i < n; ++i) {
+            sum += static_cast<double>(MLAS_Half2Float(Input[i]));
+        }
+
+        mean_val = static_cast<float>(sum / static_cast<double>(n));
+        const __m256 vmean = _mm256_set1_ps(mean_val);
+        __m256 vvariance = _mm256_setzero_ps();
+        i = 0;
+        for (; i + 8 <= n; i += 8) {
+            const __m256 vx = _mm256_cvtph_ps(
+                _mm_loadu_si128(reinterpret_cast<const __m128i*>(Input + i)));
+            const __m256 vd = _mm256_sub_ps(vx, vmean);
+            vvariance = _mm256_fmadd_ps(vd, vd, vvariance);
+        }
+
+        __m128 r = _mm_add_ps(
+            _mm256_castps256_ps128(vvariance),
+            _mm256_extractf128_ps(vvariance, 1));
+        r = _mm_add_ps(r, _mm_movehl_ps(r, r));
+        r = _mm_add_ss(r, _mm_movehdup_ps(r));
+        float variance = _mm_cvtss_f32(r);
+
+        for (; i < n; ++i) {
+            const float delta = MLAS_Half2Float(Input[i]) - mean_val;
+            variance += delta * delta;
+        }
+
+        inv_denom = 1.0f / sqrtf(variance / static_cast<float>(n) + Epsilon);
+    }
+
+    const __m256 vmean = _mm256_set1_ps(mean_val);
+    const __m256 vinv = _mm256_set1_ps(inv_denom);
+    size_t i = 0;
+
+    if (Simplified) {
+        for (; i + 8 <= n; i += 8) {
+            const __m256 vx = _mm256_cvtph_ps(
+                _mm_loadu_si128(reinterpret_cast<const __m128i*>(Input + i)));
+            const __m256 vs = _mm256_loadu_ps(Scale + i);
+            const __m256 vy = _mm256_mul_ps(_mm256_mul_ps(vx, vinv), vs);
+            const __m128i half = _mm256_cvtps_ph(vy, _MM_FROUND_TO_NEAREST_INT);
+            _mm_storeu_si128(reinterpret_cast<__m128i*>(Output + i), half);
+        }
+        for (; i < n; ++i) {
+            Output[i] = MLAS_Float2Half(MLAS_Half2Float(Input[i]) * inv_denom * Scale[i]);
+        }
+    } else if (Bias == nullptr) {
+        for (; i + 8 <= n; i += 8) {
+            const __m256 vx = _mm256_cvtph_ps(
+                _mm_loadu_si128(reinterpret_cast<const __m128i*>(Input + i)));
+            const __m256 vs = _mm256_loadu_ps(Scale + i);
+            const __m256 vy = _mm256_mul_ps(_mm256_mul_ps(_mm256_sub_ps(vx, vmean), vinv), vs);
+            const __m128i half = _mm256_cvtps_ph(vy, _MM_FROUND_TO_NEAREST_INT);
+            _mm_storeu_si128(reinterpret_cast<__m128i*>(Output + i), half);
+        }
+        for (; i < n; ++i) {
+            Output[i] = MLAS_Float2Half(
+                (MLAS_Half2Float(Input[i]) - mean_val) * inv_denom * Scale[i]);
+        }
+    } else {
+        for (; i + 8 <= n; i += 8) {
+            const __m256 vx = _mm256_cvtph_ps(
+                _mm_loadu_si128(reinterpret_cast<const __m128i*>(Input + i)));
+            const __m256 vs = _mm256_loadu_ps(Scale + i);
+            const __m256 vb = _mm256_loadu_ps(Bias + i);
+            const __m256 normalized = _mm256_mul_ps(_mm256_sub_ps(vx, vmean), vinv);
+            const __m256 vy = _mm256_fmadd_ps(normalized, vs, vb);
+            const __m128i half = _mm256_cvtps_ph(vy, _MM_FROUND_TO_NEAREST_INT);
+            _mm_storeu_si128(reinterpret_cast<__m128i*>(Output + i), half);
+        }
+        for (; i < n; ++i) {
+            Output[i] = MLAS_Float2Half(
+                (MLAS_Half2Float(Input[i]) - mean_val) * inv_denom * Scale[i] + Bias[i]);
+        }
+    }
+
+    if (MeanOut != nullptr) {
+        *MeanOut = mean_val;
+    }
+    if (InvStdDevOut != nullptr) {
+        *InvStdDevOut = inv_denom;
+    }
+}
+
 #endif  // MLAS_TARGET_AMD64 || MLAS_TARGET_IX86

onnxruntime/core/mlas/lib/mlasi.h

@@ -766,6 +766,18 @@ typedef void(MLASCALL MLAS_LAYERNORM_F32_KERNEL)(
     bool Simplified
 );
 
+typedef void(MLASCALL MLAS_LAYERNORM_F16_KERNEL)(
+    const uint16_t* Input,
+    const float* Scale,
+    const float* Bias,
+    uint16_t* Output,
+    float* MeanOut,
+    float* InvStdDevOut,
+    size_t NormSize,
+    float Epsilon,
+    bool Simplified
+);
+
 typedef
 void
 (MLASCALL MLAS_QLINEAR_BINARY_OP_S8_KERNEL)(
@@ -1513,6 +1525,7 @@ MlasReorderOutputNchwBlock16Avx512F(
 
 #if defined(MLAS_TARGET_AMD64) || defined(MLAS_TARGET_IX86)
     MLAS_LAYERNORM_F32_KERNEL MlasLayerNormKernelAvx2;
+    MLAS_LAYERNORM_F16_KERNEL MlasLayerNormKernelF16Avx2;
 #endif
 }
 
@@ -1963,6 +1976,7 @@ MLAS_COMPUTE_TANH_FP16_KERNEL* TanhFP16KernelRoutine = nullptr;
     MLAS_CAST_F16_TO_F32_KERNEL* CastF16ToF32Kernel;
     MLAS_CAST_F32_TO_F16_KERNEL* CastF32ToF16Kernel;
     MLAS_LAYERNORM_F32_KERNEL* LayerNormF32Kernel{nullptr};
+    MLAS_LAYERNORM_F16_KERNEL* LayerNormF16Kernel{nullptr};
 
     const MLAS_ROPE_DISPATCH* RopeDispatch{nullptr};
     const MLAS_HGEMM_DISPATCH* HGemmDispatch{nullptr};

onnxruntime/core/mlas/lib/platform.cpp

@@ -573,6 +573,9 @@ Return Value:
                 // TODO(vraspar): check if this really goes here or if there are other platform reqs that we need to fulfill
                 this->LutGenKernel = &MlasLutGenKernelAvx2;
                 this->LayerNormF32Kernel = &MlasLayerNormKernelAvx2;
+                if ((Cpuid1[2] & (1u << 29)) != 0) {
+                    this->LayerNormF16Kernel = &MlasLayerNormKernelF16Avx2;
+                }
 
                 //
                 // Check if the processor supports Hybrid core architecture.
@@ -718,6 +721,9 @@ Return Value:
 
             if (((Cpuid1[2] & 0x1000) != 0) && ((Cpuid7[1] & 0x20) != 0)) {
                 this->LayerNormF32Kernel = &MlasLayerNormKernelAvx2;
+                if ((Cpuid1[2] & (1u << 29)) != 0) {
+                    this->LayerNormF16Kernel = &MlasLayerNormKernelF16Avx2;
+                }
             }
 #endif  // MLAS_TARGET_IX86
 

onnxruntime/core/providers/cpu/nn/layer_norm_impl.cc

@@ -190,6 +190,21 @@ void ComputeJob(
     const ptrdiff_t task_idx, const int64_t norm_size, const int64_t broadcast_param,
     const float* scale_float_ptr, const float* bias_float_ptr, float epsilon, bool simplified,
     MLFloat16* Y_data, U* mean_data, U* inv_std_dev_data, AllocatorPtr alloc) {
+  const int64_t scale_offset = LAYER_NORM_SCALE_BIAS_OFFSET(broadcast_param, task_idx, norm_size);
+  const ptrdiff_t input_offset = SafeInt<ptrdiff_t>(task_idx) * norm_size;
+  if constexpr (std::is_same_v<U, float>) {
+    if (MlasLayerNormF16(
+            reinterpret_cast<const uint16_t*>(X_data + input_offset),
+            scale_float_ptr + scale_offset,
+            (simplified || bias_float_ptr == nullptr) ? nullptr : bias_float_ptr + scale_offset,
+            reinterpret_cast<uint16_t*>(Y_data + input_offset),
+            mean_data ? &mean_data[task_idx] : nullptr,
+            inv_std_dev_data ? &inv_std_dev_data[task_idx] : nullptr,
+            static_cast<size_t>(norm_size), epsilon, simplified)) {
+      return;
+    }
+  }
+
   ComputeJobNarrow(
       X_data, scale_data, bias_data, task_idx, norm_size, broadcast_param,
       scale_float_ptr, bias_float_ptr, epsilon, simplified, Y_data, mean_data, inv_std_dev_data, alloc);

onnxruntime/test/contrib_ops/layer_norm_op_test.cc

@@ -221,6 +221,67 @@ TEST(LayerNormTest, LayerNorm_Scale_Bias_Float16Input) {
             kOpenVINOExecutionProvider, kNnapiExecutionProvider, kCoreMLExecutionProvider, kWebGpuExecutionProvider});
 }
 
+TEST(LayerNormTest, LayerNorm_Float16_OptimizedRowWithStats) {
+  auto cpu_ep = DefaultCpuExecutionProvider();
+  if (!cpu_ep) {
+    GTEST_SKIP() << "CPU EP not available in this build.";
+  }
+
+  constexpr size_t row_count = 2;
+  constexpr size_t norm_size = 31;
+  constexpr float epsilon = 1e-5f;
+  std::vector<float> input(row_count * norm_size);
+  std::vector<float> scale(norm_size);
+  std::vector<float> bias(norm_size);
+  for (size_t i = 0; i < norm_size; ++i) {
+    scale[i] = 0.5f + static_cast<float>(i % 7) * 0.125f;
+    bias[i] = static_cast<float>(static_cast<int>(i % 5) - 2) * 0.25f;
+  }
+
+  std::vector<float> expected(row_count * norm_size);
+  std::vector<float> expected_mean(row_count);
+  std::vector<float> expected_inv_std(row_count);
+  for (size_t row = 0; row < row_count; ++row) {
+    double sum = 0.0;
+    for (size_t i = 0; i < norm_size; ++i) {
+      const float value = static_cast<float>(static_cast<int>((row * 11 + i * 3) % 17) - 8) * 0.25f;
+      input[row * norm_size + i] = value;
+      sum += static_cast<double>(value);
+    }
+
+    const float mean = static_cast<float>(sum / static_cast<double>(norm_size));
+    float variance = 0.0f;
+    for (size_t i = 0; i < norm_size; ++i) {
+      const float delta = input[row * norm_size + i] - mean;
+      variance += delta * delta;
+    }
+
+    const float inv_std = 1.0f / std::sqrt(variance / static_cast<float>(norm_size) + epsilon);
+    expected_mean[row] = mean;
+    expected_inv_std[row] = inv_std;
+    for (size_t i = 0; i < norm_size; ++i) {
+      expected[row * norm_size + i] =
+          (input[row * norm_size + i] - mean) * inv_std * scale[i] + bias[i];
+    }
+  }
+
+  OpTester test("LayerNormalization", 17);
+  test.AddAttribute<int64_t>("axis", 1);
+  test.AddAttribute<float>("epsilon", epsilon);
+  test.AddInput<MLFloat16>("X", {static_cast<int64_t>(row_count), static_cast<int64_t>(norm_size)},
+                           ToFloat16(input));
+  test.AddInput<MLFloat16>("Scale", {static_cast<int64_t>(norm_size)}, ToFloat16(scale), true);
+  test.AddInput<MLFloat16>("Bias", {static_cast<int64_t>(norm_size)}, ToFloat16(bias), true);
+  test.AddOutput<MLFloat16>("Y", {static_cast<int64_t>(row_count), static_cast<int64_t>(norm_size)},
+                            ToFloat16(expected));
+  test.AddOutput<float>("Mean", {static_cast<int64_t>(row_count), 1}, expected_mean);
+  test.AddOutput<float>("InvStdDev", {static_cast<int64_t>(row_count), 1}, expected_inv_std);
+  test.SetOutputAbsErr("Y", 1e-3f);
+  test.SetOutputAbsErr("Mean", 1e-5f);
+  test.SetOutputAbsErr("InvStdDev", 1e-4f);
+  test.ConfigEp(std::move(cpu_ep)).RunWithConfig();
+}
+
 TEST(LayerNormTest, LayerNorm_Scale_Bias_Float16ScaleBiasOutput) {
   OpTester test("LayerNormalization");
   test.AddAttribute<float>("epsilon", 1e-05f);