tritonparse 0.3.1.dev20251023071548__tar.gz → 0.3.1.dev20251025071457__tar.gz

{tritonparse-0.3.1.dev20251023071548/tritonparse.egg-info → tritonparse-0.3.1.dev20251025071457}/PKG-INFO

@@ -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.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/tritonparse/__main__.py

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

{tritonparse-0.3.1.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/tritonparse/context_manager.py

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

{tritonparse-0.3.1.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/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-0.3.1.dev20251025071457/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-0.3.1.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/tritonparse/ir_parser.py

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

{tritonparse-0.3.1.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/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-0.3.1.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/tritonparse/reproducer/cli.py

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

{tritonparse-0.3.1.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/tritonparse/reproducer/function_extractor.py

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

{tritonparse-0.3.1.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/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-0.3.1.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/tritonparse/reproducer/orchestrator.py

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

{tritonparse-0.3.1.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/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-0.3.1.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/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-0.3.1.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/tritonparse/reproducer/types.py

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

{tritonparse-0.3.1.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/tritonparse/reproducer/utils.py

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

{tritonparse-0.3.1.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/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-0.3.1.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/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-0.3.1.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/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-0.3.1.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/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-0.3.1.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/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-0.3.1.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/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-0.3.1.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/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 → tritonparse-0.3.1.dev20251025071457/tritonparse.egg-info}/PKG-INFO

@@ -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.dev20251023071548 → tritonparse-0.3.1.dev20251025071457}/tritonparse.egg-info/SOURCES.txt

@@ -35,6 +35,7 @@ tritonparse/common.py
 tritonparse/context_manager.py
 tritonparse/event_diff.py
 tritonparse/extract_source_mappings.py
+tritonparse/ir_analysis.py
 tritonparse/ir_parser.py
 tritonparse/mapper.py
 tritonparse/shared_vars.py