Support int16 in vision quantized_layer_norm (#24094)

Summary:

The vision quantized_layer_norm kernel only handled uint8 and int8; int16 input aborted at its ET_CHECK_MSG. It also had the same precision problems as the HiFi and generic kernels:
- int32 accumulators, which overflow for int16 because one squared value can reach 2^30.
- An E[x^2] - E[x]^2 variance, which cancels catastrophically when the mean is large compared to the spread.

This kernel is registered for Cadence Vision builds through functions_vision.yaml. This change closes the last int16 gap among the Cadence layer norm kernels.

Changes:
- Add an int16 (Short) branch to both overloads.
- Accumulate in int64 for 16-bit types; 8-bit keeps int32 in the inner loop.
- Widen to the accumulator type before squaring, so the product cannot overflow int32.
- Apply the zero-point correction in int64 and compute the variance numerator N*sum(v^2) - sum(v)^2 exactly in integers. The per-row mean and variance stay in double as before.
- Declare both kernels in vision/operators/operators.h.

Reviewed By: DrJessop

Differential Revision: D126953450

Author: mvartani-meta

backends/cadence/vision/operators/op_quantized_layer_norm.cpp

@@ -10,6 +10,8 @@
 #include <executorch/runtime/kernel/kernel_includes.h>
 
 #include <cmath>
+#include <cstdint>
+#include <type_traits>
 
 using ::executorch::aten::Tensor;
 using ::executorch::runtime::getLeadingDims;
@@ -49,29 +51,37 @@ void quantized_layer_norm_per_tensor_(
     const T* x = in_data + i * last_dim;
     T* y = out_data + i * last_dim;
 
-    // compute sum and squared sum. The fp32 sum can be approximated as:
-    // (X_1 - in_zero_point) * in_scale + (X_2 - in_zero_point) * in_scale + ...
-    // (X_N - in_zero_point) * in_scale.
-    int32_t sum = 0;
-    int32_t sq_sum = last_dim * input_zero_point * input_zero_point;
+    // 16-bit squares overflow an int32 sum after two elements.
+    using acc_t = std::conditional_t<sizeof(T) == 1, int32_t, int64_t>;
+    acc_t sum = 0;
+    acc_t sq_sum = 0;
     for (size_t j = 0; j < last_dim; ++j) {
-      int32_t val = x[j];
+      const int32_t val = x[j];
       sum += val;
-      sq_sum += val * val;
+      sq_sum += static_cast<acc_t>(val) * val;
     }
-    sq_sum -= (2 * sum * input_zero_point);
-    sum -= (last_dim * input_zero_point);
 
-    float mean = (input_scale * sum) / last_dim;
-    float variance =
-        (sq_sum * input_scale * input_scale) / last_dim - mean * mean;
+    // Shift raw sums to zero-point-centered sums, then form the variance
+    // numerator N * sum(v^2) - sum(v)^2 exactly in integers: doing
+    // E[x^2] - E[x]^2 in fp32 cancels catastrophically for 16-bit inputs.
+    const int64_t n = static_cast<int64_t>(last_dim);
+    const int64_t zp = input_zero_point;
+    const int64_t centered_sum = static_cast<int64_t>(sum) - n * zp;
+    const int64_t centered_sq_sum =
+        static_cast<int64_t>(sq_sum) - 2 * zp * sum + n * zp * zp;
+    const int64_t var_numerator =
+        n * centered_sq_sum - centered_sum * centered_sum;
+
+    const double n_d = static_cast<double>(n);
+    float mean = input_scale * centered_sum / n_d;
+    float variance = var_numerator / (n_d * n_d) * (input_scale * input_scale);
     float inv_std = 1.0f / std::sqrt(variance + eps);
 
     // y = (x - mean) / std * kGamma + kBeta
     for (int j = 0; j < last_dim; ++j) {
       // y[j] = (x[j] - mean) / std * kGamma + kBeta;
       // Since X is quantized, we dequantize it, compute fp32 result, and
-      // quantize the result to an int8/uint8 value.
+      // quantize the result back to T.
       float val = kernels::dequantize<T>(x[j], input_scale, input_zero_point);
 
       val = (val - mean) * inv_std * weight_data[j] + bias_data[j];
@@ -144,6 +154,17 @@ void quantized_layer_norm_out(
         output_scale,
         output_zero_point,
         out);
+  } else if (input.scalar_type() == executorch::aten::ScalarType::Short) {
+    quantized_layer_norm_<int16_t>(
+        input,
+        in_scale,
+        in_zero_point,
+        weight,
+        bias,
+        eps,
+        output_scale,
+        output_zero_point,
+        out);
   } else {
     ET_CHECK_MSG(
         false,
@@ -186,6 +207,17 @@ void quantized_layer_norm_per_tensor_out(
         output_scale,
         output_zero_point,
         out);
+  } else if (input.scalar_type() == executorch::aten::ScalarType::Short) {
+    quantized_layer_norm_per_tensor_<int16_t>(
+        input,
+        in_scale,
+        in_zero_point,
+        weight,
+        bias,
+        eps,
+        output_scale,
+        output_zero_point,
+        out);
   } else {
     ET_CHECK_MSG(
         false,

backends/cadence/vision/operators/operators.h

@@ -96,6 +96,32 @@ void quantized_conv2d_nhwc_out(
     const ::executorch::aten::Tensor& out_shift,
     ::executorch::aten::Tensor& out);
 
+void quantized_layer_norm_out(
+    ::executorch::runtime::KernelRuntimeContext& ctx,
+    const ::executorch::aten::Tensor& input,
+    const ::executorch::aten::Tensor& in_scale,
+    const ::executorch::aten::Tensor& in_zero_point,
+    const ::executorch::aten::IntArrayRef normalized_shape,
+    const ::executorch::aten::Tensor& weight,
+    const ::executorch::aten::Tensor& bias,
+    double eps,
+    double output_scale,
+    int64_t output_zero_point,
+    ::executorch::aten::Tensor& out);
+
+void quantized_layer_norm_per_tensor_out(
+    ::executorch::runtime::KernelRuntimeContext& ctx,
+    const ::executorch::aten::Tensor& input,
+    double in_scale,
+    int64_t in_zero_point,
+    const ::executorch::aten::IntArrayRef normalized_shape,
+    const ::executorch::aten::Tensor& weight,
+    const ::executorch::aten::Tensor& bias,
+    double eps,
+    double output_scale,
+    int64_t output_zero_point,
+    ::executorch::aten::Tensor& out);
+
 } // namespace native
 } // namespace vision
 } // namespace impl