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

tritonparse/__main__.py

@@ -1,3 +1,5 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
 from .cli import main
 
 

tritonparse/context_manager.py

@@ -1,3 +1,5 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
 import os
 import shutil
 import tempfile

tritonparse/event_diff.py

@@ -1,3 +1,5 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
 import json
 from collections import defaultdict
 from typing import Any, Dict, List, Tuple

tritonparse/ir_analysis.py

@@ -0,0 +1,74 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+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/ir_parser.py

@@ -1,3 +1,5 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
 import logging
 import os
 import re

tritonparse/mapper.py

@@ -1,3 +1,5 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
 import logging
 from collections import defaultdict
 from typing import Any, Dict, List, Tuple

tritonparse/reproducer/cli.py

@@ -1,3 +1,5 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
 import argparse
 
 from tritonparse.reproducer.types import KernelImportMode

tritonparse/reproducer/function_extractor.py

@@ -1,3 +1,5 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
 """
 Function extractor for reproducer utility functions.
 

tritonparse/reproducer/ingestion/ndjson.py

@@ -1,3 +1,5 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Tuple
 

tritonparse/reproducer/orchestrator.py

@@ -1,3 +1,5 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
 from pathlib import Path
 from typing import Optional
 

tritonparse/reproducer/placeholder_replacer.py

@@ -1,3 +1,5 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
 from abc import ABC
 
 from typing import Any, Dict, Protocol

tritonparse/reproducer/templates/loader.py

@@ -1,3 +1,5 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
 from importlib.resources import files as pkg_files
 from pathlib import Path
 from typing import List

tritonparse/reproducer/types.py

@@ -1,3 +1,5 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
 from enum import Enum
 
 

tritonparse/reproducer/utils.py

@@ -1,3 +1,5 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
 import importlib
 import importlib.util
 import json

tritonparse/shared_vars.py

@@ -1,3 +1,4 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
 # We'd like to sperate structured logging module and tritonparse module as much as possible. So, put the shared variables here.
 import os
 

tritonparse/sourcemap_utils.py

@@ -1,5 +1,10 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+import logging
 from typing import Any, Dict, List
 
+logger = logging.getLogger("SourceMapping")
+
 
 def get_file_extension(filename: str) -> str:
     """
@@ -70,3 +75,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/tools/decompress_bin_ndjson.py

@@ -1,4 +1,6 @@
 #!/usr/bin/env python3
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
 """
 Script to decompress .bin.ndjson files back to regular .ndjson format.
 

tritonparse/tools/format_fix.py

@@ -1,4 +1,6 @@
 #!/usr/bin/env python3
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
 """
 Format fix script for tritonparse project.
 

tritonparse/tools/load_tensor.py

@@ -1,4 +1,6 @@
 #!/usr/bin/env python3
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
 """
 Simple tensor loading utility for tritonparse saved tensors.
 Usage:

tritonparse/tools/prettify_ndjson.py

@@ -1,4 +1,6 @@
 #!/usr/bin/env python3
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
 """
 Convert an NDJSON file to a prettified JSON file.
 

tritonparse/trace_processor.py

@@ -1,3 +1,5 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
 import gzip
 import json
 import logging
@@ -6,6 +8,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 +16,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 +87,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 +300,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.dev20251025071457.dist-info}/METADATA

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

