tritonparse 0.3.1.dev20251027071529__py3-none-any.whl → 0.3.1.dev20251028071524__py3-none-any.whl

tritonparse/cli.py

@@ -70,7 +70,7 @@ def main():
     elif args.func == "reproduce":
         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,

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/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/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/reproducer/orchestrator.py

@@ -29,7 +29,7 @@ def reproduce(
 
     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.

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"
                 )

{tritonparse-0.3.1.dev20251027071529.dist-info → tritonparse-0.3.1.dev20251028071524.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tritonparse
-Version: 0.3.1.dev20251027071529
+Version: 0.3.1.dev20251028071524
 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.dist-info → tritonparse-0.3.1.dev20251028071524.dist-info}/RECORD

@@ -1,28 +1,28 @@
 tritonparse/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tritonparse/__main__.py,sha256=RXbkALBewcb1xlJBnsQl9IaBRUNln7U8NuRZKT8UdIk,117
-tritonparse/cli.py,sha256=w_sCIv1NlcGDzoflsYuTvtQis92_YuZEC51ZlTVsYCU,2517
+tritonparse/cli.py,sha256=Z3nz_rGYXF6NKmI3LsTomWqM51sJ8Tp-ybxRq7l20BI,2569
 tritonparse/common.py,sha256=MJo9bVCgSKkwXpEoUkUczPo_5jOYpJgXLq4UsWYqN3c,13924
 tritonparse/context_manager.py,sha256=OdMn11qbApYL2c9IlbUpcT27r04ZSa4DfvrY2mLA958,2243
 tritonparse/event_diff.py,sha256=USCjfjYr-7Ie-EfZgtCFMZMA1KRzFRDe7yDFy98zYI4,4962
 tritonparse/extract_source_mappings.py,sha256=Z6UxFj2cCE5NCWLQTYPKqUpLfbYhqP8xgCl5mvud9KI,1451
-tritonparse/ir_analysis.py,sha256=YObFsRKkpdjkwL3lIeWRXU6DY13fGVqdh8li4PiCpoQ,2482
-tritonparse/ir_parser.py,sha256=Qtb8kA6IMBtbLbABEGSDFdGQQI_GsPd1iC2GHbEviaI,9379
+tritonparse/ir_analysis.py,sha256=MoOXuHsUGZ705R4JnXmlsrBn9gJdLO1Dnf0L5AxcaBM,2551
+tritonparse/ir_parser.py,sha256=MH4RwoNZMBdWUxkFyEhemJ7Aa7-asoba66b06bGPNsk,13237
 tritonparse/mapper.py,sha256=QBCUMHM9pu3x3ahFp0wyXJbmv9TFGVPdkcLULok1E-k,4205
 tritonparse/shared_vars.py,sha256=RifXq55KisHgspYAmGcaCWY6ZHX8iejFHvwIewvcWZE,707
 tritonparse/source_type.py,sha256=nmYEQS8rfkIN9BhNhQbkmEvKnvS-3zAxRGLY4TaZdi8,1676
 tritonparse/sourcemap_utils.py,sha256=uI02n5Sgnlx7Nc15QAX5N6_tZZMips0PyJuo1n3eouY,2654
 tritonparse/structured_logging.py,sha256=L1xkkCx8Jr9YQbM0Kgtf2g6L3aWMkYOEeFFEOSo8Lkk,60306
 tritonparse/tp_logger.py,sha256=vXzY7hMDmVnRBGBhIjFZe3nHZzG5NKKPONGUszJhGgU,242
-tritonparse/trace_processor.py,sha256=cJHagBOSHFpvUKYLxrvtgoFODDZ5CEFkU_xZTKsD9s8,12981
+tritonparse/trace_processor.py,sha256=aQPqlnpTtWoGzHYv4BXWUH4nCeUQGSK3o-fj0LD9I0c,14147
 tritonparse/utils.py,sha256=Jnlptcd79llSDev-_1XyyOnv2izUqv0PEL74A8GF2tc,4565
 tritonparse/reproducer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tritonparse/reproducer/cli.py,sha256=wk0K8qJhvP9gty2EBMH3WEc3TSFcszvNq3JbfYu_sTw,1577
 tritonparse/reproducer/function_extractor.py,sha256=kQr10JKHy8EvAN7ic4Azjz6TYe-udBW2DVmbQ--c1pc,6643
