Add debug handle to inductor provenance tracking (pytorch#161110)

Summary:
Pull Request resolved: pytorch#161110

Use debug handle on kernel names to distinguish different calls to the same kernel.

Previous kernel name: kernel_name

New kernel name: kernel_name:debug_handle

We add the debug handle to the tlparse artifacts: `inductor_provenance_tracking_node_mappings` and `inductor_provenance_tracking_kernel_stack_traces`.

We also add debug handles in the comments of the generated code so we can map to them in the provenance tracking highlighter tool: pytorch/tlparse#134

Example output code:

```
        # Topologically Sorted Source Nodes: [y, z], Original ATen: [aten.addmm, aten.gelu]
        # [Provenance debug handles] triton_poi_fused_addmm_gelu_2:3
        stream0 = get_raw_stream(0)
        triton_poi_fused_addmm_gelu_2.run(buf4, primals_5, 300, stream=stream0)
```

The debug handles will also be used by downstream profilers such as zoomer.

Test Plan:
```
buck run mode/opt fbcode//caffe2/test/inductor:provenance_tracing
```

Rollback Plan:

Reviewed By: angelayi

Differential Revision: D78994959

Author: yushangdi

test/inductor/test_provenance_tracing.py

@@ -8,7 +8,6 @@
 import re
 import shutil
 import tempfile
-import unittest
 import zipfile
 from pathlib import Path
 
@@ -19,11 +18,11 @@
     create_kernel_information_json,
     create_mapping_pre_post_grad_nodes,
     create_node_mapping_kernel_to_post_grad,
+    reset_inductor_kernel_provenance_debug_handle,
 )
 from torch._inductor.fx_passes.post_grad import post_grad_passes
 from torch._inductor.test_case import run_tests, TestCase
 from torch._inductor.virtualized import V
-from torch.testing._internal.inductor_utils import HAS_GPU
 from torch.testing._internal.triton_utils import requires_cuda_and_triton
 
 
