Support eq.Scalar

Differential Revision: D86891707

Author: cccclai

backends/qualcomm/_passes/replace_inplace_copy.py

@@ -466,6 +466,7 @@ def define_tensor(
             tensor_source_node, target_build_node
         )
         dtype = self.get_data_type(tensor, quant_configs)
+        print(f"tensor_name: {tensor_name}, tensor_type: {tensor_type}, dtype: {dtype}")
         if isinstance(tensor, torch._subclasses.fake_tensor.FakeTensor):
             tensor_wrapper = PyQnnWrapper.TensorWrapper(
                 tensor_name,

backends/qualcomm/builders/node_visitor.py

@@ -8,15 +8,22 @@
 import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper
 
 import torch
+from executorch.exir.dialects._ops import ops as exir_ops
 
 from .node_visitor import NodeVisitor
 from .node_visitor_manager import register_node_visitor
 from .qnn_constants import OpElementWiseEqual, QNN_OP_PACKAGE_NAME_QTI_AISW
-
+from executorch.backends.qualcomm.utils.constants import (
+    QCOM_QUANT_ATTRS,
+    QCOM_QUANT_MAX,
+    QCOM_QUANT_MIN,
+    QCOM_SCALE,
+    QCOM_ZERO_POINT,
+)
 
 @register_node_visitor
 class Equal(NodeVisitor):
-    target = ["aten.eq.Tensor"]
+    target = ["aten.eq.Tensor", "aten.eq.Scalar"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)
@@ -37,11 +44,43 @@ def define_node(
         output_tensors = [output_tensor_wrapper]
 
         input_tensors = []
-        for index in range(2):
-            input_node = self.get_node(node.args[index])
-            input_tensor = self.get_tensor(input_node, node)
-            tensor_type = PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE
-
+        for index, arg in enumerate(node.args):
+            if isinstance(arg, torch.fx.Node):
+                # Normal tensor input
+                input_node = self.get_node(arg)
+                input_tensor = self.get_tensor(input_node, node)
+                tensor_type = PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE
+            else:
+                assert index == 1, f"eq op arg at index 1 has to be int, but the type is {type(arg)}"
+                assert isinstance(arg, int), f"eq op arg {arg} has to be int , but the type is {type(arg)}"
+                print(f"arg is {arg}, type is {type(arg)}")
+                # Handle scalar input (e.g., int or float)
+                scalar = arg
+                scalar_value = float(scalar)
+                input_tensor = torch.tensor(
+                    scalar_value, dtype=torch.int32
+                )
+                input_node = torch.fx.Node(
+                    node.graph,
+                    node.name + "_runtime_scalar",
+                    "call_function",
+                    exir_ops.edge.aten.scalar_tensor.default,
+                    (),  # args
+                    {},  # kwargs
+                )
+                tensor_type = PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC
+                if quant_attrs := node.args[index ^ 1].meta.get(QCOM_QUANT_ATTRS):
+                    quant_attrs = quant_attrs.copy()
+                    quant_range = (
+                        quant_attrs[QCOM_QUANT_MAX] - quant_attrs[QCOM_QUANT_MIN]
+                    )
+                    quant_attrs[QCOM_ZERO_POINT] = (
+                        0 if scalar >= 0 else quant_attrs[QCOM_QUANT_MAX]
+                    )
+                    quant_attrs[QCOM_SCALE] = (
+                        scalar / quant_range if scalar >= 0 else -scalar / quant_range
+                    )
+                    input_node.meta[QCOM_QUANT_ATTRS] = quant_attrs
             input_tensor_wrapper = self.define_tensor(
                 input_node,
                 node,

backends/qualcomm/builders/op_eq.py

@@ -944,6 +944,17 @@ def forward(self, x):
         return x == self.constant
 
 
+class EqualFromInplaceCopyDecomp(torch.nn.Module):
+    def __init__(self, hidden_size=4):
+        super().__init__()
+        # a small state tensor
+        self.register_buffer("h", torch.zeros((1, hidden_size)))
+
+    def forward(self, x):
+        self.h[0] = x
+        return self.h[0]
+
+
 class ExpandCopy(torch.nn.Module):
     def __init__(self):
         super().__init__()

backends/qualcomm/tests/models.py

@@ -765,6 +765,19 @@ def test_qnn_backend_equal(self):
                     test[QCOM_MODULE], test[QCOM_SAMPLE_INPUTS]
                 )
 
+    def test_qnn_backend_equal_debug(self):
+        test_comb = [
+            {
+                QCOM_MODULE: EqualFromInplaceCopyDecomp(),  # noqa: F405
+                QCOM_SAMPLE_INPUTS: (torch.tensor([1.0, 2.0, 3.0, 4.0]), ),
+            },
+        ]
+        for i, test in enumerate(test_comb):
+            with self.subTest(i=i):
+                self.lower_module_and_test_output(
+                    test[QCOM_MODULE], test[QCOM_SAMPLE_INPUTS]
+                )
+
     def test_qnn_backend_expand(self):
         modules = [ExpandAs(), ExpandCopy()]  # noqa: F405
         sample_inputs = [
@@ -2842,6 +2855,26 @@ def test_qnn_backend_equal(self):
                 )
                 self.lower_module_and_test_output(module, test[QCOM_SAMPLE_INPUTS])
 
+    def test_qnn_backend_equal_debug(self):
+        test_comb = [
+            {
+                QCOM_MODULE: EqualFromInplaceCopyDecomp(),  # noqa: F405
+                QCOM_SAMPLE_INPUTS: (torch.tensor([1.0, 2.0, 3.0, 4.0]), ),
+            },
+        ]
+        for i, test in enumerate(test_comb):
+            with self.subTest(i=i):
+                module = self.get_qdq_module(
+                    test[QCOM_MODULE], test[QCOM_SAMPLE_INPUTS]
+                )
+
+                print("quantized module")
+                module.print_readable()
+
+                self.lower_module_and_test_output(
+                    module, test[QCOM_SAMPLE_INPUTS]
+                )
+
     def test_qnn_backend_expand(self):
         modules = [ExpandAs(), ExpandCopy()]  # noqa: F405
         sample_inputs = [

backends/qualcomm/tests/test_qnn_delegate.py

@@ -0,0 +1,103 @@
+# 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.
+
+# pyre-strict
+
+from typing import Set
+
+import torch
+from executorch.exir.dialects._ops import ops
+from torch.export import ExportedProgram
+
+
+def _is_index_put(node: torch.fx.Node) -> bool:
+    """Check if a node is an index_put operation."""
+    return node.op == "call_function" and node.target in (
+        torch.ops.aten.index_put.default,
+        ops.edge.aten.index_put.default,
+    )
+
+
+def _is_safe_to_reinplace(
+    node: torch.fx.Node,
+    later_nodes: Set[torch.fx.Node],
+    inputs: Set[torch.fx.Node],
+    mutable_inputs: Set[torch.fx.Node],
+) -> bool:
+    # This node is used later in the graph so we can't reinplace it
+    # There is probably a faster way to do this but this works for now.
+    if node in later_nodes:
+        return False
+    # If its not an input then we can reinplace it
+    if node not in inputs:
+        return True
+    # If its a mutable input then we can reinplace it
+    elif node in mutable_inputs:
+        return True
+    else:  # input but not mutable input
+        return False
+
+
+def _is_mutable_user_input(
+    node: torch.fx.Node, exported_program: ExportedProgram
+) -> bool:
+    return (
+        node.target in exported_program.graph_signature.user_inputs_to_mutate.values()
+    )
+
+
+def _is_mutable_buffer(node: torch.fx.Node, exported_program: ExportedProgram) -> bool:
+    if node.target not in exported_program.graph_signature.inputs_to_buffers:
+        return False
+    buf = exported_program.graph_signature.inputs_to_buffers[node.target]
+    return buf in exported_program.graph_signature.buffers_to_mutate.values()
+
+
+def functionalize_pass(ep: ExportedProgram) -> ExportedProgram:
+    """
+    Pass that loops over nodes in an exported program and collects the first argument
+    of every call_function node that is a view_copy operation.
+
+    Args:
+        exported_program: The ExportedProgram to analyze
+
+    Returns:
+        Set of nodes that are first arguments to view_copy operations
+    """
+    seen_nodes: Set[torch.fx.Node] = set()
+    # Get all placeholders
+    inputs = set()
+    for node in ep.graph.nodes:
+        if node.op == "placeholder":
+            inputs.add(node)
+    # Get all inputs that we could potentially mutate
+    mutable_nodes = set(
+        [
+            node
+            for node in inputs
+            if _is_mutable_user_input(node, ep) or _is_mutable_buffer(node, ep)
+        ]
+    )
+
+    results = set()
+    for node in reversed(ep.graph.nodes):
+        if _is_index_put(node):
+            # Check if this index_put node is safe to inplace
+            # The first argument is the base tensor being indexed into
+            first_arg = node.args[0]
+            if _is_safe_to_reinplace(first_arg, seen_nodes, inputs, mutable_nodes):
+                # This index_put is safe to reinplace
+                with ep.graph.inserting_before(node):
+                    new_node = ep.graph.call_function(
+                        ops.edge.aten.index_put_.default, args=node.args
+                    )
+                    new_node.meta["val"] = node.meta["val"]
+                    node.replace_all_uses_with(new_node)
+                    ep.graph.erase_node(node)
+                results.add(first_arg)
+        elif node.op == "call_function":
+            seen_nodes.update(node.all_input_nodes)
+    return ep