Be more strict with elided arithmetic op shapes.

gonnet · xnnpack-bot · commit 658b15fbf92f · 2025-11-13T06:40:18.000-08:00
PiperOrigin-RevId: 831383616
diff --git a/src/subgraph.c b/src/subgraph.c
@@ -2812,33 +2812,16 @@ static enum xnn_status optimize_common_subgraphs_static_reshapes(
     return xnn_status_success;
   }
 
-  // Set the shape of the static-shaped value.
-  struct xnn_shape new_shape;
-  if (node->type == xnn_node_type_static_reshape) {
-    // Replace the old shape with the new shape, filling any gaps from the input
-    // shape.
-    new_shape = node->params.static_reshape.new_shape;
-    XNN_RETURN_IF_ERROR(xnn_shape_fill_gaps(&input_value->shape, &new_shape));
-  } else if (node->type == xnn_node_type_static_expand_dims) {
-    const struct xnn_shape* new_dims = &node->params.static_reshape.new_shape;
-    new_shape.num_dims = input_value->shape.num_dims + new_dims->num_dims;
-    for (uint32_t idx_new = 0, idx_old = 0, k = 0; k < new_shape.num_dims;
-         k++) {
-      if (idx_new < new_dims->num_dims && new_dims->dim[idx_new] == k) {
-        new_shape.dim[k] = 1;
-        idx_new++;
-      } else {
-        new_shape.dim[k] = input_value->shape.dim[idx_old++];
-      }
-    }
-  }
-
   // If the input is a static value, apply the new shape to it directly.
   bool elide = true;
   if (xnn_value_is_static(input_value->allocation_type)) {
-    input_value->shape = new_shape;
-  } else {
-    elide = xnn_shape_match(&new_shape, &input_value->shape);
+    input_value->shape = output_value->shape;
+  }
+
+  // Otherwise, if the new shape is the old shape, do away with the reshape
+  // entirely.
+  else {
+    elide = xnn_shape_match(&input_value->shape, &output_value->shape);
   }
 
   if (elide) {
@@ -3427,6 +3410,12 @@ static enum xnn_status optimize_common_subgraphs_binary_const_noop(
     }
   }
 
+  // If we don't know the shape of the constant, then we can't really guarantee
+  // anything.
+  if ((const_value->flags & XNN_VALUE_FLAG_SHAPE_IS_STATIC) == 0) {
+    return xnn_status_success;
+  }
+
   const enum xnn_binary_operator binary_operator = node->binary_operator;
   const bool const_is_zero = (const_value->flags & XNN_VALUE_FLAG_IS_ZERO) != 0;
   const bool const_is_one = (const_value->flags & XNN_VALUE_FLAG_IS_ONE) != 0;
