Support int16 in vision quantized_layer_norm (#24094)

backends/cadence/generic/operators/op_quantized_layer_norm.cpp

@@ -11,6 +11,7 @@
 #include <cmath>
 #include <cstdint>
 #include <cstdlib>
+#include <type_traits>
 
 #include <executorch/backends/cadence/generic/kernels/kernels.h>
 #include <executorch/backends/cadence/generic/operators/cadence_type_util.h>
@@ -61,29 +62,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 (size_t 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 = dequantize<T>(x[j], input_scale, input_zero_point);
 
       val = (val - mean) * inv_std * weight_data[j] + bias_data[j];
@@ -149,7 +158,7 @@ Tensor& quantized_layer_norm_out(
 
   ScalarType dtype = input.scalar_type();
   switch (dtype) {
-    ET_FORALL_CADENCE_QUANTIZED_TYPES(typed_quantized_layer_norm)
+    ET_FORALL_CADENCE_QUANTIZED_TYPES_WITH_INT16(typed_quantized_layer_norm)
     default:
       ET_DCHECK_MSG(
           false, "Unhandled dtype %s", torch::executor::toString(dtype));
@@ -188,7 +197,7 @@ Tensor& quantized_layer_norm_per_tensor_out(
 
   ScalarType dtype = input.scalar_type();
   switch (dtype) {
-    ET_FORALL_CADENCE_QUANTIZED_TYPES(typed_quantized_layer_norm)
+    ET_FORALL_CADENCE_QUANTIZED_TYPES_WITH_INT16(typed_quantized_layer_norm)
     default:
       ET_DCHECK_MSG(
           false, "Unhandled dtype %s", torch::executor::toString(dtype));

backends/cadence/generic/operators/tests/BUCK

@@ -4,6 +4,25 @@ load("@fbsource//xplat/on_device_ai/Assistant/Jarvis/build:cxx_wrapper.bzl", "ja
 
 oncall("executorch")
 
+jarvis_wrapper.cxx_test(
+    name = "test_op_quantized_layer_norm",
+    srcs = [
+        "test_op_quantized_layer_norm.cpp",
+    ],
+    compatible_backends = ["generic"],
+    labels = [ci.linux(ci.mode("fbsource//arvr/mode/platform010/dev"))],
+    platforms = CXX,
+    visibility = [
+        "fbsource//xplat/executorch/backends/cadence/...",
+        "fbcode//executorch/backends/cadence/...",
+    ],
+    deps = [
+        "fbsource//xplat/executorch/backends/cadence/generic/operators:op_quantized_layer_norm",
+        "fbsource//xplat/executorch/kernels/test:gtest_utils",
+        "fbsource//xplat/executorch/runtime/core/exec_aten/testing_util:tensor_util",
+    ],
+)
+
 jarvis_wrapper.cxx_test(
     name = "test_op_sdpa_bitwise_mask_gen",
     srcs = [

backends/cadence/generic/operators/tests/CMakeLists.txt

@@ -48,7 +48,7 @@ endif()
 # checks with EXPECT_TENSOR_EQ, mirroring the BUCK tests.
 add_executable(
   cadence_generic_op_tests
-  test_op_sdpa_bitwise_mask_gen.cpp
+  test_op_quantized_layer_norm.cpp test_op_sdpa_bitwise_mask_gen.cpp
   ${EXECUTORCH_ROOT}/runtime/core/exec_aten/testing_util/tensor_util.cpp
 )
 target_compile_definitions(cadence_generic_op_tests PRIVATE GTEST_HAS_PTHREAD=0)

backends/cadence/generic/operators/tests/test_op_quantized_layer_norm.cpp

@@ -0,0 +1,161 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <gtest/gtest.h>
+
+#include <executorch/kernels/test/TestUtil.h>
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
+#include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
+
+#include <executorch/backends/cadence/generic/operators/op_quantized_layer_norm.h>
+
+namespace impl {
+namespace generic {
+namespace native {
+namespace {
+
+using ::executorch::aten::IntArrayRef;
+using ::executorch::aten::ScalarType;
+using ::executorch::aten::Tensor;
+using ::executorch::runtime::testing::TensorFactory;
+
+// Scales are powers of two and inputs are symmetric around their mean, so
+// every row normalizes to exactly +/-1 (eps is negligible).
+constexpr double kEps = 1e-9;
+
+class GenericQuantizedLayerNormTest : public OperatorTest {};
+
+TEST_F(GenericQuantizedLayerNormTest, Int16LargeMagnitudePerTensor) {
+  TensorFactory<ScalarType::Short> tf_short;
+  TensorFactory<ScalarType::Float> tf_float;
+  const std::vector<int64_t> normalized_shape{8};
+
+  // Per-row sum of squares is 8 * 30000^2, well past int32 range.
+  const Tensor input = tf_short.make(
+      {2, 8},
+      {30000,
+       -30000,
+       30000,
+       -30000,
+       30000,
+       -30000,
+       30000,
+       -30000,
+       -30000,
+       30000,
+       -30000,
+       30000,
+       -30000,
+       30000,
+       -30000,
+       30000});
+  const Tensor weight = tf_float.ones({8});
+  const Tensor bias = tf_float.zeros({8});
+  Tensor out = tf_short.zeros({2, 8});
+
+  quantized_layer_norm_per_tensor_out(
+      context_,
+      input,
+      /*in_scale=*/1.0 / 1024,
+      /*in_zero_point=*/0,
+      IntArrayRef(normalized_shape.data(), normalized_shape.size()),
+      weight,
+      bias,
+      kEps,
+      /*output_scale=*/1.0 / 16384,
+      /*output_zero_point=*/0,
+      out);
+
+  const Tensor expected = tf_short.make(
+      {2, 8},
+      {16384,
+       -16384,
+       16384,
+       -16384,
+       16384,
+       -16384,
+       16384,
+       -16384,
+       -16384,
+       16384,
+       -16384,
+       16384,
+       -16384,
+       16384,
+       -16384,
+       16384});
+  EXPECT_TENSOR_EQ(out, expected);
+}
+
+TEST_F(GenericQuantizedLayerNormTest, Int16LargeMeanSmallSpreadWithZeroPoints) {
+  TensorFactory<ScalarType::Short> tf_short;
+  TensorFactory<ScalarType::Float> tf_float;
+  TensorFactory<ScalarType::Long> tf_long;
+  const std::vector<int64_t> normalized_shape{8};
+
+  // Centered values are 21000 +/- 4: the mean dwarfs the spread, which breaks
+  // an fp32 E[x^2] - E[x]^2 variance.
+  const Tensor input = tf_short.make(
+      {1, 8}, {20004, 19996, 20004, 19996, 20004, 19996, 20004, 19996});
+  const Tensor in_scale = tf_float.make({1}, {1.0f / 64});
+  const Tensor in_zero_point = tf_long.make({1}, {-1000});
+  const Tensor weight = tf_float.full({8}, 2.0f);
+  const Tensor bias = tf_float.full({8}, 0.5f);
+  Tensor out = tf_short.zeros({1, 8});
+
+  quantized_layer_norm_out(
+      context_,
+      input,
+      in_scale,
+      in_zero_point,
+      IntArrayRef(normalized_shape.data(), normalized_shape.size()),
+      weight,
+      bias,
+      kEps,
+      /*output_scale=*/1.0 / 4096,
+      /*output_zero_point=*/100,
+      out);
+
+  // (+1 * 2 + 0.5) * 4096 + 100 and (-1 * 2 + 0.5) * 4096 + 100.
+  const Tensor expected = tf_short.make(
+      {1, 8}, {10340, -6044, 10340, -6044, 10340, -6044, 10340, -6044});
+  EXPECT_TENSOR_EQ(out, expected);
+}
+
+TEST_F(GenericQuantizedLayerNormTest, Int8WithZeroPoint) {
+  TensorFactory<ScalarType::Char> tf_char;
+  TensorFactory<ScalarType::Float> tf_float;
+  const std::vector<int64_t> normalized_shape{4};
+
+  const Tensor input = tf_char.make({1, 4}, {10, 30, 10, 30});
+  const Tensor weight = tf_float.ones({4});
+  const Tensor bias = tf_float.zeros({4});
+  Tensor out = tf_char.zeros({1, 4});
+
+  quantized_layer_norm_per_tensor_out(
+      context_,
+      input,
+      /*in_scale=*/1.0 / 8,
+      /*in_zero_point=*/20,
+      IntArrayRef(normalized_shape.data(), normalized_shape.size()),
+      weight,
+      bias,
+      kEps,
+      /*output_scale=*/1.0 / 64,
+      /*output_zero_point=*/0,
+      out);
+
+  const Tensor expected = tf_char.make({1, 4}, {-64, 64, -64, 64});
+  EXPECT_TENSOR_EQ(out, expected);
+}
+
+} // namespace
+} // namespace native
+} // namespace generic
+} // namespace impl

backends/cadence/hifi/operators/op_quantized_layer_norm.cpp

@@ -11,7 +11,9 @@
 #include <executorch/runtime/kernel/kernel_includes.h>
 #include <algorithm>
 #include <cmath>
+#include <cstdint>
 #include <tuple>
+#include <type_traits>
 
 using ::executorch::aten::IntArrayRef;
 using ::executorch::aten::ScalarType;
@@ -55,32 +57,41 @@ void quantized_layer_norm_per_tensor_(
     const T* __restrict__ x = in_data + i * last_dim;
     T* __restrict__ 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;
 #pragma simd
     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 = XT_DIV_S(XT_MUL_S(input_scale, sum), last_dim);
-    float variance =
-        XT_DIV_S(
-            XT_MUL_S(sq_sum, XT_MUL_S(input_scale, input_scale)), last_dim) -
-        XT_MUL_S(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 float n_f = static_cast<float>(n);
+    float mean =
+        XT_DIV_S(XT_MUL_S(input_scale, static_cast<float>(centered_sum)), n_f);
+    float variance = XT_MUL_S(
+        XT_DIV_S(static_cast<float>(var_numerator), XT_MUL_S(n_f, n_f)),
+        XT_MUL_S(input_scale, input_scale));
     float inv_std = XT_RECIP_S(XT_SQRT_S(XT_ADD_S(variance, (float)eps)));
 
     // y = (x - mean) / std * kGamma + kBeta
 #pragma simd
     for (size_t j = 0; j < last_dim; ++j) {
       // 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 = dequantize<T>(x[j], input_scale, input_zero_point);
       val = (val - mean) * inv_std * weight_data[j] + bias_data[j];
       y[j] = quantize<T>(val, output_inv_scale, output_zero_point);
@@ -147,7 +158,7 @@ void quantized_layer_norm_out(
 
   ScalarType dtype = input.scalar_type();
   switch (dtype) {
-    ET_FORALL_CADENCE_QUANTIZED_TYPES(typed_quantized_layer_norm)
+    ET_FORALL_CADENCE_QUANTIZED_TYPES_WITH_INT16(typed_quantized_layer_norm)
     default:
       ET_DCHECK_MSG(
           false, "Unhandled dtype %s", torch::executor::toString(dtype));
@@ -185,7 +196,7 @@ void quantized_layer_norm_per_tensor_out(
 
   ScalarType dtype = input.scalar_type();
   switch (dtype) {
-    ET_FORALL_CADENCE_QUANTIZED_TYPES(typed_quantized_layer_norm)
+    ET_FORALL_CADENCE_QUANTIZED_TYPES_WITH_INT16(typed_quantized_layer_norm)
     default:
       ET_DCHECK_MSG(
           false, "Unhandled dtype %s", torch::executor::toString(dtype));

backends/cadence/hifi/operators/operators.h

@@ -62,6 +62,32 @@ void quantized_relu_out(
     const ::executorch::aten::Tensor& out_shift,
     ::executorch::aten::Tensor& output);
 
+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);
+
 void quantized_linear_out(
     ::executorch::runtime::KernelRuntimeContext& ctx,
     const ::executorch::aten::Tensor& in,