tritonparse 0.3.2.dev20251210071601__py3-none-any.whl

tritonparse/reproducer/cli.py

@@ -0,0 +1,72 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+import argparse
+
+from tritonparse.reproducer.types import KernelImportMode
+
+
+def _add_reproducer_args(parser: argparse.ArgumentParser) -> None:
+    """Add common arguments for the reproducer to a parser."""
+    parser.add_argument("input", help="Path to the ndjson/ndjson.gz log file")
+    parser.add_argument(
+        "--line",
+        type=int,
+        default=0,
+        help=(
+            "The line index (0-based) of the launch event in the input file to reproduce. "
+            "Defaults to 0 (first launch event). Mutually exclusive with --kernel/--launch-id."
+        ),
+    )
+    parser.add_argument(
+        "--kernel",
+        type=str,
+        default=None,
+        help=(
+            "Kernel name (exact match, case-sensitive) to reproduce. "
+            "Use with --launch-id to specify which launch of the kernel. "
+            "Mutually exclusive with --line."
+        ),
+    )
+    parser.add_argument(
+        "--launch-id",
+        type=int,
+        default=0,
+        help=(
+            "0-based launch index for the kernel specified by --kernel. "
+            "Defaults to 0 (first launch). Only used when --kernel is provided."
+        ),
+    )
+    parser.add_argument(
+        "--out-dir",
+        default="repro_output",
+        help=(
+            "Directory to save the reproducer script and context JSON. Defaults to "
+            "'repro_output/<kernel_name>/' if not provided."
+        ),
+    )
+    parser.add_argument(
+        "--template",
+        default="example",
+        help=(
+            "Template name (builtin, without .py) or a filesystem path to a .py file. "
+            "Defaults to 'example'."
+        ),
+    )
+    parser.add_argument(
+        "--kernel-import",
+        type=KernelImportMode,
+        choices=list(KernelImportMode),
+        default=KernelImportMode.DEFAULT,
+        help=(
+            "Kernel import strategy:\n"
+            "  default: Import kernel from original file (current behavior)\n"
+            "  copy: Embed kernel source code directly in reproducer\n"
+            "  override-ttir: Use TTIR from compilation event (bypass Python frontend)\n"
+            "Defaults to 'default'."
+        ),
+    )
+    parser.add_argument(
+        "--use-fbcode",
+        action="store_true",
+        help=("Use fbcode to setup repro environment."),
+    )

tritonparse/reproducer/consolidated_result.py

@@ -0,0 +1,52 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# pyre-strict
+
+"""
+ConsolidatedResult: Data structures for multi-file call graph analysis results.
+
+This module defines the output structure for the MultiFileCallGraphAnalyzer,
+containing all extracted functions, imports, edges, and statistics.
+"""
+
+from dataclasses import dataclass
+
+from tritonparse.reproducer.ast_analyzer import Edge
+from tritonparse.reproducer.import_info import ImportInfo
+
+
+@dataclass
+class AnalysisStats:
+    """Statistics about the multi-file analysis."""
+
+    total_files_analyzed: int
+    total_functions_found: int
+    total_imports: int
+    external_imports: int
+    internal_imports: int
+
+
+@dataclass
+class ConsolidatedResult:
+    """Consolidated analysis results across all files."""
+
+    # All dependent functions with their source code
+    functions: dict[str, str]  # qualified_name -> source_code
+
+    # Mapping of function to file
+    function_locations: dict[str, str]  # qualified_name -> file_path
+
+    # Short names for standalone file generation (just the function name)
+    function_short_names: dict[str, str]  # qualified_name -> short_name
+
+    # All required imports, deduplicated and organized
+    imports: list[ImportInfo]
+
+    # Call graph edges across all files
+    edges: list[Edge]
+
+    # Files analyzed
+    analyzed_files: set[str]
+
+    # Statistics
+    stats: AnalysisStats

tritonparse/reproducer/function_extractor.py

