tritonparse 0.3.1.dev20251027071529__tar.gz → 0.3.1.dev20251029071541__tar.gz

{tritonparse-0.3.1.dev20251027071529/tritonparse.egg-info → tritonparse-0.3.1.dev20251029071541}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tritonparse
-Version: 0.3.1.dev20251027071529
+Version: 0.3.1.dev20251029071541
 Summary: TritonParse: A Compiler Tracer, Visualizer, and mini-Reproducer Generator for Triton Kernels
 Author-email: Yueming Hao <yhao@meta.com>
 License-Expression: BSD-3-Clause
@@ -134,7 +134,7 @@ from tritonparse.reproducer.orchestrator import reproduce
 
 result = reproduce(
     input_path="./parsed_output/trace.ndjson.gz",
-    line_index=1,           # Which launch event (1-based)
+    line_index=0,           # 0-based index (first event is 0)
     out_dir="repro_output"
 )
 ```

{tritonparse-0.3.1.dev20251027071529 → tritonparse-0.3.1.dev20251029071541}/README.md

@@ -116,7 +116,7 @@ from tritonparse.reproducer.orchestrator import reproduce
 
 result = reproduce(
     input_path="./parsed_output/trace.ndjson.gz",
-    line_index=1,           # Which launch event (1-based)
+    line_index=0,           # 0-based index (first event is 0)
     out_dir="repro_output"
 )
 ```

{tritonparse-0.3.1.dev20251027071529 → tritonparse-0.3.1.dev20251029071541}/tests/test_tritonparse.py

@@ -187,6 +187,70 @@ class TestTritonparseCPU(unittest.TestCase):
 
         assert convert(nested_structure) == expected_nested
 
+    def test_loc_alias_parsing(self):
+        """Test parsing of location aliases in TTIR/TTGIR"""
+        from tritonparse.ir_parser import extract_loc_definitions
+
+        # Test case 1: Bare #loc reference (no number)
+        ir_with_bare_loc = """
+module {
+  #loc = loc("/tmp/test.py":10:5)
+  #loc13 = loc("x_ptr"(#loc))
+  func @kernel(%arg0: !tt.ptr<f32> loc(#loc13)) {
+    return loc(#loc)
+  }
+}
+"""
+        locs = extract_loc_definitions(ir_with_bare_loc)
+        # Main #loc should be stored with "" key
+        assert "" in locs, "Main #loc not found"
+        assert locs[""]["file"] == "/tmp/test.py"
+        assert locs[""]["line"] == 10
+        # Alias #loc13 should resolve to same location
+        assert "13" in locs, "#loc13 not found"
+        assert locs["13"]["file"] == "/tmp/test.py"
+        assert locs["13"]["line"] == 10
+        assert locs["13"]["alias_name"] == "x_ptr"
+        assert locs["13"]["alias_of"] == ""
+
+        # Test case 2: Named alias with numbered reference
+        ir_with_numbered_alias = """
+#loc = loc("/tmp/test.py":5:0)
+#loc2 = loc("/tmp/test.py":20:28)
+#loc16 = loc("pid"(#loc2))
+%0 = tt.get_program_id x : i32 loc(#loc16)
+"""
+        locs = extract_loc_definitions(ir_with_numbered_alias)
+        assert "2" in locs
+        assert locs["2"]["line"] == 20
+        assert "16" in locs
+        assert locs["16"]["file"] == "/tmp/test.py"
+        assert locs["16"]["line"] == 20
+        assert locs["16"]["alias_name"] == "pid"
+        assert locs["16"]["alias_of"] == "2"
+
+        # Test case 3: Simple alias (no name)
+        ir_with_simple_alias = """
+#loc = loc("/tmp/test.py":1:1)
+#loc1 = loc("/tmp/test.py":15:10)
+#loc20 = loc(#loc1)
+%1 = arith.constant 0 : i32 loc(#loc20)
+"""
+        locs = extract_loc_definitions(ir_with_simple_alias)
+        assert "1" in locs
+        assert "20" in locs
+        assert locs["20"]["file"] == "/tmp/test.py"
+        assert locs["20"]["line"] == 15
+        assert locs["20"]["alias_of"] == "1"
+        assert "alias_name" not in locs["20"]
+
+        # Test case 4: Definition line tracking
+        assert "def_line" in locs[""]
+        assert "def_line" in locs["1"]
+        assert "def_line" in locs["20"]
+
+        print("✓ All loc alias parsing tests passed")
+
 
 class TestTritonparseCUDA(unittest.TestCase):
     """CUDA tests (require GPU)"""

