Use scalar fast path in optimized layer_norm for small tensors (pytorch#18636)

Differential Revision: D98795281

Pull Request resolved: pytorch#18636

Author: xiaodong705

kernels/optimized/cpu/op_native_layer_norm.cpp

@@ -72,6 +72,24 @@ void layer_norm(
   const bool gamma_null = gamma_data == nullptr;
   const bool beta_null = beta_data == nullptr;
 
+  // For small normalized dimensions, fall back to the portable scalar
+  // implementation since SIMD vectorization setup/tail-handling overhead
+  // exceeds the benefit for small N.
+  constexpr size_t kSmallNThreshold = 256;
+  if (N < kSmallNThreshold) {
+    layer_norm_scalar<CTYPE>(
+        input_data,
+        gamma_data,
+        beta_data,
+        out_data,
+        mean_data,
+        rstd_data,
+        M,
+        N,
+        eps);
+    return;
+  }
+
   for (size_t i = 0; i < M; ++i) {
     const CTYPE* src_ptr = input_data + i * N;
     CTYPE* dst_ptr = out_data + i * N;

kernels/portable/cpu/op_native_layer_norm.cpp

@@ -8,7 +8,6 @@
 #include <c10/util/irange.h>
 
 #include <executorch/kernels/portable/cpu/util/normalization_ops_util.h>
-#include <executorch/kernels/portable/cpu/vec_ops.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
 #include <cmath>
 #include <tuple>
@@ -54,41 +53,21 @@ void layer_norm(
   }
 
   const CTYPE* input_data = input.const_data_ptr<CTYPE>();
-  const CTYPE* weight_data;
-  if (weight.has_value()) {
-    weight_data = weight.value().const_data_ptr<CTYPE>();
-  } else {
-    weight_data = nullptr;
-  }
-  const CTYPE* bias_data;
-  if (bias.has_value()) {
-    bias_data = bias.value().const_data_ptr<CTYPE>();
-  } else {
-    bias_data = nullptr;
-  }
-
-  const CTYPE ct_normalized = static_cast<CTYPE>(normalized);
-  for (const auto i : c10::irange(leading)) {
-    const CTYPE* x = input_data + i * normalized;
-    CTYPE* y = out_data + i * normalized;
-
-    // compute E[X] and Var[x] = E[x^2] - E[x]^2
-    float sum = reduce_add(x, ct_normalized);
-    float sq_sum = vec_powerf(x, ct_normalized);
-    float mean_value = sum / ct_normalized;
-    float variance = sq_sum / ct_normalized - mean_value * mean_value;
-    float std = std::sqrt(variance + eps);
-
-    // Calculate the elements of output
-    for (const auto j : c10::irange(normalized)) {
-      CTYPE w = weight_data ? weight_data[j] : static_cast<CTYPE>(1);
-      CTYPE b = bias_data ? bias_data[j] : static_cast<CTYPE>(0);
-      y[j] = (x[j] - mean_value) / std * w + b;
-    }
-
-    mean_data[i] = mean_value;
-    rstd_data[i] = 1.0 / std;
-  }
+  const CTYPE* weight_data =
+      weight.has_value() ? weight.value().const_data_ptr<CTYPE>() : nullptr;
+  const CTYPE* bias_data =
+      bias.has_value() ? bias.value().const_data_ptr<CTYPE>() : nullptr;
+
+  layer_norm_scalar<CTYPE>(
+      input_data,
+      weight_data,
+      bias_data,
+      out_data,
+      mean_data,
+      rstd_data,
+      leading,
+      normalized,
+      eps);
 }
 
 } // namespace

kernels/portable/cpu/util/normalization_ops_util.h

@@ -9,10 +9,54 @@
 #pragma once
 
 #include <executorch/runtime/kernel/kernel_includes.h>
+#include <cmath>
+#include <numeric>
 
 namespace torch {
 namespace executor {
 
+/**
+ * Scalar layer_norm computation over M rows of N elements each.
+ * Computes mean/variance in float, normalizes with (x - mean) / std * gamma +
+ * beta. Caller must handle M==0 and N==0 edge cases before calling.
+ */
+template <typename CTYPE>
+inline void layer_norm_scalar(
+    const CTYPE* input_data,
+    const CTYPE* weight_data, // nullable
+    const CTYPE* bias_data, // nullable
+    CTYPE* out_data,
+    CTYPE* mean_data,
+    CTYPE* rstd_data,
+    size_t M,
+    size_t N,
+    float eps) {
+  for (size_t i = 0; i < M; ++i) {
+    const CTYPE* x = input_data + i * N;
+    CTYPE* y = out_data + i * N;
+
+    // compute E[X] and Var[x] = E[x^2] - E[x]^2
+    float sum = std::accumulate(x, x + N, 0.0f);
+    float sq_sum = 0;
+    for (size_t j = 0; j < N; ++j) {
+      sq_sum += static_cast<float>(x[j]) * x[j];
+    }
+    float mean_value = sum / N;
+    float variance = sq_sum / N - mean_value * mean_value;
+    float std = std::sqrt(variance + eps);
+
+    // Calculate the elements of output
+    for (size_t j = 0; j < N; ++j) {
+      CTYPE w = weight_data ? weight_data[j] : static_cast<CTYPE>(1);
+      CTYPE b = bias_data ? bias_data[j] : static_cast<CTYPE>(0);
+      y[j] = (x[j] - mean_value) / std * w + b;
+    }
+
+    mean_data[i] = mean_value;
+    rstd_data[i] = 1.0 / std;
+  }
+}
+
 bool check_batch_norm_args(
     const Tensor& in,
     const std::optional<Tensor>& weight,