Code Review and adding some comments

Author: winskuo-quic

backends/qualcomm/_passes/lpai_partition_fallback_support.py

@@ -155,10 +155,11 @@ def get_unsupported_nodes(self, graph_module: torch.fx.GraphModule):
             compiler_specs=self.compiler_specs,
             skip_node_id_set=self.skip_node_id_set,
             skip_node_op_set=self.skip_node_op_set,
+            is_qnn_partitioner=False,
+            phase="LpaiPartitionFallbackSupport",
         )
 
         unsupported_nodes = []
-
         for node in graph_module.graph.nodes:
             if node.op != "call_function":
                 continue
@@ -263,20 +264,28 @@ def handle_back_to_back_nodes(self, graph_module: torch.fx.GraphModule):
                 continue
             if node.target in q_ops or node.target in dq_ops:
                 continue
-            if node.meta.get(QCOM_FALLBACK_NODE):
+
+            if node.meta.get(QCOM_FALLBACK_NODE) and QCOM_QUANT_ATTRS in node.meta:
+
+                # Skip node like get_attr
+                input_call_func_nodes = [
+                    input_node
+                    for input_node in node.all_input_nodes
+                    if input_node.op == "call_function"
+                ]
                 assert all(
-                    input_node.target in dq_ops for input_node in node.all_input_nodes
+                    input_node.target in dq_ops for input_node in input_call_func_nodes
                 ), (
-                    f"Expected every input of fall back node {node.name} to be "
+                    f"Expected every call_function input of fp fall back node {node.name} to be "
                     f"a dequantize op (insert_partition_qdq should have wrapped "
                     f"each input with a Q1/DQ1 pair), got: "
-                    f"{[(n.name, n.target) for n in node.all_input_nodes]}"
+                    f"{[(n.name, n.target) for n in input_call_func_nodes]}"
                 )
