Fix OVERRIDE_TTIR reproducer with stub kernel + ir_override (#376)

Summary:
Pull Request resolved: #376

The existing OVERRIDE_TTIR mode was broken: it skipped defining the kernel
function (causing NameError), only worked for autotuned kernels, and discarded
constexpr values.

This rewrites OVERRIDE_TTIR to generate a stub triton.jit function (same name
and params, `pass` body) wrapped with triton.autotune carrying the captured
constexpr values, compile params, and ir_override pointing to the captured TTGIR.
The stub compiles trivially through the frontend, then ir_override replaces the
compilation output with the real kernel's IR.

This eliminates the need to copy kernel source code and its transitive
dependencies (which was flaky due to complex imports, generated code, and
compiled extensions).

Key changes:
- New stub_generator.py: generates stub functions and extracts constexpr values
- Fixed _replace_kernel_import for OVERRIDE_TTIR: generates stub + autotune config
- Fixed _replace_kernel_invocation: filters constexpr/compile params (autotune provides them)
- Save captured IR files from compilation event's file_content to captured_irs/
- Reuses _is_constexpr_annotation from function_extractor.py (no duplication)
- lru_cache on extract_params_from_source to avoid redundant AST parses

Reviewed By: FindHao

Differential Revision: D99893169

fbshipit-source-id: ce1c42a9afcbce08cb853ce84f5c934551449911

Author: srivatsan-ramesh

tritonparse/reproducer/orchestrator.py

@@ -11,7 +11,11 @@
 )
 from tritonparse.reproducer.templates.loader import load_template_code
 from tritonparse.reproducer.types import KernelImportMode
-from tritonparse.reproducer.utils import determine_output_paths, format_python_code
+from tritonparse.reproducer.utils import (
+    determine_output_paths,
+    format_python_code,
+    save_captured_irs,
+)
 from tritonparse.shared_vars import is_fbcode
 from tritonparse.tools.prettify_ndjson import load_ndjson, save_prettified_json
 from tritonparse.tp_logger import logger
@@ -102,6 +106,10 @@ def reproduce(
         )
         save_prettified_json(context_bundle.raw_comp_event, comp_json_path)
 
+    # Save IR files for OVERRIDE_TTIR mode
+    if kernel_import == KernelImportMode.OVERRIDE_TTIR:
+        save_captured_irs(context_bundle.raw_comp_event, out_py_path.parent)
+
     logger.debug("Loading reproducer template.")
     template_code = load_template_code(template)
 

tritonparse/reproducer/placeholder_replacer.py

@@ -15,6 +15,11 @@
     match_autotune_config,
 )
 from tritonparse.reproducer.ingestion.ndjson import ContextBundle
