tritonparse 0.2.4.dev20251013071533__py3-none-any.whl → 0.3.0__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,16 @@ 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"
+            "  override-ttir: Use TTIR from compilation event (bypass Python frontend)\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))
@@ -43,6 +46,15 @@ def reproduce(
         out_dir, context_bundle.kernel_info.function_name
     )
     save_prettified_json(context_bundle.raw_launch_event, temp_json_path)
+
+    # Save compilation event JSON if using OVERRIDE_TTIR mode
+    comp_json_path = None
+    if kernel_import == KernelImportMode.OVERRIDE_TTIR:
+        comp_json_path = (
+            temp_json_path.parent / f"{temp_json_path.stem}_compilation.json"
+        )
+        save_prettified_json(context_bundle.raw_comp_event, comp_json_path)
+
     logger.debug("Loading reproducer template.")
     template_code = load_template_code(template)
 
@@ -51,7 +63,11 @@ 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,
+        comp_json_filename=comp_json_path.name if comp_json_path else None,
     )
 
     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,
@@ -76,6 +77,9 @@ class DefaultPlaceholderReplacer(PlaceholderReplacer):
         super().__init__()
         # Register all default handlers
         self.register("{{JSON_FILE_NAME_PLACEHOLDER}}", self._replace_json_filename)
