tritonparse 0.3.2.dev20251210071601__py3-none-any.whl

tritonparse/tools/decompress_bin_ndjson.py

@@ -0,0 +1,120 @@
+#!/usr/bin/env python3
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+"""
+Script to decompress .bin.ndjson files back to regular .ndjson format.
+
+The .bin.ndjson format stores each JSON record as a separate gzip member,
+concatenated in sequence within a single binary file. This script uses
+gzip.open() which automatically handles member concatenation to read
+the compressed file and write out the original NDJSON format.
+
+Usage:
+    python decompress_bin_ndjson.py trace.bin.ndjson
+"""
+
+import argparse
+import gzip
+import sys
+from pathlib import Path
+
+
+def decompress_bin_ndjson(input_file: str, output_file: str = None) -> None:
+    """
+    Decompress a .bin.ndjson file to regular .ndjson format.
+
+    Args:
+        input_file: Path to the .bin.ndjson file
+        output_file: Path for the output .ndjson file (optional)
+    """
+    input_path = Path(input_file)
+
+    # Validate input file
+    if not input_path.exists():
+        print(f"Error: Input file '{input_file}' does not exist", file=sys.stderr)
+        return
+
+    if not input_path.suffix.endswith(".bin.ndjson"):
+        print(f"Warning: Input file '{input_file}' doesn't have .bin.ndjson extension")
+
+    # Determine output file path
+    if output_file is None:
+        if input_path.name.endswith(".bin.ndjson"):
+            # Replace .bin.ndjson with .ndjson
+            output_file = str(input_path.with_suffix("").with_suffix(".ndjson"))
+        else:
+            # Add .decompressed.ndjson suffix
+            output_file = str(input_path.with_suffix(".decompressed.ndjson"))
+
+    output_path = Path(output_file)
+
+    try:
+        line_count = 0
+        # Because we use NDJSON format, each line is a complete JSON record.
+        # It is guruanteed here https://github.com/meta-pytorch/tritonparse/blob/
+        # c8dcc2a94ac10ede4342dba7456f6ebd8409b95d/tritonparse/structured_logging.py#L320
+        with gzip.open(input_path, "rt", encoding="utf-8") as compressed_file:
+            with open(output_path, "w", encoding="utf-8") as output:
+                for line in compressed_file:
+                    # gzip.open automatically handles member concatenation
+                    # Each line is already a complete JSON record with newline
+                    output.write(line)
+                    line_count += 1
+
+        # Get file sizes for comparison
+        input_size = input_path.stat().st_size
+        output_size = output_path.stat().st_size
+        compression_ratio = (
+            (1 - input_size / output_size) * 100 if output_size > 0 else 0
+        )
+
+        print(f"Successfully decompressed '{input_file}' to '{output_file}'")
+        print(f"  Input size:  {input_size:,} bytes")
+        print(f"  Output size: {output_size:,} bytes")
+        print(f"  Compression ratio: {compression_ratio:.1f}%")
+        print(f"  Records processed: {line_count:,}")
+
+    except gzip.BadGzipFile as e:
+        print(f"Error: Invalid gzip format in '{input_file}': {e}", file=sys.stderr)
+    except UnicodeDecodeError as e:
+        print(f"Error: Unicode decode error in '{input_file}': {e}", file=sys.stderr)
+    except Exception as e:
+        print(f"Error: Failed to decompress '{input_file}': {e}", file=sys.stderr)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Decompress .bin.ndjson files to regular .ndjson format",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  %(prog)s trace.bin.ndjson
+  %(prog)s trace.bin.ndjson -o output.ndjson
+  %(prog)s /logs/dedicated_log_triton_trace_user_.bin.ndjson
+        """,
+    )
+
+    parser.add_argument("input_file", help="Input .bin.ndjson file to decompress")
+
+    parser.add_argument(
+        "-o",
+        "--output",
+        help="Output .ndjson file path (default: replace .bin.ndjson with .ndjson)",
+    )
+
+    parser.add_argument(
+        "-v", "--verbose", action="store_true", help="Enable verbose output"
+    )
+
+    args = parser.parse_args()
+
+    if args.verbose:
+        print(f"Decompressing: {args.input_file}")
+        if args.output:
+            print(f"Output file: {args.output}")
+
+    decompress_bin_ndjson(args.input_file, args.output)
+
+
+if __name__ == "__main__":
+    main()

tritonparse/tools/disasm.py

@@ -0,0 +1,81 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+import re
+import subprocess
+
+# Regex patterns for nvdisasm output
+NVDISASM_FNAME_RE = re.compile(r"^\s*\.global\s+(\w+)")
+
+
+def path_to_nvdisasm():
+    from triton import knobs
+
+    return knobs.nvidia.nvdisasm.path
+
+
+def is_nvdisasm_available():
+    try:
+        if path_to_nvdisasm():
+            return True
+        else:
+            return False
+    except RuntimeError:
+        return False
+
+
+def extract(file_path):
+    """Extract SASS from CUBIN using nvdisasm.
+
+    nvdisasm output is much cleaner than cuobjdump:
+    - Single line per instruction (no encoding lines)
+    - Labels are already symbolized (.L_x_0 instead of addresses)
+    - Source line information is included
+    - No need for complex address remapping
+
+    nvdisasm Documentation:
+    https://docs.nvidia.com/cuda/cuda-binary-utilities/index.html
+    """
+    nvdisasm = path_to_nvdisasm()
+    args = [nvdisasm, "-c", "-gp", "-g", "-gi", file_path]
+    sass_str = subprocess.check_output(args)
+    sass_lines = sass_str.splitlines()
+    line_idx = 0
+
+    while line_idx < len(sass_lines):
+        line = sass_lines[line_idx].decode()
+
+        # Find function definition (.global function_name)
+        while NVDISASM_FNAME_RE.match(line) is None:
+            line_idx += 1
+            if line_idx >= len(sass_lines):
+                return None
+            line = sass_lines[line_idx].decode()
+
+        # Extract function name
+        match = NVDISASM_FNAME_RE.match(line)
+        if match is None:
+            return None
+        fname = match.group(1)
+        ret = f"Function:{fname}\n"
+
+        # Find the actual start of function content (.text.kernel_name:)
+        text_section_pattern = f".text.{fname}:"
+        line_idx += 1
+        while line_idx < len(sass_lines):
+            line = sass_lines[line_idx].decode().strip()
+            if line == text_section_pattern:
+                line_idx += 1  # Move past the .text.kernel_name: line
+                break
+            line_idx += 1
+
+        # Process all lines until next .headerflags or end of file
+        while line_idx < len(sass_lines):
+            line = sass_lines[line_idx].decode().rstrip()
+
+            # Stop if we encounter next function's headerflags
+            if line.strip().startswith(".headerflags"):
+                break
+            ret += line + "\n"
+            line_idx += 1
+
+        ret += "\n"
+        return ret

tritonparse/tools/extract_irs.py

@@ -0,0 +1,244 @@
+#!/usr/bin/env python3
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+"""
+Extract IR files from NDJSON trace logs.
+
+This script extracts intermediate representation (IR) files from a Triton trace NDJSON file.
+For compilation events, it extracts the IR files (ttir, ttgir, llir, ptx, etc.) contained in
+the file_content field and saves them as individual files.
+
+Example:
+    Extract IRs from line 0 (first line) of the NDJSON file:
+        python extract_irs.py -i logs.ndjson --line 0 -o output_folder
+
+    Extract from line 5:
+        python extract_irs.py -i logs.ndjson --line 5 -o ./irs
+
+Usage:
+    python extract_irs.py -i <input.ndjson> --line <line_number> -o <output_folder>
+"""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+
+def read_ndjson_line(file_path: Path, line_number: int) -> Optional[Dict[str, Any]]:
+    """
+    Read a specific line from an NDJSON file (0-based indexing).
+
+    Args:
+        file_path: Path to the NDJSON file
+        line_number: Line number to read (0-based, where 0 = first line)
+
+    Returns:
+        Parsed JSON object from the specified line, or None if line doesn't exist
+
+    Raises:
+        FileNotFoundError: If the input file doesn't exist
+        json.JSONDecodeError: If the line contains invalid JSON
+    """
+    if not file_path.exists():
+        raise FileNotFoundError(f"File not found: {file_path}")
+
+    try:
+        with open(file_path, "r", encoding="utf-8") as f:
+            for current_line_num, line in enumerate(f):
+                if current_line_num == line_number:
+                    line = line.strip()
+                    if not line:
+                        print(f"Warning: Line {line_number} is empty", file=sys.stderr)
+                        return None
+                    return json.loads(line)
+
+        print(
+            f"Error: Line {line_number} not found in file (file has fewer lines)",
+            file=sys.stderr,
+        )
+        return None
+
+    except json.JSONDecodeError as e:
+        print(f"Error: Invalid JSON on line {line_number}: {e}", file=sys.stderr)
+        raise
+
+
+def extract_irs(
+    json_obj: Dict[str, Any], output_dir: Path, kernel_name: Optional[str] = None
+) -> int:
+    """
+    Extract IR files from a JSON object and save them to the output directory.
+
+    Args:
+        json_obj: Parsed JSON object from the NDJSON file
+        output_dir: Directory to save the extracted IR files
+        kernel_name: Optional kernel name to use for file naming (overrides metadata.name)
+
+    Returns:
+        Number of files extracted
+
+    Raises:
+        ValueError: If the JSON object is not a compilation event or missing required fields
+    """
+    # Validate that this is a compilation event
+    event_type = json_obj.get("event_type")
+    if event_type != "compilation":
+        raise ValueError(f"Not a compilation event (event_type: {event_type})")
+
+    payload = json_obj.get("payload")
+    if not payload:
+        raise ValueError("Missing 'payload' field in JSON object")
+
+    # Get file_content
+    file_content = payload.get("file_content")
+    if not file_content:
+        raise ValueError("Missing 'file_content' field in payload")
+
+    # Determine kernel name
+    if kernel_name is None:
+        metadata = payload.get("metadata", {})
+        kernel_name = metadata.get("name", "kernel")
+
+    # Create output directory if it doesn't exist
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Extract each IR file
+    files_extracted = 0
+    for file_key, content in file_content.items():
+        # Determine file extension from the key
+        # file_key is typically like "embedding_forward_kernel.ttir"
+        # We want to extract just the extension
+        if "." in file_key:
+            extension = file_key.split(".")[-1]
+        else:
+            extension = "txt"
+
+        # Create output filename
+        output_filename = f"{kernel_name}.{extension}"
+        output_path = output_dir / output_filename
+
+        # Write content to file
+        try:
+            with open(output_path, "w", encoding="utf-8") as f:
+                f.write(content)
+            print(f"Extracted: {output_path}")
+            files_extracted += 1
+        except OSError as e:
+            print(f"Error writing file {output_path}: {e}", file=sys.stderr)
+
+    # Optionally extract Python source code
+    python_source = payload.get("python_source")
+    if python_source and isinstance(python_source, dict):
+        source_code = python_source.get("code")
+        if source_code:
+            output_path = output_dir / f"{kernel_name}_source.py"
+            try:
+                with open(output_path, "w", encoding="utf-8") as f:
+                    # Add header comment with file path and line range
+                    file_path_info = python_source.get("file_path", "unknown")
+                    start_line = python_source.get("start_line", "?")
+                    end_line = python_source.get("end_line", "?")
+                    f.write(f"# Source: {file_path_info}\n")
+                    f.write(f"# Lines: {start_line}-{end_line}\n\n")
+                    f.write(source_code)
+                print(f"Extracted Python source: {output_path}")
+                files_extracted += 1
+            except OSError as e:
+                print(
+                    f"Error writing Python source file {output_path}: {e}",
+                    file=sys.stderr,
+                )
+
+    return files_extracted
+
+
+def main():
+    """Main function to handle command line arguments and orchestrate IR extraction."""
+    parser = argparse.ArgumentParser(
+        description="Extract IR files from Triton trace NDJSON logs",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  Extract IRs from line 0 (first line):
+    python extract_irs.py -i logs.ndjson --line 0 -o output_folder
+  
+  Extract from line 5:
+    python extract_irs.py -i logs.ndjson --line 5 -o ./irs
+  
+  Specify custom kernel name:
+    python extract_irs.py -i logs.ndjson --line 0 -o ./irs --kernel-name my_kernel
+        """,
+    )
+
+    parser.add_argument(
+        "-i", "--input", type=str, required=True, help="Path to the input NDJSON file"
+    )
+
+    parser.add_argument(
+        "--line",
+        type=int,
+        required=True,
+        help="Line number to extract (0-based indexing, where 0 = first line)",
+    )
+
+    parser.add_argument(
+        "-o",
+        "--output",
+        type=str,
+        required=True,
+        help="Output directory to save extracted IR files",
+    )
+
+    parser.add_argument(
+        "--kernel-name",
+        type=str,
+        help="Custom kernel name for output files (default: use metadata.name from JSON)",
+    )
+
+    args = parser.parse_args()
+
+    # Validate line number
+    if args.line < 0:
+        print(
+            f"Error: Line number must be non-negative (got {args.line})",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+    # Convert to Path objects
+    input_path = Path(args.input)
+    output_dir = Path(args.output)
+
+    try:
+        # Read the specified line
+        print(f"Reading line {args.line} from {input_path}...")
+        json_obj = read_ndjson_line(input_path, args.line)
+
+        if json_obj is None:
+            print("Error: Failed to read JSON from specified line", file=sys.stderr)
+            sys.exit(1)
+
+        # Extract IRs
+        print(f"Extracting IRs to {output_dir}...")
+        num_files = extract_irs(json_obj, output_dir, args.kernel_name)
+
+        print(f"\nSuccess! Extracted {num_files} file(s) to {output_dir}")
+
+    except FileNotFoundError as e:
+        print(f"Error: {e}", file=sys.stderr)
+        sys.exit(1)
+    except ValueError as e:
+        print(f"Error: {e}", file=sys.stderr)
+        sys.exit(1)
+    except json.JSONDecodeError as e:
+        print(f"Error: Failed to parse JSON - {e}", file=sys.stderr)
+        sys.exit(1)
+    except Exception as e:
+        print(f"Unexpected error: {e}", file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

tritonparse/tools/format_fix.py

@@ -0,0 +1,151 @@
+#!/usr/bin/env python3
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+"""
+Format fix script for tritonparse project.
+
+This script runs all linter tools to format and fix code issues:
+- usort: Import sorting
+- ruff: Linting only
+- black: Code formatting
+
+Usage:
+    python -m tritonparse.tools.format_fix [options]
+
+Options:
+    --check-only    Only check for issues, don't fix them
+    --verbose       Verbose output
+    --help          Show this help message
+"""
+
+import argparse
+import subprocess
+import sys
+
+
+def run_command(cmd: list[str], verbose: bool = False) -> bool:
+    """Run a command and return success status."""
+    if verbose:
+        print(f"Running: {' '.join(cmd)}")
+
+    try:
+        result = subprocess.run(cmd, capture_output=True, text=True, check=False)
+
+        if result.returncode != 0:
+            if verbose:
+                print(f"Command failed with return code {result.returncode}")
+                if result.stdout:
+                    print("STDOUT:", result.stdout)
+                if result.stderr:
+                    print("STDERR:", result.stderr)
+            return False
+
+        if verbose and result.stdout:
+            print(result.stdout)
+
+        return True
+    except Exception as e:
+        if verbose:
+            print(f"Error running command: {e}")
+        return False
+
+
+def run_usort(check_only: bool = False, verbose: bool = False) -> bool:
+    """Run usort for import sorting."""
+    cmd = ["usort"]
+
+    if check_only:
+        cmd.extend(["check", "."])
+    else:
+        cmd.extend(["format", "."])
+
+    return run_command(cmd, verbose)
+
+
+def run_ruff_check(check_only: bool = False, verbose: bool = False) -> bool:
+    """Run ruff for linting only."""
+    cmd = ["ruff", "check", "."]
+
+    if check_only:
+        cmd.append("--diff")
+    else:
+        cmd.append("--fix")
+
+    return run_command(cmd, verbose)
+
+
+def run_black(check_only: bool = False, verbose: bool = False) -> bool:
+    """Run black for code formatting."""
+    cmd = ["black"]
+
+    if check_only:
+        cmd.extend(["--check", "--diff", "."])
+    else:
+        cmd.append(".")
+
+    return run_command(cmd, verbose)
+
+
+def main():
+    """Main function."""
+    parser = argparse.ArgumentParser(
+        description="Format fix script for tritonparse project",
+        epilog="""
+Examples:
+    # Fix all formatting issues
+    python -m tritonparse.tools.format_fix
+    
+    # Check for issues without fixing
+    python -m tritonparse.tools.format_fix --check-only
+    
+    # Verbose output
+    python -m tritonparse.tools.format_fix --verbose
+        """,
+    )
+
+    parser.add_argument(
+        "--check-only",
+        action="store_true",
+        help="Only check for issues, don't fix them",
+    )
+    parser.add_argument("--verbose", action="store_true", help="Verbose output")
+
+    args = parser.parse_args()
+
+    # Run formatters on the entire project
+    success = True
+
+    # 1. Run usort for import sorting
+    print("Running usort for import sorting...")
+    if not run_usort(args.check_only, args.verbose):
+        print("❌ usort failed")
+        success = False
+    else:
+        print("✅ usort completed")
+
+    # 2. Run ruff for linting only
+    print("Running ruff for linting...")
+    if not run_ruff_check(args.check_only, args.verbose):
+        print("❌ ruff linting failed")
+        success = False
+    else:
+        print("✅ ruff linting completed")
+
+    # 3. Run black for code formatting
+    print("Running black for code formatting...")
+    if not run_black(args.check_only, args.verbose):
+        print("❌ black failed")
+        success = False
+    else:
+        print("✅ black completed")
+
+    if success:
+        print("\n🎉 All formatting tools completed successfully!")
+        return 0
+    else:
+        print("\n❌ Some formatting tools failed. Please check the output above.")
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tritonparse/tools/load_tensor.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+"""
+Simple tensor loading utility for tritonparse saved tensors.
+Usage:
+import tritonparse.tools.load_tensor as load_tensor
+tensor = load_tensor.load_tensor(tensor_file_path, device)
+"""
+
+import gzip
+import hashlib
+import io
+from pathlib import Path
+from typing import Union
+
+import torch
+
+
+def load_tensor(tensor_file_path: Union[str, Path], device: str = None) -> torch.Tensor:
+    """
+    Load a tensor from its file path and verify its integrity using the hash in the filename.
+
+    Args:
+        tensor_file_path (str | Path): Direct path to the tensor file. Supports both:
+                               - .bin.gz: gzip-compressed tensor (hash is of uncompressed data)
+                               - .bin: uncompressed tensor (for backward compatibility)
+        device (str, optional): Device to load the tensor to (e.g., 'cuda:0', 'cpu').
+                               If None, keeps the tensor on its original device.
+
+    Returns:
+        torch.Tensor: The loaded tensor (moved to the specified device if provided)
+
+    Raises:
+        FileNotFoundError: If the tensor file doesn't exist
+        RuntimeError: If the tensor cannot be loaded
+        ValueError: If the computed hash doesn't match the filename hash
+    """
+    blob_path = Path(tensor_file_path)
+
+    if not blob_path.exists():
+        raise FileNotFoundError(f"Tensor blob not found: {blob_path}")
+
+    # Detect compression by file extension
+    is_compressed = blob_path.name.endswith(".bin.gz")
+
+    # Read file contents (decompress if needed)
+    try:
+        with open(blob_path, "rb") as f:
+            file_obj = gzip.GzipFile(fileobj=f, mode="rb") if is_compressed else f
+            file_contents = file_obj.read()
+    except (OSError, gzip.BadGzipFile) as e:
+        if is_compressed:
+            raise RuntimeError(f"Failed to decompress gzip file {blob_path}: {str(e)}")
+        else:
+            raise RuntimeError(f"Failed to read file {blob_path}: {str(e)}")
+
+    # Extract expected hash from filename
+    # abc123.bin.gz -> abc123 or abc123.bin -> abc123
+    expected_hash = blob_path.name.removesuffix(".bin.gz" if is_compressed else ".bin")
+
+    # Compute hash of uncompressed data
+    computed_hash = hashlib.blake2b(file_contents).hexdigest()
+
+    # Verify hash matches filename
+    if computed_hash != expected_hash:
+        raise ValueError(
+            f"Hash verification failed: expected '{expected_hash}' but computed '{computed_hash}'"
+        )
+
+    try:
+        # Load the tensor from memory buffer
+        tensor = torch.load(io.BytesIO(file_contents), map_location=device)
+        return tensor
+    except Exception as e:
+        raise RuntimeError(f"Failed to load tensor from {blob_path}: {str(e)}")