Add function extractor module for reproducer utility functions

This commit introduces a new module, `function_extractor.py`, which extracts utility functions from `utils.py` and `load_tensor.py` using the `inspect` module. The extracted functions are intended for use in the reproducer template, allowing for the generation of standalone code. Additionally, the `placeholder_replacer.py` has been updated to include a new placeholder for utility functions, which will be replaced with the extracted code during the reproduction process.

Changes include:
- New `extract_utility_functions` function to gather necessary utility functions.
- Integration of utility function extraction into the placeholder replacement logic.
- Minor updates to `utils.py` for function renaming and additional utility functions.

This enhancement aims to streamline the reproducer generation process by ensuring all required utility functions are readily available in the generated code.

Author: FindHao

tritonparse/reproducer/function_extractor.py

@@ -0,0 +1,123 @@
+"""
+Function extractor for reproducer utility functions.
+
+This module extracts utility functions from utils.py and load_tensor.py
+using the inspect module, and generates standalone code for reproducers.
+"""
+
+import importlib.util
+import inspect
+from pathlib import Path
+
+
+def extract_utility_functions() -> str:
+    """
+    Extract all utility functions needed for the reproducer template.
+
+    Returns:
+        str: Complete Python code including imports and all utility functions.
+    """
+    # Import the modules
+    utils_module = _import_module_from_path(
+        "tritonparse.reproducer.utils", Path(__file__).parent / "utils.py"
+    )
+    load_tensor_module = _import_module_from_path(
+        "tritonparse.tools.load_tensor",
+        Path(__file__).parent.parent / "tools" / "load_tensor.py",
+    )
+
+    # Functions to extract from utils.py
+    utils_functions = [
+        "_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",
+    ]
+
+    # Functions to extract from load_tensor.py
+    load_tensor_functions = [
+        "load_tensor",
+    ]
+
+    # Extract all function source code
+    extracted_code = []
+
+    # Add required imports
+    imports = _generate_imports()
+    extracted_code.append(imports)
+
+    # Add TRITON_KERNELS_CUSTOM_TYPES constant
+    constant_code = inspect.getsource(utils_module).split("\n")
+    for i, line in enumerate(constant_code):
+        if line.startswith("TRITON_KERNELS_CUSTOM_TYPES"):
+            # Find the end of this statement
+            j = i
+            while j < len(constant_code) and not constant_code[j].rstrip().endswith(
+                ")"
+            ):
+                j += 1
+            constant_lines = constant_code[i : j + 1]
+            extracted_code.append("\n".join(constant_lines))
+            break
+
+    extracted_code.append("")
+
+    # Extract load_tensor function
+    for func_name in load_tensor_functions:
+        func = getattr(load_tensor_module, func_name)
+        source = inspect.getsource(func)
+        extracted_code.append(source)
+
+    # Extract utils functions
+    for func_name in utils_functions:
+        func = getattr(utils_module, func_name)
+        source = inspect.getsource(func)
+        extracted_code.append(source)
+
+    return "\n\n".join(extracted_code)
+
+
+def _import_module_from_path(module_name: str, file_path: Path):
+    """
+    Import a module from a file path.
+
+    Args:
+        module_name: Name for the module
+        file_path: Path to the Python file
+
+    Returns:
+        The imported module
+    """
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+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/reproducer/placeholder_replacer.py

@@ -2,6 +2,7 @@
 
 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 @@ def __init__(self):
         )
         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 @@ def _replace_kernel_import(
         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: