Merge branch 'main' into fixBugPrint

Author: yangw-dev

.ci/scripts/benchmark_tooling/analyze_benchmark_stability.py

@@ -221,7 +221,7 @@ jobs:
                       model.use_sdpa_with_kv_cache=true \
                       backend.xnnpack.enabled=true \
                       backend.xnnpack.extended_ops=true \
-                      base.preq_mode="8da4w_output_8da8w" \
+                      base.preq_mode="preq_8da4w_out_8da8w" \
                       base.preq_group_size=32 \
                       export.max_seq_length=2048 \
                       export.max_context_length=2048 \
@@ -247,7 +247,7 @@ jobs:
                       base.params="${DOWNLOADED_PATH}/params.json" \
                       quantization.use_qat=true \
                       base.use_lora=16 \
-                      base.preq_mode="8da4w_output_8da8w" \
+                      base.preq_mode="preq_8da4w_out_8da8w" \
                       base.preq_group_size=32 \
                       base.preq_embedding_quantize=\'8,0\' \
                       model.use_sdpa_with_kv_cache=true \

.github/workflows/android-perf.yml

@@ -19,6 +19,7 @@
     op_dynamic_dequantize_ops,
     op_dynamic_quantize_ops,
     op_elu,
+    op_exp,
     op_floor,
     op_gelu,
     op_hardswish,

backends/xnnpack/operators/__init__.py

@@ -0,0 +1,52 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Dict
+
+import torch
+from executorch.backends.xnnpack.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import (
+    XNNExp,
+    XNNGraph,
+    XNode,
+)
+from executorch.backends.xnnpack.utils.utils import get_input_node
+
+
+@register_node_visitor
+class ExpVisitor(NodeVisitor):
+    target = "aten.exp.default"
+
+    def __init__(self, *args) -> None:
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        xnn_graph: XNNGraph,
+        vals_to_ids: Dict[torch.fx.Node, int],
+        debug_handle: int,
+    ) -> None:
+        self.define_nodes_tensor_inputs_outputs(node, xnn_graph, vals_to_ids)
+
+        # input
+        input_id = vals_to_ids[get_input_node(node, 0)]
+
+        # output
+        output_id = vals_to_ids[node]
+
+        ser_node = XNode(
+            xnode_union=XNNExp(
+                input_id=input_id,
+                output_id=output_id,
+                flags=0,
+            ),
+            debug_handle=debug_handle,
+        )
+        xnn_graph.xnodes.append(ser_node)

backends/xnnpack/operators/op_exp.py

@@ -25,10 +25,11 @@
     ConstantPadConfig,
     DeQuantizedPerTensorConfig,
     DivConfig,
+    # EluConfig,
+    ExpConfig,
     FloorConfig,
     GeluConfig,
     HardswishConfig,
-    # EluConfig,
     HardtanhConfig,
     LeakyReLUConfig,
     LogConfig,
@@ -80,6 +81,7 @@
     ClampConfig,
     DivConfig,
     # EluConfig, # Waiting for PyTorch Pin Update
+    ExpConfig,
     FloorConfig,
     GeluConfig,
     HardtanhConfig,

backends/xnnpack/partition/config/__init__.py

@@ -348,6 +348,13 @@ def get_original_aten(self) -> Optional[torch._ops.OpOverload]:
         return torch.ops.aten.upsample_bilinear2d.vec
 
 
+class ExpConfig(GenericNodePartitionerConfig):
+    target_name = "exp.default"
+
+    def supported_precision_types(self) -> List[ConfigPrecisionType]:
+        return [ConfigPrecisionType.FP32]
+
+
 class FloorConfig(GenericNodePartitionerConfig):
     target_name = "floor.default"
 

backends/xnnpack/partition/config/generic_node_configs.py

@@ -67,6 +67,7 @@
     exir_ops.edge.aten.log.default,
     exir_ops.edge.aten.gelu.default,
     exir_ops.edge.aten.tanh.default,
+    exir_ops.edge.aten.exp.default,
 ]
 
 SUPPORTED_MODULES = [

backends/xnnpack/partition/configs.py

@@ -1162,6 +1162,36 @@ Error defineArgMaxPooling2dNode(
   return Error::Ok;
 }
 
+/*
+Define serialized exp node into the subgraph, using the remapped ids
+to map the serialized ids, to the new ids generated when defining the
+tensor value
+*/
+Error defineExpNode(
+    xnn_subgraph_t subgraph_ptr,
+    const std::unordered_map<uint32_t, uint32_t>& remapped_ids,
+    const NodePtr node,
+    const fb_xnnpack::XNNGraph* graph) noexcept {
+  MAYBE_UNUSED(graph);
+
+  auto graph_node = node->xnode_union_as_XNNExp();
+
+  xnn_status status = xnn_define_exp(
+      subgraph_ptr,
+      remapped_ids.at(graph_node->input_id()),
+      remapped_ids.at(graph_node->output_id()),
+      graph_node->flags());
+
+  ET_CHECK_OR_RETURN_ERROR(
+      status == xnn_status_success,
+      Internal,
+      "Failed to create exp node %i with code: %s",
+      node->debug_handle(),
+      xnn_status_to_string(status));
+
+  return Error::Ok;
+}
+
 /*
 Define serialized tanh node into the subgraph, using the remapped ids
 to map the serialized ids, to the new ids generated when defining the
@@ -1733,6 +1763,7 @@ DefineNodeFunc getDefineNodeFunc(fb_xnnpack::XNodeUnion nodeType) {
     _DEFINE(Clamp)
     _DEFINE(LeakyReLU)
     _DEFINE(ELU)
+    _DEFINE(Exp)
     _DEFINE(Abs)
     _DEFINE(Floor)
     _DEFINE(PReLU)

backends/xnnpack/runtime/XNNCompiler.cpp

@@ -155,6 +155,7 @@ union XNodeUnion {
   XNNLog: _XNNNode1x1,
   XNNGelu: _XNNNode1x1,
   XNNTanh: _XNNNode1x1,
+  XNNExp: _XNNNode1x1,
 }
 
 union XValueUnion {

backends/xnnpack/serialization/runtime_schema.fbs

@@ -151,6 +151,7 @@ union XNodeUnion {
   XNNLog: _XNNNode1x1,
   XNNGelu: _XNNNode1x1,
   XNNTanh: _XNNNode1x1,
+  XNNExp: _XNNNode1x1,
 }
 
 union XValueUnion {