Leaky ReLU operator in Subgraph API

PiperOrigin-RevId: 315817869

7 files changed, 113 insertions, 0 deletions

diff --git a/BUILD.bazel b/BUILD.bazel
index 268969fda..2029c89cb 100644
--- a/BUILD.bazel
+++ b/BUILD.bazel
@@ -86,6 +86,7 @@ SUBGRAPH_SRCS = [
     "src/subgraph/fully-connected.c",
     "src/subgraph/floor.c",
     "src/subgraph/hardswish.c",
+    "src/subgraph/leaky-relu.c",
     "src/subgraph/max-pooling-2d.c",
     "src/subgraph/maximum2.c",
     "src/subgraph/minimum2.c",
diff --git a/CMakeLists.txt b/CMakeLists.txt
index ca1cf995d..5cc7f42da 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -193,6 +193,7 @@ SET(XNNPACK_SUBGRAPH_SRCS
   src/subgraph/fully-connected.c
   src/subgraph/floor.c
   src/subgraph/hardswish.c
+  src/subgraph/leaky-relu.c
   src/subgraph/max-pooling-2d.c
   src/subgraph/maximum2.c
   src/subgraph/minimum2.c
diff --git a/include/xnnpack.h b/include/xnnpack.h
index 96a438622..646bda905 100644
--- a/include/xnnpack.h
+++ b/include/xnnpack.h
@@ -862,6 +862,21 @@ enum xnn_status xnn_define_hardswish(
   uint32_t output_id,
   uint32_t flags);
 
+/// Define a Leaky ReLU Node and add it to a Subgraph.
+///
+/// @param subgraph - a Subgraph object that will own the created Node.
+/// @param negative_slope - scale factor for negative input 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.
+/// @param flags - binary features of the Leaky ReLU Node. No supported flags are currently defined.
+enum xnn_status xnn_define_leaky_relu(
+  xnn_subgraph_t subgraph,
+  float negative_slope,
+  uint32_t input_id,
+  uint32_t output_id,
+  uint32_t flags);
+
 /// Define a Negate Node and add it to a Subgraph.
 ///
 /// @param subgraph - a Subgraph object that will own the created Node.
diff --git a/src/runtime.c b/src/runtime.c
index 394735bce..96757c577 100644
--- a/src/runtime.c
+++ b/src/runtime.c
@@ -394,6 +394,21 @@ enum xnn_status xnn_create_runtime_v2(
         runtime->opdata[i].inputs[0] = node->inputs[0];
         runtime->opdata[i].outputs[0] = node->outputs[0];
         break;
+      case xnn_node_type_leaky_relu:
+        status = xnn_create_leaky_relu_nc_f32(
+          values[node->inputs[0]].shape.dim[values[node->inputs[0]].shape.num_dims - 1] /* channels */,
+          values[node->inputs[0]].shape.dim[values[node->inputs[0]].shape.num_dims - 1] /* input stride */,
+          values[node->inputs[0]].shape.dim[values[node->inputs[0]].shape.num_dims - 1] /* output stride */,
+          node->params.leaky_relu.negative_slope,
+          node->flags,
+          &runtime->opdata[i].operator_object);
+        if (status != xnn_status_success) {
+          goto error;
+        }
+        runtime->opdata[i].batch_size = product_non_channel_dims(&values[node->inputs[0]].shape);
+        runtime->opdata[i].inputs[0] = node->inputs[0];
+        runtime->opdata[i].outputs[0] = node->outputs[0];
+        break;
       case xnn_node_type_max_pooling_2d:
         status = xnn_create_max_pooling2d_nhwc_f32(
           node->params.pooling_2d.padding_top,
@@ -861,6 +876,16 @@ enum xnn_status xnn_setup_runtime(
           runtime->blobs[opdata->outputs[0]].data,
           runtime->threadpool);
         break;
+      case xnn_operator_type_leaky_relu_nc_f32:
+        assert(runtime->blobs[opdata->inputs[0]].data != NULL);
+        assert(runtime->blobs[opdata->outputs[0]].data != NULL);
+        status = xnn_setup_leaky_relu_nc_f32(
+          opdata->operator_object,
+          opdata->batch_size,
+          runtime->blobs[opdata->inputs[0]].data,
+          runtime->blobs[opdata->outputs[0]].data,
+          runtime->threadpool);
+        break;
       case xnn_operator_type_max_pooling_nhwc_f32:
         assert(runtime->blobs[opdata->inputs[0]].data != NULL);
         assert(runtime->blobs[opdata->outputs[0]].data != NULL);
diff --git a/src/subgraph-strings.c b/src/subgraph-strings.c
index 27e641fa7..d72673ba5 100644
--- a/src/subgraph-strings.c
+++ b/src/subgraph-strings.c
@@ -48,6 +48,8 @@ const char* xnn_node_type_to_string(enum xnn_node_type type) {
       return "Floor";
     case xnn_node_type_hardswish:
       return "HardSwish";
+    case xnn_node_type_leaky_relu:
+      return "Leaky ReLU";
     case xnn_node_type_maximum2:
       return "Maximum2";
     case xnn_node_type_minimum2:
diff --git a/src/subgraph/leaky-relu.c b/src/subgraph/leaky-relu.c
new file mode 100644
index 000000000..b4d2b0bdc
--- /dev/null
+++ b/src/subgraph/leaky-relu.c
@@ -0,0 +1,65 @@
+// Copyright 2020 Google LLC
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include <math.h>
+#include <stddef.h>
+#include <stdint.h>
+
+#include <xnnpack.h>
+#include <xnnpack/log.h>
+#include <xnnpack/params.h>
+#include <xnnpack/subgraph.h>
+
+
+enum xnn_status xnn_define_leaky_relu(
+  xnn_subgraph_t subgraph,
+  float negative_slope,
+  uint32_t input_id,
+  uint32_t output_id,
+  uint32_t flags)
+{
+  if (!xnn_params.initialized) {
+    xnn_log_error("failed to define %s operator: XNNPACK is not initialized",
+      xnn_node_type_to_string(xnn_node_type_leaky_relu));
+    return xnn_status_uninitialized;
+  }
+
+  if (!isfinite(negative_slope)) {
+    xnn_log_error(
+      "failed to create %s operator with %f negative slope: finite number expected",
+      xnn_node_type_to_string(xnn_node_type_leaky_relu),
+      negative_slope);
+    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",
+      xnn_node_type_to_string(xnn_node_type_leaky_relu), input_id);
+    return xnn_status_invalid_parameter;
+  }
+
+  if (output_id >= subgraph->num_values) {
+    xnn_log_error(
+      "failed to define %s operator with output ID #%" PRIu32 ": invalid Value ID",
+      xnn_node_type_to_string(xnn_node_type_leaky_relu), output_id);
+    return xnn_status_invalid_parameter;
+  }
+
+  struct xnn_node* node = xnn_subgraph_new_node(subgraph);
+  if (node == NULL) {
+    return xnn_status_out_of_memory;
+  }
+
+  node->type = xnn_node_type_leaky_relu;
+  node->params.leaky_relu.negative_slope = negative_slope;
+  node->num_inputs = 1;
+  node->inputs[0] = input_id;
+  node->num_outputs = 1;
+  node->outputs[0] = output_id;
+  node->flags = flags;
+
+  return xnn_status_success;
+}
diff --git a/src/xnnpack/subgraph.h b/src/xnnpack/subgraph.h
index 93188ec21..6688a8025 100644
--- a/src/xnnpack/subgraph.h
+++ b/src/xnnpack/subgraph.h
@@ -83,6 +83,7 @@ enum xnn_node_type {
   xnn_node_type_fully_connected,
   xnn_node_type_floor,
   xnn_node_type_hardswish,
+  xnn_node_type_leaky_relu,
   xnn_node_type_max_pooling_2d,
   xnn_node_type_maximum2,
   xnn_node_type_minimum2,
@@ -161,6 +162,9 @@ struct xnn_node {
       uint32_t dilation_width;
     } pooling_2d;
     struct {
+      float negative_slope;
+    } leaky_relu;
+    struct {
       size_t pre_paddings[XNN_MAX_TENSOR_DIMS];
       size_t post_paddings[XNN_MAX_TENSOR_DIMS];
       uint32_t padding_value;