{tritonparse-0.3.1.dev20251027071529 → tritonparse-0.3.1.dev20251029071541}/tritonparse/cli.py

@@ -68,12 +68,20 @@ def main():
         }
         unified_parse(**parse_args)
     elif args.func == "reproduce":
+        replacer = None
+        if args.use_fbcode:
+            from tritonparse.fb.reproducer.replacer import FBCodePlaceholderReplacer
+
+            replacer = FBCodePlaceholderReplacer()
+            print(f"Using FBCode placeholder replacer for template: {args.template}")
+
         reproduce(
             input_path=args.input,
-            line_index=args.line,
+            line_index=args.line - 1,  # Convert 1-based line number to 0-based index
             out_dir=args.out_dir,
             template=args.template,
             kernel_import=args.kernel_import,
+            replacer=replacer,
         )
     else:
         raise RuntimeError(f"Unknown command: {args.func}")

{tritonparse-0.3.1.dev20251027071529 → tritonparse-0.3.1.dev20251029071541}/tritonparse/ir_analysis.py

@@ -42,7 +42,7 @@ def process_amd_gcn_bufferops(
 ) -> dict[str, int]:
     ir_content = load_ir_contents(key, file_content, file_path)
     # TODO: Add atomics
-    io_keys = ["global_load_", "global_store_", "buffer_load_", "buffer_store_"]
+    io_keys = ["global_load", "global_store", "buffer_load", "buffer_store"]
     return process_amd_bufferop(ir_content, io_keys)
 
 
@@ -66,9 +66,12 @@ def _generate_ir_analysis(entry: str):
         gcn_bufferops_info = process_amd_gcn_bufferops(
             amdgcn_key, file_content, file_path
         )
+        io_counts = {}
         # NDJSON format requires a newline at the end of each line
         if ttgir_bufferops_info:
-            ir_analysis["amd_ttgir_bufferops_count"] = ttgir_bufferops_info
+            io_counts["amd_ttgir_bufferops_count"] = ttgir_bufferops_info
         if gcn_bufferops_info:
-            ir_analysis["amd_gcn_bufferops_count"] = gcn_bufferops_info
-    return {"ir_analysis": ir_analysis}
+            io_counts["amd_gcn_bufferops_count"] = gcn_bufferops_info
+        if io_counts:
+            ir_analysis["io_counts"] = io_counts
+    return ir_analysis

{tritonparse-0.3.1.dev20251027071529 → tritonparse-0.3.1.dev20251029071541}/tritonparse/ir_parser.py

@@ -10,7 +10,8 @@ logger = logging.getLogger("SourceMapping")
 
 # the definition of the #loc directive. they are in the bottom of the IR files
 # Example:#loc2 = loc("/tmp/torchinductor_yhao/yp/abcdef.py":20:28)
-LOC_PATTERN = re.compile(r'#loc(\d*) = loc\("([^"]+)":(\d+):(\d+)\)')
+# Note: This should only match numbered locs like #loc1, #loc2, not bare #loc
+LOC_PATTERN = re.compile(r'#loc(\d+) = loc\("([^"]+)":(\d+):(\d+)\)')
 
 # the reference to the #loc directive. they are in the end of lines of the IR files
 # Example: loc(#loc2)
@@ -33,6 +34,17 @@ AMDGCN_LOC_PATTERN = re.compile(
 )
 
 
