WebNN: Define pad operator in mojo

The pad operator inflates the tensor with constant or mirrored
values on the edges.

This CL moves the validation steps of pad to //components/,
implements CreatePadOperation() for creating pad mojo operation
and adds operator validation on service side.

Bug: 1273291
Change-Id: I569998d76c06c6f47030b157c8ffc270567403f1
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/4940504
Reviewed-by: Rafael Cintron <rafael.cintron@microsoft.com>
Reviewed-by: ningxin hu <ningxin.hu@intel.com>
Reviewed-by: Jiewei Qian <qjw@chromium.org>
Reviewed-by: Alex Gough <ajgo@chromium.org>
Commit-Queue: Mingming1 Xu <mingming1.xu@intel.com>
Cr-Commit-Position: refs/heads/main@{#1212525}

components/ml/webnn/graph_validation_utils.cc

@@ -380,6 +380,41 @@ base::expected<Operand, std::string> ValidateConv2dAndInferOutput(
   return Operand(input.data_type, std::move(output_shape));
 }
 
+base::expected<Operand, std::string> ValidatePadAndInferOutput(
+    const Operand& input,
+    base::span<const uint32_t> beginning_padding,
+    base::span<const uint32_t> ending_padding) {
+  // Validate the beginning_padding and ending_padding.
+  const auto input_shape = input.dimensions;
+  auto input_rank = input_shape.size();
+  if (beginning_padding.size() != input_rank) {
+    return base::unexpected(
+        "The length of beginningPadding must be "
+        "equal to the rank of the input tensor.");
+  }
+  if (ending_padding.size() != input_rank) {
+    return base::unexpected(
+        "The length of endingPadding must be "
+        "equal to the rank of the input tensor.");
+  }
+
+  // Infer the output.
+  // Each dimension of the output tensor can be calculated as follow:
+  // input_size = input_shape[i];
+  // output_size = beginning_padding + input_size + ending_padding.
+  std::vector<uint32_t> output_shape(input_rank);
+  for (size_t i = 0; i < input_rank; ++i) {
+    auto checked_output_size = base::MakeCheckedNum<uint32_t>(input_shape[i]) +
+                               beginning_padding[i] + ending_padding[i];
+    if (!checked_output_size.AssignIfValid(&output_shape[i])) {
+      return base::unexpected(base::StringPrintf(
+          "The padding of dimension (%zu) is too large.", i));
+    }
+  }
+
+  return Operand(input.data_type, std::move(output_shape));
+}
+
 base::expected<Operand, std::string> ValidatePool2dAndInferOutput(
     const Operand& input,
     const Pool2dAttributes& attributes) {

components/ml/webnn/graph_validation_utils.h

@@ -210,6 +210,13 @@ base::expected<Operand, std::string> ValidateConv2dAndInferOutput(
     const Operand& filter,
     const Conv2dAttributes& attributes);
 
+// Validate and infer output information of pad operator defined in
+// WebIDL here https://www.w3.org/TR/webnn/#api-mlgraphbuilder-pad
+base::expected<Operand, std::string> ValidatePadAndInferOutput(
+    const Operand& input,
+    base::span<const uint32_t> beginning_padding,
+    base::span<const uint32_t> ending_padding);
+
 // Validate and infer output information of 2-D pooling operator defined in
 // WebIDL here https://www.w3.org/TR/webnn/#api-mlgraphbuilder-pool2d
 base::expected<Operand, std::string> ValidatePool2dAndInferOutput(

services/webnn/dml/graph_impl.cc

@@ -96,6 +96,8 @@ std::string OpTagToString(Operation::Tag tag) {
       return "conv2d";
     case Operation::Tag::kElementWiseBinary:
       return "element-wise binary";
+    case Operation::Tag::kPad:
+      return "pad";
     case Operation::Tag::kPool2d:
       return "pool2d";
     case Operation::Tag::kResample2d:

services/webnn/public/mojom/webnn_graph.mojom

@@ -146,6 +146,41 @@ struct ElementWiseBinary  {
   uint64 output_operand;
 };
 
+// Specifies the different ways to pad a tensor. The padding value is only
+// specified when the mode is "constant".
+struct ConstantPadding {
+  float value = 0;
+};
+struct EdgePadding {};
+struct ReflectionPadding {};
+struct SymmetricPadding {};
+
+union PaddingMode {
+  ConstantPadding constant;
+  EdgePadding edge;
+  ReflectionPadding reflection;
+  SymmetricPadding symmetric;
+};
+
+// Represents a pad operation which inflates the input tensor with constant or
+// mirrored values on the edges.
+struct Pad {
+  // The id of input operand is used to get the `Operand` description from
+  // `GraphInfo.id_to_operand_map`.
+  uint64 input_operand_id;
+  // The id of output operand is used to get the `Operand` description from
+  // `GraphInfo.id_to_operand_map`.
+  uint64 output_operand_id;
+  // The number of padding values to add at the beginning of each input
+  // dimension. The array length should be equal to the rank of input tensor.
+  array<uint32> beginning_padding;
+  // The number of padding values to add at the ending of each input
+  // dimension. The array length should be equal to the rank of input tensor.
+  array<uint32> ending_padding;
+
+  PaddingMode mode;
+};
+
 // Represents an average or max pooling operation across all the elements with
 // moving window over the input tensor.
 // This struct also contains the attributes of pool2d operator, but the
@@ -317,6 +352,7 @@ union Operation {
   Concat concat;
   Conv2d conv2d;
   ElementWiseBinary element_wise_binary;
+  Pad pad;
   Pool2d pool2d;
   Relu relu;
   Resample2d resample2d;

services/webnn/webnn_graph_impl.cc

@@ -385,6 +385,28 @@ bool ValidateGemm(const IdToOperandMap& id_to_operand_map,
   return true;
 }
 
+bool ValidatePad(const IdToOperandMap& id_to_operand_map,
+                 const mojom::PadPtr& pad) {
+  auto* input = GetMojoOperand(id_to_operand_map, pad->input_operand_id);
+  auto* output = GetMojoOperand(id_to_operand_map, pad->output_operand_id);
+  if (!input || !output || output == input) {
+    // The pad operator is invalid.
+    return false;
+  }
+
+  auto validated_output =
+      ValidatePadAndInferOutput(ConvertToComponentOperand(input),
+                                pad->beginning_padding, pad->ending_padding);
+  if (!validated_output.has_value()) {
+    return false;
+  }
+  if (validated_output != ConvertToComponentOperand(output)) {
+    return false;
+  }
+
+  return true;
+}
+
 bool ValidatePool2d(const IdToOperandMap& id_to_operand_map,
                     const mojom::Pool2dPtr& pool2d) {
   auto* input = GetMojoOperand(id_to_operand_map, pool2d->input_operand_id);
@@ -628,6 +650,8 @@ bool ValidateOperation(const IdToOperandMap& id_to_operand_map,
     case mojom::Operation::Tag::kElementWiseBinary:
       return ValidateElementWiseBinary(id_to_operand_map,
                                        operation->get_element_wise_binary());
+    case mojom::Operation::Tag::kPad:
+      return ValidatePad(id_to_operand_map, operation->get_pad());
     case mojom::Operation::Tag::kPool2d:
       return ValidatePool2d(id_to_operand_map, operation->get_pool2d());
     case mojom::Operation::Tag::kResample2d:

services/webnn/webnn_graph_impl_unittest.cc

@@ -919,6 +919,113 @@ TEST_F(WebNNGraphImplTest, GemmTest) {
   }
 }
 
+struct PadTester {
+  OperandInfo input;
+  std::vector<uint32_t> beginning_padding;
+  std::vector<uint32_t> ending_padding;
+  mojom::PaddingMode::Tag mode = mojom::PaddingMode::Tag::kConstant;
+  float value = 0;
+  OperandInfo output;
+  bool expected;
+
+  void Test() {
+    // Build the graph with mojo type.
+    GraphInfoBuilder builder;
+    uint64_t input_operand_id =
+        builder.BuildInput("input", input.dimensions, input.type);
+    uint64_t output_operand_id =
+        builder.BuildOutput("output", output.dimensions, output.type);
+    builder.BuildPad(input_operand_id, output_operand_id, beginning_padding,
+                     ending_padding, mode, value);
+    EXPECT_EQ(WebNNGraphImpl::ValidateGraph(builder.GetGraphInfo()), expected);
+  }
+};
+
+TEST_F(WebNNGraphImplTest, PadTest) {
+  {
+    // Test pad with default options, beginningPadding = {1, 2} and
+    // endingPadding = {1, 2}.
+    PadTester{.input = {.type = mojom::Operand::DataType::kFloat32,
+                        .dimensions = {2, 3}},
+              .beginning_padding = {1, 2},
+              .ending_padding = {1, 2},
+              .output = {.type = mojom::Operand::DataType::kFloat32,
+                         .dimensions = {4, 7}},
+              .expected = true}
+        .Test();
+  }
+  {
+    // Test pad with mode = "edge", beginningPadding = {1, 2} and
+    // endingPadding = {1, 2}.
+    PadTester{.input = {.type = mojom::Operand::DataType::kFloat32,
+                        .dimensions = {2, 3}},
+              .beginning_padding = {1, 2},
+              .ending_padding = {1, 2},
+              .mode = mojom::PaddingMode::Tag::kEdge,
+              .output = {.type = mojom::Operand::DataType::kFloat32,
+                         .dimensions = {4, 7}},
+              .expected = true}
+        .Test();
+  }
+  {
+    // Test pad with value = 1, beginningPadding = {1, 2} and
+    // endingPadding = {1, 2}.
+    PadTester{.input = {.type = mojom::Operand::DataType::kFloat32,
+                        .dimensions = {2, 3}},
+              .beginning_padding = {1, 2},
+              .ending_padding = {1, 2},
+              .value = 1,
+              .output = {.type = mojom::Operand::DataType::kFloat32,
+                         .dimensions = {4, 7}},
+              .expected = true}
+        .Test();
+  }
+  {
+    // Test the invalid graph when the length of beginningPadding is not equal
+    // to the input rank.
+    PadTester{.input = {.type = mojom::Operand::DataType::kFloat32,
+                        .dimensions = {2, 3}},
+              .beginning_padding = {1},
+              .ending_padding = {1, 2},
+              .output = {.type = mojom::Operand::DataType::kFloat32,
+                         .dimensions = {4, 7}},
+              .expected = false}
+        .Test();
+  }
+  {
+    // Test the invalid graph when the length of endingPadding is not equal to
+    // the input rank.
+    PadTester{.input = {.type = mojom::Operand::DataType::kFloat32,
+                        .dimensions = {2, 3}},
+              .beginning_padding = {1, 0},
+              .ending_padding = {1, 2, 0},
+              .output = {.type = mojom::Operand::DataType::kFloat32,
+                         .dimensions = {4, 7}},
+              .expected = false}
+        .Test();
+  }
+  {
+    // Test the invalid graph when the padding of one dimension is too large.
+    PadTester{.input = {.type = mojom::Operand::DataType::kFloat32,
+                        .dimensions = {2, 3}},
+              .beginning_padding = {2294967295, 0},
+              .ending_padding = {3294967295, 2},
+              .output = {.type = mojom::Operand::DataType::kFloat32,
+                         .dimensions = {1294967294, 5}},
+              .expected = false}
+        .Test();
+  }
+  {
+    // Test the invalid graph when the input is as same as output.
+    GraphInfoBuilder builder;
+    uint64_t input_operand_id =
+        builder.BuildInput("input", {2, 3}, mojom::Operand::DataType::kFloat32);
+    builder.BuildPad(input_operand_id, input_operand_id, {1, 1}, {1, 1},
+                     mojom::PaddingMode::Tag::kConstant, 0);
+    EXPECT_FALSE(WebNNGraphImpl::ValidateGraph(builder.GetGraphInfo()));
+  }
+}
+
 struct Pool2dTester {
   OperandInfo input;
   struct Pool2dAttributes {

services/webnn/webnn_test_utils.cc

@@ -63,6 +63,40 @@ uint64_t GraphInfoBuilder::BuildOutput(const std::string& name,
   return operand_id;
 }
 
+void GraphInfoBuilder::BuildPad(uint64_t input_operand_id,
+                                uint64_t output_operand_id,
+                                const std::vector<uint32_t>& beginning_padding,
+                                const std::vector<uint32_t>& ending_padding,
+                                mojom::PaddingMode::Tag mode,
+                                float value) {
+  mojom::PadPtr pad = mojom::Pad::New();
+  pad->input_operand_id = input_operand_id;
+  pad->output_operand_id = output_operand_id;
+  pad->beginning_padding = beginning_padding;
+  pad->ending_padding = ending_padding;
+  switch (mode) {
+    case mojom::PaddingMode::Tag::kConstant: {
+      auto constant_padding = mojom::ConstantPadding::New();
+      constant_padding->value = value;
+      pad->mode = mojom::PaddingMode::NewConstant(std::move(constant_padding));
+      break;
+    }
+    case mojom::PaddingMode::Tag::kEdge:
+      pad->mode = mojom::PaddingMode::NewEdge(mojom::EdgePadding::New());
+      break;
+    case mojom::PaddingMode::Tag::kReflection:
+      pad->mode =
+          mojom::PaddingMode::NewReflection(mojom::ReflectionPadding::New());
+      break;
+    case mojom::PaddingMode::Tag::kSymmetric:
+      pad->mode =
+          mojom::PaddingMode::NewSymmetric(mojom::SymmetricPadding::New());
+      break;
+  }
+
+  graph_info_->operations.push_back(mojom::Operation::NewPad(std::move(pad)));
+}
+
 void GraphInfoBuilder::BuildSplit(
     uint64_t input_operand_id,
     const std::vector<uint64_t>& output_operand_ids,

services/webnn/webnn_test_utils.h

@@ -121,6 +121,13 @@ class GraphInfoBuilder final {
                               uint64_t rhs_operand,
                               uint64_t output_operand);
 
+  void BuildPad(uint64_t input_operand_id,
+                uint64_t output_operand_id,
+                const std::vector<uint32_t>& beginning_padding,
+                const std::vector<uint32_t>& ending_padding,
+                mojom::PaddingMode::Tag mode,
+                float value);
+
   // A `Pool2dAttributes` type should have the following members:
   // struct Pool2dAttributes {
   //   std::vector<uint32_t> window_dimensions;

third_party/blink/renderer/modules/ml/webnn/ml_graph_builder.cc

@@ -1178,17 +1178,12 @@ MLOperand* MLGraphBuilder::pad(const MLOperand* input,
                                const Vector<uint32_t>& ending_padding,
                                const MLPadOptions* options,
                                ExceptionState& exception_state) {
-  const auto input_rank = input->Dimensions().size();
-  if (beginning_padding.size() != input_rank) {
-    exception_state.ThrowDOMException(DOMExceptionCode::kDataError,
-                                      "The length of beginningPadding must be "
-                                      "equal to the rank of the input tensor.");
-    return nullptr;
-  }
-  if (ending_padding.size() != input_rank) {
-    exception_state.ThrowDOMException(DOMExceptionCode::kDataError,
-                                      "The length of endingPadding must be "
-                                      "equal to the rank of the input tensor.");
+  auto validated_output = webnn::ValidatePadAndInferOutput(
+      ConvertToComponentOperand(input), beginning_padding, ending_padding);
+  if (!validated_output.has_value()) {
+    exception_state.ThrowDOMException(
+        DOMExceptionCode::kDataError,
+        String::FromUTF8(validated_output.error()));
     return nullptr;
   }
 
@@ -1199,28 +1194,13 @@ MLOperand* MLGraphBuilder::pad(const MLOperand* input,
         "constant.");
   }
 
-  // Each dimension of the output tensor can be calculated as follow:
-  // output_size = beginning_padding + input_size + ending_padding.
-  Vector<uint32_t> output_shape(input_rank);
-  for (wtf_size_t i = 0; i < input_rank; i++) {
-    auto checked_output_size =
-        base::MakeCheckedNum<uint32_t>(input->Dimensions()[i]) +
-        beginning_padding[i] + ending_padding[i];
-    if (!checked_output_size.AssignIfValid(&output_shape[i])) {
-      exception_state.ThrowDOMException(
-          DOMExceptionCode::kDataError,
-          String::Format("The padding of dimension (%u) is too large.", i));
-      return nullptr;
-    }
-  }
-
   auto* pad = MakeGarbageCollected<MLPadOperator>(this, beginning_padding,
                                                   ending_padding, options);
   // According to WebNN spec
   // https://www.w3.org/TR/webnn/#api-mlgraphbuilder-pad, the output
   // tensor of pad has the same type as its input.
   auto output = MLOperand::ValidateAndCreateOutput(
-      this, input->Type(), std::move(output_shape), pad);
+      this, input->Type(), Vector<uint32_t>(validated_output->dimensions), pad);
   if (!output.has_value()) {
     exception_state.ThrowDOMException(DOMExceptionCode::kDataError,
                                       output.error());