@@ -94,11 +93,12 @@ class TestProvenanceTracingArtifact(TestCase):
     corresponding "inductor triton kernel node" is expected.
     """
 
-    def _check_provenance_tracing_artifact(self, filepath, expected_data):
+    def _check_provenance_tracing_kernel_to_post_grad(self, filepath, expected_data):
         self.assertTrue(filepath.is_dir())
-        filename = Path(filepath) / "inductor_generated_kernel_to_post_grad_nodes.json"
+        filename = Path(filepath) / "inductor_provenance_tracking_node_mappings.json"
         with open(filename) as f:
             actual_data = json.load(f)
+        actual_data = actual_data["cppCodeToPost"]
         # check that the generated provenance tracing artifact is expected
         self.assertEqual(sorted(actual_data.items()), sorted(expected_data.items()))
 
@@ -116,10 +116,11 @@ def _test_triton_kernel_to_post_grad_tracing(self, device):
         c = torch.randn(10, 30, device=device)
         example_inputs = (a, b, c)
 
-        model = Model()
+        model = Model().to(device)
         filepath = None
 
         for backend in ["aot_inductor", "inductor"]:
+            reset_inductor_kernel_provenance_debug_handle()
             try:
                 with config.patch(
                     {
@@ -142,28 +143,12 @@ def _test_triton_kernel_to_post_grad_tracing(self, device):
                     self.assertTrue(m)
                     filepath = Path(m.group(1))
                     if device == "cuda":
-                        expected_data = {
-                            "triton_poi_fused_mul_0": ["mul"],
-                            "triton_poi_fused_addmm_gelu_1": [
-                                "mul_3",
-                                "mul_1",
-                                "add_tensor",
-                                "add",
-                                "erf",
-                                "mul_2",
-                            ],
-                        }
-                        if backend == "aot_inductor":
-                            expected_data["aoti_torch_cuda_mm_out"] = ["mm_default"]
-                        else:
-                            expected_data["extern_kernels.mm"] = ["mm_default"]
-                        self._check_provenance_tracing_artifact(filepath, expected_data)
                         expected_mapping = [
                             (
                                 "cppCodeToPost",
                                 {
-                                    "triton_poi_fused_mul_0": ["mul"],
-                                    "triton_poi_fused_addmm_gelu_1": [
+                                    "triton_poi_fused_mul_0:1": ["mul"],
+                                    "triton_poi_fused_addmm_gelu_1:2": [
                                         "mul_3",
                                         "mul_1",
                                         "add_tensor",
@@ -176,13 +161,13 @@ def _test_triton_kernel_to_post_grad_tracing(self, device):
                             (
                                 "postToCppCode",
                                 {
-                                    "mul": ["triton_poi_fused_mul_0"],
-                                    "mul_3": ["triton_poi_fused_addmm_gelu_1"],
-                                    "mul_1": ["triton_poi_fused_addmm_gelu_1"],
-                                    "add_tensor": ["triton_poi_fused_addmm_gelu_1"],
-                                    "add": ["triton_poi_fused_addmm_gelu_1"],
-                                    "erf": ["triton_poi_fused_addmm_gelu_1"],
-                                    "mul_2": ["triton_poi_fused_addmm_gelu_1"],
+                                    "mul": ["triton_poi_fused_mul_0:1"],
+                                    "mul_3": ["triton_poi_fused_addmm_gelu_1:2"],
+                                    "mul_1": ["triton_poi_fused_addmm_gelu_1:2"],
+                                    "add_tensor": ["triton_poi_fused_addmm_gelu_1:2"],
+                                    "add": ["triton_poi_fused_addmm_gelu_1:2"],
+                                    "erf": ["triton_poi_fused_addmm_gelu_1:2"],
+                                    "mul_2": ["triton_poi_fused_addmm_gelu_1:2"],
                                 },
                             ),
                             (
@@ -208,15 +193,19 @@ def _test_triton_kernel_to_post_grad_tracing(self, device):
                             ),
                         ]
                         if backend == "aot_inductor":
-                            expected_mapping[0][1]["aoti_torch_cuda_mm_out"] = [
+                            expected_mapping[0][1]["aoti_torch_cuda_mm_out:3"] = [
                                 "mm_default"
                             ]
                             expected_mapping[1][1]["mm_default"] = [
-                                "aoti_torch_cuda_mm_out"
+                                "aoti_torch_cuda_mm_out:3"
                             ]
                         else:
-                            expected_mapping[0][1]["extern_kernels.mm"] = ["mm_default"]
-                            expected_mapping[1][1]["mm_default"] = ["extern_kernels.mm"]
+                            expected_mapping[0][1]["extern_kernels.mm:3"] = [
+                                "mm_default"
+                            ]
+                            expected_mapping[1][1]["mm_default"] = [
+                                "extern_kernels.mm:3"
+                            ]
                         self._check_provenance_tracking_node_mappings(
                             filepath, expected_mapping
                         )
@@ -225,9 +214,9 @@ def _test_triton_kernel_to_post_grad_tracing(self, device):
                         # check the inductor kernel to post grad nodes mapping is expected for cpu
                         if backend == "aot_inductor":
                             expected_data = {
-                                "cpp_fused_mul_0": ["mul"],
-                                "aoti_torch_cpu_addmm_out": ["addmm"],
-                                "cpp_fused_gelu_1": [
+                                "cpp_fused_mul_0:1": ["mul"],
+                                "aoti_torch_cpu_addmm_out:3": ["addmm"],
+                                "cpp_fused_gelu_1:2": [
                                     "mul_3",
                                     "mul_1",
                                     "add",
@@ -238,17 +227,19 @@ def _test_triton_kernel_to_post_grad_tracing(self, device):
                         else:
                             # backend == "inductor"
                             expected_data = {
-                                "cpp_fused_mul_0": ["mul"],
-                                "cpp_fused_gelu_1": [
+                                "cpp_fused_mul_0:1": ["mul"],
+                                "cpp_fused_gelu_1:2": [
                                     "mul_3",
                                     "mul_1",
                                     "add",
                                     "erf",
                                     "mul_2",
                                 ],
-                                "extern_kernels.addmm": ["addmm"],
+                                "extern_kernels.addmm:3": ["addmm"],
                             }
-                        self._check_provenance_tracing_artifact(filepath, expected_data)
+                        self._check_provenance_tracing_kernel_to_post_grad(
+                            filepath, expected_data
+                        )
 
             finally:
                 if filepath:
@@ -258,7 +249,6 @@ def _test_triton_kernel_to_post_grad_tracing(self, device):
     def test_triton_kernel_to_post_grad_tracing_cuda(self):
         self._test_triton_kernel_to_post_grad_tracing(device="cuda")
 
-    @unittest.skipIf(HAS_GPU, "the test is only for cpu")
     def test_triton_kernel_to_post_grad_tracing_cpu(self):
         self._test_triton_kernel_to_post_grad_tracing(device="cpu")
 
@@ -274,6 +264,7 @@ def test_triton_kernel_to_post_grad_tracing_extern_kernel(self):
         filepath = None
 
         for backend in ["aot_inductor", "inductor"]:
+            reset_inductor_kernel_provenance_debug_handle()
             try:
                 with config.patch(
                     {
@@ -297,15 +288,17 @@ def test_triton_kernel_to_post_grad_tracing_extern_kernel(self):
                     filepath = Path(m.group(1))
                     if backend == "inductor":
                         expected_data = {
-                            "extern_kernels.addmm": ["addmm"],
+                            "extern_kernels.addmm:1": ["addmm"],
                         }
                     else:
                         # backend = aot_inductor
                         expected_data = {
-                            "aoti_torch_cuda_addmm_out": ["addmm"],
-                            "triton_poi_fused_0": ["_tensor_constant1"],
+                            "aoti_torch_cuda_addmm_out:2": ["addmm"],
+                            "triton_poi_fused_0:1": ["_tensor_constant1"],
                         }
-                    self._check_provenance_tracing_artifact(filepath, expected_data)
+                    self._check_provenance_tracing_kernel_to_post_grad(
+                        filepath, expected_data
+                    )
             finally:
                 if filepath:
                     shutil.rmtree(filepath)
@@ -319,6 +312,7 @@ def _test_pt_tracing_combo_kernel(self, backend):
         example_inputs = (a, b, c)
 
         model = Model2()
+        reset_inductor_kernel_provenance_debug_handle()
 
         with config.patch(
             {
@@ -342,8 +336,8 @@ def _test_pt_tracing_combo_kernel(self, backend):
             m = re.match(r"WARNING.* debug trace: (.*)", cm.output[0])
             self.assertTrue(m)
             filepath = Path(m.group(1)).resolve()
-            expected_data = {"triton_poi_fused_0": ["relu", "sigmoid", "tanh"]}
-            self._check_provenance_tracing_artifact(filepath, expected_data)
+            expected_data = {"triton_poi_fused_0:1": ["relu", "sigmoid", "tanh"]}
+            self._check_provenance_tracing_kernel_to_post_grad(filepath, expected_data)
 
     @requires_cuda_and_triton
     def test_triton_kernel_to_post_grad_tracing_combo_kernel(self):
@@ -556,25 +550,28 @@ def test_tlparse_kernel_stack_traces(self):
         example_inputs = (x, a, b, c)
 
         expected = {
-            "triton_poi_fused_addmm_relu_sigmoid_threshold_backward_0": [
+            "triton_poi_fused_addmm_relu_sigmoid_threshold_backward_0:1": [
                 "x = self.sigmoid(x)",
                 "x = self.fc1(x)",
                 "x = self.relu(x)",
             ],
-            "triton_poi_fused_mul_1": [
+            "triton_poi_fused_mul_1:2": [
                 "d = a * 3.14",
             ],
-            "triton_poi_fused_addmm_gelu_2": [
+            "triton_poi_fused_addmm_gelu_2:3": [
                 "z = torch.nn.functional.gelu(y)",
                 "y = torch.addmm(c, d, b)",
             ],
-            "extern_kernels.mm": [
+            "extern_kernels.mm:4": [
                 "x = self.fc1(x)",
+            ],
+            "extern_kernels.mm:5": [
                 "y = torch.addmm(c, d, b)",
             ],
         }
 
         with self._setup_provenance_capture() as payload_buffer:
+            reset_inductor_kernel_provenance_debug_handle()
             compiled = torch.compile(model)
             compiled(*example_inputs)
             payload_content = payload_buffer.getvalue().strip()
@@ -623,6 +620,7 @@ def test_kernel_information_generation(self):
         with tempfile.TemporaryDirectory() as temp_dir:
             ep = torch.export.export(model, inputs, strict=False)
             pt2_file = os.path.join(temp_dir, "model.pt2")
+            reset_inductor_kernel_provenance_debug_handle()
             torch._inductor.aoti_compile_and_package(ep, package_path=pt2_file)
 
             # Extract and check kernel_information.json exists in the package
@@ -646,7 +644,7 @@ def test_kernel_information_generation(self):
                 kernel_info = json.load(f)
 
             expected = {
-                "triton_poi_fused_addmm_relu_sigmoid_0": {
+                "triton_poi_fused_addmm_relu_sigmoid_0:1": {
                     "stack_traces": [
                         "x = self.sigmoid(x)",
                         "x = self.fc1(x)",
@@ -655,14 +653,14 @@ def test_kernel_information_generation(self):
                     "post_grad_nodes": ["sigmoid", "relu", "add_tensor_1"],
                     "pre_grad_nodes": ["sigmoid", "relu", "linear"],
                 },
-                "triton_poi_fused_mul_1": {
+                "triton_poi_fused_mul_1:2": {
                     "stack_traces": [
                         "d = a * 3.14",
                     ],
                     "post_grad_nodes": ["mul"],
                     "pre_grad_nodes": ["mul"],
                 },
-                "triton_poi_fused_addmm_gelu_2": {
+                "triton_poi_fused_addmm_gelu_2:3": {
                     "stack_traces": [
                         "z = torch.nn.functional.gelu(y)",
                         "y = torch.addmm(c, d, b)",
@@ -677,13 +675,19 @@ def test_kernel_information_generation(self):
                     ],
                     "pre_grad_nodes": ["gelu", "addmm"],
                 },
-                "aoti_torch_cuda_mm_out": {
+                "aoti_torch_cuda_mm_out:4": {
                     "stack_traces": [
                         "x = self.fc1(x)",
+                    ],
+                    "post_grad_nodes": ["mm_default_1"],
+                    "pre_grad_nodes": ["linear"],
+                },
+                "aoti_torch_cuda_mm_out:5": {
+                    "stack_traces": [
                         "y = torch.addmm(c, d, b)",
                     ],
-                    "post_grad_nodes": ["mm_default_1", "mm_default"],
-                    "pre_grad_nodes": ["linear", "addmm"],
+                    "post_grad_nodes": ["mm_default"],
+                    "pre_grad_nodes": ["addmm"],
                 },
             }
 

torch/_inductor/codegen/cpp.py

@@ -5392,10 +5392,6 @@ def define_kernel(self, src_code, nodes, kernel_args=None):
             else ""
         )
         kernel_name = "_".join(["cpp", fused_name, wrapper.next_kernel_suffix()])
-        # below add provenance tracing info for cpu CppKernel types
-        if config.trace.provenance_tracking_level != 0:
-            set_kernel_post_grad_provenance_tracing(nodes, kernel_name)
-
         kernel_decl_name = kernel_name if V.graph.cpp_wrapper else "kernel"
         src_code = src_code.replace(str(Placeholder.KERNEL_NAME), kernel_decl_name)
         src_code = src_code.replace(str(Placeholder.DESCRIPTIVE_NAME), kernel_name)
@@ -5434,7 +5430,15 @@ def flush(self):
             kernel_name = self.define_kernel(
                 src_code, self.kernel_group.scheduled_nodes
             )
-            self.kernel_group.call_kernel(V.graph.wrapper_code, kernel_name)
+            # below add provenance tracing info for cpu CppKernel types
+            debug_handle: Optional[int] = None
+            if config.trace.provenance_tracking_level != 0:
+                debug_handle = set_kernel_post_grad_provenance_tracing(
+                    self.kernel_group.scheduled_nodes, kernel_name
+                )
+            self.kernel_group.call_kernel(
+                V.graph.wrapper_code, kernel_name, debug_handle=debug_handle
+            )
         self.reset_kernel_group()
         self._set_flush_status(False)
 
@@ -5509,10 +5513,14 @@ def codegen_group(self, name=None) -> str:
             code.splice(self.loops_code)
         return code.getvalue()
 
-    def call_kernel(self, wrapper, kernel_name):
+    def call_kernel(self, wrapper, kernel_name, debug_handle: Optional[int] = None):
         _, call_args, arg_types = self.args.cpp_argdefs()
         wrapper.generate_kernel_call(
-            kernel_name, call_args, triton=False, arg_types=arg_types
+            kernel_name,
+            call_args,
+            triton=False,
+            arg_types=arg_types,
+            debug_handle=debug_handle,
         )
 
 

torch/_inductor/codegen/cpp_wrapper_cpu.py

@@ -1219,6 +1219,7 @@ def generate_c_shim_extern_kernel_call(
         device: str,
         *,
         debug_args: Optional[list[str]] = None,
+        debug_handle: Optional[int] = None,
     ) -> None:
         """debug_args kwarg allows CppWrapperCpuArrayRef to pass in wrapped arguments in
         place of args while preserving debug printer output."""
@@ -1235,14 +1236,16 @@ def generate_c_shim_extern_kernel_call(
         ]
         with debug_printer_manager:
             shim_fn = self.get_c_shim_func_name(kernel, device)
+            self.write_provenance_debug_handle(shim_fn, debug_handle)
             shim_fn_codes = (
                 f"AOTI_TORCH_ERROR_CODE_CHECK({shim_fn}({', '.join(args)}));"
             )
             if enable_kernel_profile:
+                debug_handle_str = "" if debug_handle is None else f":{debug_handle}"
                 shim_fn_codes = textwrap.dedent(
                     f"""
                     {{
-                      RAIIAtenRecordFunctionHandle record_{shim_fn}_("{shim_fn}", nullptr);
+                      RAIIAtenRecordFunctionHandle record_{shim_fn}_("{shim_fn}{debug_handle_str}", nullptr);
                       {shim_fn_codes}
                     }}
                     """
@@ -1338,6 +1341,7 @@ def _generate_extern_kernel_out_helper(
         out_view: Optional[str],
         args: list[str],
         device: str,
+        debug_handle: Optional[int] = None,
     ) -> None:
         if out_view:
             out_name = f"{out}_as_strided"
@@ -1346,7 +1350,9 @@ def _generate_extern_kernel_out_helper(
         else:
             args.insert(0, out)
 
-        self.generate_c_shim_extern_kernel_call(kernel, args, device)
+        self.generate_c_shim_extern_kernel_call(
+            kernel, args, device, debug_handle=debug_handle
+        )
 
     def generate_scatter_fallback(
         self,