From abd8e60edf09db5f5ba8e7fa2f1fcab0ae0807e1 Mon Sep 17 00:00:00 2001
From: Reilly Grant <[email protected]>
Date: Fri, 20 Mar 2026 17:26:44 -0700
Subject: [PATCH] [M147] Allow redundant reduction axes, and out of bounds reduction axes

TFlite allows the same axis to be specified in a reduction multiple times, and expects it to be treated like set, i.e. these should be deduplicated. Because of the way reshaping works, it can be difficult at delegation or subgraph creation time to determine if this is happening, so we can't just not delegate such ops.

TFlite also allows specifying reduction of a "scalar" with a reduction axis list of {0} (not {}. This is a super annoying behavior, but we need to handle it, because we can't determine if this is happening at delegation time. I think it is reasonable to simply allow out of bounds reduction axes, and just ignore them. This basically treats that reduction axis as an implied dimension of extent 1, which a lot of other things do already (e.g. binary elementwise ops).

As a result of this, I realized that the sum => mean rewrite we do currently is not safe from reshaping. The graph might be equivalent to a mean at the time of construction, but become not so after reshaping. (This would almost certainly be a bug in the client code, but it's a bug that we should not silently fix for the client.)

In addition, I think the sum => mean rewrite probably has negligible impact on performance. I locally modified the layer norm benchmark to use a sum + multiply to implement the two means, and the performance impact is actually an improvement (though very small, and that doesn't really make sense):
```
name                                                                time/op       time/op     vs base
FP32LayerNorm/M:128/N:256/K:512/NormMask:1/process_time/real_time   73.64m ± 2%   74.00m ± 1%       ~ (p=0.394 n=6)
FP32LayerNorm/M:128/N:256/K:512/NormMask:2/process_time/real_time   69.82m ± 1%   69.31m ± 1%  -0.74% (p=0.041 n=6)
FP32LayerNorm/M:128/N:256/K:512/NormMask:3/process_time/real_time   69.82m ± 2%   69.29m ± 1%       ~ (p=0.485 n=6)
FP32LayerNorm/M:128/N:256/K:512/NormMask:4/process_time/real_time   62.28m ± 1%   61.31m ± 1%  -1.56% (p=0.009 n=6)
FP32LayerNorm/M:128/N:256/K:512/NormMask:5/process_time/real_time   61.80m ± 1%   61.77m ± 1%       ~ (p=0.699 n=6)
FP32LayerNorm/M:128/N:256/K:512/NormMask:6/process_time/real_time   60.46m ± 2%   59.42m ± 2%  -1.72% (p=0.041 n=6)
FP32LayerNorm/M:128/N:256/K:512/NormMask:7/process_time/real_time   60.01m ± 2%   59.42m ± 3%       ~ (p=0.240 n=6)
geomean                                                             65.21m        64.71m       -0.76%
```

To fix this issue, I've replaced this rewrite with a `widen_fp16_accumulators` rewrite, that leaves the subgraph mostly intact, but changes the types of the intermediate tensors to fp32, and inserts a convert to fp16 after the division.

(Cherry-picked from commit de3504fd8cfcedf194cd0ae43afb37cdff824aa2.)

PiperOrigin-RevId: 884165426
Bug: 492350403
Change-Id: I74b0d03c6ce57674ca9d16e9f93e7f9f3a37108d
---

diff --git a/src/operators/reduce-nd.c b/src/operators/reduce-nd.c
index fb34c75..818db69 100644
--- a/src/operators/reduce-nd.c
+++ b/src/operators/reduce-nd.c
@@ -142,12 +142,12 @@
     return xnn_status_unsupported_parameter;
   }
 
-  if (num_reduction_axes > num_input_dims) {
+  if (num_reduction_axes > XNN_MAX_TENSOR_DIMS) {
     xnn_log_error(
         "failed to reshape %s operator with %zu reduction axes: the number of "
-        "reduction axes must not exceed the number of input dimensions %zu",
+        "reduction axes must not exceed %d",
         xnn_operator_type_to_string_v2(reduce_op), num_reduction_axes,
-        num_input_dims);
+        XNN_MAX_TENSOR_DIMS);
     return xnn_status_invalid_parameter;
   }
 
@@ -163,19 +163,6 @@
   assert(num_input_dims <= XNN_MAX_TENSOR_DIMS);
   memcpy(normalized_input_shape, input_shape, num_input_dims * sizeof(size_t));
 
-  for (size_t i = 0; i < num_reduction_axes; i++) {
-    const int64_t signed_num_input_dims = (int64_t)num_input_dims;
-    if (signed_num_input_dims <= reduction_axes[i] ||
-        reduction_axes[i] < -signed_num_input_dims) {
-      xnn_log_error(
-          "failed to reshape %s operator with #%zu reduction axis of %" PRIi64
-          ": the index is out of bounds for a %zuD input shape",
-          xnn_operator_type_to_string_v2(reduce_op), i, reduction_axes[i],
-          num_input_dims);
-      return xnn_status_invalid_parameter;
-    }
-  }
-
   size_t normalized_reduction_axes[XNN_MAX_TENSOR_DIMS];
   assert(num_reduction_axes <= XNN_MAX_TENSOR_DIMS);
   for (int i = 0; i < num_reduction_axes; i++) {
@@ -186,15 +173,20 @@
   qsort(normalized_reduction_axes, num_reduction_axes, sizeof(size_t),
         cmp_value_size_t);
 
-  for (size_t i = 1; i < num_reduction_axes; i++) {
-    if (normalized_reduction_axes[i] <= normalized_reduction_axes[i - 1]) {
-      xnn_log_error(
-          "failed to reshape %s operator with #%zu reduction axis of %" PRIi64
-          ": the reduction axes must be unique",
-          xnn_operator_type_to_string_v2(reduce_op), i, reduction_axes[i]);
-      return xnn_status_invalid_parameter;
+  // Remove duplicate reduction axes.
+  int i = 0;
+  // The array is sorted, we're done if an axis is bigger than num_input_dims.
+  for (int j = 1;
+       j < num_reduction_axes && normalized_reduction_axes[j] < num_input_dims;
+       ++j) {
+    // Shift non-duplicate elements forward.
+    if (normalized_reduction_axes[i] != normalized_reduction_axes[j]) {
+      normalized_reduction_axes[++i] = normalized_reduction_axes[j];
     }
   }
+  num_reduction_axes = i + 1;
+
+  assert(num_reduction_axes <= num_input_dims);
 
   xnn_normalize_reduction(
     &num_reduction_axes, normalized_reduction_axes,
diff --git a/src/subgraph.c b/src/subgraph.c
index 2d785ec..9682395 100644
--- a/src/subgraph.c
+++ b/src/subgraph.c
@@ -2455,11 +2455,24 @@
   return xnn_status_success;
 }
 
+static void convert_static_value_to_fp32(struct xnn_value* value) {
+  assert(xnn_value_is_static(value->allocation_type));
+  if (value->flags & XNN_VALUE_FLAG_NEEDS_CLEANUP) {
+    xnn_release_memory(value->data);
+  }
+  value->data = xnn_allocate_memory(sizeof(float));
+  float data = get_scalar_value_as_float(value);
+  memcpy(value->data, &data, sizeof(float));
+  value->flags |= XNN_VALUE_FLAG_NEEDS_CLEANUP;
+  value->datatype = xnn_datatype_fp32;
+}
+
 // Replace `mul(reduce_sum(x), 1/n)`, `div(reduce_sum(x), n)`  or
 // `mul(reduce_sum_squared(x), 1/n)`, `div(reduce_sum_squared(x), n)`
 // with `reduce_mean(x)` or `reduce_mean_squared(x)`, respectively.
-static enum xnn_status optimize_common_subgraphs_scaled_sum_to_mean(
-    xnn_subgraph_t subgraph, uint32_t node_id, size_t* changes) {
+static enum xnn_status widen_fp16_accumulators(xnn_subgraph_t subgraph,
+                                               uint32_t node_id,
+                                               size_t* changes) {
   struct xnn_node* node = &subgraph->nodes[node_id];
 
   if (node->type != xnn_node_type_binary_elementwise ||
@@ -2468,78 +2481,66 @@
     return xnn_status_success;
   }
 
-  struct xnn_value* reduce_value = &subgraph->values[node->inputs[0]];
+  struct xnn_value* reduced_value = &subgraph->values[node->inputs[0]];
   struct xnn_value* arg_value = &subgraph->values[node->inputs[1]];
   if (xnn_shape_multiply_all_dims(&arg_value->shape) != 1 ||
       !xnn_value_is_static(arg_value->allocation_type)) {
-    if (xnn_shape_multiply_all_dims(&reduce_value->shape) == 1 &&
-        xnn_value_is_static(reduce_value->allocation_type)) {
-      swap_value_pointers(&reduce_value, &arg_value);
+    if (xnn_shape_multiply_all_dims(&reduced_value->shape) == 1 &&
+        xnn_value_is_static(reduced_value->allocation_type)) {
+      swap_value_pointers(&reduced_value, &arg_value);
     } else {
       return xnn_status_success;
     }
   }
 
   // Check that one of the args is a sum or sum2 reduction.
-  if (!(reduce_value->datatype == xnn_datatype_fp16 ||
-        reduce_value->datatype == xnn_datatype_fp32) ||
-      reduce_value->producer == XNN_INVALID_NODE_ID) {
+  if (reduced_value->datatype != xnn_datatype_fp16 ||
+      reduced_value->producer == XNN_INVALID_NODE_ID) {
     return xnn_status_success;
   }
-  struct xnn_node* reduce_node = &subgraph->nodes[reduce_value->producer];
+
+  if (reduced_value->num_consumers > 1 || arg_value->num_consumers > 1) {
+    // Don't rewrite if we might modify an unrelated consumer.
+    return xnn_status_success;
+  }
+
+  struct xnn_node* reduce_node = &subgraph->nodes[reduced_value->producer];
   const enum xnn_node_type reduce_node_type = reduce_node->type;
   if (!(reduce_node_type == xnn_node_type_static_sum ||
         reduce_node_type == xnn_node_type_static_sum_squared)) {
     return xnn_status_success;
   }
 
-  // Check that the other arg is the product of the dimensions of
-  // the reduction axes, or its inverse.
-  struct xnn_value* input_value = &subgraph->values[reduce_node->inputs[0]];
-  const float arg_as_float = get_scalar_value_as_float(arg_value);
-  size_t num_reduced_dims = 1;
-  for (size_t k = 0; k < reduce_node->params.reduce.num_reduction_axes; k++) {
-    num_reduced_dims *= xnn_shape_get_dim(
-        &input_value->shape, reduce_node->params.reduce.reduction_axes[k]);
-  }
-  if (!num_reduced_dims) {
-    return xnn_status_success;
-  }
-  const enum xnn_binary_operator binary_operator = node->binary_operator;
-  float expected_arg = (binary_operator == xnn_binary_multiply)
-                           ? 1.0f / num_reduced_dims
-                           : (float)num_reduced_dims;
-  if (arg_value->datatype == xnn_datatype_fp16) {
-    expected_arg = xnn_float16_to_float(xnn_float16_from_float(expected_arg));
-  }
-  if (arg_as_float != expected_arg) {
-    return xnn_status_success;
-  }
+  // Rewrite the internal values to this subgraph to be fp32.
+  reduced_value->datatype = xnn_datatype_fp32;
+  convert_static_value_to_fp32(arg_value);