Replace debug handle with from_node to trace operator transformation

Summary:
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. For simplify the progress, we are trying to reuse the debug handle infrastructure by generating debug handle from from_node info via hasing.

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

@@ -723,11 +723,22 @@ 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.assertTrue(
+                CUSTOM_KEY in node.meta
+                and NUMERIC_DEBUG_HANDLE_KEY in node.meta[CUSTOM_KEY],
+                f"Node {node} doesn't have debug handle",
+            )
+
+        bfs_trace_with_node_process(model, _assert_node_has_debug_handle)
 
-        def _extract_debug_handles_from_node(node: torch.fx.Node) -> None:
+    def _extract_debug_handles(self, model) -> dict[str, int]:
+        debug_handle_map: dict[str, int] = {}
+
+        def _extract_debug_handles_from_node(node):
             nonlocal debug_handle_map
             if (
                 CUSTOM_KEY in node.meta
@@ -738,32 +749,55 @@ def _extract_debug_handles_from_node(node: torch.fx.Node) -> None:
                 ]
 
         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 (
+                CUSTOM_KEY in node.meta
+                and NUMERIC_DEBUG_HANDLE_KEY in node.meta[CUSTOM_KEY]
+            ):
+                prev_decomp_op = str(node.meta.get("nn_module_stack"))
+                debug_handle = node.meta[CUSTOM_KEY][NUMERIC_DEBUG_HANDLE_KEY]
+                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)
+        # 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 +810,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 +835,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 +859,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)