Replace debug handle with from_node to trace operator transformation (pytorch#11532)

Summary:
Pull Request resolved: pytorch#11532

X-link: pytorch/ao#2339

This diff replace the debug handle  with `from_node` infrastructure, which is a first class citizen in exported program and used to trace the node-level transformation by recording every ancestor of given node. N6213836 is a demonstration of how `from_node` infra records the node transformation after unlifting and re-exporting exported graph.

For simplify the progress, we are trying to reuse the debug handle infrastructure by generating debug handle with hashing their greatest ancestor's node.

After this change user no longer need to invoke `generate_numeric_debug_handle` for debugging. Also the original pipeline will still work under current scenario.

Reviewed By: jerryzh168

Differential Revision: D76168997

Author: Gasoonjia

backends/xnnpack/test/quantizer/test_pt2e_quantization.py

@@ -7,7 +7,7 @@
 # pyre-unsafe
 
 from collections import Counter
-from typing import Dict, Tuple
+from typing import Tuple
 
 import torch
 from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import (
@@ -36,13 +36,13 @@
 from torchao.quantization.pt2e import (
     allow_exported_model_train_eval,
     compare_results,
-    CUSTOM_KEY,
     extract_results_from_loggers,
-    generate_numeric_debug_handle,
-    NUMERIC_DEBUG_HANDLE_KEY,
+    FROM_NODE_KEY,
     prepare_for_propagation_comparison,
 )
 
+from torchao.quantization.pt2e._numeric_debugger import _generate_debug_handle_from_node
+
 from torchao.quantization.pt2e.graph_utils import bfs_trace_with_node_process
 from torchao.quantization.pt2e.quantize_pt2e import (
     convert_pt2e,
@@ -723,47 +723,72 @@ def test_save_load(self) -> None:
 instantiate_parametrized_tests(TestQuantizePT2E)
 
 
+# TODO: deduplicate with TestNumericDebugger under torchao
 class TestNumericDebugger(TestCase):
-    def _extract_debug_handles(self, model) -> Dict[str, int]:
-        debug_handle_map: Dict[str, int] = {}
+    def _assert_each_node_has_debug_handle(self, model) -> None:
+        def _assert_node_has_debug_handle(node):
+            self.assertIn(
+                FROM_NODE_KEY,
+                node.meta,
+                f"Node {node} doesn't have from_node info",
+            )
+
+        bfs_trace_with_node_process(model, _assert_node_has_debug_handle)
+
+    def _extract_debug_handles(self, model) -> dict[str, int]:
+        debug_handle_map: dict[str, int] = {}
 
-        def _extract_debug_handles_from_node(node: torch.fx.Node) -> None:
+        def _extract_debug_handles_from_node(node):
             nonlocal debug_handle_map
-            if (
-                CUSTOM_KEY in node.meta
-                and NUMERIC_DEBUG_HANDLE_KEY in node.meta[CUSTOM_KEY]
-            ):
-                debug_handle_map[str(node)] = node.meta[CUSTOM_KEY][
-                    NUMERIC_DEBUG_HANDLE_KEY
-                ]
+            if (dh := _generate_debug_handle_from_node(node)) is not None:
+                debug_handle_map[str(node)] = dh
 
         bfs_trace_with_node_process(model, _extract_debug_handles_from_node)
+
         return debug_handle_map
 
-    def _assert_each_node_has_debug_handle(self, model) -> None:
-        def _assert_node_has_debug_handle(node: torch.fx.Node) -> None:
-            self.assertTrue(
-                CUSTOM_KEY in node.meta
-                and NUMERIC_DEBUG_HANDLE_KEY in node.meta[CUSTOM_KEY],
-                f"Node {node} doesn't have debug handle",
-            )
+    def _extract_debug_handles_with_prev_decomp_op(self, model) -> dict[str, int]:
+        prev_decomp_op_to_debug_handle_map: dict[str, int] = {}
 
-        bfs_trace_with_node_process(model, _assert_node_has_debug_handle)
+        def _extract_debug_handles_with_prev_decomp_op_from_node(node):
+            nonlocal prev_decomp_op_to_debug_handle_map
+            if FROM_NODE_KEY in node.meta:
+                prev_decomp_op = str(node.meta.get("nn_module_stack"))
+                debug_handle = _generate_debug_handle_from_node(node)
+                if prev_decomp_op not in prev_decomp_op_to_debug_handle_map:
+                    prev_decomp_op_to_debug_handle_map[prev_decomp_op] = debug_handle
+                else:
+                    assert (
+                        prev_decomp_op_to_debug_handle_map[prev_decomp_op]
+                        == debug_handle
+                    ), f"Node {node} has different debug handle {debug_handle}"
+                    "than previous node sharing the same decomp op {prev_decomp_op}"
+
+        bfs_trace_with_node_process(
+            model, _extract_debug_handles_with_prev_decomp_op_from_node
+        )
+        return prev_decomp_op_to_debug_handle_map
 
-    def test_quantize_pt2e_preserve_handle(self) -> None:
+    def test_quantize_pt2e_preserve_handle(self):
         m = TestHelperModules.Conv2dThenConv1d()
         example_inputs = m.example_inputs()
         ep = export_for_training(m, example_inputs, strict=True)
-        generate_numeric_debug_handle(ep)
         m = ep.module()
 
         quantizer = XNNPACKQuantizer().set_global(
             get_symmetric_quantization_config(is_per_channel=False)
         )
-        m = prepare_pt2e(m, quantizer)  # pyre-ignore[6]
+        m = prepare_pt2e(m, quantizer)
         debug_handle_map = self._extract_debug_handles(m)
+        node_name_equip_with_output_observer = [
+            "conv2d",
+            "conv1d",
+            "squeeze",
+        ]
         res_counter = Counter(debug_handle_map.values())
-        repeated_debug_handle_ids = [1, 2, 3]
+        repeated_debug_handle_ids = [
+            debug_handle_map[n_name] for n_name in node_name_equip_with_output_observer
+        ]
         # 3 ids were repeated because we copy over the id from node to its output observer
         # torch.ops.aten.conv2d.default, torch.ops.aten.squeeze.dim and torch.ops.aten.conv1d.default
         for dh_id in repeated_debug_handle_ids:
@@ -776,22 +801,23 @@ def test_quantize_pt2e_preserve_handle(self) -> None:
         res_counter = Counter(debug_handle_map.values())
         # same set of ids where repeated, because we copy over the id from observer/fake_quant to
         # dequantize node
-        repeated_debug_handle_ids = [1, 2, 3]
+        repeated_debug_handle_ids = [
+            debug_handle_map[n_name] for n_name in node_name_equip_with_output_observer
+        ]
         for dh_id in repeated_debug_handle_ids:
             self.assertEqual(res_counter[dh_id], 2)
 
-    def test_extract_results_from_loggers(self) -> None:
+    def test_extract_results_from_loggers(self):
         m = TestHelperModules.Conv2dThenConv1d()
         example_inputs = m.example_inputs()
         ep = export_for_training(m, example_inputs, strict=True)
-        generate_numeric_debug_handle(ep)
         m = ep.module()
-        m_ref_logger = prepare_for_propagation_comparison(m)  # pyre-ignore[6]
+        m_ref_logger = prepare_for_propagation_comparison(m)
 
         quantizer = XNNPACKQuantizer().set_global(
             get_symmetric_quantization_config(is_per_channel=False)
         )
-        m = prepare_pt2e(m, quantizer)  # pyre-ignore[6]
+        m = prepare_pt2e(m, quantizer)
         m(*example_inputs)
         m = convert_pt2e(m)
         m_quant_logger = prepare_for_propagation_comparison(m)
@@ -800,29 +826,22 @@ def test_extract_results_from_loggers(self) -> None:
         m_quant_logger(*example_inputs)
         ref_results = extract_results_from_loggers(m_ref_logger)
         quant_results = extract_results_from_loggers(m_quant_logger)
-        comparison_results = compare_results(
-            ref_results,
-            quant_results,  # pyre-ignore[6]
-        )
+        comparison_results = compare_results(ref_results, quant_results)
         for node_summary in comparison_results.values():
             if len(node_summary.results) > 0:
-                self.assertGreaterEqual(
-                    node_summary.results[0].sqnr,
-                    35,  # pyre-ignore[6]
-                )
+                self.assertGreaterEqual(node_summary.results[0].sqnr, 35)
 
-    def test_extract_results_from_loggers_list_output(self) -> None:
+    def test_extract_results_from_loggers_list_output(self):
         m = TestHelperModules.Conv2dWithSplit()
         example_inputs = m.example_inputs()
         ep = export_for_training(m, example_inputs, strict=True)
-        generate_numeric_debug_handle(ep)
         m = ep.module()
-        m_ref_logger = prepare_for_propagation_comparison(m)  # pyre-ignore[6]
+        m_ref_logger = prepare_for_propagation_comparison(m)
 
         quantizer = XNNPACKQuantizer().set_global(
             get_symmetric_quantization_config(is_per_channel=False)
         )
-        m = prepare_pt2e(m, quantizer)  # pyre-ignore[6]
+        m = prepare_pt2e(m, quantizer)
         m(*example_inputs)
         m = convert_pt2e(m)
         m_quant_logger = prepare_for_propagation_comparison(m)
@@ -831,15 +850,12 @@ def test_extract_results_from_loggers_list_output(self) -> None:
         m_quant_logger(*example_inputs)
         ref_results = extract_results_from_loggers(m_ref_logger)
         quant_results = extract_results_from_loggers(m_quant_logger)
-        comparison_results = compare_results(
-            ref_results,
-            quant_results,  # pyre-ignore[6]
-        )
+        comparison_results = compare_results(ref_results, quant_results)
         for node_summary in comparison_results.values():
             if len(node_summary.results) > 0:
                 sqnr = node_summary.results[0].sqnr
                 if isinstance(sqnr, list):
                     for sqnr_i in sqnr:
                         self.assertGreaterEqual(sqnr_i, 35)
                 else:
-                    self.assertGreaterEqual(sqnr, 35)  # pyre-ignore[6]
+                    self.assertGreaterEqual(sqnr, 35)