@@ -0,0 +1,228 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+"""
+Function extractor for reproducer utility functions.
+
+This module extracts utility functions from utils.py and load_tensor.py
+using AST parsing, and generates standalone code for reproducers.
+"""
+
+import ast
+from pathlib import Path
+
+
+def extract_utility_functions() -> str:
+    """
+    Extract all utility functions needed for the reproducer template.
+
+    Uses AST parsing to extract functions and constants from source files
+    without importing them (avoiding potential side effects).
+
+    Returns:
+        str: Complete Python code including imports and all utility functions.
+    """
+    # Prepare file paths
+    base_dir = Path(__file__).parent
+    utils_path = base_dir / "utils.py"
+    load_tensor_path = base_dir.parent / "tools" / "load_tensor.py"
+
+    # Parse source files
+    utils_tree, utils_lines = _parse_source_file(utils_path)
+    load_tensor_tree, load_tensor_lines = _parse_source_file(load_tensor_path)
+
+    # Define what to extract (in dependency order)
+    utils_function_names = [
+        "_get_triton_tensor_types",
+        "create_args_from_json_file",
+        "create_args_from_json",
+        "_apply_stride_and_offset",
+        "_create_base_tensor",
+        "_create_tensor",
+        "_create_arg_from_info",
+    ]
+
+    load_tensor_function_names = [
+        "load_tensor",
+    ]
+
+    # Extract content
+    extracted_parts = []
+
+    # Add required imports
+    extracted_parts.append(_generate_imports())
+
+    # Extract constant
+    constant = _extract_assignment(
+        utils_tree, utils_lines, "TRITON_KERNELS_CUSTOM_TYPES"
+    )
+    if constant:
+        extracted_parts.append(constant)
+
+    # Extract TRITON_DTYPE_MAP constant
+    dtype_map = _extract_assignment(utils_tree, utils_lines, "TRITON_DTYPE_MAP")
+    if dtype_map:
+        extracted_parts.append(dtype_map)
+
+    # Extract load_tensor functions
+    extracted_parts.extend(
+        _extract_functions(
+            load_tensor_tree, load_tensor_lines, load_tensor_function_names
+        )
+    )
+
+    # Extract utils functions
+    extracted_parts.extend(
+        _extract_functions(utils_tree, utils_lines, utils_function_names)
+    )
+
+    # Combine all parts
+    return "\n\n".join(extracted_parts)
+
+
+def _parse_source_file(file_path: Path) -> tuple[ast.Module, list[str]]:
+    """
+    Parse a Python source file and return its AST and source lines.
+
+    Args:
+        file_path: Path to the Python source file
+
+    Returns:
+        tuple: (AST tree, list of source code lines)
+
+    Raises:
+        FileNotFoundError: If the source file doesn't exist
+        SyntaxError: If the source file has syntax errors
+    """
+    try:
+        source_code = file_path.read_text(encoding="utf-8")
+        tree = ast.parse(source_code, filename=str(file_path))
+    except FileNotFoundError as e:
+        raise FileNotFoundError(f"Source file not found: {file_path}") from e
+    except SyntaxError as e:
+        raise SyntaxError(f"Failed to parse {file_path}: {e}") from e
+
+    lines = source_code.splitlines()
+    return tree, lines
+
+
+def _extract_assignment(
+    tree: ast.Module, lines: list[str], var_name: str
+) -> str | None:
+    """
+    Extract a module-level assignment statement by variable name.
+
+    Args:
+        tree: AST tree of the source file
+        lines: Source code lines
+        var_name: Name of the variable to extract
+
+    Returns:
+        Complete assignment statement source code, or None if not found
+
+    Example:
+        Extracts:
+        TRITON_KERNELS_CUSTOM_TYPES = (
+            importlib.util.find_spec("triton_kernels") is not None
+            and importlib.util.find_spec("triton_kernels.tensor") is not None
+        )
+    """
+    # Search only at module level
+    for node in tree.body:
+        if isinstance(node, ast.Assign):
+            for target in node.targets:
+                if isinstance(target, ast.Name) and target.id == var_name:
+                    # Found it! Extract source code using line numbers
+                    start_line = node.lineno - 1  # Convert to 0-based index
+                    end_line = node.end_lineno  # Already suitable for slicing
+                    assignment_lines = lines[start_line:end_line]
+                    return "\n".join(assignment_lines)
+    return None
+
+
+def _extract_function(tree: ast.Module, lines: list[str], func_name: str) -> str | None:
+    """
+    Extract a function definition by name, including decorators.
+
+    Args:
+        tree: AST tree of the source file
+        lines: Source code lines
+        func_name: Name of the function to extract
+
+    Returns:
+        Complete function source code including decorators, or None if not found
+
+    Example:
+        Extracts:
+        @lru_cache(maxsize=1)
+        def _get_triton_tensor_types():
+            '''Docstring'''
+            ...
+    """
+    # Walk the entire tree (handles nested functions if needed)
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef) and node.name == func_name:
+            # If function has decorators, start from the first decorator
+            if node.decorator_list:
+                start_line = node.decorator_list[0].lineno - 1
+            else:
+                start_line = node.lineno - 1
+
+            end_line = node.end_lineno
+            func_lines = lines[start_line:end_line]
+            return "\n".join(func_lines)
+    return None
+
+
+def _extract_functions(
+    tree: ast.Module, lines: list[str], func_names: list[str]
+) -> list[str]:
+    """
+    Extract multiple functions from a source file.
+
+    Args:
+        tree: AST tree of the source file
+        lines: Source code lines
+        func_names: List of function names to extract
+
+    Returns:
+        List of function source codes in the same order as func_names
+
+    Raises:
+        ValueError: If any function is not found
+    """
+    extracted = []
+    for func_name in func_names:
+        func_source = _extract_function(tree, lines, func_name)
+        if func_source is None:
+            raise ValueError(
+                f"Function '{func_name}' not found in source. "
+                f"Available functions might have been renamed or removed."
+            )
+        extracted.append(func_source)
+    return extracted
+
+
+def _generate_imports() -> str:
+    """
+    Generate the import statements needed for the extracted functions.
+
+    Returns:
+        str: Import statements as a single string
+    """
+    imports = [
+        "import gzip",
+        "import hashlib",
+        "import importlib",
+        "import importlib.util",
+        "import io",
+        "import json",
+        "import logging",
+        "import sys",
+        "from functools import lru_cache",
+        "from pathlib import Path",
+        "from typing import Union",
+        "",
+        "import torch",
+        "import triton.language as tl",
+    ]
+    return "\n".join(imports)

