tritonparse 0.2.4.dev20251014071550__py3-none-any.whl → 0.3.1.dev20251015071527__py3-none-any.whl

tritonparse/reproducer/cli.py

@@ -40,6 +40,7 @@ def _add_reproducer_args(parser: argparse.ArgumentParser) -> None:
             "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

@@ -46,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)
 
@@ -58,6 +67,7 @@ def reproduce(
         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

@@ -77,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(
@@ -92,6 +95,65 @@ 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:
@@ -106,6 +168,9 @@ class DefaultPlaceholderReplacer(PlaceholderReplacer):
                 "# 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}")
 
@@ -146,6 +211,9 @@ class DefaultPlaceholderReplacer(PlaceholderReplacer):
             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}")
 

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

@@ -6,10 +6,13 @@ 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
+    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.
+    """
 
-    # Future modes can be added here:
-    # OVERRIDE_TTIR = "override-ttir"  # Use TTIR from compilation event with monkeypatch
+    DEFAULT = "default"
+    COPY = "copy"
+    OVERRIDE_TTIR = "override-ttir"

{tritonparse-0.2.4.dev20251014071550.dist-info → tritonparse-0.3.1.dev20251015071527.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tritonparse
-Version: 0.2.4.dev20251014071550
+Version: 0.3.1.dev20251015071527
 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.dev20251014071550.dist-info → tritonparse-0.3.1.dev20251015071527.dist-info}/RECORD

@@ -15,14 +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=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/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
@@ -31,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.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,,
+tritonparse-0.3.1.dev20251015071527.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
+tritonparse-0.3.1.dev20251015071527.dist-info/METADATA,sha256=CPxG95kX-c5EcK739gu_ky-k0TxDY1pAL9SaXOKBNDw,8278
+tritonparse-0.3.1.dev20251015071527.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tritonparse-0.3.1.dev20251015071527.dist-info/entry_points.txt,sha256=wEXdaieDoRRCCdhEv2p_C68iytnaXU_2pwt5CqjfbWY,56
+tritonparse-0.3.1.dev20251015071527.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
+tritonparse-0.3.1.dev20251015071527.dist-info/RECORD,,