tritonparse 0.2.4.dev20251012071503__tar.gz → 0.2.4.dev20251014071550__tar.gz

{tritonparse-0.2.4.dev20251012071503 → tritonparse-0.2.4.dev20251014071550}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tritonparse
-Version: 0.2.4.dev20251012071503
+Version: 0.2.4.dev20251014071550
 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.2.4.dev20251012071503 → tritonparse-0.2.4.dev20251014071550}/tritonparse/cli.py

@@ -73,6 +73,7 @@ def main():
             line_index=args.line,
             out_dir=args.out_dir,
             template=args.template,
+            kernel_import=args.kernel_import,
         )
     else:
         raise RuntimeError(f"Unknown command: {args.func}")

{tritonparse-0.2.4.dev20251012071503 → tritonparse-0.2.4.dev20251014071550}/tritonparse/reproducer/cli.py

@@ -1,5 +1,7 @@
 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."""
@@ -29,3 +31,15 @@ def _add_reproducer_args(parser: argparse.ArgumentParser) -> None:
             "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"
+            "Defaults to 'default'."
+        ),
+    )

{tritonparse-0.2.4.dev20251012071503 → tritonparse-0.2.4.dev20251014071550}/tritonparse/reproducer/orchestrator.py

@@ -7,6 +7,7 @@ from tritonparse.reproducer.placeholder_replacer import (
     PlaceholderReplacer,
 )
 from tritonparse.reproducer.templates.loader import load_template_code
+from tritonparse.reproducer.types import KernelImportMode
 from tritonparse.reproducer.utils import determine_output_paths
 
 from tritonparse.tools.prettify_ndjson import load_ndjson, save_prettified_json
@@ -19,6 +20,7 @@ def reproduce(
     out_dir: str,
     template: str,
     replacer: Optional[PlaceholderReplacer] = None,
+    kernel_import: KernelImportMode = KernelImportMode.DEFAULT,
 ) -> dict[str, Path]:
     """
     Generate a reproducer script from NDJSON trace file.
@@ -29,6 +31,7 @@ def reproduce(
         out_dir: Output directory for reproducer files.
         template: Template name to use for the reproducer.
         replacer: Optional custom PlaceholderReplacer instance. If None, uses DefaultPlaceholderReplacer.
+        kernel_import: Kernel import mode (DEFAULT or COPY).
     """
     logger.debug(f"Building bundle from {input_path} at line {line_index}")
     events = load_ndjson(Path(input_path))
@@ -51,7 +54,10 @@ def reproduce(
     if replacer is None:
         replacer = DefaultPlaceholderReplacer()
     final_code = replacer.replace(
-        template_code, context_bundle, temp_json_path=temp_json_path
+        template_code,
+        context_bundle,
+        temp_json_path=temp_json_path,
+        kernel_import=kernel_import,
     )
 
     out_py_path.write_text(final_code, encoding="utf-8")

{tritonparse-0.2.4.dev20251012071503 → tritonparse-0.2.4.dev20251014071550}/tritonparse/reproducer/placeholder_replacer.py

@@ -3,6 +3,7 @@ from abc import ABC
 from typing import Any, Dict, Protocol
 
 from tritonparse.reproducer.ingestion.ndjson import ContextBundle
+from tritonparse.reproducer.types import KernelImportMode
 from tritonparse.reproducer.utils import (
     _generate_import_statements,
     _generate_invocation_snippet,
@@ -95,15 +96,58 @@ class DefaultPlaceholderReplacer(PlaceholderReplacer):
         self, code: str, context_bundle: ContextBundle, **kwargs
     ) -> str:
         """Replace the kernel sys.path placeholder."""
-        sys_stmt, _ = _generate_import_statements(context_bundle.kernel_info)
-        return code.replace("# {{KERNEL_SYSPATH_PLACEHOLDER}}", sys_stmt)
+        kernel_import = kwargs.get("kernel_import", KernelImportMode.DEFAULT)
+
+        if kernel_import == KernelImportMode.DEFAULT:
+            sys_stmt, _ = _generate_import_statements(context_bundle.kernel_info)
+            return code.replace("# {{KERNEL_SYSPATH_PLACEHOLDER}}", sys_stmt)
+        elif kernel_import == KernelImportMode.COPY:
+            comment = (
+                "# Kernel sys.path setup skipped - kernel source code embedded below"
+            )
+            return code.replace("# {{KERNEL_SYSPATH_PLACEHOLDER}}", comment)
+        else:
+            raise ValueError(f"Unknown kernel_import mode: {kernel_import}")
 
     def _replace_kernel_import(
         self, code: str, context_bundle: ContextBundle, **kwargs
     ) -> str:
         """Replace the kernel import placeholder."""
-        _, import_statement = _generate_import_statements(context_bundle.kernel_info)
-        return code.replace("# {{KERNEL_IMPORT_PLACEHOLDER}}", import_statement)
+        kernel_import = kwargs.get("kernel_import", KernelImportMode.DEFAULT)
+
+        if kernel_import == KernelImportMode.DEFAULT:
+            _, import_statement = _generate_import_statements(
+                context_bundle.kernel_info
+            )
+            return code.replace("# {{KERNEL_IMPORT_PLACEHOLDER}}", import_statement)
+        elif kernel_import == KernelImportMode.COPY:
+            source_code = context_bundle.kernel_info.source_code
+            func_name = context_bundle.kernel_info.function_name
+
+            if not source_code or not source_code.strip():
+                raise ValueError("Kernel source code is empty, cannot use 'copy' mode")
+            if not func_name:
+                raise ValueError(
+                    "Cannot determine kernel function name for 'copy' mode"
+                )
+
+            # Add common imports needed for most Triton kernels
+            import_lines = [
+                "import torch",
+                "import numpy as np",
+                "import triton",
+                "import triton.language as tl",
+                "",
+            ]
+
+            # Combine: imports + kernel source code + alias
+            embedded_code = "\n".join(import_lines)
+            embedded_code += "\n" + source_code
+            embedded_code += f"\n\n# Use kernel function directly\nimported_kernel_function = {func_name}"
+
+            return code.replace("# {{KERNEL_IMPORT_PLACEHOLDER}}", embedded_code)
+        else:
+            raise ValueError(f"Unknown kernel_import mode: {kernel_import}")
 
     def _replace_kernel_invocation(
         self, code: str, context_bundle: ContextBundle, **kwargs

tritonparse-0.2.4.dev20251014071550/tritonparse/reproducer/types.py

@@ -0,0 +1,15 @@
+from enum import Enum
+
+
+class KernelImportMode(str, Enum):
+    """
+    Kernel import strategy for reproducer generation.
+
+    Inherits from str to allow direct string comparison and use in argparse.
+    """
+
+    DEFAULT = "default"  # Import kernel from original file (current behavior)
+    COPY = "copy"  # Embed kernel source code directly in reproducer
+
+    # Future modes can be added here:
+    # OVERRIDE_TTIR = "override-ttir"  # Use TTIR from compilation event with monkeypatch

{tritonparse-0.2.4.dev20251012071503 → tritonparse-0.2.4.dev20251014071550}/tritonparse.egg-info/PKG-INFO

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

@@ -57,6 +57,7 @@ tritonparse/reproducer/__init__.py
 tritonparse/reproducer/cli.py
 tritonparse/reproducer/orchestrator.py
 tritonparse/reproducer/placeholder_replacer.py
+tritonparse/reproducer/types.py
 tritonparse/reproducer/utils.py
 tritonparse/reproducer/ingestion/ndjson.py
 tritonparse/reproducer/templates/__init__.py