@@ -0,0 +1,41 @@
+tritonparse/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tritonparse/__main__.py,sha256=RXbkALBewcb1xlJBnsQl9IaBRUNln7U8NuRZKT8UdIk,117
+tritonparse/cli.py,sha256=w_sCIv1NlcGDzoflsYuTvtQis92_YuZEC51ZlTVsYCU,2517
+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/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/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/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/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
+tritonparse/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tritonparse/tools/decompress_bin_ndjson.py,sha256=Gn5foDIlxBN5D5wmcdrEmwvxo3_wRlH8ih2U2Ys3RdM,4199
+tritonparse/tools/disasm.py,sha256=c4HmNNoPPeXPQBQkPVcMaHwDHbHNZNxuqXn4UIIs1Z0,2434
+tritonparse/tools/format_fix.py,sha256=ISalg_N_L7Xktag3mLr-G9T6Opxv793s1WG6A9wUtXk,3922
+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.dev20251025071457.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
+tritonparse-0.3.1.dev20251025071457.dist-info/METADATA,sha256=AOVS0JPzUdhi5-59slIrr7oVO0Ofy_8xNQW75pbLj04,8278
+tritonparse-0.3.1.dev20251025071457.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tritonparse-0.3.1.dev20251025071457.dist-info/entry_points.txt,sha256=wEXdaieDoRRCCdhEv2p_C68iytnaXU_2pwt5CqjfbWY,56
+tritonparse-0.3.1.dev20251025071457.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
+tritonparse-0.3.1.dev20251025071457.dist-info/RECORD,,

tritonparse-0.3.1.dev20251023071548.dist-info/RECORD

@@ -1,40 +0,0 @@
-tritonparse/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tritonparse/__main__.py,sha256=wu5N2wk8mvBgyvr2ghmQf4prezAe0_i-p123VVreyYc,62
-tritonparse/cli.py,sha256=w_sCIv1NlcGDzoflsYuTvtQis92_YuZEC51ZlTVsYCU,2517
-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_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/structured_logging.py,sha256=L1xkkCx8Jr9YQbM0Kgtf2g6L3aWMkYOEeFFEOSo8Lkk,60306
-tritonparse/tp_logger.py,sha256=vXzY7hMDmVnRBGBhIjFZe3nHZzG5NKKPONGUszJhGgU,242
-tritonparse/trace_processor.py,sha256=brQBt26jdB6-quJXP5-warp2j31JSjOOFJa5ayiUZ5k,12963
-tritonparse/utils.py,sha256=Jnlptcd79llSDev-_1XyyOnv2izUqv0PEL74A8GF2tc,4565
-tritonparse/reproducer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tritonparse/reproducer/cli.py,sha256=MqYuuAP-uAIWWLQxixwyyBHJGSaQdG3xXGaVjERTLX8,1522
-tritonparse/reproducer/function_extractor.py,sha256=pB5b52Xlk9-fe-Gs-z_Rn2HkZd1bC9qHTw1JZc_epBM,6588
-tritonparse/reproducer/orchestrator.py,sha256=ZGdsiOZg2Bg6o0cqZSCGx_v6mdzR4yZX7S3SVQjZ9-c,3182
-tritonparse/reproducer/placeholder_replacer.py,sha256=TvqQIrOubmBHJ_pl0ZvkpP4dIDiMYxUll1q9MIbzZco,9552
-tritonparse/reproducer/types.py,sha256=AfVl83zoJZQ58JJoplCcMC51gK-M-OKcafatYEIGgW0,509
-tritonparse/reproducer/utils.py,sha256=LHmkM9EEAn9BwNyhHuCEp1tm7omeILbhkAcOmN-iLrk,16544
-tritonparse/reproducer/ingestion/ndjson.py,sha256=pEujTl5xXW2E2DEW8ngxXQ8qP9oawb90wBVTWHDs1jk,7372
-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=HqjfThdDVg7q2bYWry78sIaVRkUpkcA8KQDt83YrlVE,1920
-tritonparse/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tritonparse/tools/decompress_bin_ndjson.py,sha256=kpt7DM_sSA334F1X45xdkP2OR9LuB27Pc50EkGr6CPM,4144
-tritonparse/tools/disasm.py,sha256=c4HmNNoPPeXPQBQkPVcMaHwDHbHNZNxuqXn4UIIs1Z0,2434
-tritonparse/tools/format_fix.py,sha256=Ol0Sjui8D7OzHwbamAfGnq8V5Y63uwNaFTKSORN5HkQ,3867
-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,,