tritonparse 0.3.1.dev20251023071548__py3-none-any.whl → 0.3.1.dev20251024071516__py3-none-any.whl

tritonparse/ir_analysis.py

@@ -0,0 +1,72 @@
+import logging
+
+from .sourcemap_utils import load_ir_contents
+
+
+logger = logging.getLogger("IRAnalysis")
+
+
+def process_amd_bufferop(ir_content: str, io_keys: list[str]) -> dict[str, int]:
+    def make_key(prefix: str) -> str:
+        return f"{prefix}_count"
+
+    io_keys = [(make_key(prefix), prefix) for prefix in io_keys]
+    output: dict[str, int] = {}
+    for dict_key, _ in io_keys:
+        output[dict_key] = 0
+    if ir_content:
+        for line in ir_content.split("\n"):
+            for dict_key, code_key in io_keys:
+                if code_key in line:
+                    output[dict_key] += 1
+    return output
+
+
+def process_amd_ttgir_bufferops(
+    key: str,
+    file_content: dict[str, str],
+    file_path: dict[str, str],
+) -> dict[str, int]:
+    ir_content = load_ir_contents(key, file_content, file_path)
+    # TODO: Add atomics
+    io_keys = ["tt.load", "tt.store", "amdgpu.buffer_load", "amdgpu.buffer_store"]
+    return process_amd_bufferop(ir_content, io_keys)
+
+
+def process_amd_gcn_bufferops(
+    key: str,
+    file_content: dict[str, str],
+    file_path: dict[str, str],
+) -> 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_"]
+    return process_amd_bufferop(ir_content, io_keys)
+
+
+def _generate_ir_analysis(entry: str):
+    payload = entry.setdefault("payload", {})
+    file_content = payload.get("file_content", {})
+    file_path = payload.get("file_path", {})
+
+    # Find the IR file keys
+    ttgir_key = next((k for k in file_content if k.endswith(".ttgir")), None)
+    amdgcn_key = next((k for k in file_content if k.endswith(".amdgcn")), None)
+    # Skip if no IR files found
+    if not (ttgir_key or amdgcn_key):
+        logger.debug("No AMD IR found")
+        return {}
+    ir_analysis = {}
+    if amdgcn_key:
+        ttgir_bufferops_info = process_amd_ttgir_bufferops(
+            ttgir_key, file_content, file_path
+        )
+        gcn_bufferops_info = process_amd_gcn_bufferops(
+            amdgcn_key, file_content, file_path
+        )
+        # NDJSON format requires a newline at the end of each line
+        if ttgir_bufferops_info:
+            ir_analysis["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}

tritonparse/sourcemap_utils.py

@@ -1,5 +1,8 @@
+import logging
 from typing import Any, Dict, List
 
+logger = logging.getLogger("SourceMapping")
+
 
 def get_file_extension(filename: str) -> str:
     """
@@ -70,3 +73,22 @@ def _to_ranges(indices: List[int]) -> List[Dict[str, int]]:
 
     ranges.append({"start": start, "end": end})
     return ranges
+
+
+def load_ir_contents(
+    key: str,
+    file_content: dict[str, str],
+    file_path: dict[str, str],
+):
+    if not key:
+        return {}
+    logger.debug(f"Processing {key}")
+    ir_content = file_content.get(key, None)
+    if not ir_content:
+        ir_file_path = file_path.get(key, None)
+        if not ir_file_path:
+            logger.warning(f"No content found for {key}")
+            return {}
+        with open(ir_file_path, "r") as f:
+            ir_content = f.read()
+    return ir_content

tritonparse/trace_processor.py

@@ -6,6 +6,7 @@ from collections import defaultdict
 from typing import Any, Dict, List
 
 from .event_diff import _generate_launch_diff
+from .ir_analysis import _generate_ir_analysis
 
 from .ir_parser import (
     extract_code_locations,
@@ -13,7 +14,7 @@ from .ir_parser import (
     extract_ptx_amdgcn_mappings,
 )
 from .mapper import create_bidirectional_mapping, create_python_mapping
-from .sourcemap_utils import get_file_extension
+from .sourcemap_utils import get_file_extension, load_ir_contents
 
 logger = logging.getLogger("SourceMapping")
 
@@ -84,19 +85,9 @@ def process_ir(
     file_path: Dict[str, str],
     other_mappings: List[Any] = None,
 ):
-    # Generate source mappings for each IR type
-    # the key should be the full file name with extension for the IR files
-    if not key:
-        return {}
-    logger.debug(f"Processing {key}")
-    ir_content = file_content.get(key, None)
+    ir_content = load_ir_contents(key, file_content, file_path)
     if not ir_content:
-        ir_file_path = file_path.get(key, None)
-        if not ir_file_path:
-            logger.warning(f"No content found for {key}")
-            return {}
-        with open(ir_file_path, "r") as f:
-            ir_content = f.read()
+        return {}
     mapping = generate_source_mappings(ir_content, key.split(".")[1], other_mappings)
     logger.debug(f"Generated source mapping for {key}")
     return mapping
@@ -307,6 +298,13 @@ def parse_single_file(
                 json.dumps(launch_event, separators=(",", ":")) + "\n"
             )
 
+        if compilation_event:
+            ir_analysis_event = _generate_ir_analysis(compilation_event)
+            if ir_analysis_event:
+                all_output_lines[output_file].append(
+                    json.dumps(ir_analysis_event, separators=(",", ":")) + "\n"
+                )
+
         if compilation_event and launches_with_indices:
             sames, diffs, launch_index_map = _generate_launch_diff(
                 launches_with_indices

{tritonparse-0.3.1.dev20251023071548.dist-info → tritonparse-0.3.1.dev20251024071516.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tritonparse
-Version: 0.3.1.dev20251023071548
+Version: 0.3.1.dev20251024071516
 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

{tritonparse-0.3.1.dev20251023071548.dist-info → tritonparse-0.3.1.dev20251024071516.dist-info}/RECORD

@@ -5,14 +5,15 @@ tritonparse/common.py,sha256=MJo9bVCgSKkwXpEoUkUczPo_5jOYpJgXLq4UsWYqN3c,13924
 tritonparse/context_manager.py,sha256=R6zKoixi8aXF2yK8pN1EUVPPmldhhXMdEK_PgKZjaX4,2188
 tritonparse/event_diff.py,sha256=yOD6uNxLJroatfx2nEGr-erw24ObOrHU9P6V5pzr8do,4907
 tritonparse/extract_source_mappings.py,sha256=Z6UxFj2cCE5NCWLQTYPKqUpLfbYhqP8xgCl5mvud9KI,1451
+tritonparse/ir_analysis.py,sha256=ERu6EWotqcW99ZGNEH-mxGyYej45yt0IXx0OCab_bfw,2427
 tritonparse/ir_parser.py,sha256=1j1tP9jpUN7wH3e01bKUkUPgTMlNXUdp8LKRCC-WTro,9324
 tritonparse/mapper.py,sha256=prrczfi13P7Aa042OrEBsmRF1HW3jDhwxicANgPkWIM,4150
 tritonparse/shared_vars.py,sha256=fCAW24mx9nENYoNbTy-tZjiN5-w6oGTO_av-Pw1J1TY,653
 tritonparse/source_type.py,sha256=nmYEQS8rfkIN9BhNhQbkmEvKnvS-3zAxRGLY4TaZdi8,1676
-tritonparse/sourcemap_utils.py,sha256=qsQmTDuEe9yuUVyxSHRbjTR38gi0hvJEijnPkrJVAV4,2037
+tritonparse/sourcemap_utils.py,sha256=f36lKjc23qx4aOgpHOOo2lxPF4rwgSrOcULZqyUKxn0,2599
 tritonparse/structured_logging.py,sha256=L1xkkCx8Jr9YQbM0Kgtf2g6L3aWMkYOEeFFEOSo8Lkk,60306
 tritonparse/tp_logger.py,sha256=vXzY7hMDmVnRBGBhIjFZe3nHZzG5NKKPONGUszJhGgU,242
-tritonparse/trace_processor.py,sha256=brQBt26jdB6-quJXP5-warp2j31JSjOOFJa5ayiUZ5k,12963
+tritonparse/trace_processor.py,sha256=DMY03eqQbM1oUbQuMkfHV_rlaKmW4ggJ048FMCbOH3c,12926
 tritonparse/utils.py,sha256=Jnlptcd79llSDev-_1XyyOnv2izUqv0PEL74A8GF2tc,4565
 tritonparse/reproducer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tritonparse/reproducer/cli.py,sha256=MqYuuAP-uAIWWLQxixwyyBHJGSaQdG3xXGaVjERTLX8,1522
@@ -32,9 +33,9 @@ tritonparse/tools/format_fix.py,sha256=Ol0Sjui8D7OzHwbamAfGnq8V5Y63uwNaFTKSORN5H
 tritonparse/tools/load_tensor.py,sha256=94-TiSYlpXJx4MPmGK1ovmZlTt56Q_B3KQeCPaA6Cnw,2734
 tritonparse/tools/prettify_ndjson.py,sha256=r2YlHwFDTHgML7KljRmMsHaDg29q8gOQAgyDKWJhxRM,11062
 tritonparse/tools/readme.md,sha256=w6PWYfYnRgoPArLjxG9rVrpcLUkoVMGuRlbpF-o0IQM,110
-tritonparse-0.3.1.dev20251023071548.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
-tritonparse-0.3.1.dev20251023071548.dist-info/METADATA,sha256=QCqUUpoXTPPeUN3mhNWN_fsipg9W3a1Pq5GiHln1oRM,8278
-tritonparse-0.3.1.dev20251023071548.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tritonparse-0.3.1.dev20251023071548.dist-info/entry_points.txt,sha256=wEXdaieDoRRCCdhEv2p_C68iytnaXU_2pwt5CqjfbWY,56
-tritonparse-0.3.1.dev20251023071548.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
-tritonparse-0.3.1.dev20251023071548.dist-info/RECORD,,
+tritonparse-0.3.1.dev20251024071516.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
+tritonparse-0.3.1.dev20251024071516.dist-info/METADATA,sha256=zSXJHPi5BEV2zS_aiZBVSEGQXY8_xLp1f6YoC2NHFgQ,8278
+tritonparse-0.3.1.dev20251024071516.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tritonparse-0.3.1.dev20251024071516.dist-info/entry_points.txt,sha256=wEXdaieDoRRCCdhEv2p_C68iytnaXU_2pwt5CqjfbWY,56
+tritonparse-0.3.1.dev20251024071516.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
+tritonparse-0.3.1.dev20251024071516.dist-info/RECORD,,