+        self.register(
+            "# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}", self._replace_ir_override_setup
+        )
         self.register("# {{KERNEL_SYSPATH_PLACEHOLDER}}", self._replace_kernel_syspath)
         self.register("# {{KERNEL_IMPORT_PLACEHOLDER}}", self._replace_kernel_import)
         self.register(
@@ -91,19 +95,127 @@ class DefaultPlaceholderReplacer(PlaceholderReplacer):
             raise ValueError("temp_json_path is required for JSON filename replacement")
         return code.replace("{{JSON_FILE_NAME_PLACEHOLDER}}", temp_json_path.name)
 
+    def _replace_ir_override_setup(
+        self, code: str, context_bundle: ContextBundle, **kwargs
+    ) -> str:
+        """Replace the IR override setup placeholder."""
+        kernel_import = kwargs.get("kernel_import", KernelImportMode.DEFAULT)
+
+        if kernel_import != KernelImportMode.OVERRIDE_TTIR:
+            return code.replace("# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}", "")
+
+        comp_json_filename = kwargs.get("comp_json_filename")
+        if not comp_json_filename:
+            raise ValueError("comp_json_filename is required for OVERRIDE_TTIR mode")
+
+        setup_code = f'''
+def create_ttir_tempfile():
+    """Extract TTIR from compilation event and create temporary file."""
+    script_dir = Path(__file__).resolve().parent
+    comp_json_file = script_dir / "{comp_json_filename}"
+    
+    with open(comp_json_file, 'r') as f:
+        comp_data = json.load(f)
+    
+    # Extract TTIR content
+    kernel_name = comp_data['payload']['metadata']['name']
+    ttir_key = f"{{kernel_name}}.ttir"
+    ttir_content = comp_data['payload']['file_content'][ttir_key]
+    
+    # Create temporary file
+    temp_file = tempfile.NamedTemporaryFile(
+        mode='w', 
+        suffix='.ttir', 
+        delete=False,
+        prefix=f'{{kernel_name}}_'
+    )
+    temp_file.write(ttir_content)
+    temp_file.close()
+    return temp_file.name
+
+
+# Monkeypatch triton.autotune to use our TTIR
+_ttir_file = create_ttir_tempfile()
+_original_autotune = None
+
+def _patched_autotune(configs, key=None, **kwargs):
+    """Patched autotune that uses our TTIR file."""
+    import triton
+    # Replace configs with our single config using ir_override
+    new_configs = [triton.Config(kwargs={{}}, ir_override=_ttir_file)]
+    # Call original autotune with our config
+    return _original_autotune(new_configs, key=[], **kwargs)
+
+# Apply the monkeypatch before importing the kernel
+import triton
+_original_autotune = triton.autotune
+triton.autotune = _patched_autotune
+'''
+
+        return code.replace("# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}", setup_code)
+
     def _replace_kernel_syspath(
         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)
+        elif kernel_import == KernelImportMode.OVERRIDE_TTIR:
+            comment = "# Kernel sys.path setup skipped - using IR override mode"
+            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)
+        elif kernel_import == KernelImportMode.OVERRIDE_TTIR:
+            comment = "# Kernel import skipped - using IR override mode with TTIR"
+            return code.replace("# {{KERNEL_IMPORT_PLACEHOLDER}}", comment)
+        else:
+            raise ValueError(f"Unknown kernel_import mode: {kernel_import}")
 
     def _replace_kernel_invocation(
         self, code: str, context_bundle: ContextBundle, **kwargs

tritonparse/reproducer/templates/example.py

@@ -6,6 +6,7 @@ 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
@@ -16,6 +17,8 @@ from typing import Union
 
 import torch
 
+# {{IR_OVERRIDE_SETUP_PLACEHOLDER}}
+
 # {{KERNEL_SYSPATH_PLACEHOLDER}}
 
 # {{KERNEL_IMPORT_PLACEHOLDER}}

tritonparse/reproducer/types.py

@@ -0,0 +1,18 @@
+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.
+
+    Attributes:
+        DEFAULT: Import kernel from original file (current behavior).
+        COPY: Embed kernel source code directly in reproducer.
+        OVERRIDE_TTIR: Use TTIR from compilation event with monkeypatch.
+    """
+
+    DEFAULT = "default"
+    COPY = "copy"
+    OVERRIDE_TTIR = "override-ttir"

{tritonparse-0.2.4.dev20251013071533.dist-info → tritonparse-0.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tritonparse
-Version: 0.2.4.dev20251013071533
+Version: 0.3.0
 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.dev20251013071533.dist-info → tritonparse-0.3.0.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,13 +15,14 @@ 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=MqYuuAP-uAIWWLQxixwyyBHJGSaQdG3xXGaVjERTLX8,1522
+tritonparse/reproducer/orchestrator.py,sha256=rZyA7TNVNgVuQkLlte3sPi1tQZ0ZAXGiFd5IJRyaD9Q,3180
+tritonparse/reproducer/placeholder_replacer.py,sha256=YPspknFkoZ1WLHQBrJefSlp4RerbdX9LpBe-QSCmMRs,9026
+tritonparse/reproducer/types.py,sha256=AfVl83zoJZQ58JJoplCcMC51gK-M-OKcafatYEIGgW0,509
 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
-tritonparse/reproducer/templates/example.py,sha256=gOYKKKptvzyI4mq4eUygi1NrIEQkChm4d1bJXed1w5U,13904
+tritonparse/reproducer/templates/example.py,sha256=QAyykFvmd1nI4ZvCjkJfrpAs2XBsMnFOK6WzhXxQ1uI,13963
 tritonparse/reproducer/templates/loader.py,sha256=HqjfThdDVg7q2bYWry78sIaVRkUpkcA8KQDt83YrlVE,1920
 tritonparse/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tritonparse/tools/decompress_bin_ndjson.py,sha256=kpt7DM_sSA334F1X45xdkP2OR9LuB27Pc50EkGr6CPM,4144
@@ -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.dev20251013071533.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
-tritonparse-0.2.4.dev20251013071533.dist-info/METADATA,sha256=C1nbmBxgnz-P7TCLT9d0JB1DVPu7TVp_1g08vB6mS4c,8278
-tritonparse-0.2.4.dev20251013071533.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tritonparse-0.2.4.dev20251013071533.dist-info/entry_points.txt,sha256=wEXdaieDoRRCCdhEv2p_C68iytnaXU_2pwt5CqjfbWY,56
-tritonparse-0.2.4.dev20251013071533.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
-tritonparse-0.2.4.dev20251013071533.dist-info/RECORD,,
+tritonparse-0.3.0.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
+tritonparse-0.3.0.dist-info/METADATA,sha256=YoDJfNP6fBgH4ypV7Fvzmy-en3Ax5mIuM660u6S3gnc,8260
+tritonparse-0.3.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tritonparse-0.3.0.dist-info/entry_points.txt,sha256=wEXdaieDoRRCCdhEv2p_C68iytnaXU_2pwt5CqjfbWY,56
+tritonparse-0.3.0.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
+tritonparse-0.3.0.dist-info/RECORD,,