+# alias loc definitions in TTGIR/TTIR
+# Example: #loc16 = loc("pid"(#loc2))
+# Example: #loc13 = loc("x_ptr"(#loc)) - bare #loc without number
+ALIAS_WITH_NAME_PATTERN = re.compile(
+    r'#loc(\d+)\s*=\s*loc\("([^"]+)"\s*\(\s*#loc(\d*)\s*\)\s*\)'
+)
+
+# Example: #loc20 = loc(#loc16)
+ALIAS_SIMPLE_PATTERN = re.compile(r"#loc(\d+)\s*=\s*loc\(\s*#loc(\d*)\s*\)")
+
+
 def extract_loc_definitions(ir_content: str) -> Dict[str, Dict[str, Any]]:
     """
     Extracts location definitions from the given IR content.
@@ -50,9 +62,10 @@ def extract_loc_definitions(ir_content: str) -> Dict[str, Dict[str, Any]]:
     """
     locations = {}
     # The first #loc directive is a special case. It locates at the top of the IR files
+    # Store it with empty string "" as key to avoid conflict with #loc1
     main_match = re.search(r'#loc = loc\("([^"]+)":(\d+):(\d+)\)', ir_content)
     if main_match:
-        locations["1"] = {
+        locations[""] = {
             "file": main_match.group(1),
             "line": int(main_match.group(2)),
             "column": int(main_match.group(3)),
@@ -61,6 +74,84 @@ def extract_loc_definitions(ir_content: str) -> Dict[str, Dict[str, Any]]:
     for loc_id, filename, line, col in LOC_PATTERN.findall(ir_content):
         key = loc_id
         locations[key] = {"file": filename, "line": int(line), "column": int(col)}
+
+    # Handle alias-style loc definitions that reference another #loc
+    # Build alias map first: alias_id -> target_id
+    alias_map: Dict[str, str] = {}
+    for m in ALIAS_WITH_NAME_PATTERN.finditer(ir_content):
+        alias_id, _name, target_id = m.groups()
+        # Empty target_id means bare #loc, map to "" (main loc key)
+        alias_map[alias_id] = target_id or ""
+    for m in ALIAS_SIMPLE_PATTERN.finditer(ir_content):
+        alias_id, target_id = m.groups()
+        # Empty target_id means bare #loc, map to "" (main loc key)
+        alias_map[alias_id] = target_id or ""
+
+    # Build definition line map and alias name map by scanning lines
+    def_line_map: Dict[str, int] = {}
+    alias_name_map: Dict[str, str] = {}
+    main_loc_line: int = 0
+    for i, line in enumerate(ir_content.split("\n"), start=1):
+        if m := ALIAS_WITH_NAME_PATTERN.search(line):
+            alias_id, name, target_id = m.groups()
+            def_line_map[alias_id] = i
+            alias_name_map[alias_id] = name
+            # ensure alias map is populated even if only found in line scan
+            # Empty target_id means bare #loc, map to "" (main loc key)
+            alias_map.setdefault(alias_id, target_id or "")
+        elif m := ALIAS_SIMPLE_PATTERN.search(line):
+            alias_id, target_id = m.groups()
+            def_line_map[alias_id] = i
+            # Empty target_id means bare #loc, map to "" (main loc key)
+            alias_map.setdefault(alias_id, target_id or "")
+        if m2 := LOC_PATTERN.search(line):
+            base_id, _fn, _ln, _col = m2.groups()
+            def_line_map[base_id] = i
+        if re.search(r'#loc\s*=\s*loc\("[^"]+":\d+:\d+\)', line):
+            # main #loc = loc("file":line:col) without id
+            main_loc_line = main_loc_line or i
+
+    # Resolve aliases to base locations (file/line/column)
+    resolving_stack = set()
+
+    def resolve_alias(current_id: str) -> Dict[str, Any]:
+        # Already a concrete location
+        if current_id in locations:
+            return locations[current_id]
+        # Detect cycles
+        if current_id in resolving_stack:
+            return {}
+        resolving_stack.add(current_id)
+        parent_id = alias_map.get(current_id)
+        result: Dict[str, Any] = {}
+        if parent_id is not None:
+            base = resolve_alias(parent_id)
+            if base:
+                # copy to avoid sharing the same dict by reference
+                result = {
+                    "file": base.get("file"),
+                    "line": base.get("line"),
+                    "column": base.get("column"),
+                }
+                locations[current_id] = result
+        resolving_stack.remove(current_id)
+        return result
+
+    # Resolve aliases and attach alias metadata
+    for alias_id, target_id in alias_map.items():
+        if alias_id not in locations:
+            resolve_alias(alias_id)
+        if alias_id in locations:
+            locations[alias_id]["alias_of"] = target_id
+            if alias_id in alias_name_map:
+                locations[alias_id]["alias_name"] = alias_name_map[alias_id]
+
+    # Attach definition line metadata
+    for k, v in def_line_map.items():
+        if k in locations:
+            locations[k]["def_line"] = v
+    if main_loc_line and "" in locations:
+        locations[""]["def_line"] = main_loc_line
     return locations
 
 

{tritonparse-0.3.1.dev20251027071529 → tritonparse-0.3.1.dev20251029071541}/tritonparse/reproducer/cli.py

@@ -46,3 +46,8 @@ def _add_reproducer_args(parser: argparse.ArgumentParser) -> None:
             "Defaults to 'default'."
         ),
     )
