tritonparse 0.2.4.dev20251012071503__py3-none-any.whl → 0.2.4.dev20251014071550__py3-none-any.whl

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/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/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/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/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.dist-info → tritonparse-0.2.4.dev20251014071550.dist-info}/METADATA

@@ -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.dist-info → tritonparse-0.2.4.dev20251014071550.dist-info}/RECORD

@@ -1,6 +1,6 @@
 tritonparse/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tritonparse/__main__.py,sha256=wu5N2wk8mvBgyvr2ghmQf4prezAe0_i-p123VVreyYc,62
-tritonparse/cli.py,sha256=ODfLvbMrrRm5JX6-3XDCu9f59gkJ00sg0R_uzNoVNgw,2471
+tritonparse/cli.py,sha256=w_sCIv1NlcGDzoflsYuTvtQis92_YuZEC51ZlTVsYCU,2517
 tritonparse/common.py,sha256=9coQbzpyHWAAdv6lx2YQTualiIH49ULJuZTA7VB_V7A,13946
 tritonparse/context_manager.py,sha256=MqAWI8uD1uCeCbmkcEWPejc5_IeY---BdajhsSR9T_E,2221
 tritonparse/event_diff.py,sha256=yOD6uNxLJroatfx2nEGr-erw24ObOrHU9P6V5pzr8do,4907
@@ -15,9 +15,10 @@ tritonparse/tp_logger.py,sha256=vXzY7hMDmVnRBGBhIjFZe3nHZzG5NKKPONGUszJhGgU,242
 tritonparse/trace_processor.py,sha256=brQBt26jdB6-quJXP5-warp2j31JSjOOFJa5ayiUZ5k,12963
 tritonparse/utils.py,sha256=Jnlptcd79llSDev-_1XyyOnv2izUqv0PEL74A8GF2tc,4565
 tritonparse/reproducer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tritonparse/reproducer/cli.py,sha256=bhtjD3k8pr7l2R2wmoleL-pGer2YndhUaLGnZq4rRBQ,948
-tritonparse/reproducer/orchestrator.py,sha256=VliooEC1Z7kXXGiNnbFkDT48tBQ2FFKH0qRs3jlbVmY,2530
-tritonparse/reproducer/placeholder_replacer.py,sha256=x9ddhIXVcYoEf6mBpByECPdZeA44rAXuo6a1cpy3F6c,4500
+tritonparse/reproducer/cli.py,sha256=xV3HZdNiWF44JCfJOGKHpoF8TErqbm4ohuLAK0pXcRw,1432
+tritonparse/reproducer/orchestrator.py,sha256=AtfR78HAsqDYMF-C00-jH6jVWs3aQ0bt74RBCakoqyc,2767
+tritonparse/reproducer/placeholder_replacer.py,sha256=0dtnfuaVqj_aX2fH4omMPFeIz6673Vc4yeQcZCLr5UA,6425
+tritonparse/reproducer/types.py,sha256=3cwVzWe4wRJXtR_KCta8hYtIwNQOxL5SJDocLnMQNgg,485
 tritonparse/reproducer/utils.py,sha256=UTclw48vH49g6Z2ljJL5DOZ6Rl4UDudyr0PeUySa3p8,13857
 tritonparse/reproducer/ingestion/ndjson.py,sha256=pEujTl5xXW2E2DEW8ngxXQ8qP9oawb90wBVTWHDs1jk,7372
 tritonparse/reproducer/templates/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -30,9 +31,9 @@ tritonparse/tools/format_fix.py,sha256=Ol0Sjui8D7OzHwbamAfGnq8V5Y63uwNaFTKSORN5H
 tritonparse/tools/load_tensor.py,sha256=94-TiSYlpXJx4MPmGK1ovmZlTt56Q_B3KQeCPaA6Cnw,2734
 tritonparse/tools/prettify_ndjson.py,sha256=r2YlHwFDTHgML7KljRmMsHaDg29q8gOQAgyDKWJhxRM,11062
 tritonparse/tools/readme.md,sha256=w6PWYfYnRgoPArLjxG9rVrpcLUkoVMGuRlbpF-o0IQM,110
-tritonparse-0.2.4.dev20251012071503.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
-tritonparse-0.2.4.dev20251012071503.dist-info/METADATA,sha256=bMh41VgVdUUuA-MtwkeY-UzrIXtIr9U6Gvpz0gAB7tk,8278
-tritonparse-0.2.4.dev20251012071503.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tritonparse-0.2.4.dev20251012071503.dist-info/entry_points.txt,sha256=wEXdaieDoRRCCdhEv2p_C68iytnaXU_2pwt5CqjfbWY,56
-tritonparse-0.2.4.dev20251012071503.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
-tritonparse-0.2.4.dev20251012071503.dist-info/RECORD,,
+tritonparse-0.2.4.dev20251014071550.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
+tritonparse-0.2.4.dev20251014071550.dist-info/METADATA,sha256=NRX3WAhjKHGcj7dw72sUFiiKzuAhiqgWDwRxaKvoXiM,8278
+tritonparse-0.2.4.dev20251014071550.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tritonparse-0.2.4.dev20251014071550.dist-info/entry_points.txt,sha256=wEXdaieDoRRCCdhEv2p_C68iytnaXU_2pwt5CqjfbWY,56
+tritonparse-0.2.4.dev20251014071550.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
+tritonparse-0.2.4.dev20251014071550.dist-info/RECORD,,