-tritonparse/reproducer/orchestrator.py,sha256=h55xXEVYtIr33GvwCuvM03bsYc01_yPa5X_B-3w71kE,3237
+tritonparse/reproducer/orchestrator.py,sha256=GotBOJjrShN1oCFc_xTMXn8WWT1Jlfap5qcM21dKBpM,3259
 tritonparse/reproducer/placeholder_replacer.py,sha256=ARPZAa9A3Fyit_dIclOKe1JzFgUPBFdHvfy3z20x2E8,9607
 tritonparse/reproducer/types.py,sha256=86wql3NaGgpkOzx0gDFb5qexNjKExzhL0uIwGU7grrw,564
 tritonparse/reproducer/utils.py,sha256=yFS1Mg2IhRgW-1UNfqjWH5gRSqc8Wbn5Ykre8L-EWcU,16599
-tritonparse/reproducer/ingestion/ndjson.py,sha256=QX_4XyK_vhSOFyJtfDxQNwzXqVXolrfk4cefzlShFps,7427
+tritonparse/reproducer/ingestion/ndjson.py,sha256=7amSwpbtG-od1-pW18Nm9AiaFc3Etd0-UETXwiYCmgw,7443
 tritonparse/reproducer/templates/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tritonparse/reproducer/templates/example.py,sha256=jR3c8_d7fAFJYaj1DuUuthnI4Xd-_606bWDRdUPMNyo,785
 tritonparse/reproducer/templates/loader.py,sha256=x14KHXkovOIcXFKii3Jx4XjpEhXqUMqp575qAffi370,1975
@@ -33,9 +33,9 @@ tritonparse/tools/format_fix.py,sha256=ISalg_N_L7Xktag3mLr-G9T6Opxv793s1WG6A9wUt
 tritonparse/tools/load_tensor.py,sha256=7-LbpboKDNJFBLNhiKS3enoqRlVAb55OjPc70PwHXAw,2789
 tritonparse/tools/prettify_ndjson.py,sha256=kR8hmBCv-iJeuzpi2_6CZv9T4_edRQbBOSOPpMm6wrw,11117
 tritonparse/tools/readme.md,sha256=w6PWYfYnRgoPArLjxG9rVrpcLUkoVMGuRlbpF-o0IQM,110
-tritonparse-0.3.1.dev20251027071529.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
-tritonparse-0.3.1.dev20251027071529.dist-info/METADATA,sha256=lpDRFaDfU5veptgFOboAjJLvvVPZfd9BGhDan4JLQgQ,8278
-tritonparse-0.3.1.dev20251027071529.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tritonparse-0.3.1.dev20251027071529.dist-info/entry_points.txt,sha256=wEXdaieDoRRCCdhEv2p_C68iytnaXU_2pwt5CqjfbWY,56
-tritonparse-0.3.1.dev20251027071529.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
-tritonparse-0.3.1.dev20251027071529.dist-info/RECORD,,
+tritonparse-0.3.1.dev20251028071524.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
+tritonparse-0.3.1.dev20251028071524.dist-info/METADATA,sha256=6PkpYMi1Qjf4Lar46WHsMPBo6dQts6i6n3IcwmijYeg,8282
+tritonparse-0.3.1.dev20251028071524.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tritonparse-0.3.1.dev20251028071524.dist-info/entry_points.txt,sha256=wEXdaieDoRRCCdhEv2p_C68iytnaXU_2pwt5CqjfbWY,56
+tritonparse-0.3.1.dev20251028071524.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
+tritonparse-0.3.1.dev20251028071524.dist-info/RECORD,,