Check that zero point and scale of inputs and outputs match for Concatenate operator.

The initial checks within VisitConcatenationNode all pass as xnn_define_concatenate is not called when checking if a node may be delegated or not. It is during the second pass that the check fails, causing delegation to fail.

INFO: Created TensorFlow Lite XNNPACK delegate for CPU.
ERROR: failed to delegate CONCATENATION node #31
ERROR: failed to delegate CONCATENATION node #36
ERROR: Node number 55 (TfLiteXNNPackDelegate) failed to prepare.
PiperOrigin-RevId: 518874585

Author: alankelly

tensorflow/lite/delegates/xnnpack/concatenation_tester.cc

@@ -170,16 +170,16 @@ std::vector<char> ConcatenationTester::CreateTfLiteModel(
 
   std::vector<flatbuffers::Offset<Tensor>> tensors;
   for (size_t i = 0; i < NumInputs(); i++) {
-    tensors.push_back(
-        CreateTensor(builder,
-                     builder.CreateVector<int32_t>(InputShape(i).data(),
-                                                   InputShape(i).size()),
-                     tensor_type,
-                     /*buffer=*/0, /*name=*/0,
-                     CreateQuantizationParameters(
-                         builder, /*min=*/0, /*max=*/0,
-                         builder.CreateVector<float>({/*scale=*/1.0f}),
-                         builder.CreateVector<int64_t>({/*zero_point=*/0}))));
+    tensors.push_back(CreateTensor(
+        builder,
+        builder.CreateVector<int32_t>(InputShape(i).data(),
+                                      InputShape(i).size()),
+        tensor_type,
+        /*buffer=*/0, /*name=*/0,
+        CreateQuantizationParameters(
+            builder, /*min=*/0, /*max=*/0,
+            builder.CreateVector<float>({input_scales_[i]}),
+            builder.CreateVector<int64_t>({input_zero_points_[i]}))));
   }
 
   tensors.push_back(CreateTensor(
@@ -189,8 +189,8 @@ std::vector<char> ConcatenationTester::CreateTfLiteModel(
       /*buffer=*/0, /*name=*/0,
       CreateQuantizationParameters(
           builder, /*min=*/0, /*max=*/0,
-          builder.CreateVector<float>({/*scale=*/1.0f}),
-          builder.CreateVector<int64_t>({/*zero_point=*/0}))));
+          builder.CreateVector<float>({output_scale_}),
+          builder.CreateVector<int64_t>({output_zero_point_}))));
 
   std::vector<int32_t> op_inputs;
   for (size_t i = 0; i < NumInputs(); i++) {

tensorflow/lite/delegates/xnnpack/concatenation_tester.h

@@ -53,13 +53,46 @@ class ConcatenationTester {
       }
     }
     input_shapes_ = shapes;
+    input_scales_.resize(shapes.size(), 1);
+    output_scale_ = 1.f;
+    input_zero_points_.resize(shapes.size(), 0);
+    output_zero_point_ = 0;
     return *this;
   }
 
   inline std::vector<int32_t> InputShape(size_t i) const {
     return input_shapes_[i];
   }
 
+  inline ConcatenationTester& InputScales(const std::vector<float> scales) {
+    input_scales_ = scales;
+    return *this;
+  }
+
+  inline float InputScale(size_t i) const { return input_scales_[i]; }
+
+  inline ConcatenationTester& InputZeroPoint(
+      const std::vector<int32_t> zero_points) {
+    input_zero_points_ = zero_points;
+    return *this;
+  }
+
+  inline float InputZeroPoint(size_t i) const { return input_zero_points_[i]; }
+
+  inline ConcatenationTester& OutputScale(float scale) {
+    output_scale_ = scale;
+    return *this;
+  }
+
+  inline float OutputScale() const { return output_scale_; }
+
+  inline ConcatenationTester& OutputZeroPoint(int32_t zero_point) {
+    output_zero_point_ = zero_point;
+    return *this;
+  }
+
+  inline int32_t OutputZeroPoint() const { return output_scale_; }
+
   inline size_t NumInputs() const { return input_shapes_.size(); }
 
   std::vector<int32_t> OutputShape() const {
@@ -86,6 +119,10 @@ class ConcatenationTester {
   int axis_;
   std::vector<int32_t> output_shape_;
   std::vector<std::vector<int32_t>> input_shapes_;
+  std::vector<float> input_scales_;
+  float output_scale_;
+  std::vector<int32_t> input_zero_points_;
+  int32_t output_zero_point_;
 };
 
 }  // namespace xnnpack

