xmos-ai-tools 1.1.2.dev216__py3-none-macosx_11_0_arm64.whl → 1.1.2.dev236__py3-none-macosx_11_0_arm64.whl

xmos_ai_tools/runtime/include/tensorflow/lite/kernels/internal/reference/integer_ops/add.h

@@ -16,6 +16,7 @@ limitations under the License.
 #define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_INTEGER_OPS_ADD_H_
 
 #include <algorithm>
+#include <cstddef>
 #include <limits>
 
 #include "tensorflow/lite/kernels/internal/common.h"
@@ -35,12 +36,35 @@ inline void CheckArithmeticParams(const ArithmeticParams& params) {
   TFLITE_DCHECK_LE(-params.input2_offset, std::numeric_limits<int8_t>::max());
 }
 
+// TODO: b/270589088 - move to a more appropriate file (b/270589088#comment2)
 template <typename T>
-inline void ElementWise(
-    int size, const ArithmeticParams& params, const T* input1_data,
-    const T* input2_data, T* output_data,
-    void (*check_arithmetic_params)(const ArithmeticParams&),
-    T (*binary_func)(T, T, const ArithmeticParams&)) {
+void BroadcastInput1(int size, const ArithmeticParams& params,
+                     const T* input1_data, const T* input2_data, T* output_data,
+                     void (*check_arithmetic_params)(const ArithmeticParams&),
+                     T (*binary_func)(T, T, const ArithmeticParams&)) {
+  CheckArithmeticParams(params);
+  for (int i = 0; i < size; ++i) {
+    output_data[i] = binary_func(input1_data[0], input2_data[i], params);
+  }
+}
+
+template <typename T>
+void BroadcastInput2(int size, const ArithmeticParams& params,
+                     const T* input1_data, const T* input2_data, T* output_data,
+                     void (*check_arithmetic_params)(const ArithmeticParams&),
+                     T (*binary_func)(T, T, const ArithmeticParams&)) {
+  CheckArithmeticParams(params);
+  for (int i = 0; i < size; ++i) {
+    output_data[i] = binary_func(input1_data[i], input2_data[0], params);
+  }
+}
+
+// TODO: b/270589088 - move to a more appropriate file (b/270589088#comment2)
+template <typename T>
+void ElementWise(int size, const ArithmeticParams& params, const T* input1_data,
+                 const T* input2_data, T* output_data,
+                 void (*check_arithmetic_params)(const ArithmeticParams&),
+                 T (*binary_func)(T, T, const ArithmeticParams&)) {
   CheckArithmeticParams(params);
   for (int i = 0; i < size; ++i) {
     output_data[i] = binary_func(input1_data[i], input2_data[i], params);
@@ -48,19 +72,66 @@ inline void ElementWise(
 }
 
 template <typename T>
-inline void BroadcastBinaryFunction6DSlow(
+inline void BroadcastAddRecursiveDimensions(
+    const ArithmeticParams& params, int dimension, size_t* input1_offset_p,
+    size_t* input2_offset_p, size_t* output_offset,
+    size_t* compressed_input1_stride, size_t* compressed_input2_stride,
+    size_t* compressed_output_shape, const T* input1_data, const T* input2_data,
+    T* output_data, void (*check_arithmetic_params)(const ArithmeticParams&),
+    T (*binary_func)(T, T, const ArithmeticParams&)) {
+  if (dimension > 0) {
+    for (size_t c = 0; c < compressed_output_shape[dimension]; ++c) {
+      size_t input1_offset_c = *input1_offset_p;
+      size_t input2_offset_c = *input2_offset_p;
+      BroadcastAddRecursiveDimensions(
+          params, dimension - 1, &input1_offset_c, &input2_offset_c,
+          output_offset, compressed_input1_stride, compressed_input2_stride,
+          compressed_output_shape, input1_data, input2_data, output_data,
+          check_arithmetic_params, binary_func);
+      *input1_offset_p += compressed_input1_stride[dimension];
+      *input2_offset_p += compressed_input2_stride[dimension];
+    }
+  } else {
+    TFLITE_DCHECK(dimension == 0);
+    bool input1_is_broadcast = compressed_input1_stride[dimension] == 0;
+    bool input2_is_broadcast = compressed_input2_stride[dimension] == 0;
+    TFLITE_DCHECK(!(input1_is_broadcast && input2_is_broadcast));
+    const T* input1_data_ptr = input1_data + *input1_offset_p;
+    const T* input2_data_ptr = input2_data + *input2_offset_p;
+    T* output_data_ptr = output_data + *output_offset;
+    if (input1_is_broadcast) {
+      // input1 is broadcast.
+      BroadcastInput1<T>(compressed_output_shape[dimension], params,
+                         input1_data_ptr, input2_data_ptr, output_data_ptr,
+                         check_arithmetic_params, binary_func);
+      *input2_offset_p += compressed_output_shape[dimension];
+    } else if (input2_is_broadcast) {
+      // input2 is broadcast.
+      BroadcastInput2<T>(compressed_output_shape[dimension], params,
+                         input1_data_ptr, input2_data_ptr, output_data_ptr,
+                         check_arithmetic_params, binary_func);
+      *input1_offset_p += compressed_output_shape[dimension];
+    } else {
+      // Add element-wise.
+      ElementWise<T>(compressed_output_shape[dimension], params,
+                     input1_data_ptr, input2_data_ptr, output_data_ptr,
+                     check_arithmetic_params, binary_func);
+      *input1_offset_p += compressed_output_shape[dimension];
+      *input2_offset_p += compressed_output_shape[dimension];
+    }
+    *output_offset += compressed_output_shape[dimension];
+  }
+}
+
+// TODO: b/270589088 - move to a more appropriate file. (b/270589088#comment2)
+template <typename T>
+void BroadcastBinaryFunction6DSlow(
     const ArithmeticParams& params, const RuntimeShape& input1_shape,
     const T* input1_data, const RuntimeShape& input2_shape,
-    const T* input2_data, const RuntimeShape& output_shape,
-    T* output_data,
+    const T* input2_data, const RuntimeShape& output_shape, T* output_data,
     void (*check_arithmetic_params)(const ArithmeticParams&),
     T (*binary_func)(T, T, const ArithmeticParams&)) {
-  NdArrayDesc<6> desc1;
-  NdArrayDesc<6> desc2;
-  NdArrayDescsForElementwiseBroadcast(input1_shape, input2_shape, &desc1,
-                                      &desc2);
-  const RuntimeShape extended_output_shape =
-      RuntimeShape::ExtendedShape(6, output_shape);
+  constexpr int kMaxBroadcastDim = 6;
 
   // In Tensorflow, the dimensions are canonically named (batch_number, row,
   // col, channel), with extents (batches, height, width, depth), with the
@@ -73,65 +144,25 @@ inline void BroadcastBinaryFunction6DSlow(
   // We name our variables by their Tensorflow convention, but generate C code
   // nesting loops such that the innermost loop has the smallest stride for the
   // best cache behavior.
-  size_t input1_offset_a = 0;
-  size_t input2_offset_a = 0;
-  size_t output_offset_a = 0;
-  for (int a = 0; a < extended_output_shape.Dims(0); ++a) {
-    size_t input1_offset_d = input1_offset_a;
-    size_t input2_offset_d = input2_offset_a;
-    size_t output_offset_d = output_offset_a;
-    for (int d = 0; d < extended_output_shape.Dims(1); ++d) {
-      size_t input1_offset_b = input1_offset_d;
-      size_t input2_offset_b = input2_offset_d;
-      size_t output_offset_b = output_offset_d;
-      for (int b = 0; b < extended_output_shape.Dims(2); ++b) {
-        size_t input1_offset_y = input1_offset_b;
-        size_t input2_offset_y = input2_offset_b;
-        size_t output_offset_y = output_offset_b;
-        for (int y = 0; y < extended_output_shape.Dims(3); ++y) {
-          size_t input1_offset_x = input1_offset_y;
-          size_t input2_offset_x = input2_offset_y;
-          size_t output_offset_x = output_offset_y;
-          for (int x = 0; x < extended_output_shape.Dims(4); ++x) {
-            size_t input1_offset_c = input1_offset_x;
-            size_t input2_offset_c = input2_offset_x;
-            size_t output_offset_c = output_offset_x;
-            for (int c = 0; c < extended_output_shape.Dims(5); ++c) {
-              output_data[output_offset_c] =
-                  binary_func(input1_data[input1_offset_c],
-                              input2_data[input2_offset_c], params);
-              input1_offset_c += desc1.strides[5];
-              input2_offset_c += desc2.strides[5];
-              ++output_offset_c;
-            }
-            input1_offset_x += desc1.strides[4];
-            input2_offset_x += desc2.strides[4];
-            output_offset_x += extended_output_shape.Dims(5);
-          }
-          input1_offset_y += desc1.strides[3];
-          input2_offset_y += desc2.strides[3];
-          output_offset_y +=
-              extended_output_shape.Dims(4) * extended_output_shape.Dims(5);
-        }
-        input1_offset_b += desc1.strides[2];
-        input2_offset_b += desc2.strides[2];
-        output_offset_b += extended_output_shape.Dims(3) *
-                           extended_output_shape.Dims(4) *
-                           extended_output_shape.Dims(5);
-      }
-      input1_offset_d += desc1.strides[1];
-      input2_offset_d += desc2.strides[1];
-      output_offset_d +=
-          extended_output_shape.Dims(2) * extended_output_shape.Dims(3) *
-          extended_output_shape.Dims(4) * extended_output_shape.Dims(5);
-    }
-    input1_offset_a += desc1.strides[0];
-    input2_offset_a += desc2.strides[0];
-    output_offset_a +=
-        extended_output_shape.Dims(1) * extended_output_shape.Dims(2) *
-        extended_output_shape.Dims(3) * extended_output_shape.Dims(4) *
-        extended_output_shape.Dims(5);
+  size_t compressed_input1_stride[kMaxBroadcastDim];
+  size_t compressed_input2_stride[kMaxBroadcastDim];
+  size_t compressed_output_shape[kMaxBroadcastDim];
+  bool broadcastable_shape = ReduceDimensionsForBroadcast<kMaxBroadcastDim>(
+      input1_shape, input2_shape, compressed_input1_stride,
+      compressed_input2_stride, compressed_output_shape);
+  // Skip broadcasting for degenerate shapes.
+  if (!broadcastable_shape) {
+    return;
   }
+
+  size_t input1_offset = 0;
+  size_t input2_offset = 0;
+  size_t output_offset = 0;
+  BroadcastAddRecursiveDimensions(
+      params, kMaxBroadcastDim - 1, &input1_offset, &input2_offset,
+      &output_offset, compressed_input1_stride, compressed_input2_stride,
+      compressed_output_shape, input1_data, input2_data, output_data,
+      check_arithmetic_params, binary_func);
 }
 
 template <typename T>

xmos_ai_tools/runtime/include/tensorflow/lite/kernels/internal/reference/integer_ops/conv.h

@@ -63,8 +63,10 @@ inline void ConvPerChannel(
   const int filter_width = filter_shape.Dims(2);
   const int filter_input_depth = filter_shape.Dims(3);
   const int groups = input_depth / filter_input_depth;
+  TFLITE_DCHECK_NE(groups, 0);
   TFLITE_DCHECK_EQ(input_depth % filter_input_depth, 0);
   const int filters_per_group = output_depth / groups;
+  TFLITE_DCHECK_NE(filters_per_group, 0);
   const int output_height = output_shape.Dims(1);
   const int output_width = output_shape.Dims(2);
   for (int batch = 0; batch < batches; ++batch) {

xmos_ai_tools/runtime/include/tensorflow/lite/kernels/internal/reference/integer_ops/mean.h

@@ -1,10 +1,10 @@
-/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
 
-    http://www.apache.org/licenses/LICENSE-2.0
+http://www.apache.org/licenses/LICENSE-2.0
 
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
@@ -15,65 +15,4 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_INTEGER_OPS_MEAN_H_
 #define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_INTEGER_OPS_MEAN_H_
 
-#include <algorithm>
-
-#include "tensorflow/lite/kernels/internal/common.h"
-
-namespace tflite {
-namespace reference_integer_ops {
-
-template <typename integer_type>
-inline void Mean(const tflite::MeanParams& op_params, int32_t multiplier,
-                 int32_t shift, const RuntimeShape& unextended_input_shape,
-                 const integer_type* input_data, int32_t input_zero_point,
-                 const RuntimeShape& unextended_output_shape,
-                 integer_type* output_data, int32_t output_zero_point) {
-  // Current implementation only supports dimension equals 4 and simultaneous
-  // reduction over width and height.
-  TFLITE_CHECK_EQ(unextended_input_shape.DimensionsCount(), 4);
-  TFLITE_CHECK_LE(unextended_output_shape.DimensionsCount(), 4);
-  const RuntimeShape input_shape =
-      RuntimeShape::ExtendedShape(4, unextended_input_shape);
-  const RuntimeShape output_shape =
-      RuntimeShape::ExtendedShape(4, unextended_output_shape);
-  const int output_batch = output_shape.Dims(0);
-  const int output_height = output_shape.Dims(1);
-  const int output_width = output_shape.Dims(2);
-  const int output_depth = output_shape.Dims(3);
-  const int input_height = input_shape.Dims(1);
-  const int input_width = input_shape.Dims(2);
-  const int num_elements_in_axis = input_width * input_height;
-
-  TFLITE_CHECK_EQ(op_params.axis_count, 2);
-  TFLITE_CHECK((op_params.axis[0] == 1 && op_params.axis[1] == 2) ||
-               (op_params.axis[0] == 2 && op_params.axis[1] == 1));
-  TFLITE_CHECK_EQ(output_height, 1);
-  TFLITE_CHECK_EQ(output_width, 1);
-
-  static constexpr int32_t kMinInt = std::numeric_limits<integer_type>::min();
-  static constexpr int32_t kMaxInt = std::numeric_limits<integer_type>::max();
-
-  for (int out_b = 0; out_b < output_batch; ++out_b) {
-    for (int out_d = 0; out_d < output_depth; ++out_d) {
-      int32_t acc = 0;
-      for (int in_h = 0; in_h < input_height; ++in_h) {
-        for (int in_w = 0; in_w < input_width; ++in_w) {
-          acc += input_data[Offset(input_shape, out_b, in_h, in_w, out_d)] -
-                 input_zero_point;
-        }
-      }
-      acc = MultiplyByQuantizedMultiplier(acc, multiplier, shift);
-      acc = acc > 0 ? (acc + num_elements_in_axis / 2) / num_elements_in_axis
-                    : (acc - num_elements_in_axis / 2) / num_elements_in_axis;
-      acc += output_zero_point;
-      acc = std::min(std::max(acc, kMinInt), kMaxInt);
-      output_data[Offset(output_shape, out_b, 0, 0, out_d)] =
-          static_cast<integer_type>(acc);
-    }
-  }
-}
-
-}  // namespace reference_integer_ops
-}  // namespace tflite
-
 #endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_INTEGER_OPS_MEAN_H_

xmos_ai_tools/runtime/include/tensorflow/lite/kernels/internal/reference/integer_ops/mul.h

@@ -24,6 +24,9 @@ limitations under the License.
 namespace tflite {
 namespace reference_integer_ops {
 
+// Maximum dimension supported by the broadcast mul operation.
+constexpr int kMaxMulBroadcastDim = 6;
+
 template <typename InputType, typename OutputType>
 void MulElementwise(int size, const ArithmeticParams& params,
                     const InputType* input1_data, const InputType* input2_data,
@@ -88,46 +91,104 @@ inline void Mul(const ArithmeticParams& params,
 }
 
 template <typename T>
-inline void BroadcastMul4DSlow(
+inline void BroadcastMul6DSlow(
     const ArithmeticParams& params, const RuntimeShape& input1_shape,
     const T* input1_data, const RuntimeShape& input2_shape,
     const T* input2_data, const RuntimeShape& output_shape, T* output_data) {
-  ruy::profiler::ScopeLabel label("BroadcastMul4DSlow");
+  ruy::profiler::ScopeLabel label("BroadcastMul6DSlow");
 
-  NdArrayDesc<4> desc1;
-  NdArrayDesc<4> desc2;
+  NdArrayDesc<kMaxMulBroadcastDim> desc1;
+  NdArrayDesc<kMaxMulBroadcastDim> desc2;
   // The input shapes are extended as part of NdArrayDesc initialization.
   NdArrayDescsForElementwiseBroadcast(input1_shape, input2_shape, &desc1,
                                       &desc2);
   const RuntimeShape extended_output_shape =
-      RuntimeShape::ExtendedShape(4, output_shape);
-
-  for (int b = 0; b < extended_output_shape.Dims(0); ++b) {
-    for (int y = 0; y < extended_output_shape.Dims(1); ++y) {
-      for (int x = 0; x < extended_output_shape.Dims(2); ++x) {
-        for (int c = 0; c < extended_output_shape.Dims(3); ++c) {
-          const int32_t input1_val =
-              params.input1_offset +
-              input1_data[SubscriptToIndex(desc1, b, y, x, c)];
-          const int32_t input2_val =
-              params.input2_offset +
-              input2_data[SubscriptToIndex(desc2, b, y, x, c)];
-          const int32_t unclamped_result =
-              params.output_offset +
-              MultiplyByQuantizedMultiplier(input1_val * input2_val,
-                                            params.output_multiplier,
-                                            params.output_shift);
-          const int32_t clamped_output = std::min(
-              params.quantized_activation_max,
-              std::max(params.quantized_activation_min, unclamped_result));
-          output_data[Offset(extended_output_shape, b, y, x, c)] =
-              static_cast<T>(clamped_output);
+      RuntimeShape::ExtendedShape(kMaxMulBroadcastDim, output_shape);
+  // Cache output shape dimensions.
+  int32_t extended_output_shape_dims[kMaxMulBroadcastDim];
+  std::memcpy(extended_output_shape_dims, extended_output_shape.DimsData(),
+              sizeof(extended_output_shape_dims));
+
+  size_t input1_offset_a = 0;
+  size_t input2_offset_a = 0;
+  size_t output_offset_a = 0;
+  for (int a = 0; a < extended_output_shape_dims[0]; ++a) {
+    size_t input1_offset_d = input1_offset_a;
+    size_t input2_offset_d = input2_offset_a;
+    size_t output_offset_d = output_offset_a;
+    for (int d = 0; d < extended_output_shape_dims[1]; ++d) {
+      size_t input1_offset_b = input1_offset_d;
+      size_t input2_offset_b = input2_offset_d;
+      size_t output_offset_b = output_offset_d;
+      for (int b = 0; b < extended_output_shape_dims[2]; ++b) {
+        size_t input1_offset_y = input1_offset_b;
+        size_t input2_offset_y = input2_offset_b;
+        size_t output_offset_y = output_offset_b;
+        for (int y = 0; y < extended_output_shape_dims[3]; ++y) {
+          size_t input1_offset_x = input1_offset_y;
+          size_t input2_offset_x = input2_offset_y;
+          size_t output_offset_x = output_offset_y;
+          for (int x = 0; x < extended_output_shape_dims[4]; ++x) {
+            size_t input1_offset_c = input1_offset_x;
+            size_t input2_offset_c = input2_offset_x;
+            size_t output_offset_c = output_offset_x;
+            for (int c = 0; c < extended_output_shape_dims[5]; ++c) {
+              const int32_t input1_val =
+                  params.input1_offset + input1_data[input1_offset_c];
+              const int32_t input2_val =
+                  params.input2_offset + input2_data[input2_offset_c];
+              const int32_t unclamped_result =
+                  params.output_offset +
+                  MultiplyByQuantizedMultiplier(input1_val * input2_val,
+                                                params.output_multiplier,
+                                                params.output_shift);
+              const int32_t clamped_output = std::min(
+                  params.quantized_activation_max,
+                  std::max(params.quantized_activation_min, unclamped_result));
+              output_data[output_offset_c] = static_cast<T>(clamped_output);
+              input1_offset_c += desc1.strides[5];
+              input2_offset_c += desc2.strides[5];
+              ++output_offset_c;
+            }
+            input1_offset_x += desc1.strides[4];
+            input2_offset_x += desc2.strides[4];
+            output_offset_x += extended_output_shape_dims[5];
+          }
+          input1_offset_y += desc1.strides[3];
+          input2_offset_y += desc2.strides[3];
+          output_offset_y +=
+              extended_output_shape_dims[4] * extended_output_shape_dims[5];
         }
+        input1_offset_b += desc1.strides[2];
+        input2_offset_b += desc2.strides[2];
+        output_offset_b += extended_output_shape_dims[3] *
+                           extended_output_shape_dims[4] *
+                           extended_output_shape_dims[5];
       }
+      input1_offset_d += desc1.strides[1];
+      input2_offset_d += desc2.strides[1];
+      output_offset_d +=
+          extended_output_shape_dims[2] * extended_output_shape_dims[3] *
+          extended_output_shape_dims[4] * extended_output_shape_dims[5];
     }
+    input1_offset_a += desc1.strides[0];
+    input2_offset_a += desc2.strides[0];
+    output_offset_a +=
+        extended_output_shape_dims[1] * extended_output_shape_dims[2] *
+        extended_output_shape_dims[3] * extended_output_shape_dims[4] *
+        extended_output_shape_dims[5];
   }
 }
 
+template <typename T>
+inline void BroadcastMul4DSlow(
+    const ArithmeticParams& params, const RuntimeShape& input1_shape,
+    const T* input1_data, const RuntimeShape& input2_shape,
+    const T* input2_data, const RuntimeShape& output_shape, T* output_data) {
+  BroadcastMul6DSlow(params, input1_shape, input1_data, input2_shape,
+                     input2_data, output_shape, output_data);
+}
+
 }  // namespace reference_integer_ops
 }  // namespace tflite
 #endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_INTEGER_OPS_MUL_H_