tritonparse 0.3.1.dev20251019071438__tar.gz → 0.3.1.dev20251021071528__tar.gz

{tritonparse-0.3.1.dev20251019071438/tritonparse.egg-info → tritonparse-0.3.1.dev20251021071528}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tritonparse
-Version: 0.3.1.dev20251019071438
+Version: 0.3.1.dev20251021071528
 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.dev20251021071528/tritonparse/reproducer/function_extractor.py

@@ -0,0 +1,220 @@
+"""
+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 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",
+    ]
+    return "\n".join(imports)

{tritonparse-0.3.1.dev20251019071438 → tritonparse-0.3.1.dev20251021071528}/tritonparse/reproducer/placeholder_replacer.py

@@ -2,6 +2,7 @@ from abc import ABC
 
 from typing import Any, Dict, Protocol
 
+from tritonparse.reproducer.function_extractor import extract_utility_functions
 from tritonparse.reproducer.ingestion.ndjson import ContextBundle
 from tritonparse.reproducer.types import KernelImportMode
 from tritonparse.reproducer.utils import (
@@ -82,6 +83,9 @@ class DefaultPlaceholderReplacer(PlaceholderReplacer):
         )
         self.register("# {{KERNEL_SYSPATH_PLACEHOLDER}}", self._replace_kernel_syspath)
         self.register("# {{KERNEL_IMPORT_PLACEHOLDER}}", self._replace_kernel_import)
+        self.register(
+            "# {{UTILITY_FUNCTIONS_PLACEHOLDER}}", self._replace_utility_functions
+        )
         self.register(
             "# {{KERNEL_INVOCATION_PLACEHOLDER}}", self._replace_kernel_invocation
         )
@@ -217,6 +221,13 @@ triton.autotune = _patched_autotune
         else:
             raise ValueError(f"Unknown kernel_import mode: {kernel_import}")
 
+    def _replace_utility_functions(
+        self, code: str, context_bundle: ContextBundle, **kwargs
+    ) -> str:
+        """Replace the utility functions placeholder with extracted functions."""
+        utility_code = extract_utility_functions()
+        return code.replace("# {{UTILITY_FUNCTIONS_PLACEHOLDER}}", utility_code)
+
     def _replace_kernel_invocation(
         self, code: str, context_bundle: ContextBundle, **kwargs
     ) -> str:

tritonparse-0.3.1.dev20251021071528/tritonparse/reproducer/templates/example.py

@@ -0,0 +1,30 @@
+"""
+This file is automatically generated by TritonParse reproducer.
+It contains a smallest testing example for a Triton kernel.
+"""
+
+import torch
+
+# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}
+
+# {{KERNEL_SYSPATH_PLACEHOLDER}}
+
+# {{KERNEL_IMPORT_PLACEHOLDER}}
+
+# {{UTILITY_FUNCTIONS_PLACEHOLDER}}
+
+
+if __name__ == "__main__":
+    script_dir = Path(__file__).resolve().parent  # noqa: F821
+    json_file = script_dir / "{{JSON_FILE_NAME_PLACEHOLDER}}"
+    grid, args_dict = create_args_from_json_file(str(json_file))  # noqa: F821
+
+    print("Generated kernel arguments dictionary:")
+    for name, arg in args_dict.items():
+        print(f"  {name}: {arg}")
+    print(f"Grid: {grid}")
+
+    # {{KERNEL_INVOCATION_PLACEHOLDER}}
+
+    torch.cuda.synchronize()
+    print("Kernel execution finished.")

tritonparse-0.3.1.dev20251019071438/tritonparse/reproducer/templates/example.py → tritonparse-0.3.1.dev20251021071528/tritonparse/reproducer/utils.py

@@ -1,27 +1,16 @@
-"""
-This file is automatically generated by TritonParse reproducer.
-It contains a smallest testing example for a Triton kernel.
-"""
-
-import gzip
-import hashlib
 import importlib
 import importlib.util
-import io
 import json
 import logging
 import sys
+from datetime import datetime
 from functools import lru_cache
 from pathlib import Path
-from typing import Union
 
 import torch
 
-# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}
-
-# {{KERNEL_SYSPATH_PLACEHOLDER}}
-
-# {{KERNEL_IMPORT_PLACEHOLDER}}
+from tritonparse.tools.load_tensor import load_tensor
+from tritonparse.tp_logger import logger
 
 TRITON_KERNELS_CUSTOM_TYPES = (
     importlib.util.find_spec("triton_kernels") is not None
@@ -31,15 +20,6 @@ TRITON_KERNELS_CUSTOM_TYPES = (
 
 @lru_cache(maxsize=1)
 def _get_triton_tensor_types():
-    """
-    Import and cache Triton custom tensor types.
-
-    Returns:
-        tuple: (Tensor, Storage, StridedLayout) classes from triton_kernels.tensor.
-
-    Raises:
-        ImportError: If the optional module 'triton_kernels.tensor' is not available.
-    """
     mod = importlib.import_module("triton_kernels.tensor")
     return (
         mod.Tensor,
@@ -48,66 +28,16 @@ def _get_triton_tensor_types():
     )
 
 
-def load_tensor(tensor_file_path: Union[str, Path], device: str = None) -> torch.Tensor:
+def create_args_from_json_file(json_path):
     """