@@ -3453,32 +3442,29 @@ static enum xnn_status optimize_common_subgraphs_binary_const_noop(
            (const_is_zero &&
             (binary_operator == xnn_binary_add ||
              (binary_operator == xnn_binary_subtract && const_is_rhs)))) {
-    if (short_circuit(subgraph, input_value->id, node->outputs[0])) {
-      xnn_log_info("Elided spurious %s[#%u](v%03u, %s).",
-                   binary_operator == xnn_binary_multiply ? "mul"
-                   : binary_operator == xnn_binary_divide ? "div"
-                   : binary_operator == xnn_binary_add    ? "add"
-                                                          : "sub",
-                   node->id, input_value->id, const_is_zero ? "0.0" : "1.0");
-      xnn_node_clear(node);
-      (*changes)++;
-    } else if (input_value->flags & XNN_VALUE_FLAG_SHAPE_IS_STATIC &&
-               const_value->flags & XNN_VALUE_FLAG_SHAPE_IS_STATIC) {
-      // If this node cannot be elided, and both input shapes are static, then
-      // try to replace it with a `copy` or `broadcast` of the input value.
-      struct xnn_shape* output_shape = &subgraph->values[node->outputs[0]].shape;
-      XNN_RETURN_IF_ERROR(
-          xnn_shape_binary_broadcast(&input_value->shape, &const_value->shape,
-                                     output_shape),
-          "Incompatible input shapes for %s[#%u](v%03u, %s).",
-          binary_operator == xnn_binary_multiply ? "mul"
-          : binary_operator == xnn_binary_divide ? "div"
-          : binary_operator == xnn_binary_add    ? "add"
-                                                 : "sub",
-          node->id, input_value->id, const_is_zero ? "0.0" : "1.0");
-
-      // If the output shape matches the input shape, just copy the input.
-      if (xnn_shape_match(&input_value->shape, output_shape)) {
+    // We can safely elide this operation if we know it will not change the
+    // shape of the output, e.g. if the constant is a scalar or the shapes are
+    // static and equal.
+    if ((xnn_shape_multiply_all_dims(&const_value->shape) == 1 &&
+         const_value->shape.num_dims <= input_value->shape.num_dims) ||
+        ((input_value->flags & XNN_VALUE_FLAG_SHAPE_IS_STATIC) &&
+         xnn_shape_match(&input_value->shape, &output_value->shape))) {
+      // If the node be elided (not load-bearing), then just remove it.
+      if (short_circuit(subgraph, input_value->id, node->outputs[0])) {
+        xnn_log_info("Elided spurious %s[#%u](v%03u, %s).",
+                     binary_operator == xnn_binary_multiply ? "mul"
+                     : binary_operator == xnn_binary_divide ? "div"
+                     : binary_operator == xnn_binary_add    ? "add"
+                                                            : "sub",
+                     node->id, input_value->id, const_is_zero ? "0.0" : "1.0");
+        xnn_node_clear(node);
+        (*changes)++;
+      }
+
+      // Otherwise, replace it with a copy.
+      else {
+        // If the constant is a scalar, then it won't affect the shape of the
+        // output.
         XNN_RETURN_IF_ERROR(xnn_define_copy(subgraph, input_value->id,
                                             node->outputs[0], node->flags),
                             "Failed to create new `Copy` node.");
@@ -3492,26 +3478,6 @@ static enum xnn_status optimize_common_subgraphs_binary_const_noop(
             node->id, input_value->id, const_is_zero ? "0.0" : "1.0", node_id,
             input_value->id);
       }
-
-      // Otherwise, we need to broadcast the input to the output shape.
-      else {
-        XNN_RETURN_IF_ERROR(
-            xnn_define_static_broadcast(subgraph, output_shape->num_dims,
-                                        output_shape->dim, input_value->id,
-                                        node->outputs[0],
-                                        node->flags),
-            "Failed to create new `Broadcast` node.");
-        node = move_last_node_to(subgraph, node_id);
-        xnn_log_info(
-            "Replaced spurious %s[#%u](v%03u, %s) with "
-            "static_broadcast[#%u](v%03u).",
-            binary_operator == xnn_binary_multiply ? "mul"
-            : binary_operator == xnn_binary_divide ? "div"
-            : binary_operator == xnn_binary_add    ? "add"
-                                                   : "sub",
-            node->id, input_value->id, const_is_zero ? "0.0" : "1.0", node_id,
-            input_value->id);
-      }
       (*changes)++;
     }
   }
@@ -3637,9 +3603,55 @@ static enum xnn_status optimize_common_subgraphs_iter(
                                       XNN_VALUE_FLAG_SHAPE_IS_STATIC) != 0;
     }
     if (all_input_shapes_are_static) {
-      for (int k = 0; k < node->num_outputs; k++) {
-        subgraph->values[node->outputs[k]].flags |=
-            XNN_VALUE_FLAG_SHAPE_IS_STATIC;
+      switch (node->type) {
+        case xnn_node_type_unary_elementwise:
+          subgraph->values[node->outputs[0]].shape =
+              subgraph->values[node->inputs[0]].shape;
+          subgraph->values[node->outputs[0]].flags |=
+              XNN_VALUE_FLAG_SHAPE_IS_STATIC;
+          break;
+
+        case xnn_node_type_binary_elementwise:
+          xnn_shape_binary_broadcast(&subgraph->values[node->inputs[0]].shape,
+                                     &subgraph->values[node->inputs[1]].shape,
+                                     &subgraph->values[node->outputs[0]].shape);
+          subgraph->values[node->outputs[0]].flags |=
+              XNN_VALUE_FLAG_SHAPE_IS_STATIC;
+          break;
+
+        case xnn_node_type_static_transpose:
+          // Apply the transpose to the output shape.
+          for (int k = 0; k < node->params.transpose.num_dims; k++) {
+            subgraph->values[node->outputs[0]].shape.dim[k] =
+                subgraph->values[node->inputs[0]]
+                    .shape.dim[node->params.transpose.perm[k]];
+          }
+          subgraph->values[node->outputs[0]].flags |=
+              XNN_VALUE_FLAG_SHAPE_IS_STATIC;
+          break;
+
+        case xnn_node_type_static_expand_dims: {
+          const struct xnn_shape* new_dims =
+              &node->params.static_reshape.new_shape;
+          subgraph->values[node->outputs[0]].shape.num_dims =
+              subgraph->values[node->inputs[0]].shape.num_dims +
+              new_dims->num_dims;
+          for (uint32_t idx_new = 0, idx_in = 0, k = 0;
+               k < subgraph->values[node->outputs[0]].shape.num_dims; k++) {
+            if (idx_new < new_dims->num_dims && new_dims->dim[idx_new] == k) {
+              subgraph->values[node->outputs[0]].shape.dim[k] = 1;
+              idx_new++;
+            } else {
+              subgraph->values[node->outputs[0]].shape.dim[k] =
+                  subgraph->values[node->inputs[0]].shape.dim[idx_in++];
+            }
+          }
+          subgraph->values[node->outputs[0]].flags |=
+              XNN_VALUE_FLAG_SHAPE_IS_STATIC;
+        } break;
+
+        default:
+          break;
       }
     }
 
diff --git a/test/subgraph/rewrites.cc b/test/subgraph/rewrites.cc
@@ -19,7 +19,6 @@
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
-#include "include/experimental.h"
 #include "include/xnnpack.h"
 #include "src/subgraph/subgraph-utils.h"
 #include "src/xnnpack/buffer.h"
@@ -844,7 +843,7 @@ TEST_P(RewriteArithmeticTest, ElidesNoOpStaticShapeMul) {
         // Add a scalar static tensor with the value `1.0`.
         uint32_t static_one_value_id;
         std::tie(static_one_tensor, static_one_value_id) =
-            add_static_tensor<float>(rng, subgraph, /*shape=*/{1}, 1.0, 1.0);
+            add_static_tensor<float>(rng, subgraph, /*shape=*/{}, 1.0, 1.0);
 
         // Add the binary `multiply` op with the constant 1.0.
         auto inputs =
@@ -934,7 +933,7 @@ TEST_P(RewriteArithmeticTest, ElidesNoOpStaticShapeDiv) {
         // Add a scalar static tensor with the value `1.0`.
         uint32_t static_one_value_id;
         std::tie(static_one_tensor, static_one_value_id) =
-            add_static_tensor<float>(rng, subgraph, /*shape=*/{1}, 1.0, 1.0);
+            add_static_tensor<float>(rng, subgraph, /*shape=*/{}, 1.0, 1.0);
 
         // Add the binary `divide` op by the constant 1.0.
         subgraph.AddBinary(xnn_binary_divide, /*params=*/nullptr,
@@ -1022,7 +1021,7 @@ TEST_P(RewriteArithmeticTest, ElidesNoOpStaticShapeAdd) {
         // Add a scalar static tensor with the value `0.0`.
         uint32_t static_zero_value_id;
         std::tie(static_zero_tensor, static_zero_value_id) =
-            add_static_tensor<float>(rng, subgraph, /*shape=*/{1}, 0.0, 0.0);
+            add_static_tensor<float>(rng, subgraph, /*shape=*/{}, 0.0, 0.0);
 
         // Add the binary `add` op with the constant 0.0.
         auto inputs =
@@ -1112,7 +1111,7 @@ TEST_P(RewriteArithmeticTest, ElidesNoOpStaticShapeSub) {
         // Add a scalar static tensor with the value `0.0`.
         uint32_t static_zero_value_id;
         std::tie(static_zero_tensor, static_zero_value_id) =
-            add_static_tensor<float>(rng, subgraph, /*shape=*/{1}, 0.0, 0.0);
+            add_static_tensor<float>(rng, subgraph, /*shape=*/{}, 0.0, 0.0);
 
         // Add the binary `subtract` op with the constant 0.0.
         subgraph.AddBinary(xnn_binary_subtract, /*params=*/nullptr,
@@ -1247,7 +1246,7 @@ TEST_P(RewriteArithmeticTest, ElidesNoOpChainOfStaticShapeMulZeroAdd) {
         // Add a scalar static tensor with the value `0.0`.
         uint32_t static_zero_value_id;
         std::tie(static_zero_tensor, static_zero_value_id) =
-            add_static_tensor<float>(rng, subgraph, /*shape=*/{1}, 0.0, 0.0);
+            add_static_tensor<float>(rng, subgraph, /*shape=*/{}, 0.0, 0.0);
 
         // Add the binary `multiply` op with the constant 0.0.
         uint32_t dynamic_zero_value_id =
@@ -1361,7 +1360,7 @@ TEST_P(RewriteArithmeticTest, ElidesNoOpChainOfStaticShapeDivOneMul) {
         // Add a scalar static tensor with the value `1.0`.
         uint32_t static_one_value_id;
         std::tie(static_one_tensor, static_one_value_id) =
-            add_static_tensor<float>(rng, subgraph, /*shape=*/{1}, 1.0, 1.0);
+            add_static_tensor<float>(rng, subgraph, /*shape=*/{}, 1.0, 1.0);
 
         // Add the static `1.0` to the absolute value of the inputs to make sure
         // they are non-negative