+    parser.add_argument(
+        "--use-fbcode",
+        action="store_true",
+        help=("Use fbcode to setup repro environment."),
+    )

{tritonparse-0.3.1.dev20251027071529 → tritonparse-0.3.1.dev20251029071541}/tritonparse/reproducer/ingestion/ndjson.py

@@ -40,7 +40,7 @@ def get_launch_and_compilation_events(
 
     Args:
         events: List of parsed event dictionaries.
-        line_index: Index of the launch event to process.
+        line_index: 0-based index of the launch event to process.
 
     Returns:
         Tuple of (launch_event, compilation_event).
@@ -179,7 +179,7 @@ def build_context_bundle(
 
     Args:
         events: List of parsed event dictionaries.
-        line_index: Index of the launch event to process.
+        line_index: 0-based index of the launch event to process.
 
     Returns:
         ContextBundle containing all information needed to reproduce the kernel launch.

{tritonparse-0.3.1.dev20251027071529 → tritonparse-0.3.1.dev20251029071541}/tritonparse/reproducer/orchestrator.py

@@ -23,13 +23,13 @@ def reproduce(
     template: str,
     replacer: Optional[PlaceholderReplacer] = None,
     kernel_import: KernelImportMode = KernelImportMode.DEFAULT,
-) -> dict[str, Path]:
+) -> dict[str, str]:
     """
     Generate a reproducer script from NDJSON trace file.
 
     Args:
         input_path: Path to the NDJSON trace file.
-        line_index: Line index of the launch event to reproduce.
+        line_index: 0-based index of the launch event to reproduce in the events list.
         out_dir: Output directory for reproducer files.
         template: Template name to use for the reproducer.
         replacer: Optional custom PlaceholderReplacer instance. If None, uses DefaultPlaceholderReplacer.
@@ -45,7 +45,7 @@ def reproduce(
         f"Built context bundle for kernel: {context_bundle.kernel_info.function_name}"
     )
     out_py_path, temp_json_path = determine_output_paths(
-        out_dir, context_bundle.kernel_info.function_name
+        out_dir, context_bundle.kernel_info.function_name, template
     )
     save_prettified_json(context_bundle.raw_launch_event, temp_json_path)
 

{tritonparse-0.3.1.dev20251027071529 → tritonparse-0.3.1.dev20251029071541}/tritonparse/reproducer/placeholder_replacer.py

@@ -76,21 +76,39 @@ class DefaultPlaceholderReplacer(PlaceholderReplacer):
     - # {{KERNEL_INVOCATION_PLACEHOLDER}}: Replaced with kernel invocation code
     """
 
+    KERNEL_NAME_PLACEHOLDER = "{{KERNEL_NAME_PLACEHOLDER}}"
+    JSON_FILE_NAME_PLACEHOLDER = "{{JSON_FILE_NAME_PLACEHOLDER}}"
+    IR_OVERRIDE_SETUP_PLACEHOLDER = "# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}"
+    KERNEL_SYSPATH_PLACEHOLDER = "# {{KERNEL_SYSPATH_PLACEHOLDER}}"
+    KERNEL_IMPORT_PLACEHOLDER = "# {{KERNEL_IMPORT_PLACEHOLDER}}"
+    UTILITY_FUNCTIONS_PLACEHOLDER = "# {{UTILITY_FUNCTIONS_PLACEHOLDER}}"
+    KERNEL_INVOCATION_PLACEHOLDER = "# {{KERNEL_INVOCATION_PLACEHOLDER}}"
+
     def __init__(self):
         super().__init__()
         # Register all default handlers
-        self.register("{{JSON_FILE_NAME_PLACEHOLDER}}", self._replace_json_filename)
+        self.register(self.JSON_FILE_NAME_PLACEHOLDER, self._replace_json_filename)
         self.register(
-            "# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}", self._replace_ir_override_setup
+            self.IR_OVERRIDE_SETUP_PLACEHOLDER, self._replace_ir_override_setup
         )
-        self.register("# {{KERNEL_SYSPATH_PLACEHOLDER}}", self._replace_kernel_syspath)
-        self.register("# {{KERNEL_IMPORT_PLACEHOLDER}}", self._replace_kernel_import)
+        self.register(self.KERNEL_SYSPATH_PLACEHOLDER, self._replace_kernel_syspath)
+        self.register(self.KERNEL_IMPORT_PLACEHOLDER, self._replace_kernel_import)
         self.register(
-            "# {{UTILITY_FUNCTIONS_PLACEHOLDER}}", self._replace_utility_functions
+            self.UTILITY_FUNCTIONS_PLACEHOLDER, self._replace_utility_functions
         )
         self.register(
-            "# {{KERNEL_INVOCATION_PLACEHOLDER}}", self._replace_kernel_invocation
+            self.KERNEL_INVOCATION_PLACEHOLDER, self._replace_kernel_invocation
         )
+        self.register(self.KERNEL_NAME_PLACEHOLDER, self._replace_kernel_name)
+
+    def _replace_kernel_name(
+        self, code: str, context_bundle: ContextBundle, **kwargs
+    ) -> str:
+        """Replace the kernel name placeholder."""
+        kernel_name = context_bundle.kernel_info.function_name
+        if not kernel_name:
+            raise ValueError("Kernel function name is not available")
+        return code.replace(self.KERNEL_NAME_PLACEHOLDER, kernel_name)
 
     def _replace_json_filename(
         self, code: str, context_bundle: ContextBundle, **kwargs
@@ -99,7 +117,7 @@ class DefaultPlaceholderReplacer(PlaceholderReplacer):
         temp_json_path = kwargs.get("temp_json_path")
         if temp_json_path is None:
             raise ValueError("temp_json_path is required for JSON filename replacement")
-        return code.replace("{{JSON_FILE_NAME_PLACEHOLDER}}", temp_json_path.name)
+        return code.replace(self.JSON_FILE_NAME_PLACEHOLDER, temp_json_path.name)
 
     def _replace_ir_override_setup(
         self, code: str, context_bundle: ContextBundle, **kwargs
@@ -108,7 +126,7 @@ class DefaultPlaceholderReplacer(PlaceholderReplacer):
         kernel_import = kwargs.get("kernel_import", KernelImportMode.DEFAULT)
 
         if kernel_import != KernelImportMode.OVERRIDE_TTIR:
-            return code.replace("# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}", "")
+            return code.replace(self.IR_OVERRIDE_SETUP_PLACEHOLDER, "")
 
         comp_json_filename = kwargs.get("comp_json_filename")
         if not comp_json_filename:
@@ -158,7 +176,7 @@ _original_autotune = triton.autotune
 triton.autotune = _patched_autotune
 '''
 
-        return code.replace("# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}", setup_code)
+        return code.replace(self.IR_OVERRIDE_SETUP_PLACEHOLDER, setup_code)
 
     def _replace_kernel_syspath(
         self, code: str, context_bundle: ContextBundle, **kwargs
@@ -168,15 +186,15 @@ triton.autotune = _patched_autotune
 
         if kernel_import == KernelImportMode.DEFAULT:
             sys_stmt, _ = _generate_import_statements(context_bundle.kernel_info)
-            return code.replace("# {{KERNEL_SYSPATH_PLACEHOLDER}}", sys_stmt)
+            return code.replace(self.KERNEL_SYSPATH_PLACEHOLDER, sys_stmt)
         elif kernel_import == KernelImportMode.COPY:
             comment = (
                 "# Kernel sys.path setup skipped - kernel source code embedded below"
             )
-            return code.replace("# {{KERNEL_SYSPATH_PLACEHOLDER}}", comment)
+            return code.replace(self.KERNEL_SYSPATH_PLACEHOLDER, comment)
         elif kernel_import == KernelImportMode.OVERRIDE_TTIR:
             comment = "# Kernel sys.path setup skipped - using IR override mode"
-            return code.replace("# {{KERNEL_SYSPATH_PLACEHOLDER}}", comment)
+            return code.replace(self.KERNEL_SYSPATH_PLACEHOLDER, comment)
         else:
             raise ValueError(f"Unknown kernel_import mode: {kernel_import}")
 
@@ -190,7 +208,7 @@ triton.autotune = _patched_autotune
             _, import_statement = _generate_import_statements(
                 context_bundle.kernel_info
             )
-            return code.replace("# {{KERNEL_IMPORT_PLACEHOLDER}}", import_statement)
+            return code.replace(self.KERNEL_IMPORT_PLACEHOLDER, import_statement)
         elif kernel_import == KernelImportMode.COPY:
             source_code = context_bundle.kernel_info.source_code
             func_name = context_bundle.kernel_info.function_name
@@ -216,10 +234,10 @@ triton.autotune = _patched_autotune
             embedded_code += "\n" + source_code
             embedded_code += f"\n\n# Use kernel function directly\nimported_kernel_function = {func_name}"
 
-            return code.replace("# {{KERNEL_IMPORT_PLACEHOLDER}}", embedded_code)
+            return code.replace(self.KERNEL_IMPORT_PLACEHOLDER, embedded_code)
         elif kernel_import == KernelImportMode.OVERRIDE_TTIR:
             comment = "# Kernel import skipped - using IR override mode with TTIR"
-            return code.replace("# {{KERNEL_IMPORT_PLACEHOLDER}}", comment)
+            return code.replace(self.KERNEL_IMPORT_PLACEHOLDER, comment)
         else:
             raise ValueError(f"Unknown kernel_import mode: {kernel_import}")
 
@@ -228,7 +246,7 @@ triton.autotune = _patched_autotune
     ) -> str:
         """Replace the utility functions placeholder with extracted functions."""
         utility_code = extract_utility_functions()
-        return code.replace("# {{UTILITY_FUNCTIONS_PLACEHOLDER}}", utility_code)
+        return code.replace(self.UTILITY_FUNCTIONS_PLACEHOLDER, utility_code)
 
     def _replace_kernel_invocation(
         self, code: str, context_bundle: ContextBundle, **kwargs
@@ -237,4 +255,4 @@ triton.autotune = _patched_autotune
         source_code = context_bundle.kernel_info.source_code
         pos_args, kw_args = _parse_kernel_signature(source_code)
         invocation_snippet = _generate_invocation_snippet(pos_args, kw_args)
-        return code.replace("# {{KERNEL_INVOCATION_PLACEHOLDER}}", invocation_snippet)
+        return code.replace(self.KERNEL_INVOCATION_PLACEHOLDER, invocation_snippet)

{tritonparse-0.3.1.dev20251027071529 → tritonparse-0.3.1.dev20251029071541}/tritonparse/reproducer/templates/example.py

@@ -14,7 +14,7 @@ import torch
 # {{UTILITY_FUNCTIONS_PLACEHOLDER}}
 
 
-if __name__ == "__main__":
+def launch_kernel():
     script_dir = Path(__file__).resolve().parent  # noqa: F821
     json_file = script_dir / "{{JSON_FILE_NAME_PLACEHOLDER}}"
     grid, args_dict = create_args_from_json_file(str(json_file))  # noqa: F821
@@ -28,3 +28,7 @@ if __name__ == "__main__":
 
     torch.cuda.synchronize()
     print("Kernel execution finished.")
+
+
+if __name__ == "__main__":
+    launch_kernel()

tritonparse-0.3.1.dev20251029071541/tritonparse/reproducer/templates/tritonbench.py

@@ -0,0 +1,103 @@
+# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+
+from pathlib import Path
+from typing import Any, Callable, Dict, Optional, Tuple
+
+import torch
+from tritonbench.utils.triton_op import (
+    BenchmarkOperator,
+    register_benchmark,
+    REGISTERED_X_VALS,
+)
+
+
+imported_kernel_function: Optional[Callable[[Tuple[int], Dict[str, Any]], None]] = None
+
+# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}
+
+# {{KERNEL_SYSPATH_PLACEHOLDER}}
+
+# {{KERNEL_IMPORT_PLACEHOLDER}}
+
+# {{UTILITY_FUNCTIONS_PLACEHOLDER}}
+
+assert imported_kernel_function is not None, "imported_kernel_function is missing"
+
+KERNEL_NAME = "{{KERNEL_NAME_PLACEHOLDER}}"
+REPRO_CONTEXT_FILE_NAME = "{{JSON_FILE_NAME_PLACEHOLDER}}"
+
+
+def _get_launch_kernel_args() -> Tuple[Tuple[int], Dict[str, Any]]:
+    script_dir = Path(__file__).resolve().parent  # noqa: F821
+    json_file = script_dir / REPRO_CONTEXT_FILE_NAME
+
+    grid, args_dict = create_args_from_json_file(json_file)  # noqa: F821, F841
+
+    print("Recorded kernel arguments dictionary:")
+    for name, arg in args_dict.items():
+        if isinstance(arg, torch.Tensor):
+            print(
+                f"  {name}: Tensor:  {arg.shape} {arg.dtype} stride: {arg.stride()}, is_contiguous: {arg.is_contiguous()}"
+            )
+        else:
+            print(f"  {name}: {arg}")
+    print(f"Grid: {grid}")
+
+    return tuple(grid), args_dict
+
+
+grid, args_dict = _get_launch_kernel_args()
+
+
+def _launch_kernel(grid: tuple[int], args_dict: dict[str, Any]):
+    try:
+        assert grid is not None
+        assert args_dict is not None
+
+        # {{KERNEL_INVOCATION_PLACEHOLDER}}
+
+    except Exception as e:
+        print(f"Error: {e}")
+        print("Failed to launch kernel!")
+
+
+# HACK: @register_x_val doesn't allow us to pass `operator_name`` as a parameter
+tensor_args = {k: v for k, v in args_dict.items() if isinstance(v, torch.Tensor)}
+x_vals_label = ", ".join(tensor_args.keys())
+REGISTERED_X_VALS[KERNEL_NAME] = x_vals_label
+
+
+class Operator(BenchmarkOperator):
+    @register_benchmark(operator_name=KERNEL_NAME)
+    def run_kernel(self, grid, args_dict):
+        return lambda: _launch_kernel(grid, args_dict)
+
+    def get_input_iter(self):
+        yield {"grid": grid, "args_dict": args_dict}
+
+    def get_x_val(self, example_inputs):
+        tensors_shapes = [
+            tuple(v.shape)
+            for v in example_inputs["args_dict"].values()
+            if isinstance(v, torch.Tensor)
+        ]
+        return tuple(tensors_shapes)
+
+
+if __name__ == "__main__":
+    print("do_benchmark...")
+
+    args = [
+        "--benchmark-name",
+        KERNEL_NAME,
+    ]
+
+    from tritonbench.utils.parser import get_parser
+
+    parser = get_parser(args)
+    tb_args, extra_args = parser.parse_known_args(args)
+    bench = Operator(tb_args, extra_args)
+    bench.run()
+
+    print(bench.output)
+    print("Benchmark completed successfully!")

{tritonparse-0.3.1.dev20251027071529 → tritonparse-0.3.1.dev20251029071541}/tritonparse/reproducer/utils.py

@@ -327,7 +327,7 @@ def _create_arg_from_info(arg_info):
         return None
 
 
-def determine_output_paths(out_dir: str, kernel_name: str):
+def determine_output_paths(out_dir: str, kernel_name: str, template: str):
     """
     Determine output file paths for reproducer script and context data.
 
@@ -342,7 +342,12 @@ def determine_output_paths(out_dir: str, kernel_name: str):
     output_directory = Path(out_dir) / kernel_name
     output_directory.mkdir(parents=True, exist_ok=True)
 
-    out_py_path = output_directory / f"repro_{timestamp}.py"
+    filename_parts = ["repro"]
+    if template != "example":
+        filename_parts.append(template.replace(".", "_"))
+    filename_parts.append(timestamp)
+    filename = "_".join(filename_parts) + ".py"
+    out_py_path = output_directory / filename
     temp_json_path = output_directory / f"repro_context_{timestamp}.json"
 
     return out_py_path, temp_json_path

{tritonparse-0.3.1.dev20251027071529 → tritonparse-0.3.1.dev20251029071541}/tritonparse/trace_processor.py

@@ -71,12 +71,38 @@ def generate_source_mappings(
             }
         elif loc_id in loc_defs:
             info = loc_defs[loc_id]
-            mappings[str(ln)] = {
+            entry = {
                 "file": info["file"],
                 "line": info["line"],
                 "column": info["column"],
                 f"{ir_type}_line": ln,
             }
+            # Propagate alias metadata if present
+            if "alias_name" in info:
+                entry["alias_name"] = info["alias_name"]
+            if "alias_of" in info:
+                entry["loc_id"] = loc_id
+            mappings[str(ln)] = entry
+
+    # Add separate entries for loc definition lines
+    for loc_id, info in loc_defs.items():
+        if "def_line" not in info:
+            continue
+        def_ln = info["def_line"]
+        # Only create mapping if this line doesn't already have one
+        if str(def_ln) not in mappings:
+            entry = {
+                "file": info["file"],
+                "line": info["line"],
+                "column": info["column"],
+                f"{ir_type}_line": def_ln,
+                "kind": "loc_def",
+            }
+            if "alias_name" in info:
+                entry["alias_name"] = info["alias_name"]
+            if "alias_of" in info:
+                entry["loc_id"] = loc_id
+            mappings[str(def_ln)] = entry
 
     return mappings
 
@@ -301,8 +327,13 @@ def parse_single_file(
             )
 
         if compilation_event:
-            ir_analysis_event = _generate_ir_analysis(compilation_event)
-            if ir_analysis_event:
+            ir_analysis = _generate_ir_analysis(compilation_event)
+            if ir_analysis:
+                ir_analysis_event = {
+                    "event_type": "ir_analysis",
+                    "hash": _kernel_hash,
+                    "ir_analysis": ir_analysis,
+                }
                 all_output_lines[output_file].append(
                     json.dumps(ir_analysis_event, separators=(",", ":")) + "\n"
                 )