-    Load a tensor from its file path and verify its integrity using the hash in the filename.
+    Load and parse a reproducer JSON file.
 
     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.
+        json_path (str): Path to the JSON file describing the kernel launch.
 
     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
+        tuple[list, dict]: Grid specification list and map of argument name to value.
     """
-    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)}")
-
-
-def create_args_from_json_file(json_path):
     with open(json_path, "r") as f:
         data = json.load(f)
     return create_args_from_json(data)
@@ -118,7 +48,7 @@ def create_args_from_json(data):
     Parse a reproducer JSON and build kernel grid and argument dictionary.
 
     Args:
-        json_path (str): Path to the JSON file describing the kernel launch.
+        data (dict | list): JSON data describing the kernel launch.
 
     Returns:
         tuple[list, dict]: Grid specification list and map of argument name to value.
@@ -197,6 +127,15 @@ def _apply_stride_and_offset(tensor, shape, stride, storage_offset):
 
 
 def _create_base_tensor(arg_info) -> torch.Tensor:
+    """
+    Create a base tensor without stride/offset modifications.
+
+    Args:
+        arg_info (dict): Argument information including dtype, shape, device, etc.
+
+    Returns:
+        torch.Tensor: The created base tensor
+    """
     if arg_info.get("blob_path"):
         return load_tensor(arg_info.get("blob_path"), arg_info.get("device"))
 
@@ -210,6 +149,9 @@ def _create_base_tensor(arg_info) -> torch.Tensor:
 
     shape = arg_info.get("shape", [])
     device = arg_info.get("device", "cpu")
+    # Normalize cuda device to cuda:0
+    if isinstance(device, str) and device.startswith("cuda"):
+        device = "cuda:0"
 
     # Extract statistical information if available
     mean = arg_info.get("mean")
@@ -301,6 +243,15 @@ def _create_base_tensor(arg_info) -> torch.Tensor:
 
 
 def _create_tensor(arg_info) -> torch.Tensor:
+    """
+    Create a tensor with stride and storage offset if needed.
+
+    Args:
+        arg_info (dict): Argument information including dtype, shape, stride, etc.
+
+    Returns:
+        torch.Tensor: The created tensor with applied stride/offset
+    """
     tensor = _create_base_tensor(arg_info)
 
     # Apply stride and storage offset if needed
@@ -374,17 +325,135 @@ def _create_arg_from_info(arg_info):
         return None
 
 
-if __name__ == "__main__":
-    script_dir = Path(__file__).resolve().parent
-    json_file = script_dir / "{{JSON_FILE_NAME_PLACEHOLDER}}"
-    grid, args_dict = create_args_from_json_file(str(json_file))
+def determine_output_paths(out_dir: str, kernel_name: str):
+    """
+    Determine output file paths for reproducer script and context data.
+
+    Args:
+        out_dir: Output directory path. If empty, uses default location.
+        kernel_name: Name of the kernel for default directory naming.
+
+    Returns:
+        Tuple of (python_script_path, json_context_path) as Path objects.
+    """
+    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+    output_directory = Path(out_dir) / kernel_name
+    output_directory.mkdir(parents=True, exist_ok=True)
 
-    print("Generated kernel arguments dictionary:")
-    for name, arg in args_dict.items():
-        print(f"  {name}: {arg}")
-    print(f"Grid: {grid}")
+    out_py_path = output_directory / f"repro_{timestamp}.py"
+    temp_json_path = output_directory / f"repro_context_{timestamp}.json"
 
-    # {{KERNEL_INVOCATION_PLACEHOLDER}}
+    return out_py_path, temp_json_path
 
-    torch.cuda.synchronize()
-    print("Kernel execution finished.")
+
+def _generate_import_statements(kernel_info) -> tuple[str, str]:
+    """
+    Generate (sys.path insertion statement, import statement) for the kernel.
+
+    Strategy:
+    - Always add the kernel file's parent directory to sys.path.
+    - If the filename (without .py) is a valid identifier, import using that
+      module name: `from <stem> import <func> as imported_kernel_function`.
+    - Otherwise, fall back to dynamic import via importlib.util and bind
+      `imported_kernel_function` from the loaded module.
+    """
+    file_path = Path(kernel_info.file_path)
+    function_name = kernel_info.function_name
+
+    if not file_path or not function_name:
+        raise ValueError("Kernel file path or function name missing from context.")
+
+    # Always add the file's parent directory to sys.path
+    sys_stmt = (
+        "import sys; p = r'" + str(file_path.parent) + "';\n"
+        "if p not in sys.path: sys.path.insert(0, p)"
+    )
+
+    module_name = file_path.with_suffix("").name
+    if module_name.isidentifier():
+        import_stmt = (
+            f"from {module_name} import {function_name} as imported_kernel_function"
+        )
+        logger.debug("Generated direct import statement: %s", import_stmt)
+        return sys_stmt, import_stmt
+
+    # Fallback: dynamic import when filename is not a valid identifier
+    import_stmt = (
+        "import importlib.util\n"
+        f"_spec = importlib.util.spec_from_file_location('kernel_mod', r'{str(file_path)}')\n"
+        "_mod = importlib.util.module_from_spec(_spec)\n"
+        "_spec.loader.exec_module(_mod)\n"
+        f"imported_kernel_function = getattr(_mod, '{function_name}')"
+    )
+    logger.debug("Generated dynamic import for file: %s", file_path)
+    return sys_stmt, import_stmt
+
+
+def _parse_kernel_signature(kernel_source_code: str) -> tuple[list[str], list[str]]:
+    """
+    Parses a Triton kernel's source code to distinguish positional args
+    from keyword args (those with default values).
+    """
+    signature_lines = []
+    in_signature = False
+    for line in kernel_source_code.splitlines():
+        # Mark beginning of signature when function definition is found
+        if line.strip().startswith("def "):
+            in_signature = True
+        if in_signature:
+            # Strip comments and leading/trailing whitespace
+            clean_line = line.split("#")[0].strip()
+            signature_lines.append(clean_line)
+            # Stop capturing after the signature ends
+            if "):" in line:
+                break
+
+    full_signature = "".join(signature_lines)
+    # Extract content between the first '(' and the last '):'
+    try:
+        params_str = full_signature[
+            full_signature.find("(") + 1 : full_signature.rfind("):")
+        ]
+    except IndexError as exc:
+        raise ValueError("Could not parse kernel signature.") from exc
+
+    # Clean up and split the parameters string
+    params = [p.strip() for p in params_str.replace("\n", "").split(",") if p.strip()]
+
+    positional_args = []
+    keyword_args = []
+
+    for param in params:
+        if "=" in param:
+            # Keyword arguments have a default value
+            arg_name = param.split("=")[0].strip()
+            keyword_args.append(arg_name)
+        else:
+            # Positional arguments do not have a default value
+            arg_name = param.split(":")[0].strip()
+            positional_args.append(arg_name)
+
+    logger.debug("Parsed positional args: %s", positional_args)
+    logger.debug("Parsed keyword args: %s", keyword_args)
+    return positional_args, keyword_args
+
+
+def _generate_invocation_snippet(
+    positional_args: list[str], keyword_args: list[str]
+) -> str:
+    """Generates a single-line Python code snippet for kernel invocation."""
+    # Prepare positional args for direct injection into the call
+    pos_args_str = ", ".join([f'args_dict["{arg}"]' for arg in positional_args])
+
+    # Prepare keyword args for direct injection
+    kw_args_str = ", ".join([f'{arg}=args_dict["{arg}"]' for arg in keyword_args])
+
+    # Combine them, ensuring proper comma separation
+    all_args = []
+    if pos_args_str:
+        all_args.append(pos_args_str)
+    if kw_args_str:
+        all_args.append(kw_args_str)
+
+    # Create the single-line call
+    return f"imported_kernel_function[tuple(grid)]({', '.join(all_args)})"

{tritonparse-0.3.1.dev20251019071438 → tritonparse-0.3.1.dev20251021071528/tritonparse.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tritonparse
-Version: 0.3.1.dev20251019071438
+Version: 0.3.1.dev20251021071528
 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.dev20251019071438 → tritonparse-0.3.1.dev20251021071528}/tritonparse.egg-info/SOURCES.txt

@@ -52,6 +52,7 @@ tritonparse.egg-info/requires.txt
 tritonparse.egg-info/top_level.txt
 tritonparse/reproducer/__init__.py
 tritonparse/reproducer/cli.py
+tritonparse/reproducer/function_extractor.py
 tritonparse/reproducer/orchestrator.py
 tritonparse/reproducer/placeholder_replacer.py
 tritonparse/reproducer/types.py