Add xnn_define_static_constant_pad_v2 and deprecate xnn_define_static_constant_pad

`xnn_define_static_constant_pad_v2` takes `num_padding_dims` as a parameter, instead of inferring it from the input value's shape as `xnn_define_static_constant_pad` does.

This is necessary because we infer the rank of values in the subgraph. If the rank inference changes the rank of the input value after the padding node is defined, the padding node is unlikely to behave as intended. This change makes the intended behavior explicit.

See #9775 for more context.

PiperOrigin-RevId: 890668243

Author: dsharletg

include/xnnpack.h

@@ -1370,16 +1370,27 @@ XNN_DEPRECATED enum xnn_status xnn_define_squared_difference(
 /// Define a Constant Pad Node with static padding specification and add it to a Subgraph.
 ///
 /// @param subgraph - a Subgraph object that will own the created Node.
+/// @param num_padding_dims - number of dimensions to pad.
 /// @param pre_paddings - number of padding elements to insert before input elements for every dimension. This array
-///                       must have as many elements as the number of dimensions in the input tensor.
+///                       must have @num_padding_dims elements.
 /// @param post_paddings - number of padding elements to insert after input elements for every dimension. This array
-///                        must have as many elements as the number of dimensions in the input tensor.
+///                        must have @num_padding_dims elements.
 /// @param padding_value - constant value used to initialize padding elements.
 /// @param input_id - Value ID for the input tensor. The input tensor must be defined in the @a subgraph.
 /// @param output_id - Value ID for the output tensor. The output tensor must be defined in the @a subgraph, and its
 ///                    shape must match the shape of the input tensor with padding.
 /// @param flags - binary features of the Constant Pad Node. No supported flags are currently defined.
-enum xnn_status xnn_define_static_constant_pad(
+enum xnn_status xnn_define_static_constant_pad_v2(
+  xnn_subgraph_t subgraph,
+  size_t num_padding_dims,
+  const size_t* pre_paddings,
+  const size_t* post_paddings,
+  float padding_value,
+  uint32_t input_id,
+  uint32_t output_id,
+  uint32_t flags);
+
+XNN_DEPRECATED enum xnn_status xnn_define_static_constant_pad(
   xnn_subgraph_t subgraph,
   const size_t* pre_paddings,
   const size_t* post_paddings,

src/subgraph/deprecated.c

@@ -1,8 +1,11 @@
+#include <inttypes.h>
 #include <math.h>
 #include <stddef.h>
 #include <stdint.h>
 
 #include "include/xnnpack.h"
+#include "src/xnnpack/log.h"
+#include "src/xnnpack/node-type.h"
 #include "src/xnnpack/subgraph.h"
 
 enum xnn_status xnn_define_add2(xnn_subgraph_t subgraph, float output_min,
@@ -386,3 +389,22 @@ enum xnn_status xnn_define_even_split4(xnn_subgraph_t subgraph, int32_t split_di
   const uint32_t outputs_id[4] = {output1_id, output2_id, output3_id, output4_id};
   return xnn_define_even_split(subgraph, split_dim, input_id, /*num_outputs=*/4, outputs_id, flags);
 }
+
+enum xnn_status xnn_define_static_constant_pad(
+  xnn_subgraph_t subgraph,
+  const size_t* pre_paddings,
+  const size_t* post_paddings,
+  float padding_value,
+  uint32_t input_id,
+  uint32_t output_id,
+  uint32_t flags) {
+  if (input_id >= subgraph->num_values) {
+    xnn_log_error(
+      "failed to define %s operator with input ID #%" PRIu32 ": invalid Value ID",
+      xnn_node_type_to_string(xnn_node_type_static_constant_pad), input_id);
+    return xnn_status_invalid_parameter;
+  }
+  size_t num_padding_dims = subgraph->values[input_id].shape.num_dims;
+  return xnn_define_static_constant_pad_v2(subgraph, num_padding_dims, pre_paddings,
+                                           post_paddings, padding_value, input_id, output_id, flags);
+}

src/subgraph/static-constant-pad.c

@@ -173,8 +173,9 @@ static enum xnn_status setup_constant_pad_operator(
   }
 }
 
-enum xnn_status xnn_define_static_constant_pad(
+enum xnn_status xnn_define_static_constant_pad_v2(
   xnn_subgraph_t subgraph,
+  size_t num_padding_dims,
   const size_t* pre_paddings,
   const size_t* post_paddings,
   float padding_value,
@@ -187,6 +188,13 @@ enum xnn_status xnn_define_static_constant_pad(
     return status;
   }
 
+  if (num_padding_dims > XNN_MAX_TENSOR_DIMS) {
+    xnn_log_error(
+      "failed to define %s operator with %zu dims: maximum number of dimensions is %d",
+      xnn_node_type_to_string(xnn_node_type_static_constant_pad), num_padding_dims, XNN_MAX_TENSOR_DIMS);
+    return xnn_status_invalid_parameter;
+  }
+
   if (input_id >= subgraph->num_values) {
     xnn_log_error(
       "failed to define %s operator with input ID #%" PRIu32 ": invalid Value ID",
@@ -244,9 +252,11 @@ enum xnn_status xnn_define_static_constant_pad(
     return xnn_status_out_of_memory;
   }
 
-  const size_t num_dims = subgraph->values[input_id].shape.num_dims;
-  memcpy(&node->params.static_pad.pre_paddings, pre_paddings, num_dims * sizeof(size_t));
-  memcpy(&node->params.static_pad.post_paddings, post_paddings, num_dims * sizeof(size_t));
+  node->params.static_pad.num_padding_dims = num_padding_dims;
+  memset(&node->params.static_pad.pre_paddings, 0, XNN_MAX_TENSOR_DIMS * sizeof(size_t));
+  memset(&node->params.static_pad.post_paddings, 0, XNN_MAX_TENSOR_DIMS * sizeof(size_t));
+  memcpy(&node->params.static_pad.pre_paddings, pre_paddings, num_padding_dims * sizeof(size_t));
+  memcpy(&node->params.static_pad.post_paddings, post_paddings, num_padding_dims * sizeof(size_t));
   switch (output_value->datatype) {
     case xnn_datatype_fp32:
       node->params.static_pad.padding_value = float_as_uint32(padding_value);

src/xnnpack/subgraph.h

@@ -364,6 +364,7 @@ struct xnn_node {
     struct {
       size_t pre_paddings[XNN_MAX_TENSOR_DIMS];
       size_t post_paddings[XNN_MAX_TENSOR_DIMS];
+      size_t num_padding_dims;
       uint32_t padding_value;
     } static_pad;
     struct {

test/subgraph/BUILD

@@ -342,7 +342,6 @@ xnnpack_unit_test(
         "//:node_type",
         "//:operator_h",
         "//:params",
-        "//:subgraph",
         "//:subgraph_h",
         "//:xnnpack_h",
         "//test:replicable_random_device",

test/subgraph/fusion.cc

@@ -341,8 +341,8 @@ TEST(CONSTANT_PAD_THEN_CONVOLUTION, fusion) {
   uint32_t filter_id = 2;
   uint32_t bias_id = 3;
   uint32_t output_id = 4;
-  size_t pre_paddings[4] = {0, 2, 4, 0};
-  size_t post_paddings[4] = {0, 6, 8, 0};
+  std::vector<size_t> pre_paddings = {0, 2, 4, 0};
+  std::vector<size_t> post_paddings = {0, 6, 8, 0};
   float padding_value = 0.0f;
 
   tester.AddInputTensorF32({1, 254, 254, 3}, input_id)
@@ -386,8 +386,8 @@ TEST(CONSTANT_PAD_THEN_CONVOLUTION, fusion_quantized_int8) {
   uint32_t filter_id = 2;
   uint32_t bias_id = 3;
   uint32_t output_id = 4;
-  size_t pre_paddings[4] = {0, 2, 4, 0};
-  size_t post_paddings[4] = {0, 6, 8, 0};
+  std::vector<size_t> pre_paddings = {0, 2, 4, 0};
+  std::vector<size_t> post_paddings = {0, 6, 8, 0};
   float padding_value = 0.0f;
   using qint8 = xnnpack::quantized<int8_t>;
   using qint32 = xnnpack::quantized<int32_t>;
@@ -446,8 +446,8 @@ TEST(CONSTANT_PAD_THEN_CONVOLUTION,
   uint32_t bias_id = 3;
   uint32_t output_id = 4;
   // Non-zero pre-padding in the N or C dimension.
-  size_t pre_paddings[4] = {1, 2, 4, 0};
-  size_t post_paddings[4] = {0, 6, 8, 0};
+  std::vector<size_t> pre_paddings = {1, 2, 4, 0};
+  std::vector<size_t> post_paddings = {0, 6, 8, 0};
   float padding_value = 0.0f;
 
   tester.AddInputTensorF32({1, 254, 254, 3}, input_id)
@@ -480,8 +480,8 @@ TEST(CONSTANT_PAD_THEN_CONVOLUTION, not_fused_due_to_padding_value_not_zero) {
   uint32_t filter_id = 2;
   uint32_t bias_id = 3;
   uint32_t output_id = 4;
-  size_t pre_paddings[4] = {1, 2, 4, 0};
-  size_t post_paddings[4] = {0, 6, 8, 0};
+  std::vector<size_t> pre_paddings = {1, 2, 4, 0};
+  std::vector<size_t> post_paddings = {0, 6, 8, 0};
   float padding_value = 1.0f;
 
   tester.AddInputTensorF32({1, 254, 254, 3}, input_id)
@@ -514,8 +514,8 @@ TEST(CONSTANT_PAD_THEN_DEPTHWISE_CONVOLUTION, fusion) {
   uint32_t filter_id = 2;
   uint32_t bias_id = 3;
   uint32_t output_id = 4;
-  size_t pre_paddings[4] = {0, 2, 4, 0};
-  size_t post_paddings[4] = {0, 6, 8, 0};
+  std::vector<size_t> pre_paddings = {0, 2, 4, 0};
+  std::vector<size_t> post_paddings = {0, 6, 8, 0};
   float padding_value = 0.0f;
   tester.AddInputTensorF32({1, 128, 128, 4}, input_id)
       .AddDynamicTensorF32({1, 136, 140, 4}, intermediate_id)
@@ -558,8 +558,8 @@ TEST(CONSTANT_PAD_THEN_DEPTHWISE_CONVOLUTION,
   uint32_t bias_id = 3;
   uint32_t output_id = 4;
   // Non-zero pre-padding in the N or C dimension.
-  size_t pre_paddings[4] = {1, 2, 4, 0};
-  size_t post_paddings[4] = {0, 6, 8, 0};
+  std::vector<size_t> pre_paddings = {1, 2, 4, 0};
+  std::vector<size_t> post_paddings = {0, 6, 8, 0};
   float padding_value = 0.0f;
   tester.AddInputTensorF32({1, 128, 128, 4}, input_id)
       .AddDynamicTensorF32({2, 136, 140, 4}, intermediate_id)
@@ -590,8 +590,8 @@ TEST(CONSTANT_PAD_THEN_DEPTHWISE_CONVOLUTION,
   uint32_t filter_id = 2;
   uint32_t bias_id = 3;
   uint32_t output_id = 4;
-  size_t pre_paddings[4] = {0, 2, 4, 0};
-  size_t post_paddings[4] = {0, 6, 8, 0};
+  std::vector<size_t> pre_paddings = {0, 2, 4, 0};
+  std::vector<size_t> post_paddings = {0, 6, 8, 0};
   float padding_value = 1.0f;
   tester.AddInputTensorF32({1, 128, 128, 4}, input_id)
       .AddDynamicTensorF32({1, 136, 140, 4}, intermediate_id)

test/subgraph/static-constant-pad.cc

@@ -26,8 +26,9 @@ void TestImpl(size_t rank) {
   ASSERT_EQ(xnn_status_success, xnn_initialize(nullptr /* allocator */));
 
   for (auto _ : FuzzTest(std::chrono::milliseconds(250))) {
-    std::vector<size_t> pre_padding = random_shape(rng, rank, 0, 3);
-    std::vector<size_t> post_padding = random_shape(rng, rank, 0, 3);
+    size_t padding_rank = std::uniform_int_distribution<size_t>(1, rank)(rng);
+    std::vector<size_t> pre_padding = random_shape(rng, padding_rank, 0, 3);
+    std::vector<size_t> post_padding = random_shape(rng, padding_rank, 0, 3);
     float pad_value = 1.0f;
 
     xnn_quantization_params quantization =
@@ -40,6 +41,10 @@ void TestImpl(size_t rank) {
         .AddConstantPad(pre_padding, post_padding, pad_value, 0, 1)
         .CreateRuntime();
 
+    // The test code needs the padding to be the same rank as the input/output.
+    pre_padding.resize(rank);
+    post_padding.resize(rank);
+
     for (int reshape = 0; reshape < 2; ++reshape) {
       std::vector<size_t> shape = random_shape(rng, rank);
 
@@ -48,7 +53,7 @@ void TestImpl(size_t rank) {
       input.generate([&]() { return generator(rng); });
 
       std::vector<size_t> output_shape(shape);
-      for (size_t i = 0; i < rank; ++i) {
+      for (size_t i = 0; i < padding_rank; ++i) {
         output_shape[i] += pre_padding[i] + post_padding[i];
       }
 

test/subgraph/subgraph-fp16.cc

@@ -222,8 +222,8 @@ TEST(SUBGRAPH_FP16, fully_connected_f16_weights_no_biases) {
 
 TEST(SUBGRAPH_FP16, value_both_external_output_and_input) {
   SubgraphTester tester(4);
-  std::array<size_t, 4> pre_paddings = {0, 1, 0, 0};
-  std::array<size_t, 4> post_paddings = {0, 1, 0, 0};
+  std::vector<size_t> pre_paddings = {0, 1, 0, 0};
+  std::vector<size_t> post_paddings = {0, 1, 0, 0};
   // external input[0]
   //      /
   // [constant pad]
@@ -239,7 +239,7 @@ TEST(SUBGRAPH_FP16, value_both_external_output_and_input) {
       .AddDynamicTensorF32({1, 1, 1, 3}, 1)
       .AddOutputTensorF32({1, 4, 2, 3}, 2)
       .AddOutputTensorF32({1, 4, 2, 3}, 3)
-      .AddConstantPad(pre_paddings.data(), post_paddings.data(), 0.0f, 0, 2)
+      .AddConstantPad(pre_paddings, post_paddings, 0.0f, 0, 2)
       .AddAddition(2, 1, 3)
       .Optimize()
       .RewriteForFp16();
@@ -940,14 +940,13 @@ TEST(SUBGRAPH_FP16_DYNAMIC_FULLY_CONNECTED,
   const uint32_t input2_id = 1;
   const uint32_t weights_id = 3;
   const uint32_t fully_connected_out_id = 2;
-  std::array<size_t, 4> pre_paddings = {1, 0, 0, 0};
-  std::array<size_t, 4> post_paddings = {0, 0, 0, 0};
+  std::vector<size_t> pre_paddings = {1, 0, 0, 0};
+  std::vector<size_t> post_paddings = {0, 0, 0, 0};
   tester.AddInputTensorF32({1, 5, 5, 3}, input_id)
       .AddInputTensorF32({1, 1, 1, 3}, input2_id)
       .AddOutputTensorF32({1, 5, 5, 2}, fully_connected_out_id)
       .AddDynamicTensorF32({2, 1, 1, 3}, weights_id)
-      .AddConstantPad(pre_paddings.data(), post_paddings.data(), 0.0f,
-                      input2_id, weights_id)
+      .AddConstantPad(pre_paddings, post_paddings, 0.0f, input2_id, weights_id)
       .AddFullyConnected(input_id, weights_id,
                          /*bias_id=*/XNN_INVALID_VALUE_ID,
                          fully_connected_out_id)
@@ -974,15 +973,14 @@ TEST(SUBGRAPH_FP16_DYNAMIC_FULLY_CONNECTED,
   const uint32_t weights_id = 3;
   const uint32_t bias_id = 4;
   const uint32_t fully_connected_out_id = 2;
-  std::array<size_t, 4> pre_paddings = {1, 0, 0, 0};
-  std::array<size_t, 4> post_paddings = {0, 0, 0, 0};
+  std::vector<size_t> pre_paddings = {1, 0, 0, 0};
+  std::vector<size_t> post_paddings = {0, 0, 0, 0};
   tester.AddInputTensorF32({1, 5, 5, 3}, input_id)
       .AddInputTensorF32({1, 1, 1, 3}, input2_id)
       .AddOutputTensorF32({1, 5, 5, 2}, fully_connected_out_id)
       .AddDynamicTensorF32({2, 1, 1, 3}, weights_id)
       .AddStaticTensorF32({2}, TensorType::kDense, bias_id)
-      .AddConstantPad(pre_paddings.data(), post_paddings.data(), 0.0f,
-                      input2_id, weights_id)
+      .AddConstantPad(pre_paddings, post_paddings, 0.0f, input2_id, weights_id)
       .AddFullyConnected(input_id, weights_id, bias_id, fully_connected_out_id)
       .Optimize()
       .RewriteForFp16();
@@ -1007,15 +1005,14 @@ TEST(SUBGRAPH_FP16_DYNAMIC_FULLY_CONNECTED,
   const uint32_t weights_id = 3;
   const uint32_t bias_id = 4;
   const uint32_t fully_connected_out_id = 2;
-  std::array<size_t, 4> pre_paddings = {1};
-  std::array<size_t, 4> post_paddings = {0};
+  std::vector<size_t> pre_paddings = {1};
+  std::vector<size_t> post_paddings = {0};
   tester.AddInputTensorF32({1, 5, 5, 3}, input_id)
       .AddInputTensorF32({1}, input2_id)
       .AddOutputTensorF32({1, 5, 5, 2}, fully_connected_out_id)
       .AddStaticTensorF32({2, 1, 1, 3}, TensorType::kDense, weights_id)
       .AddDynamicTensorF32({2}, bias_id)
-      .AddConstantPad(pre_paddings.data(), post_paddings.data(), 0.0f,
-                      input2_id, bias_id)
+      .AddConstantPad(pre_paddings, post_paddings, 0.0f, input2_id, bias_id)
       .AddFullyConnected(input_id, weights_id, bias_id, fully_connected_out_id)
       .Optimize()
       .RewriteForFp16();
@@ -1042,21 +1039,21 @@ TEST(SUBGRAPH_FP16_DYNAMIC_FULLY_CONNECTED,
   const uint32_t bias_id = 4;
   const uint32_t fully_connected_out_id = 5;
 
-  std::array<size_t, 4> weights_pre_paddings = {1, 0, 0, 0};
-  std::array<size_t, 4> weights_post_paddings = {0, 0, 0, 0};
-  std::array<size_t, 4> bias_pre_paddings = {1};
-  std::array<size_t, 4> bias_post_paddings = {0};
+  std::vector<size_t> weights_pre_paddings = {1, 0, 0, 0};
+  std::vector<size_t> weights_post_paddings = {0, 0, 0, 0};
+  std::vector<size_t> bias_pre_paddings = {1};
+  std::vector<size_t> bias_post_paddings = {0};
 
   tester.AddInputTensorF32({1, 5, 5, 3}, input_id)
       .AddInputTensorF32({1, 1, 1, 3}, input2_id)
       .AddInputTensorF32({1}, input3_id)
       .AddOutputTensorF32({1, 5, 5, 2}, fully_connected_out_id)
       .AddDynamicTensorF32({2, 1, 1, 3}, weights_id)
       .AddDynamicTensorF32({2}, bias_id)
-      .AddConstantPad(weights_pre_paddings.data(), weights_post_paddings.data(),
-                      0.0f, input2_id, weights_id)
-      .AddConstantPad(bias_pre_paddings.data(), bias_post_paddings.data(), 0.0f,
-                      input3_id, bias_id)
+      .AddConstantPad(weights_pre_paddings, weights_post_paddings, 0.0f,
+                      input2_id, weights_id)
+      .AddConstantPad(bias_pre_paddings, bias_post_paddings, 0.0f, input3_id,
+                      bias_id)
       .AddFullyConnected(input_id, weights_id, bias_id, fully_connected_out_id)
       .Optimize()
       .RewriteForFp16();

test/subgraph/subgraph-tester.cc

@@ -285,25 +285,14 @@ SubgraphTester& SubgraphTester::AddConcatenate(size_t axis,
   return *this;
 }
 
-SubgraphTester& SubgraphTester::AddConstantPad(const size_t* pre_paddings,
-                                               const size_t* post_paddings,
-                                               float padding_value,
-                                               uint32_t input_id,
-                                               uint32_t output_id) {
-  const xnn_status status = xnn_define_static_constant_pad(
-      subgraph_.get(), pre_paddings, post_paddings, padding_value, input_id,
-      output_id, /*flags=*/0);
-  EXPECT_EQ(status, xnn_status_success);
-  return *this;
-}
-
 SubgraphTester& SubgraphTester::AddConstantPad(
     const std::vector<size_t>& pre_paddings,
     const std::vector<size_t>& post_paddings, float padding_value,
     uint32_t input_id, uint32_t output_id) {
-  const xnn_status status = xnn_define_static_constant_pad(
-      subgraph_.get(), pre_paddings.data(), post_paddings.data(), padding_value,
-      input_id, output_id, /*flags=*/0);
+  assert(pre_paddings.size() == post_paddings.size());
+  const xnn_status status = xnn_define_static_constant_pad_v2(
+      subgraph_.get(), pre_paddings.size(), pre_paddings.data(),
+      post_paddings.data(), padding_value, input_id, output_id, /*flags=*/0);
   EXPECT_EQ(status, xnn_status_success);
   return *this;
 }

test/subgraph/subgraph-tester.h

@@ -361,11 +361,6 @@ class SubgraphTester {
   SubgraphTester& AddConcatenate(size_t axis, std::vector<uint32_t> input_ids,
                                  uint32_t output_id);
 
-  SubgraphTester& AddConstantPad(const size_t* pre_paddings,
-                                 const size_t* post_paddings,
-                                 float padding_value, uint32_t input_id,
-                                 uint32_t output_id);
-
   SubgraphTester& AddConstantPad(const std::vector<size_t>& pre_paddings,
                                  const std::vector<size_t>& post_paddings,
                                  float padding_value, uint32_t input_id,