validate runtime input channels in convolution reshape paths

src/subgraph/convolution-2d.c

@@ -682,6 +682,24 @@ enum xnn_status reshape_convolution_operator(struct xnn_operator_data* opdata,
   const size_t input_height = input_value->shape.dim[1];
   const size_t input_width = input_value->shape.dim[2];
 
+  // The indirection buffer strides across input pixels by the operator's
+  // create-time input_pixel_stride, so an NHWC input carrying fewer channels
+  // than that makes the microkernel read past the end of the input buffer.
+  const xnn_operator_t convolution_op = opdata->operator_objects[0];
+  if (convolution_op->type != xnn_operator_type_convolution_nchw_f16 &&
+      convolution_op->type != xnn_operator_type_convolution_nchw_f32) {
+    const size_t input_channels =
+        input_value->shape.dim[input_value->shape.num_dims - 1];
+    if (input_channels != convolution_op->input_pixel_stride) {
+      xnn_log_error("failed to reshape %s operator with input ID #%" PRIu32
+                    ": mismatch at channel dimension (%zu != %zu)",
+                    xnn_node_type_to_string(xnn_node_type_convolution_2d),
+                    input_id, input_channels,
+                    convolution_op->input_pixel_stride);
+      return xnn_status_invalid_parameter;
+    }
+  }
+
   size_t output_height, output_width;
   enum xnn_status status = xnn_status_invalid_state;
   const size_t old_workspace_size = opdata->workspace_size;

src/subgraph/deconvolution-2d.c

@@ -414,6 +414,21 @@ static enum xnn_status reshape_deconvolution_operator(
   const size_t batch_size = values[input_id].shape.dim[0];
   const size_t input_height = values[input_id].shape.dim[1];
   const size_t input_width = values[input_id].shape.dim[2];
+
+  // The indirection buffer strides across input pixels by the operator's
+  // create-time input_pixel_stride, so an input carrying fewer channels than
+  // that makes the microkernel read past the end of the input buffer.
+  const size_t input_channels =
+      values[input_id].shape.dim[values[input_id].shape.num_dims - 1];
+  if (input_channels != opdata->operator_objects[0]->input_pixel_stride) {
+    xnn_log_error("failed to reshape %s operator with input ID #%" PRIu32
+                  ": mismatch at channel dimension (%zu != %zu)",
+                  xnn_node_type_to_string(xnn_node_type_deconvolution_2d),
+                  input_id, input_channels,
+                  opdata->operator_objects[0]->input_pixel_stride);
+    return xnn_status_invalid_parameter;
+  }
+
   enum xnn_status status = xnn_status_invalid_state;
   const size_t old_workspace_size = opdata->workspace_size;
   size_t output_height, output_width;

src/subgraph/depthwise-convolution-2d.c

@@ -412,6 +412,25 @@ static enum xnn_status reshape_depthwise_convolution_operator(
   const size_t batch_size = input_value->shape.dim[0];
   const size_t input_height = input_value->shape.dim[1];
   const size_t input_width = input_value->shape.dim[2];
+
+  // The indirection buffer strides across input pixels by the operator's
+  // create-time input_pixel_stride, so an NHWC input carrying fewer channels
+  // than that makes the microkernel read past the end of the input buffer.
+  const xnn_operator_t convolution_op = opdata->operator_objects[0];
+  if (convolution_op->type != xnn_operator_type_convolution_nchw_f16 &&
+      convolution_op->type != xnn_operator_type_convolution_nchw_f32) {
+    const size_t input_channels =
+        input_value->shape.dim[input_value->shape.num_dims - 1];
+    if (input_channels != convolution_op->input_pixel_stride) {
+      xnn_log_error("failed to reshape %s operator with input ID #%" PRIu32
+                    ": mismatch at channel dimension (%zu != %zu)",
+                    xnn_node_type_to_string(xnn_node_type_depthwise_convolution_2d),
+                    input_id, input_channels,
+                    convolution_op->input_pixel_stride);
+      return xnn_status_invalid_parameter;
+    }
+  }
+
   enum xnn_status status = xnn_status_invalid_state;
   const size_t old_workspace_size = opdata->workspace_size;
   size_t output_height, output_width;

test/subgraph/convolution-2d.cc

@@ -375,4 +375,46 @@ TEST(Convolution2DQD8F32QC8W, test) {
   TestImpl<float, qcint8, float>(/*convert_to=*/xnn_datatype_qdint8);
 }
 
+TEST(Convolution2D, reshape_rejects_input_channel_mismatch) {
+  ASSERT_EQ(xnn_status_success, xnn_initialize(nullptr /* allocator */));
+
+  ConvolutionParams params;
+  params.padding = {0, 0, 0, 0};
+  params.kernel = {3, 3};  // >1 so it is not rewritten to fully-connected.
+  params.subsampling = {1, 1};
+  params.dilation = {1, 1};
+  params.groups = 1;
+  params.group_input_channels = 4;
+  params.group_output_channels = 8;
+
+  const uint32_t input_id = 0;
+  const uint32_t filter_id = 1;
+  const uint32_t bias_id = 2;
+  const uint32_t output_id = 3;
+
+  Tensor<float> filter(std::vector<size_t>{8, 3, 3, 4});
+  filter.fill(0.0f);
+  Tensor<float> bias(std::vector<size_t>{8});
+  bias.fill(0.0f);
+
+  SubgraphTester subgraph(4);
+  subgraph.AddInputTensor(4, xnn_datatype_fp32, input_id)
+      .AddStaticTensor(filter.extents(), filter_id, filter.base())
+      .AddStaticTensor(bias.extents(), bias_id, bias.base())
+      .AddOutputTensor(4, xnn_datatype_fp32, output_id)
+      .AddConvolution2D(params, input_id, filter_id, bias_id, output_id);
+  if (subgraph.CreateRuntime() == xnn_status_unsupported_hardware) {
+    GTEST_SKIP();
+  }
+
+  // The operator was created for groups*group_input_channels == 4 channels. An
+  // external input reshaped to a single channel must be rejected: the
+  // indirection buffer strides by the create-time channel count, so a smaller
+  // channel dim would read past the end of the input buffer.
+  std::vector<float> input(1 * 8 * 8 * 1);
+  subgraph.ReshapeExternalTensor({1, 8, 8, 1}, input.data(), input_id)
+      .ReshapeRuntime();
+  EXPECT_EQ(subgraph.Status(), xnn_status_invalid_parameter);
+}
+
 }  // namespace xnnpack

test/subgraph/deconvolution-2d.cc

@@ -341,4 +341,46 @@ TEST(Deconvolution2DQD8F32QC8W, test) {
   TestImpl<float, qcint8, float>(/*convert_to=*/xnn_datatype_qdint8);
 }
 
+TEST(Deconvolution2D, reshape_rejects_input_channel_mismatch) {
+  ASSERT_EQ(xnn_status_success, xnn_initialize(nullptr /* allocator */));
+
+  DeconvolutionParams params;
+  params.padding = {0, 0, 0, 0};
+  params.adjustment = {0, 0};
+  params.kernel = {3, 3};
+  params.upsampling = {1, 1};
+  params.dilation = {1, 1};
+  params.groups = 1;
+  params.group_input_channels = 4;
+  params.group_output_channels = 8;
+
+  const uint32_t input_id = 0;
+  const uint32_t filter_id = 1;
+  const uint32_t bias_id = 2;
+  const uint32_t output_id = 3;
+
+  Tensor<float> filter(std::vector<size_t>{8, 3, 3, 4});
+  filter.fill(0.0f);
+  Tensor<float> bias(std::vector<size_t>{8});
+  bias.fill(0.0f);
+
+  SubgraphTester subgraph(4);
+  subgraph.AddInputTensor(4, xnn_datatype_fp32, input_id)
+      .AddStaticTensor(filter.extents(), filter_id, filter.base())
+      .AddStaticTensor(bias.extents(), bias_id, bias.base())
+      .AddOutputTensor(4, xnn_datatype_fp32, output_id)
+      .AddDeconvolution2D(params, input_id, filter_id, bias_id, output_id);
+  if (subgraph.CreateRuntime() == xnn_status_unsupported_hardware) {
+    GTEST_SKIP();
+  }
+
+  // Create-time input channels are groups*group_input_channels == 4; a
+  // single-channel runtime input must be rejected before the indirection buffer
+  // strides past the end of the smaller input.
+  std::vector<float> input(1 * 8 * 8 * 1);
+  subgraph.ReshapeExternalTensor({1, 8, 8, 1}, input.data(), input_id)
+      .ReshapeRuntime();
+  EXPECT_EQ(subgraph.Status(), xnn_status_invalid_parameter);
+}
+
 }  // namespace xnnpack

test/subgraph/depthwise-convolution-2d.cc

@@ -350,4 +350,45 @@ TEST(DepthwiseConvolution2DF16F32, test) {
 }
 TEST(DepthwiseConvolution2DF32, test) { TestImpl<float, float, float>(); }
 
+TEST(DepthwiseConvolution2D, reshape_rejects_input_channel_mismatch) {
+  ASSERT_EQ(xnn_status_success, xnn_initialize(nullptr /* allocator */));
+
+  DepthwiseConvolutionParams params;
+  params.padding = {0, 0, 0, 0};
+  params.kernel = {3, 3};
+  params.subsampling = {1, 1};
+  params.dilation = {1, 1};
+  params.depth_multiplier = 1;
+  params.input_channels = 4;
+
+  const uint32_t input_id = 0;
+  const uint32_t filter_id = 1;
+  const uint32_t bias_id = 2;
+  const uint32_t output_id = 3;
+
+  Tensor<float> filter(std::vector<size_t>{1, 3, 3, 4});
+  filter.fill(0.0f);
+  Tensor<float> bias(std::vector<size_t>{4});
+  bias.fill(0.0f);
+
+  SubgraphTester subgraph(4);
+  subgraph.AddInputTensor(4, xnn_datatype_fp32, input_id)
+      .AddStaticTensor(filter.extents(), filter_id, filter.base())
+      .AddStaticTensor(bias.extents(), bias_id, bias.base())
+      .AddOutputTensor(4, xnn_datatype_fp32, output_id)
+      .AddDepthwiseConvolution2D(params, input_id, filter_id, bias_id,
+                                 output_id);
+  if (subgraph.CreateRuntime() == xnn_status_unsupported_hardware) {
+    GTEST_SKIP();
+  }
+
+  // The operator was created for input_channels == 4. A single-channel runtime
+  // input must be rejected before the indirection buffer strides past the end
+  // of the smaller input.
+  std::vector<float> input(1 * 8 * 8 * 1);
+  subgraph.ReshapeExternalTensor({1, 8, 8, 1}, input.data(), input_id)
+      .ReshapeRuntime();
+  EXPECT_EQ(subgraph.Status(), xnn_status_invalid_parameter);
+}
+
 }  // namespace xnnpack