-                for dq_node in list(node.all_input_nodes):
+                for dq_node in input_call_func_nodes:
                     q_node = dq_node.args[0]
                     if q_node.op == "placeholder":
                         raise RuntimeError(
-                            "Fallback nodes with weights is currently supported."
+                            "Fallback nodes with weights is not currently supported."
                         )
                     assert (
                         q_node.target in q_ops

backends/qualcomm/builders/node_visitor.py

@@ -118,10 +118,19 @@ def __init__(
         external_ids,
         edge_program: torch.export.ExportedProgram,
         enable_tensor_dump,
+        is_qnn_partitioner=True,
     ) -> None:
+        # is_qnn_partitioner: True for the real qnn_partitioner / qnn_preprocess
+        # flows, where external_ids is built from the program being processed so
+        # graph-input ids resolve correctly.
+        # Set False for lpai_partition_fallback_support pass, where constructor
+        # takes an older exported_program whose node objects differ from the
+        # live graph; the input_{id} naming would mismatch there, so we skip
+        # it during op validation.
         self.external_ids = external_ids or {}
         self.edge_program = edge_program
         self.enable_tensor_dump = enable_tensor_dump
+        self.is_qnn_partitioner = is_qnn_partitioner
 
     def get_node(self, node):
         """
@@ -371,9 +380,11 @@ def get_tensor_type(
         is_output = is_graph_output(node)
         # handle logic for input/output tensors
         if is_input or is_output:
-            assert (
-                node in self.external_ids
-            ), f"Node {node}, is_input: {is_input}, is_output: {is_output}, ext_ids: {self.external_ids.keys()}"
+            # For more info about existence of self.is_qnn_partitioner, check constructor for explanation.
+            assert not self.is_qnn_partitioner or node in self.external_ids, (
+                f"Node {node}, is_input: {is_input}, is_output: {is_output}, "
+                f"ext_ids: {self.external_ids.keys()}"
+            )
             if is_input:
                 return PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_APP_WRITE
             if is_output:
@@ -422,21 +433,24 @@ def get_tensor_name(
         # the input order between QNN and the original graph’s forward function may differ.
         # The `mutbuf_{id}` is utilized for mapping I/O of mutable buffer at runtime.
         # The `output_` is identified as the graph’s output at runtime to prevent confusion with per_tensor_dump.
-        if is_mutable_buffer_input(node, self.edge_program):
-            fqn = self.edge_program.graph_signature.inputs_to_buffers[node.target]
-            position_index = list(
-                self.edge_program.graph_signature.buffers_to_mutate.values()
-            ).index(fqn)
-            tensor_name = f"input_{str(self.external_ids[node])}_mutbuf_{str(position_index)}_{tensor_name}"
-        elif is_graph_input(node, self.edge_program):
-            tensor_name = f"input_{str(self.external_ids[node])}_{tensor_name}"
-        elif is_mutable_buffer_output(node, self.edge_program):
-            position_index = list(
-                self.edge_program.graph_signature.buffers_to_mutate.keys()
-            ).index(node.name)
-            tensor_name = f"output_mutbuf_{position_index}_{tensor_name}"
-        elif is_graph_output(node):
-            tensor_name = f"output_{tensor_name}"
+
+        # For more info about self.is_qnn_partitioner, check constructor for explanation.
+        if self.is_qnn_partitioner:
+            if is_mutable_buffer_input(node, self.edge_program):
+                fqn = self.edge_program.graph_signature.inputs_to_buffers[node.target]
+                position_index = list(
+                    self.edge_program.graph_signature.buffers_to_mutate.values()
+                ).index(fqn)
+                tensor_name = f"input_{str(self.external_ids[node])}_mutbuf_{str(position_index)}_{tensor_name}"
+            elif is_graph_input(node, self.edge_program):
+                tensor_name = f"input_{str(self.external_ids[node])}_{tensor_name}"
+            elif is_mutable_buffer_output(node, self.edge_program):
+                position_index = list(
+                    self.edge_program.graph_signature.buffers_to_mutate.keys()
+                ).index(node.name)
+                tensor_name = f"output_mutbuf_{position_index}_{tensor_name}"
+            elif is_graph_output(node):
+                tensor_name = f"output_{tensor_name}"
 
         # Only add qcom_tensor_name when enable tensor dump.
         # Only runs in qnn_preprocess (not op validation) since that's when

backends/qualcomm/builders/node_visitor_manager.py

@@ -53,6 +53,7 @@ def get_node_visitors(
     edge_program: torch.export.ExportedProgram,
     enable_tensor_dump=False,
     op_package_infos: List[QnnExecuTorchOpPackageInfo] = None,
+    is_qnn_partitioner=True,
 ) -> Dict[str, NodeVisitor]:
     """Create a new class instance at runtime, and put them in a dict"""
     node_to_external_map = generate_node_to_external_map(edge_program)
@@ -62,7 +63,7 @@ def get_node_visitors(
             visitor
         ), f"Expecting a callable class, but got {visitor} of type {type(visitor)}"
         node_visitors[target] = visitor(
-            node_to_external_map, edge_program, enable_tensor_dump
+            node_to_external_map, edge_program, enable_tensor_dump, is_qnn_partitioner
         )
     if op_package_infos:
         custom_ops = []
@@ -73,6 +74,7 @@ def get_node_visitors(
                     node_to_external_map,
                     edge_program,
                     enable_tensor_dump,
+                    is_qnn_partitioner,
                 )
                 node_visitors[op_package_info.custom_op_name] = custom_op_builder
                 custom_ops.append(op_package_info.custom_op_name)

backends/qualcomm/partition/qnn_partitioner.py

@@ -56,16 +56,23 @@ def __init__(
         compiler_specs,
         skip_node_id_set: set = None,
         skip_node_op_set: set = None,
+        is_qnn_partitioner=True,
+        phase="QnnPartitioner",
     ):
         option = generate_qnn_executorch_option(compiler_specs)
         python_options = flatbuffer_to_option(option)
         self.node_visitors = node_visitor_manager.get_node_visitors(
             edge_program,
             op_package_infos=python_options.op_package_options.op_package_infos,
+            is_qnn_partitioner=is_qnn_partitioner,
         )
 
         self.skip_node_op_set = skip_node_op_set
         self.skip_node_id_set = skip_node_id_set
+        self.is_qnn_partitioner = is_qnn_partitioner
+        # Label prefixed to op-support log lines so callers that reuse this
+        # checker (e.g. the LPAI fallback pass) are distinguishable in the logs.
+        self.phase = phase
         self.nodes_to_wrappers = defaultdict(dict)
         self.qnn_manager = get_current_qnn_manager(
             python_options.backend_options.backend_type, compiler_specs
@@ -77,12 +84,12 @@ def is_node_supported(self, _, node: torch.fx.Node) -> bool:
 
         if node.target in to_be_implemented_operator:
             logger.info(
-                f"{node.target.__name__} | Skipped, this op can be supported, please report an issue in https://github.com/pytorch/executorch/issues"
+                f"[{self.phase}] {node.target.__name__} | Skipped, this op can be supported, please report an issue in https://github.com/pytorch/executorch/issues"
             )
             return False
 
         if node.meta.get(QCOM_FALLBACK_NODE, False):
-            logger.info(f"{node.target.__name__} | Forced Fallback")
+            logger.info(f"[{self.phase}] {node.target.__name__} | Forced Fallback")
             return False
 
         if (
@@ -92,14 +99,14 @@ def is_node_supported(self, _, node: torch.fx.Node) -> bool:
             # bypass dequantize op for parameters & buffers
             or node.meta.get(QCOM_BYPASS_NODE, False)
         ):
-            logger.info(f"{node.target.__name__} | Forced Passed")
+            logger.info(f"[{self.phase}] {node.target.__name__} | Forced Passed")
             return True
 
         if (
             node.name in self.skip_node_id_set
             or node.target.__name__ in self.skip_node_op_set
         ):
-            logger.info(f"{node.target.__name__} | Skipped")
+            logger.info(f"[{self.phase}] {node.target.__name__} | Skipped")
             return False
 
         supported = False
@@ -121,7 +128,7 @@ def is_node_supported(self, _, node: torch.fx.Node) -> bool:
             )
 
         self.nodes_to_wrappers.clear()
-        logger.info(f"{node.target.__name__} | {supported}")
+        logger.info(f"[{self.phase}] {node.target.__name__} | {supported}")
         return supported
 
 

backends/qualcomm/tests/models.py

@@ -2325,6 +2325,17 @@ def forward(self, x):
         return torch.relu(x) + torch.sqrt(x)
 
 
+class SkipIntNode(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, y):
+        a = torch.add(x, y)
+        b = torch.mul(a, a)
+        c = torch.add(b, y)
+        return c
+
+
 class SkipMultiInput(torch.nn.Module):
 
     def __init__(self):

backends/qualcomm/tests/test_qnn_delegate.py

@@ -6796,6 +6796,19 @@ def test_qnn_backend_skip_node_id_partitioner_last_node(self):
             skip_node_id_set={"aten_add_tensor"},
         )
 
+    def test_qnn_backend_skip_node_id_partitioner_int_node(self):
+        module = SkipIntNode()  # noqa: F405
+        sample_input = (
+            torch.randint(0, 5, (2, 3, 4), dtype=torch.int32),
+            torch.randint(0, 5, (2, 3, 4), dtype=torch.int32),
+        )
+        self.lower_module_and_test_output(
+            module,
+            sample_input,
+            expected_partitions=2,
+            skip_node_id_set={"aten_mul_tensor"},
+        )
+
     def test_qnn_backend_skip_node_id_quantizer(self):
         module = SimpleModel()  # noqa: F405
         sample_input = (torch.ones(1, 32, 28, 28), torch.ones(1, 32, 28, 28))