tensorflow/lite/delegates/xnnpack/unsigned_quantized_concatenation_test.cc

@@ -324,5 +324,57 @@ TEST(UnsignedQuantizedConcatenation, 4D_of_4) {
   }
 }
 
+TEST(UnsignedQuantizedConcatenation, 2D_2_inputs_differnet_zero_points) {
+  std::unique_ptr<TfLiteDelegate, decltype(&TfLiteXNNPackDelegateDelete)>
+      xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
+                       TfLiteXNNPackDelegateDelete);
+
+  std::random_device random_device;
+  auto rng = std::mt19937(random_device());
+  auto shape_rng =
+      std::bind(std::uniform_int_distribution<int32_t>(2, 10), std::ref(rng));
+
+  for (int i = -1; i < 2; i++) {
+    // All dimensions must be the same, except for axis.
+    const std::vector<int32_t> shape1({shape_rng(), shape_rng()});
+    auto shape2 = SameShapeDifferentAxis(shape1, i, shape_rng());
+
+    // clang-format off
+    ConcatenationTester()
+        .InputShapes({shape1, shape2})
+        .Axis(i)
+        .InputZeroPoint({2, 3})
+        .OutputZeroPoint(1)
+        .Test(TensorType_UINT8, xnnpack_delegate.get());
+    // clang-format on
+  }
+}
+
+TEST(UnsignedQuantizedConcatenation, 2D_2_inputs_differnet_scales) {
+  std::unique_ptr<TfLiteDelegate, decltype(&TfLiteXNNPackDelegateDelete)>
+      xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
+                       TfLiteXNNPackDelegateDelete);
+
+  std::random_device random_device;
+  auto rng = std::mt19937(random_device());
+  auto shape_rng =
+      std::bind(std::uniform_int_distribution<int32_t>(2, 10), std::ref(rng));
+
+  for (int i = -1; i < 2; i++) {
+    // All dimensions must be the same, except for axis.
+    const std::vector<int32_t> shape1({shape_rng(), shape_rng()});
+    auto shape2 = SameShapeDifferentAxis(shape1, i, shape_rng());
+
+    // clang-format off
+    ConcatenationTester()
+        .InputShapes({shape1, shape2})
+        .Axis(i)
+        .InputScales({2, 3})
+        .OutputScale(1)
+        .Test(TensorType_UINT8, xnnpack_delegate.get());
+    // clang-format on
+  }
+}
+
 }  // namespace xnnpack
 }  // namespace tflite

tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc

@@ -2821,6 +2821,33 @@ class Subgraph {
     if (axis < 0) axis += NumDimensions(&output_tensor);
     int sum_axis = 0;
 
+    if (output_tensor.type == kTfLiteUInt8) {
+      const int32_t zero_point =
+          tensors[node->outputs->data[0]].params.zero_point;
+      const float scale = tensors[node->outputs->data[0]].params.scale;
+      for (int i = 0; i < num_inputs; i++) {
+        if (tensors[node->inputs->data[i]].params.zero_point != zero_point) {
+          TF_LITE_MAYBE_KERNEL_LOG(
+              logging_context,
+              "Mismatching quantization zero point across the %dth input "
+              "(%" PRId32 ") and the output (%" PRId32
+              ") for CONCATENATE operator #%d",
+              i, tensors[node->inputs->data[i]].params.zero_point, zero_point,
+              node_index);
+          return kTfLiteError;
+        }
+        if (tensors[node->inputs->data[i]].params.scale != scale) {
+          TF_LITE_MAYBE_KERNEL_LOG(
+              logging_context,
+              "Mismatching quantization scale across the %dth input (%f) "
+              "and the output (%f) for CONCATENATE operator #%d",
+              i, tensors[node->inputs->data[i]].params.scale, scale,
+              node_index);
+          return kTfLiteError;
+        }
+      }
+    }
+
     for (int i = 0; i < num_inputs; i++) {
       const TfLiteTensor& input_tensor = tensors[node->inputs->data[i]];
       TF_LITE_ENSURE_STATUS(CheckTensorFloat32OrQUInt8Type(