+from tritonparse.reproducer.stub_generator import (
+    _find_ir_override_file,
+    generate_stub_source,
+    get_constexpr_values,
+)
 from tritonparse.reproducer.templates.utils import (
     _disable_triton_autotune,
     get_function_source,
@@ -242,60 +247,7 @@ 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(self.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(self.IR_OVERRIDE_SETUP_PLACEHOLDER, setup_code)
+        return code.replace(self.IR_OVERRIDE_SETUP_PLACEHOLDER, "")
 
     def _replace_kernel_syspath(
         self, code: str, context_bundle: ContextBundle, **kwargs
@@ -312,7 +264,7 @@ def _replace_kernel_syspath(
             )
             return code.replace(self.KERNEL_SYSPATH_PLACEHOLDER, comment)
         elif kernel_import == KernelImportMode.OVERRIDE_TTIR:
-            comment = "# Kernel sys.path setup skipped - using IR override mode"
+            comment = "# Kernel sys.path setup skipped - using stub with IR override"
             return code.replace(self.KERNEL_SYSPATH_PLACEHOLDER, comment)
         else:
             raise ValueError(f"Unknown kernel_import mode: {kernel_import}")
@@ -425,8 +377,96 @@ def _replace_kernel_import(
 
             return code.replace(self.KERNEL_IMPORT_PLACEHOLDER, embedded_code)
         elif kernel_import == KernelImportMode.OVERRIDE_TTIR:
-            comment = "# Kernel import skipped - using IR override mode with TTIR"
-            return code.replace(self.KERNEL_IMPORT_PLACEHOLDER, comment)
+            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 'override-ttir' mode"
+                )
+            if not func_name:
+                raise ValueError(
+                    "Cannot determine kernel function name for 'override-ttir' mode"
+                )
+
+            # Build the autotune config with ir_override + constexpr values
+            constexpr_vals = get_constexpr_values(context_bundle)
+            compile_meta = context_bundle.compile
+
+            config_parts = []
+            # Constexpr kwargs
+            kwargs_repr = repr(constexpr_vals)
+            config_parts.append(f"kwargs={kwargs_repr}")
+            # Compile params
+            for param in TRITON_COMPILE_PARAMS:
+                value = compile_meta.get(param)
+                if value is not None:
+                    config_parts.append(f"{param}={value!r}")
+            # ir_override — find the best available IR file
+            config_parts.append("ir_override=_IR_OVERRIDE_FILE")
+
+            config_str = ", ".join(config_parts)
+
+            # Generate imports + IR file discovery + stub function + autotune wrapper
+            import_lines = list(_BASE_IMPORT_LINES) + [""]
+
+            # Scan original source for extra imports (e.g., tlx)
+            extra_imports = _detect_extra_imports(source_code)
+            if extra_imports:
+                import_lines.extend(extra_imports)
+                import_lines.append("")
+
+            import_lines.extend(
+                [
+                    "import os",
+                    "",
+                ]
+                + get_function_source(_find_ir_override_file, with_invocation=False)
+                + [
+                    "",
+                    "# Find the best captured IR file for ir_override",
+                    "_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))",
+                    "_IR_DIR = os.path.join(_SCRIPT_DIR, 'captured_irs')",
+                    "_IR_OVERRIDE_FILE = _find_ir_override_file(_IR_DIR) if os.path.isdir(_IR_DIR) else None",
+                    "",
+                ]
+            )
+
+            # Generate the stub function source
+            stub_code = generate_stub_source(func_name, source_code)
+
+            # Wrap with autotune carrying ir_override + constexpr values.
+            # Apply @triton.autotune programmatically so we can conditionally
+            # include ir_override only when the IR file exists.
+            # Detect indent of the import placeholder for the if/else block.
+            placeholder = self.KERNEL_IMPORT_PLACEHOLDER
+            indent = next(
+                (
+                    line[: len(line) - len(line.lstrip())]
+                    for line in code.splitlines()
+                    if line.lstrip().startswith(placeholder)
+                ),
+                "",
+            )
+            inner = indent + "    "
+
+            import_lines.extend(
+                [
+                    "# Stub kernel — body is replaced by captured IR via ir_override",
+                    stub_code,
+                    "",
+                    f"{indent}if _IR_OVERRIDE_FILE:",
+                    f"{inner}{func_name} = triton.autotune(",
+                    f"{inner}    configs=[triton.Config({config_str})],",
+                    f"{inner}    key=[],",
+                    f"{inner})({func_name})",
+                    "",
+                    f"{indent}imported_kernel_function = {func_name}",
+                ]
+            )
+
+            embedded_code = "\n".join(import_lines)
+            return code.replace(self.KERNEL_IMPORT_PLACEHOLDER, embedded_code)
         else:
             raise ValueError(f"Unknown kernel_import mode: {kernel_import}")
 
@@ -452,9 +492,46 @@ def _replace_kernel_invocation(
 
         Otherwise, passes all args plus compile params as usual.
         """
+        kernel_import = kwargs.get("kernel_import", KernelImportMode.DEFAULT)
         source_code = context_bundle.kernel_info.source_code
         pos_args, kw_args = _parse_kernel_signature(source_code)
 
+        if kernel_import == KernelImportMode.OVERRIDE_TTIR:
+            # When ir_override is active, autotune provides constexpr values
+            # and compile params — only pass non-constexpr args.
+            # When ir_override is missing (no captured_irs/), the stub runs
+            # as a plain @triton.jit and needs all args + compile params.
+            constexpr_vals = get_constexpr_values(context_bundle)
+            filtered_pos = [a for a in pos_args if a not in constexpr_vals]
+            filtered_kw = [a for a in kw_args if a not in constexpr_vals]
+            override_snippet = _generate_invocation_snippet(
+                filtered_pos, filtered_kw, compile_params=None
+            )
+            compile_params = _get_compile_params_for_invocation(
+                context_bundle.compile, kw_args
+            )
+            fallback_snippet = _generate_invocation_snippet(
+                pos_args, kw_args, compile_params
+            )
+            # Find the placeholder's indent so the if/else aligns in any template.
+            placeholder = self.KERNEL_INVOCATION_PLACEHOLDER
+            indent = next(
+                (
+                    line[: len(line) - len(line.lstrip())]
+                    for line in code.splitlines()
+                    if line.lstrip().startswith(placeholder)
+                ),
+                "    ",
+            )
+            inner = indent + "    "
+            invocation_snippet = (
+                f"if _IR_OVERRIDE_FILE:\n"
+                f"{inner}{override_snippet}\n"
+                f"{indent}else:\n"
+                f"{inner}{fallback_snippet}"
+            )
+            return code.replace(self.KERNEL_INVOCATION_PLACEHOLDER, invocation_snippet)
+
         if self.preserve_autotune:
             # Get full kernel source with decorators to find autotune config params
             func_name = context_bundle.kernel_info.function_name

tritonparse/reproducer/stub_generator.py

@@ -0,0 +1,137 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+"""
+Generate stub @triton.jit functions for IR-based reproducers.
+
+A stub kernel has the same function name and parameter signature as the original
+kernel but a trivial body (``pass``).  When used with ``ir_override`` on a
+``triton.Config``, the stub compiles through the frontend (producing a valid but
+trivial TTIR) and the compilation output is replaced by the captured IR from the
+original kernel.  This completely eliminates the need to copy the original
+kernel's source code and its transitive dependencies.
+"""
+
+import ast
+from functools import lru_cache
+from typing import Any, Dict, List, Tuple
+
+from tritonparse.reproducer.function_extractor import _is_constexpr_annotation
+from tritonparse.reproducer.ingestion.ndjson import ContextBundle
+from tritonparse.tp_logger import logger
+
+
+def generate_stub_source(kernel_name: str, source_code: str) -> str:
+    """Generate a stub ``@triton.jit`` function from the original kernel source.
+
+    Parses the original source to extract parameter names and ``tl.constexpr``
+    annotations, then produces a stub with the same signature and a ``pass``
+    body.  Does NOT include ``@triton.autotune`` — that is added separately
+    by the placeholder replacer.
+
+    Args:
+        kernel_name: The kernel function name.
+        source_code: The original kernel source code (from ``@triton.jit`` onward).
+
+    Returns:
+        Python source code defining the stub function.
+    """
+    params = extract_params_from_source(source_code)
+
+    # Build parameter list
+    param_strs: List[str] = []
+    for name, is_constexpr in params:
+        if is_constexpr:
+            param_strs.append(f"{name}: tl.constexpr")
+        else:
+            param_strs.append(name)
+
+    # Format as multi-line if many params
+    if len(param_strs) > 3:
+        params_block = "\n    " + ",\n    ".join(param_strs) + ",\n"
+    else:
+        params_block = ", ".join(param_strs)
+
+    return f"@triton.jit\ndef {kernel_name}({params_block}):\n    pass"
+
+
+def _find_ir_override_file(ir_dir: str) -> "str | None":
+    """Find the captured TTIR file for ir_override.
+
+    Returns the full path to the first ``.ttir`` file, or None.
+
+    NOTE: This function is extracted by ``get_function_source()`` and embedded
+    verbatim in generated reproducer scripts.  Keep imports minimal (only
+    ``os`` and ``logging`` are available).
+    """
+    import logging
+    import os
+
+    for name in sorted(os.listdir(ir_dir)):
+        if name.endswith(".ttir"):
+            return os.path.join(ir_dir, name)
+    logging.getLogger(__name__).warning(
+        "captured_irs/ directory exists but contains no .ttir files. "
+        "Stub kernel will run without ir_override (fallback mode)."
+    )
+    return None
+
+
+def get_constexpr_values(context_bundle: ContextBundle) -> Dict[str, Any]:
+    """Extract constexpr parameter names and their values from the context.
+
+    Args:
+        context_bundle: The context bundle from the captured launch event.
+
+    Returns:
+        Dictionary mapping constexpr parameter name to its value.
+    """
+    source_code = context_bundle.kernel_info.source_code
+    params = extract_params_from_source(source_code)
+    constexpr_names = {name for name, is_ce in params if is_ce}
+    args = context_bundle.args
+
+    result: Dict[str, Any] = {}
+    for name in constexpr_names:
+        arg_info = args.get(name)
+        if arg_info is None:
+            continue
+        if isinstance(arg_info, dict):
+            result[name] = arg_info.get("value")
+        else:
+            result[name] = arg_info
+
+    return result
+
+
+@lru_cache(maxsize=8)
+def extract_params_from_source(
+    source_code: str,
+) -> Tuple[Tuple[str, bool], ...]:
+    """Extract parameter names and constexpr annotations from kernel source.
+
+    Results are cached since the same source is parsed multiple times during
+    reproducer generation (stub generation, constexpr extraction, invocation).
+
+    Returns:
+        Tuple of ``(param_name, is_constexpr)`` pairs.
+    """
+    try:
+        tree = ast.parse(source_code)
+    except SyntaxError:
+        logger.warning("Failed to parse kernel source for param extraction")
+        return ()
+
+    for node in ast.walk(tree):
+        if not isinstance(node, ast.FunctionDef):
+            continue
+
+        params: List[Tuple[str, bool]] = []
+        all_args = list(node.args.args) + list(node.args.kwonlyargs)
+        for arg in all_args:
+            is_ce = arg.annotation is not None and _is_constexpr_annotation(
+                arg.annotation
+            )
+            params.append((arg.arg, is_ce))
+        return tuple(params)
+
+    return ()

tritonparse/reproducer/types.py

@@ -12,7 +12,7 @@ class KernelImportMode(str, Enum):
     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.
+        OVERRIDE_TTIR: Generate a stub kernel with captured IR override.
     """
 
     DEFAULT = "default"