diff --git a/kernels/quantized/cpu/op_dequantize.cpp b/kernels/quantized/cpu/op_dequantize.cpp index 3f5fca38c86..2958aa79f85 100644 --- a/kernels/quantized/cpu/op_dequantize.cpp +++ b/kernels/quantized/cpu/op_dequantize.cpp @@ -167,6 +167,25 @@ float get_scale(const Tensor& scale, size_t channel_ix) { } } +/** + * Reads one per-channel zero point. `dequantize_per_channel` accepts both Int + * and Long zero_point tensors, so the element width must be honoured here — + * reading an Int tensor as int64 reinterprets pairs of channels as one value + * and runs off the end of the buffer. + */ +int64_t get_zero_point(const Tensor& zero_points, size_t channel_ix) { + ET_CHECK_MSG( + (zero_points.scalar_type() == ScalarType::Int) || + (zero_points.scalar_type() == ScalarType::Long), + "zero_points.scalar_type() %" PRId8 " is not int or long type", + static_cast(zero_points.scalar_type())); + if (zero_points.scalar_type() == ScalarType::Int) { + return static_cast( + zero_points.const_data_ptr()[channel_ix]); + } + return zero_points.const_data_ptr()[channel_ix]; +} + bool can_use_optimized_dequantize_per_channel( const Tensor& in, const ScalarType in_dtype, @@ -208,17 +227,15 @@ void dequantize_per_channel_optimized( } const int8_t* in_data = in.const_data_ptr(); float* out_data = out.mutable_data_ptr(); - const int64_t* zero_points_data = nullptr; - if (opt_zero_points.has_value()) { - zero_points_data = opt_zero_points.value().const_data_ptr(); - } + const Tensor* zero_points = + opt_zero_points.has_value() ? &opt_zero_points.value() : nullptr; const StridesType axis_stride = in.strides()[axis]; const StridesType outer_stride = in.size(axis) * axis_stride; apply_over_unpacked_dim( [in_data, out_data, &scales, - zero_points_data, + zero_points, axis_stride, outer_stride, quant_min, @@ -227,8 +244,8 @@ void dequantize_per_channel_optimized( const int8_t* in_data_local = in_data + outer_idx * outer_stride + unpacked_dim_idx * axis_stride; const double scale = get_scale(scales, unpacked_dim_idx); - const int64_t zero_point = zero_points_data != nullptr - ? zero_points_data[unpacked_dim_idx] + const int64_t zero_point = zero_points != nullptr + ? get_zero_point(*zero_points, unpacked_dim_idx) : 0; float* out_data_local = out_data + outer_idx * outer_stride + unpacked_dim_idx * axis_stride; @@ -422,12 +439,8 @@ Tensor& dequantize_per_channel_out( dims[i] = i + 1; } } - const int64_t* zero_point_data; - if (opt_zero_points.has_value()) { - zero_point_data = opt_zero_points.value().const_data_ptr(); - } else { - zero_point_data = nullptr; - } + const Tensor* zero_point_tensor = + opt_zero_points.has_value() ? &opt_zero_points.value() : nullptr; std::optional> optional_dim_list{ executorch::aten::ArrayRef{dims, size_t(input.dim() - 1)}}; @@ -449,14 +462,14 @@ Tensor& dequantize_per_channel_out( axis == 0, "Axis must be 0 for a single dimensional tensors"); \ const std::optional dim; \ apply_over_dim( \ - [input_data_ptr, out_data_ptr, zero_point_data, &scale]( \ + [input_data_ptr, out_data_ptr, zero_point_tensor, &scale]( \ size_t numel, size_t stride, size_t base_ix) { \ for (size_t i = 0; i < numel; i++) { \ size_t current_ix = base_ix * stride + i; \ float _scale = get_scale(scale, current_ix); \ int64_t zero_point = 0; \ - if (zero_point_data != nullptr) { \ - zero_point = zero_point_data[current_ix]; \ + if (zero_point_tensor != nullptr) { \ + zero_point = get_zero_point(*zero_point_tensor, current_ix); \ } \ out_data_ptr[current_ix] = \ static_cast( \ @@ -472,8 +485,8 @@ Tensor& dequantize_per_channel_out( for (size_t channel_ix = 0; channel_ix < input.size(axis); ++channel_ix) { \ float _scale = get_scale(scale, channel_ix); \ int64_t _zero_point = 0; \ - if (zero_point_data != nullptr) { \ - _zero_point = zero_point_data[channel_ix]; \ + if (zero_point_tensor != nullptr) { \ + _zero_point = get_zero_point(*zero_point_tensor, channel_ix); \ } \ auto* out_data_ptr = out.mutable_data_ptr(); \ const auto* input_data_ptr = input.const_data_ptr(); \ diff --git a/kernels/quantized/test/op_dequantize_test.cpp b/kernels/quantized/test/op_dequantize_test.cpp index 407e04d9ff8..ee6caabcdba 100644 --- a/kernels/quantized/test/op_dequantize_test.cpp +++ b/kernels/quantized/test/op_dequantize_test.cpp @@ -317,3 +317,59 @@ TEST(OpDequantizeOutTest, DequantizePerChannel) { test_per_channel_dtype(); test_per_channel_dtype(); } + +// The schema allows an Int zero_point tensor as well as a Long one. Reading an +// Int tensor as int64 reinterprets pairs of channels as a single value and runs +// off the end of the buffer, which silently corrupted every channel. +template +void test_per_channel_int_zero_point() { + TensorFactory tf; + TensorFactory tf_double; + TensorFactory tf_int; + TensorFactory tfo; + + Tensor scale = tf_double.make({4}, {0.5, 0.75, 1, 2}); + Tensor zero_point = tf_int.make({4}, {30, 50, 60, 90}); + int64_t quant_min = 0; + int64_t quant_max = 127; + + // Multi-dimensional input, channel axis 0. + Tensor input = tf.full({4, 2}, 100); + Tensor out = tfo.zeros({4, 2}); + Tensor expected = tfo.make({4, 2}, {35, 35, 37.5, 37.5, 40, 40, 20, 20}); + dequantize_per_channel_out( + input, + scale, + zero_point, + /*axis=*/0, + quant_min, + quant_max, + DTYPE, + optional(), + out); + EXPECT_TENSOR_EQ(out, expected); + + // Single-dimensional input takes a separate branch in the kernel. + input = tf.make({4}, {100, 100, 100, 100}); + out = tfo.zeros({4}); + expected = tfo.make({4}, {35, 37.5, 40, 20}); + dequantize_per_channel_out( + input, + scale, + zero_point, + /*axis=*/0, + quant_min, + quant_max, + DTYPE, + optional(), + out); + EXPECT_TENSOR_EQ(out, expected); +} + +TEST(OpDequantizeOutTest, DequantizePerChannelIntZeroPoint) { + et_pal_init(); + test_per_channel_int_zero_point(); + test_per_channel_int_zero_point(); + test_per_channel_int_zero_point(); + test_per_channel_int_zero_point(); +}