tritonparse/reproducer/import_info.py

@@ -0,0 +1,25 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# pyre-strict
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class ImportInfo:
+    """Information about an import statement."""
+
+    # Import statement details
+    import_type: str  # "import" or "from_import"
+    module: str  # e.g., "torch.nn.functional"
+    names: list[str]  # Imported names: ["func1", "func2"]
+
+    # Resolution metadata
+    source_file: str  # File containing this import
+    resolved_path: str | None  # Resolved file path (None if external)
+    is_external: bool  # True for third-party/built-in
+    lineno: int  # Line number in source file
+
+    # Fields with defaults (must come after required fields)
+    aliases: dict[str, str] = field(default_factory=dict)  # {local_name: original_name}
+    level: int = 0  # 0 = absolute, 1 = ".", 2 = "..", etc.

tritonparse/reproducer/import_parser.py

@@ -0,0 +1,178 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# pyre-strict
+
+import ast
+
+from tritonparse.reproducer.import_info import ImportInfo
+from tritonparse.reproducer.import_resolver import ImportResolver
+
+
+class ImportParser:
+    """
+    Parse import statements from Python AST.
+
+    Extracts all import statements from Python source code using AST,
+    resolves them to file paths, and returns structured ImportInfo objects.
+    """
+
+    def __init__(self, import_resolver: ImportResolver) -> None:
+        """
+        Initialize the import parser.
+
+        Args:
+            import_resolver: ImportResolver instance for resolving import paths
+        """
+        self.import_resolver = import_resolver
+
+    def parse_imports(
+        self, tree: ast.Module, source_file: str, package: str | None = None
+    ) -> list[ImportInfo]:
+        """
+        Extract all import statements from AST.
+
+        Handles:
+        - import X
+        - import X as Y
+        - from X import Y
+        - from X import Y as Z
+        - from . import X  (relative)
+        - from .. import X (relative)
+
+        Args:
+            tree: Parsed AST module
+            source_file: File path containing the imports
+            package: Package context for relative imports (e.g., "pytorch.tritonparse")
+
+        Returns:
+            List of ImportInfo objects with resolved paths
+        """
+        imports: list[ImportInfo] = []
+
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Import):
+                # Handle: import X, Y as Z
+                imports.extend(self._parse_import_node(node, source_file))
+            elif isinstance(node, ast.ImportFrom):
+                # Handle: from X import Y, Z as W
+                imports.extend(self._parse_import_from_node(node, source_file, package))
+
+        return imports
+
+    def _parse_import_node(
+        self, node: ast.Import, source_file: str
+    ) -> list[ImportInfo]:
+        """
+        Parse 'import X' statements.
+
+        Args:
+            node: ast.Import node
+            source_file: File containing this import
+
+        Returns:
+            List of ImportInfo objects for each imported module
+        """
+        imports: list[ImportInfo] = []
+
+        for alias in node.names:
+            # Resolve the import to a file path
+            resolved_path, is_external = self.import_resolver.resolve_import(alias.name)
+
+            # Build aliases dict if alias is used
+            aliases: dict[str, str] = {}
+            if alias.asname:
+                aliases[alias.asname] = alias.name
+
+            imports.append(
+                ImportInfo(
+                    import_type="import",
+                    module=alias.name,
+                    names=[alias.name.split(".")[-1]],  # Last component as name
+                    aliases=aliases,
+                    source_file=source_file,
+                    resolved_path=resolved_path,
+                    is_external=is_external,
+                    lineno=node.lineno,
+                    level=0,  # Absolute import
+                )
+            )
+
+        return imports
+
+    def _parse_import_from_node(
+        self,
+        node: ast.ImportFrom,
+        source_file: str,
+        package: str | None = None,
+    ) -> list[ImportInfo]:
+        """
+        Parse 'from X import Y' statements.
+
+        Args:
+            node: ast.ImportFrom node
+            source_file: File containing this import
+            package: Package context for relative imports
+
+        Returns:
+            List of ImportInfo objects for each imported name
+        """
+        imports: list[ImportInfo] = []
+
+        # Get module name (may be None for relative imports like "from . import X")
+        module = node.module or ""
+        level = node.level  # 0 = absolute, 1 = ".", 2 = "..", etc.
+
+        # For relative imports, construct the full module name
+        if level > 0 and package:
+            # Build relative module name
+            # level=1: current package (from . import X)
+            # level=2: parent package (from .. import X)
+            # Formula: Remove (level-1) components from package, then append module
+            # Examples:
+            # - package="a.b.c", level=1, module=None -> "a.b.c" (current)
+            # - package="a.b.c", level=1, module="d" -> "a.b.c.d"
+            # - package="a.b.c", level=2, module=None -> "a.b" (parent)
+            # - package="a.b.c", level=2, module="d" -> "a.b.d"
+            package_parts = package.split(".")
+
+            # Remove (level-1) components
+            if level == 1:
+                parent_package = package
+            elif level <= len(package_parts):
+                parent_package = ".".join(package_parts[: -(level - 1)])
+            else:
+                parent_package = ""
+
+            # Append module if specified
+            if module:
+                full_module = f"{parent_package}.{module}" if parent_package else module
+            else:
+                full_module = parent_package
+        else:
+            full_module = module
+
+        # Resolve the module to a file path
+        resolved_path, is_external = self.import_resolver.resolve_import(full_module)
+
+        # Parse each imported name
+        for alias in node.names:
+            # Build aliases dict if alias is used
+            aliases: dict[str, str] = {}
+            if alias.asname:
+                aliases[alias.asname] = alias.name
+
+            imports.append(
+                ImportInfo(
+                    import_type="from_import",
+                    module=full_module,
+                    names=[alias.name],
+                    aliases=aliases,
+                    source_file=source_file,
+                    resolved_path=resolved_path,
+                    is_external=is_external,
+                    lineno=node.lineno,
+                    level=level,
+                )
+            )
+
+        return imports