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/mul.h

@@ -24,6 +24,9 @@ namespace tflite {
 
 namespace reference_ops {
 
+// Maximum dimension supported by the broadcast mul operation.
+constexpr int kMaxMulBroadcastDim = 6;
+
 // Element-wise mul that can often be used for inner loop of broadcast Mul as
 // well as the non-broadcast Mul.
 inline void MulElementwise(int size, const ArithmeticParams& params,
@@ -88,128 +91,174 @@ inline void Mul(const ArithmeticParams& params,
   MulElementwise(flat_size, params, input1_data, input2_data, output_data);
 }
 
-inline void BroadcastMul4DSlow(const ArithmeticParams& params,
+template <typename T, typename F>
+void BroadcastMulRecursiveDimensions(
+    const ArithmeticParams& params, int dimension, const T* input1_data,
+    const T* input2_data, T* output_data, size_t* input1_offset_p,
+    size_t* input2_offset_p, size_t* output_offset,
+    const NdArrayDesc<kMaxMulBroadcastDim>& desc1,
+    const NdArrayDesc<kMaxMulBroadcastDim>& desc2,
+    const int32_t extended_output_shape_dims[kMaxMulBroadcastDim],
+    F binary_func) {
+  if (dimension == kMaxMulBroadcastDim - 1) {
+    for (int c = 0; c < extended_output_shape_dims[dimension]; ++c) {
+      const T input1_val = input1_data[*input1_offset_p];
+      const T input2_val = input2_data[*input2_offset_p];
+      output_data[*output_offset] = binary_func(params, input1_val, input2_val);
+      *input1_offset_p += desc1.strides[dimension];
+      *input2_offset_p += desc2.strides[dimension];
+      ++(*output_offset);
+    }
+  } else {
+    for (int a = 0; a < extended_output_shape_dims[dimension]; ++a) {
+      size_t input1_offset_c = *input1_offset_p;
+      size_t input2_offset_c = *input2_offset_p;
+      BroadcastMulRecursiveDimensions(
+          params, dimension + 1, input1_data, input2_data, output_data,
+          &input1_offset_c, &input2_offset_c, output_offset, desc1, desc2,
+          extended_output_shape_dims, binary_func);
+      *input1_offset_p += desc1.strides[dimension];
+      *input2_offset_p += desc2.strides[dimension];
+    }
+  }
+}
+
+inline void BroadcastMul6DSlow(const ArithmeticParams& params,
                                const RuntimeShape& input1_shape,
                                const uint8_t* input1_data,
                                const RuntimeShape& input2_shape,
                                const uint8_t* input2_data,
                                const RuntimeShape& output_shape,
                                uint8_t* output_data) {
-  NdArrayDesc<4> desc1;
-  NdArrayDesc<4> desc2;
+  NdArrayDesc<kMaxMulBroadcastDim> desc1;
+  NdArrayDesc<kMaxMulBroadcastDim> desc2;
   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<uint8_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 = 0;
+  size_t input2_offset = 0;
+  size_t output_offset = 0;
+  BroadcastMulRecursiveDimensions(
+      params, 0, input1_data, input2_data, output_data, &input1_offset,
+      &input2_offset, &output_offset, desc1, desc2, extended_output_shape_dims,
+      [](const ArithmeticParams& params, const uint8_t input1_val,
+         const uint8_t input2_val) {
+        const int32_t offsetted_input1_val = params.input1_offset + input1_val;
+        const int32_t offsetted_input2_val = params.input2_offset + input2_val;
+        const int32_t unclamped_result =
+            params.output_offset +
+            MultiplyByQuantizedMultiplier(
+                offsetted_input1_val * offsetted_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));
+        return static_cast<uint8_t>(clamped_output);
+      });
 }
 
 template <typename T,
-          // For unquantized mul on small integers, explictly set to true.
+          // For unquantized mul on small integers, explicitly set to true.
           bool enable_for_short_integers = false>
 inline typename std::enable_if<
     !is_small_integer<T>::value || enable_for_short_integers, void>::type
-BroadcastMul4DSlow(const ArithmeticParams& params,
+BroadcastMul6DSlow(const ArithmeticParams& params,
                    const RuntimeShape& unextended_input1_shape,
                    const T* input1_data,
                    const RuntimeShape& unextended_input2_shape,
                    const T* input2_data,
                    const RuntimeShape& unextended_output_shape,
                    T* output_data) {
-  T output_activation_min;
-  T output_activation_max;
-  GetActivationParams(params, &output_activation_min, &output_activation_max);
-
-  TFLITE_DCHECK_LE(unextended_input1_shape.DimensionsCount(), 4);
-  TFLITE_DCHECK_LE(unextended_input2_shape.DimensionsCount(), 4);
-  TFLITE_DCHECK_LE(unextended_output_shape.DimensionsCount(), 4);
-  const RuntimeShape output_shape =
-      RuntimeShape::ExtendedShape(4, unextended_output_shape);
-
-  NdArrayDesc<4> desc1;
-  NdArrayDesc<4> desc2;
+  TFLITE_DCHECK_LE(unextended_input1_shape.DimensionsCount(), 6);
+  TFLITE_DCHECK_LE(unextended_input2_shape.DimensionsCount(), 6);
+  TFLITE_DCHECK_LE(unextended_output_shape.DimensionsCount(), 6);
+  NdArrayDesc<kMaxMulBroadcastDim> desc1;
+  NdArrayDesc<kMaxMulBroadcastDim> desc2;
   NdArrayDescsForElementwiseBroadcast(unextended_input1_shape,
                                       unextended_input2_shape, &desc1, &desc2);
+  const RuntimeShape extended_output_shape =
+      RuntimeShape::ExtendedShape(kMaxMulBroadcastDim, unextended_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));
 
   // In Tensorflow, the dimensions are canonically named (batch_number, row,
   // col, channel), with extents (batches, height, width, depth), with the
-  // trailing dimension changing most rapidly (channels has the smallest stride,
-  // typically 1 element).
+  // trailing dimension changing most rapidly (channels has the smallest
+  // stride, typically 1 element).
   //
   // In generated C code, we store arrays with the dimensions reversed. The
   // first dimension has smallest stride.
   //
   // 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.
-  for (int b = 0; b < output_shape.Dims(0); ++b) {
-    for (int y = 0; y < output_shape.Dims(1); ++y) {
-      for (int x = 0; x < output_shape.Dims(2); ++x) {
-        for (int c = 0; c < output_shape.Dims(3); ++c) {
-          output_data[Offset(output_shape, b, y, x, c)] =
-              ActivationFunctionWithMinMax<T>(
-                  input1_data[SubscriptToIndex(desc1, b, y, x, c)] *
-                      input2_data[SubscriptToIndex(desc2, b, y, x, c)],
-                  output_activation_min, output_activation_max);
-        }
-      }
-    }
-  }
+  // nesting loops such that the innermost loop has the smallest stride for
+  // the best cache behavior.
+  size_t input1_offset = 0;
+  size_t input2_offset = 0;
+  size_t output_offset = 0;
+  BroadcastMulRecursiveDimensions(
+      params, 0, input1_data, input2_data, output_data, &input1_offset,
+      &input2_offset, &output_offset, desc1, desc2, extended_output_shape_dims,
+      [](const ArithmeticParams& params, const T input1_val,
+         const T input2_val) {
+        T output_activation_min;
+        T output_activation_max;
+        GetActivationParams(params, &output_activation_min,
+                            &output_activation_max);
+        return ActivationFunctionWithMinMax<T>(input1_val * input2_val,
+                                               output_activation_min,
+                                               output_activation_max);
+      });
 }
 
-inline void BroadcastMul4DSlow(const ArithmeticParams& params,
+inline void BroadcastMul6DSlow(const ArithmeticParams& params,
                                const RuntimeShape& unextended_input1_shape,
                                const std::complex<float>* input1_data,
                                const RuntimeShape& unextended_input2_shape,
                                const std::complex<float>* input2_data,
                                const RuntimeShape& unextended_output_shape,
                                std::complex<float>* output_data) {
-  TFLITE_DCHECK_LE(unextended_input1_shape.DimensionsCount(), 4);
-  TFLITE_DCHECK_LE(unextended_input2_shape.DimensionsCount(), 4);
-  TFLITE_DCHECK_LE(unextended_output_shape.DimensionsCount(), 4);
-  const RuntimeShape output_shape =
-      RuntimeShape::ExtendedShape(4, unextended_output_shape);
-
-  NdArrayDesc<4> desc1;
-  NdArrayDesc<4> desc2;
+  TFLITE_DCHECK_LE(unextended_input1_shape.DimensionsCount(), 6);
+  TFLITE_DCHECK_LE(unextended_input2_shape.DimensionsCount(), 6);
+  TFLITE_DCHECK_LE(unextended_output_shape.DimensionsCount(), 6);
+
+  NdArrayDesc<kMaxMulBroadcastDim> desc1;
+  NdArrayDesc<kMaxMulBroadcastDim> desc2;
   NdArrayDescsForElementwiseBroadcast(unextended_input1_shape,
                                       unextended_input2_shape, &desc1, &desc2);
+  const RuntimeShape extended_output_shape =
+      RuntimeShape::ExtendedShape(kMaxMulBroadcastDim, unextended_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));
 
-  for (int b = 0; b < output_shape.Dims(0); ++b) {
-    for (int y = 0; y < output_shape.Dims(1); ++y) {
-      for (int x = 0; x < output_shape.Dims(2); ++x) {
-        for (int c = 0; c < output_shape.Dims(3); ++c) {
-          output_data[Offset(output_shape, b, y, x, c)] =
-              input1_data[SubscriptToIndex(desc1, b, y, x, c)] *
-              input2_data[SubscriptToIndex(desc2, b, y, x, c)];
-        }
-      }
-    }
-  }
+  size_t input1_offset = 0;
+  size_t input2_offset = 0;
+  size_t output_offset = 0;
+  BroadcastMulRecursiveDimensions(
+      params, 0, input1_data, input2_data, output_data, &input1_offset,
+      &input2_offset, &output_offset, desc1, desc2, extended_output_shape_dims,
+      [](const ArithmeticParams& params, const std::complex<float> input1_val,
+         const std::complex<float> input2_val) {
+        return input1_val * input2_val;
+      });
+}
+
+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) {
+  return BroadcastMul6DSlow(params, input1_shape, input1_data, input2_shape,
+                            input2_data, output_shape, output_data);
 }
 
 }  // namespace reference_ops

xmos_ai_tools/runtime/include/tensorflow/lite/kernels/internal/reference/reduce.h

@@ -268,11 +268,11 @@ inline bool Mean(const T* input_data, const int* input_dims,
   return true;
 }
 
-template <typename T>
 inline void Mean(const tflite::MeanParams& op_params,
                  const RuntimeShape& unextended_input_shape,
-                 const T* input_data,
-                 const RuntimeShape& unextended_output_shape, T* output_data) {
+                 const float* input_data,
+                 const RuntimeShape& unextended_output_shape,
+                 float* output_data) {
   ruy::profiler::ScopeLabel label("Mean4D");
 
   // Current implementation only supports dimension equals 4 and simultaneous
@@ -312,78 +312,21 @@ inline void Mean(const tflite::MeanParams& op_params,
   }
 }
 
-inline void Mean(const tflite::MeanParams& op_params,
-                 const RuntimeShape& unextended_input_shape,
-                 const uint8_t* input_data, int32_t input_zero_point,
-                 float input_scale, const RuntimeShape& unextended_output_shape,
-                 uint8_t* output_data, int32_t output_zero_point,
-                 float output_scale) {
-  ruy::profiler::ScopeLabel label("Mean4D/Uint8");
-
-  // 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 float 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);
-
-  constexpr int32_t kMinValue = std::numeric_limits<uint8_t>::min();
-  constexpr int32_t kMaxValue = std::numeric_limits<uint8_t>::max();
-
-  float temp = input_zero_point * input_scale / output_scale;
-  temp = temp > 0 ? temp + 0.5f : temp - 0.5f;
-  int32_t bias = output_zero_point - static_cast<int32_t>(temp);
-  double real_scale =
-      static_cast<double>(input_scale / (num_elements_in_axis * output_scale));
-
-  int32_t multiplier;
-  int shift;
-  QuantizeMultiplier(real_scale, &multiplier, &shift);
-  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)];
-        }
-      }
-      acc = MultiplyByQuantizedMultiplier(acc, multiplier, shift);
-      acc += bias;
-      acc = std::min(std::max(acc, kMinValue), kMaxValue);
-      output_data[Offset(output_shape, out_b, 0, 0, out_d)] =
-          static_cast<uint8_t>(acc);
-    }
-  }
-}
-
 // Computes the mean of elements across dimensions given in axis.
 // It does so in two stages, first calculates the sum of elements along the axis
 // then divides it by the number of element in axis for quantized values.
 template <typename T, typename U>
 inline bool QuantizedMeanOrSum(const T* input_data, int32_t input_zero_point,
-                               float input_scale, const int* input_dims,
-                               const int input_num_dims, T* output_data,
-                               int32_t output_zero_point, float output_scale,
+                               const int* input_dims, const int input_num_dims,
+                               T* output_data, int32_t output_multiplier,
+                               int output_shift, int32_t output_zero_point,
                                const int* output_dims,
                                const int output_num_dims, const int* axis,
                                const int num_axis_dimensions, bool keep_dims,
                                int* temp_index, int* resolved_axis, U* temp_sum,
                                bool compute_sum) {
+  const int32_t kMinValue = std::numeric_limits<T>::min();
+  const int32_t kMaxValue = std::numeric_limits<T>::max();
   const bool uint8_case = std::is_same<T, uint8_t>::value;
   const bool int16_case = std::is_same<T, int16_t>::value;
   if (uint8_case) {
@@ -430,40 +373,46 @@ inline bool QuantizedMeanOrSum(const T* input_data, int32_t input_zero_point,
   }
 
   // Calculate mean by dividing output_data by num of aggregated element.
-  size_t num_elements_in_axis = 1;
+  int64_t num_elements_in_axis = 1;
   for (int idx = 0; idx < num_resolved_axis; ++idx) {
     size_t current = static_cast<size_t>(input_dims[resolved_axis[idx]]);
     // Overflow prevention.
-    if (current > (std::numeric_limits<size_t>::max() / num_elements_in_axis)) {
+    if (current > static_cast<size_t>(std::numeric_limits<int64_t>::max() /
+                                      num_elements_in_axis)) {
       return false;
     }
     num_elements_in_axis *= current;
   }
 
-  if (num_elements_in_axis > 0) {
-    const float scale = input_scale / output_scale;
-    if (compute_sum) {
-      // TODO(b/116341117): Eliminate float and do this completely in 8bit.
-      const float bias = -input_zero_point * scale * num_elements_in_axis;
-      for (size_t idx = 0; idx < num_outputs; ++idx) {
-        const U value =
-            static_cast<U>(TfLiteRound(temp_sum[idx] * scale + bias)) +
-            output_zero_point;
-        output_data[idx] = static_cast<T>(value);
-      }
-    } else {
-      const float bias = -input_zero_point * scale;
-      for (size_t idx = 0; idx < num_outputs; ++idx) {
-        float float_mean = static_cast<float>(temp_sum[idx]) /
-                           static_cast<float>(num_elements_in_axis);
-        float result = TfLiteMin(
-            TfLiteRound(float_mean * scale + bias) + output_zero_point,
-            static_cast<float>(std::numeric_limits<T>::max()));
-        result = TfLiteMax(result,
-                           static_cast<float>(std::numeric_limits<T>::min()));
-        output_data[idx] = static_cast<T>(result);
-      }
-    }
+  if (num_elements_in_axis == 0) {
+    return true;
+  }
+
+  // Readapt output rescaling when calculating the mean to integrate a
+  // 1/num_elements_in_axis multiplier.
+  if (!compute_sum) {
+    TFLITE_DCHECK_GE(num_elements_in_axis, 0);
+    int shift =
+        63 - CountLeadingZeros(static_cast<uint64_t>(num_elements_in_axis));
+    // To avoid any overflow risk 'shift' should be <= 32 and to satisfy
+    // 'MultiplyByQuantizedMultiplier' pre-conditions 'output_shift - shift'
+    // should be >= -31. Clamp the value at the price of some precision loss.
+    shift = std::min(shift, 32);
+    shift = std::min(shift, 31 + output_shift);
+    output_multiplier = static_cast<int32_t>(
+        (static_cast<int64_t>(output_multiplier) << shift) /
+        num_elements_in_axis);
+    output_shift = output_shift - shift;
+  }
+
+  for (size_t idx = 0; idx < num_outputs; ++idx) {
+    const U shifted_sum =
+        static_cast<U>(temp_sum[idx] - input_zero_point * num_elements_in_axis);
+    int32_t output = MultiplyByQuantizedMultiplier(
+                         shifted_sum, output_multiplier, output_shift) +
+                     output_zero_point;
+    output = std::min(std::max(output, kMinValue), kMaxValue);
+    output_data[idx] = static_cast<T>(output);
   }
   return true;
 }
@@ -478,8 +427,8 @@ inline bool QuantizedMeanOrSumExtraArgs(
     bool keep_dims, int* temp_index, int* resolved_axis, U* temp_sum,
     bool compute_sum) {
   return QuantizedMeanOrSum<T, U>(
-      input_data, input_zero_point, input_scale, input_dims, input_num_dims,
-      output_data, output_zero_point, output_scale, output_dims,
+      input_data, input_zero_point, input_dims, input_num_dims, output_data,
+      output_multiplier, output_shift, output_zero_point, output_dims,
       output_num_dims, axis, num_axis_dimensions, keep_dims, temp_index,
       resolved_axis, temp_sum, compute_sum);
 }

xmos_ai_tools/runtime/include/tensorflow/lite/kernels/internal/reference/resize_bilinear.h

@@ -212,9 +212,14 @@ inline void ResizeBilinearInteger(
               (input_y - (1 << 10) * y0) * (input_x - (1 << 10) * x0);
           const int64_t output_20 =
               output_20_ll + output_20_lu + output_20_rl + output_20_ru;
+#if TFLITE_SINGLE_ROUNDING
+          const int64_t round = 1 << 19;
+          const T interpolation = static_cast<T>((output_20 + round) >> 20);
+#else
           const int64_t round = (output_20 > 0) ? (1 << 19) : -(1 << 19);
           const T interpolation =
               static_cast<T>((output_20 + round) / (1 << 20));
+#endif  // TFLITE_SINGLE_ROUNDING
           output_data[Offset(output_shape, b, y, x, c)] = interpolation;
         }
       }