diff --git a/backends/cadence/generic/operators/op_quantized_layer_norm.cpp b/backends/cadence/generic/operators/op_quantized_layer_norm.cpp index c5b53146916..c83a9f278ff 100644 --- a/backends/cadence/generic/operators/op_quantized_layer_norm.cpp +++ b/backends/cadence/generic/operators/op_quantized_layer_norm.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -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; + 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(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(last_dim); + const int64_t zp = input_zero_point; + const int64_t centered_sum = static_cast(sum) - n * zp; + const int64_t centered_sq_sum = + static_cast(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(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(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)); diff --git a/backends/cadence/generic/operators/tests/BUCK b/backends/cadence/generic/operators/tests/BUCK index be1ba3ea358..b7c25cd2658 100644 --- a/backends/cadence/generic/operators/tests/BUCK +++ b/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 = [ diff --git a/backends/cadence/generic/operators/tests/CMakeLists.txt b/backends/cadence/generic/operators/tests/CMakeLists.txt index f92ba3afadb..29c8fe71883 100644 --- a/backends/cadence/generic/operators/tests/CMakeLists.txt +++ b/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) diff --git a/backends/cadence/generic/operators/tests/test_op_quantized_layer_norm.cpp b/backends/cadence/generic/operators/tests/test_op_quantized_layer_norm.cpp new file mode 100644 index 00000000000..b90253cfdc8 --- /dev/null +++ b/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 + +#include +#include +#include +#include + +#include + +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 tf_short; + TensorFactory tf_float; + const std::vector 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 tf_short; + TensorFactory tf_float; + TensorFactory tf_long; + const std::vector 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 tf_char; + TensorFactory tf_float; + const std::vector 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 diff --git a/backends/cadence/hifi/operators/op_quantized_layer_norm.cpp b/backends/cadence/hifi/operators/op_quantized_layer_norm.cpp index 5510385b14f..2bb449a7d05 100644 --- a/backends/cadence/hifi/operators/op_quantized_layer_norm.cpp +++ b/backends/cadence/hifi/operators/op_quantized_layer_norm.cpp @@ -11,7 +11,9 @@ #include #include #include +#include #include +#include 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; + 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(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(last_dim); + const int64_t zp = input_zero_point; + const int64_t centered_sum = static_cast(sum) - n * zp; + const int64_t centered_sq_sum = + static_cast(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(n); + float mean = + XT_DIV_S(XT_MUL_S(input_scale, static_cast(centered_sum)), n_f); + float variance = XT_MUL_S( + XT_DIV_S(static_cast(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(x[j], input_scale, input_zero_point); val = (val - mean) * inv_std * weight_data[j] + bias_data[j]; y[j] = quantize(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)); diff --git a/backends/cadence/hifi/operators/operators.h b/backends/cadence/hifi/operators/operators.h index a109d9d9e1b..9a4c36ebb70 100644 --- a/backends/cadence/hifi/operators/operators.h +++ b/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, diff --git a/backends/cadence/hifi/operators/tests/BUCK b/backends/cadence/hifi/operators/tests/BUCK index aabe8a0b7a5..66357fea0c3 100644 --- a/backends/cadence/hifi/operators/tests/BUCK +++ b/backends/cadence/hifi/operators/tests/BUCK @@ -27,6 +27,27 @@ jarvis_wrapper.cxx_test( ) +jarvis_wrapper.cxx_test( + name = "test_op_quantized_layer_norm", + srcs = [ + "test_op_quantized_layer_norm.cpp", + ], + compatible_backends = ["hifi"], + labels = [ci.linux(ci.mode("fbsource//arvr/mode/platform010/dev"))], + platforms = CXX, + remote_execution = XTENSA_TEST_REMOTE_EXECUTION_ASIC, + visibility = [ + "fbsource//xplat/executorch/backends/cadence/...", + "fbcode//executorch/backends/cadence/...", + ], + deps = [ + "fbsource//xplat/executorch/kernels/test:gtest_utils", + "fbsource//xplat/executorch/backends/cadence/hifi/operators:op_quantized_layer_norm", + "fbsource//xplat/executorch/backends/cadence/runtime:et_pal", + "fbsource//xplat/executorch/runtime/core/exec_aten/testing_util:tensor_util", + ], +) + jarvis_wrapper.cxx_test( name = "test_op_quantize_per_tensor", srcs = [ diff --git a/backends/cadence/hifi/operators/tests/CMakeLists.txt b/backends/cadence/hifi/operators/tests/CMakeLists.txt index 0398400a1d5..c55f1b51e32 100644 --- a/backends/cadence/hifi/operators/tests/CMakeLists.txt +++ b/backends/cadence/hifi/operators/tests/CMakeLists.txt @@ -56,6 +56,7 @@ add_executable( test_op_permute_copy.cpp test_op_quantize_per_tensor.cpp test_op_quantized_conv2d_out.cpp + test_op_quantized_layer_norm.cpp test_op_quantized_matmul_out.cpp test_op_quantized_relu_out.cpp test_op_relu.cpp diff --git a/backends/cadence/hifi/operators/tests/test_op_quantized_layer_norm.cpp b/backends/cadence/hifi/operators/tests/test_op_quantized_layer_norm.cpp new file mode 100644 index 00000000000..e8dce5e1a9a --- /dev/null +++ b/backends/cadence/hifi/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 + +#include +#include +#include +#include + +#include + +namespace impl { +namespace HiFi { +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 HiFiQuantizedLayerNormTest : public OperatorTest {}; + +TEST_F(HiFiQuantizedLayerNormTest, Int16LargeMagnitudePerTensor) { + TensorFactory tf_short; + TensorFactory tf_float; + const std::vector 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(HiFiQuantizedLayerNormTest, Int16LargeMeanSmallSpreadWithZeroPoints) { + TensorFactory tf_short; + TensorFactory tf_float; + TensorFactory tf_long; + const std::vector 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(HiFiQuantizedLayerNormTest, Int8WithZeroPoint) { + TensorFactory tf_char; + TensorFactory tf_float; + const std::vector 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 HiFi +} // namespace impl diff --git a/backends/cadence/vision/operators/op_quantized_layer_norm.cpp b/backends/cadence/vision/operators/op_quantized_layer_norm.cpp index 90919c6296d..b25c45ae6fb 100644 --- a/backends/cadence/vision/operators/op_quantized_layer_norm.cpp +++ b/backends/cadence/vision/operators/op_quantized_layer_norm.cpp @@ -10,6 +10,8 @@ #include #include +#include +#include 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; + 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(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(last_dim); + const int64_t zp = input_zero_point; + const int64_t centered_sum = static_cast(sum) - n * zp; + const int64_t centered_sq_sum = + static_cast(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(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(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_( + 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_( + input, + in_scale, + in_zero_point, + weight, + bias, + eps, + output_scale, + output_zero_point, + out); } else { ET_CHECK_MSG( false, diff --git a/backends/cadence/vision/operators/operators.h b/backends/cadence/vision/operators/operators.h index 3ae870e368b..314514e0169 100644 --- a/backends/cadence/vision/operators/operators.h +++ b/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