tritonparse 0.2.4.dev20251013071533__tar.gz → 0.3.0__tar.gz

tritonparse-0.3.0/CHANGELOG.md

@@ -0,0 +1,308 @@
+# Changelog
+
+All notable changes to this project will be documented in this file.
+
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
+and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## [0.3.0] - 2025-10-14
+
+### TritonParse Release Notes (last 44 commits)
+
+- **Date range**: 2025-09-19 — 2025-10-14
+- **Scope**: Major feature release - Reproducer system, tensor storage, SASS support, enhanced context manager, CLI improvements.
+
+### Highlights
+
+- **🔄 Reproducer System (Complete)**: Full-featured standalone kernel script generation with template support, tensor reconstruction, and multiple import modes. Extract any traced kernel into a self-contained Python script for debugging, testing, and sharing.
+- **💾 TensorBlobManager**: Production-ready content-addressed tensor storage with automatic compression, deduplication, quota management, and efficient disk usage. Enables high-fidelity kernel reproduction with actual tensor data.
+- **🔧 SASS Disassembly Support**: Optional NVIDIA SASS disassembly during compilation tracing for low-level debugging and performance analysis. Toggle via `enable_sass_dump` parameter or `TRITONPARSE_DUMP_SASS` environment variable.
+- **🎯 Enhanced Context Manager**: Configurable `TritonParseManager` context manager with support for trace launch control, inductor compilation splitting, and flexible parsing parameters.
+- **⚡ CLI Modernization**: Refactored to subcommand structure (`tritonparse parse`, `tritonparse reproduce`) with unified entry point and improved argument handling.
+- **📊 Auto-enable Inductor Launch Tracing**: Automatic detection and tracing of PyTorch Inductor-compiled kernels without manual configuration.
+- **🌐 Website Improvements**: Light mode color scheme, improved stack display in Launch Analysis, and better file diff navigation.
+
+### Changes by area
+
+#### 🔄 **Reproducer System**
+- **Complete reproducer infrastructure** (PR #117-127):
+  - CLI subcommand structure: `tritonparse reproduce <ndjson_file> [options]`
+  - NDJSON ingestion layer with IR preservation
+  - Context bundle system for kernel metadata and parameters
+  - Standardized output paths: `repro_output/<kernel_name>/repro_<timestamp>.py`
+  - Template support with placeholder system for custom generation
+  - Example templates for tensor loading and kernel invocation
+  - Dynamic import generation for kernel dependencies
+  - Kernel signature parsing and integration
+  - Kernel invocation snippet generation with grid/block configuration
+- **Kernel import modes** (PR #165, #166):
+  - `--kernel-import direct`: Import kernel from source file
+  - `--kernel-import override-ttir`: Override and inject TTIR for advanced debugging
+  - Flexible kernel loading strategies for different debugging workflows
+- **Enhanced tensor handling** (PR #141):
+  - Improved tensor metadata logging (shape, dtype, stride, storage offset, device)
+  - Better tensor reconstruction quality in generated reproducers
+  - Support for non-contiguous tensors (commit 12f1d1b)
+- **Extensible placeholder system** (PR #149):
+  - Refactored placeholder replacement with class-based design
+  - Support for: `{{KERNEL_IMPORT_PLACEHOLDER}}`, `{{KERNEL_INVOCATION_PLACEHOLDER}}`, `{{KERNEL_SYSPATH_PLACEHOLDER}}`, `{{JSON_FILE_NAME_PLACEHOLDER}}`
+  - Easy extension for future template needs
+- **Documentation**: Comprehensive reproducer section in README (PR #161) and Usage Guide in Wiki
+
+#### 💾 **TensorBlobManager & Storage**
+- **Production-ready blob storage** (PR #156):
+  - Content-addressed storage using BLAKE2b hashing
+  - Automatic gzip compression for large tensors (>1MB)
+  - Two-level directory structure (`xx/hash.bin.gz`) to avoid filesystem limits
+  - Automatic deduplication: identical tensors stored only once
+  - Storage quota enforcement (default: 100GB)
+  - Per-tensor size limit (default: 10GB) to prevent OOM
+  - Real-time statistics: saved count, dedup hits, compression ratio
+  - Graceful degradation with warning logs when quota exceeded
+- **Compression support** (PR #157):
+  - Configurable compression level (default: 4)
+  - Atomic writes using temporary files + rename for safety
+  - Hash verification for data integrity
+- **Comprehensive testing** (PR #162):
+  - Unit tests for compression, deduplication, quota management
+  - Edge case handling and cleanup verification
+
+#### 🔧 **SASS Disassembly**
+- **SASS extraction support** (PR #137):
+  - New tool: `tritonparse/tools/disasm.py` for CUBIN disassembly
+  - Integration into structured logging behind opt-in flag
+  - Uses `nvdisasm -c -gp -g -gi` for detailed disassembly
+  - Parses output to find function blocks with preserved labels and source mapping
+- **Configuration**:
+  - Environment variable: `TRITONPARSE_DUMP_SASS=1`
+  - API parameter: `enable_sass_dump=True` in `structured_logging.init()`
+  - API parameter takes precedence over environment variable
+- **Robustness**:
+  - Error handling for subprocess failures, missing nvdisasm, and generic exceptions
+  - Writes marker messages instead of failing the trace
+  - Requires NVIDIA CUDA Binary Utilities (nvdisasm)
+- **CUDA testing** (PR #138):
+  - Strengthened tests to validate SASS extraction and persistence
+
+#### 🎯 **Context Manager & API**
+- **Enhanced context manager** (PR #144, #159):
+  - Added `__init__` method with configurable parameters:
+    - `enable_trace_launch`: Control trace launch logging
+    - `split_inductor_compilations`: Control inductor compilation splitting
+    - `**parse_kwargs`: Additional arguments for `unified_parse`
+  - Updated `__exit__` to pass parameters through to parsing pipeline
+  - More flexible for different use cases and workflows
+- **Split inductor compilations control**:
+  - Parameter threading through: `unified_parse()` → `oss_run()` → `parse_logs()` → `parse_single_file()`
+  - Renamed from `split_by_frame_id_and_compile_id` to `split_inductor_compilations` for clarity
+  - Default `True`: splits by frame_id, frame_compile_id, attempt_id, compiled_autograd_id
+  - When `False`: groups all inductor compilations together
+  - Follows tlparse's convention
+- **Unit tests** (commit a5338ce):
+  - Tests for enhanced context manager behavior
+  - Validation of split inductor compilation modes
+
+#### ⚡ **CLI & Entry Points**
+- **Subcommand structure** (PR #117):
+  - Refactored from single-command to modern subcommand architecture
+  - `tritonparse parse <source> [options]` - Run structured log parser
+  - `tritonparse reproduce <ndjson_file> [options]` - Generate reproducers
+  - Breaking change: old `python run.py <source>` no longer works
+  - Extract parser flags into `tritonparse.utils._add_parse_args()`
+  - Remove `unified_parse_from_cli` (programmatic `unified_parse()` remains)
+- **Unified entry point** (PR #133):
+  - Added proper CLI entry point in package configuration
+  - Unified argument handling across commands
+- **CLI entry point fix** (PR #154):
+  - Fixed `ModuleNotFoundError` for tritonparse CLI entry point
+  - Improved package installation and command availability
+
+#### 📊 **Logging & Tracing**
+- **Auto-enable Inductor Launch Tracing** (PR #142):
+  - Automatically detect and trace PyTorch Inductor-compiled kernels
+  - No manual configuration required for Inductor workflows
+  - Seamless integration with existing tracing infrastructure
+- **Kernel source path output** (commit 03bc1e1):
+  - Output `kernel_src_path` in trace metadata for better debugging
+- **NDJSON prettifier improvements** (PR #135):
+  - Renamed and inverted flag to default-filter IRs
+  - More intuitive filtering behavior
+- **Debug flag deprecation** (PR #132):
+  - Removed unused debugging flags
+  - Cleaner configuration surface
+
+#### 🌐 **Website & UI**
+- **Upgraded to Tailwind CSS v4** (commit 6c42d8a):
+  - Migrated from PostCSS plugin to `@tailwindcss/vite` for improved performance
+  - Updated CSS import syntax from `@tailwind` directives to `@import "tailwindcss"`
+  - Removed `tailwind.config.js` and `postcss.config.js` (now CSS-based configuration)
+  - Updated `shadow` class naming to v4 convention (`shadow` → `shadow-sm`)
+  - Cleaned up global CSS to prevent interference with Tailwind utility classes
+- **Upgraded all frontend dependencies**:
+  - Vite: 6.3.5 → 7.1.10
+  - React ecosystem: Updated to latest versions (React 19+)
+  - TypeScript: 5.7.2 → 5.7.3
+  - Added `@types/node` for Node.js type definitions
+  - Fixed dompurify security vulnerability (3.1.7 → 3.3.0) via npm overrides
+- **Light mode color scheme** (PR #139):
+  - Updated `index.css` to support only light mode
+  - Consistent, professional appearance
+- **Improved stack display** (PR #151):
+  - Better stack trace visualization in Launch Analysis
+  - Clearer debugging information
+- **Documentation cleanup** (PR #172):
+  - Removed redundant docs directory and screenshots
+  - Streamlined repository structure
+
+#### 🔧 **Bug Fixes & Maintenance**
+- **General bug fixes** (PR #153):
+  - Multiple stability and reliability improvements
+  - Better error handling throughout codebase
+- **Deserialization fix** (commit d4d7a20):
+  - Fixed unhandled types in deserialization
+  - More robust data loading
+- **README improvements** (PR #158, #164):
+  - Refactored and cleaned up README
+  - Fixed command typos in reproducer generation examples
+  - Clearer installation and usage instructions
+- **Test cleanup** (PR #160):
+  - Removed deprecated test for triton_kernels Tensor functionality
+  - Updated test suite for current codebase
+
+### Compatibility notes
+
+- **Breaking Change**: CLI now uses subcommand structure. Old usage `python run.py <source>` must be updated to `tritonparse parse <source>` or `python run.py parse <source>`.
+- **New Dependencies**: SASS disassembly requires NVIDIA CUDA Binary Utilities (`nvdisasm`). This is optional and only needed if `enable_sass_dump=True`.
+- **Storage**: TensorBlobManager introduces new blob storage directory structure. Default quota is 100GB; configure via `TensorBlobManager` initialization if needed.
+- **Context Manager API**: Enhanced with new parameters. Fully backward compatible with sensible defaults.
+
+### Upgrade guidance
+
+1. **Update CLI commands**: Change `python run.py <source>` to `tritonparse parse <source>` or use the new `tritonparse` command if installed via pip.
+2. **Reproducer usage**: Use `tritonparse reproduce ./parsed_output/trace.ndjson.gz --line <N> --out-dir <output>` to generate standalone kernel scripts.
+3. **SASS disassembly**: Opt-in by setting `TRITONPARSE_DUMP_SASS=1` or passing `enable_sass_dump=True` to `structured_logging.init()`. Requires `nvdisasm` in PATH.
+4. **Tensor storage**: Enable high-fidelity reproduction by using TensorBlobManager (enabled by default when `enable_trace_launch=True`).
+5. **Context manager**: Use enhanced `TritonParseManager` for more control over tracing and parsing behavior.
+
+## [0.2.3] - 2025-09-19
+
+### TritonParse Release Notes (last 15 commits)
+
+- **Date range**: 2025-09-13 — 2025-09-18
+- **Scope**: Website UI/UX, core library, CI/CD & packaging, documentation & testing.
+
+### Highlights
+- **Website File Diff tooling**: Introduced a new Diff Comparison view and File Diff page, preserved diff sessions across navigation, integrated Monaco editor, added preview mode, and shipped a round of UI polish with a URL redirect fix for File Diff navigation.
+- **Kernel Overview**: Added a tiled kernel view toggle to improve dense overviews.
+- **Core**: Added lazy-import support for Triton repo `triton_kernels` custom types, attribution check for `torch._utils_internal`, and safer file mapping cleanup in the log parser.
+- **CI/Packaging**: Refactored dependencies in `pyproject.toml`, removed a legacy Triton install script, and updated GitHub Actions workflows.
+- **Docs & tests**: Improved README guidance; added tests and example outputs; minor UI bug fix in `CopyCodeButton` SVG attributes.
+
+### Changes by area
+- **Website UI/UX**
+  - Introduce `DiffComparisonView` and `FileDiffView`; maintain diff session state; integrate Monaco editor; preview mode; UI polish and navigation fixes.
+  - Add tiled kernel view toggle in `KernelOverview`.
+
+- **Core library**
+  - Lazy-import support for `triton_kernels` custom types; extend tensor handling in tests.
+  - Add attribution check for `torch._utils_internal`.
+  - Refactor file mapping cleanup in `parse_logs`.
+
+- **CI/CD & packaging**
+  - Refactor dependencies in `pyproject.toml`; remove `.ci/install-triton-pip.sh`.
+  - Update GitHub Actions workflows; add helper for `triton_kernels` in CI.
+
+- **Docs & testing**
+  - Clarify tool purpose and installation in `README.md`.
+  - Add tests and sample outputs; small UI component fixes.
+
+### Compatibility notes
+- No breaking changes expected. `triton_kernels` support is optional via lazy import.
+
+### Upgrade guidance
+- Reinstall website dependencies if developing the UI to pick up the Monaco editor.
+
+## [0.2.0] - 2025-09-11
+
+### TritonParse Release Notes (last 27 commits)
+
+- **Date range**: 2025-07-25 — 2025-09-11
+- **Scope**: Core library, website UI/UX, performance & scalability, CI/CD & packaging, documentation & maintenance.
+
+### Highlights
+- **Website usability**: Drag-and-drop to open logs; one-click copy in code viewers; sticky, compact kernel selector; footer shows app version, localized build date, and Git short SHA; tensor arguments in Launch Analysis now display concise summaries with expandable details.
+- **Large-file parsing**: Streaming NDJSON parsing and robust gzip handling significantly reduce memory usage and improve stability for files >100 MB.
+- **Core & integrations**: Persist Inductor kernel config into `inductor_metadata` and pass to JIT hooks; ensure Inductor path invokes `jit_post_compile_hook`; new `init_with_env` for environment-based initialization; move compilation timing `times` into `metadata` for automatic frontend rendering.
+- **Releases & versioning**: Adopt setuptools-scm dynamic versioning; add Nightly PyPI publishing; enable stable publishing on tag push; fix nightly version potentially being older than stable; correct packaging license metadata.
+- **CI stability**: Ubuntu 24.04 compatibility; improved CUDA/cuDNN setup and detection; parallelize jobs; add parallel CI for pip-installed Triton; better error visibility in install scripts; upgrade libstdc++.
+
+### Changes by area
+- **Core library**
+  - Save Inductor kernel params to `inductor_metadata` and forward to JIT hooks.
+  - Manually invoke `jit_post_compile_hook` in the Inductor Triton compile path.
+  - Add `init_with_env` that reads `TRITON_TRACE_FOLDER` and `TRITON_TRACE_LAUNCH`.
+  - Move compilation `times` into `metadata` so the frontend auto-renders it.
+  - Use cached source in compile listener for stability.
+  - Refactor source-mapping pipeline into modular units for maintainability.
+
+- **Website UI/UX**
+  - Drag-and-drop to open supported log files.
+  - Copy button in code viewer panels.
+  - Sticky/collapsible/compact kernel selector in Kernel Overview; resizable compilation stack trace vertically.
+  - Launch Analysis: tensor args show concise summaries with expandable details.
+  - Footer displays version, localized build date, and Git short SHA.
+  - Streaming NDJSON parsing and improved error handling for large logs.
+
+- **Performance & scalability**
+  - Use streaming path for files >100 MB to reduce memory peaks and improve robustness.
+
+- **CI/CD & packaging**
+  - Enable setuptools-scm and nightly PyPI publishing.
+  - Publish stable releases on tag push; improve version computation and tag detection.
+  - Fix nightly version possibly lagging behind stable; add clear error on missing tags.
+  - Add parallel CI for pip-installed Triton; recommend pip installation in docs.
+  - Improve Ubuntu 24.04 setup, CUDA/cuDNN handling, and job parallelism.
+  - Increase error visibility in install scripts and upgrade libstdc++.
+  - Define lower bounds for prerequisites in `pyproject.toml`.
+
+- **Docs & maintenance**
+  - Move repository to `meta-pytorch` org; update links and guidance; add AI assistant context.
+  - Update/restore CONTRIBUTING docs to avoid breaking downstream consumers.
+
+- **Testing**
+  - Preserve test outputs when `TEST_KEEP_OUTPUT=1` to aid debugging.
+
+### Compatibility notes
+- Versioning & publishing: setuptools-scm with tag-based stable releases and nightly dev versions. Ensure `PYPI_API_TOKEN` is configured in CI if publishing is intended.
+- Data format: compilation timing `times` moved under `metadata`; update any downstream scripts that referenced the old location.
+- Build metadata: footer shows localized build date and Git short SHA; restart dev server to refresh these values.
+
+### Upgrade guidance
+- Prefer Triton from PyPI (≥ 3.4.0) and adhere to the lower bounds declared in `pyproject.toml`.
+- For deterministic build metadata in the website, set `BUILD_DATE` and `GIT_COMMIT_SHA_SHORT` in the environment when running dev/build.
+
+
+## [0.1.1] - 2025-07-25
+
+### Added
+
+- **Launch Difference Analysis**: A new `launch_diff` event is automatically generated for each kernel, providing a concise summary of how launch parameters vary across different calls. This helps to quickly identify changes in kernel arguments, grid dimensions, and other metadata.
+- **Enhanced Web UI for Launch Analysis**: The web interface now visualizes the `launch_diff` data, offering an interactive way to explore how kernel launches differ. It includes a detailed breakdown of constant vs. varying parameters and their value distributions.
+- **Kernel-Centric Event Grouping**: The parser now intelligently groups compilation and launch events by kernel, making it easier to analyze the entire lifecycle of a specific kernel.
+- **Launch Event Tracing Control**: Added an `enable_trace_launch` parameter to `tritonparse.structured_logging.init` to give users explicit control over whether to trace kernel launch events.
+- **Enhanced Logging and Testing**: Improved the structured logging initialization and expanded test coverage to verify the correctness of `launch` and `compilation` event counts.
+
+## [0.1.0] - 2025-07-21
+
+This is the initial public release of TritonParse.
+
+### Added
+
+- **Interactive Web Interface**: A rich, client-side web UI for exploring, comparing, and understanding Triton IRs. Features side-by-side code views, synchronized highlighting, and detailed metadata panels.
+- **Structured Logging Backend**: A powerful Python backend to capture detailed information from the Triton compiler and runtime, including IRs (TTIR, TTGIR, PTX, AMDGCN), metadata, timings, and Python source code, and outputs it as structured NDJSON logs.
+- **Source-to-Source Mapping**: Automatic generation of bidirectional mappings between Python code and all intermediate representations (IRs), allowing you to trace a line of Python code all the way down to the generated assembly and back.
+- **Kernel Launch Tracing**: Capability to trace each kernel launch, capturing the grid dimensions, kernel arguments (with detailed tensor information), and other runtime metadata.
+- **Flexible Log Parsing CLI**: A command-line interface (`run.py`) to parse logs from local files or directories, and from single or multiple ranks in a distributed training job.
+- **Prerequisites Documentation**: Clear requirements for Python (>=3.10), PyTorch, and Triton (>3.3.1, compiled from source).
+- **Getting Started Guide**: A step-by-step workflow for generating, parsing, and visualizing traces.
+- **Configuration via Environment Variables**: Support for `TRITON_TRACE`, `TRITON_TRACE_LAUNCH`, `TRITONPARSE_KERNEL_ALLOWLIST`, and `TRITON_TRACE_GZIP`.

{tritonparse-0.2.4.dev20251013071533 → tritonparse-0.3.0}/PKG-INFO

@@ -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 → tritonparse-0.3.0}/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.dev20251013071533 → tritonparse-0.3.0}/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-0.2.4.dev20251013071533 → tritonparse-0.3.0}/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-0.3.0/tritonparse/reproducer/placeholder_replacer.py

@@ -0,0 +1,227 @@
+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,
+    _parse_kernel_signature,
+)
+
+
+class HandlerProtocol(Protocol):
+    def __call__(
+        self, code: str, context_bundle: ContextBundle, **kwargs: Any
+    ) -> str: ...
+
+
+class PlaceholderReplacer(ABC):
+    """
+    Abstract base class for template placeholder replacement.
+
+    Subclasses should register replacement handlers in their __init__ method
+    by calling self.register(placeholder, handler_function).
+
+    Each handler function should have the signature:
+        handler(code: str, context_bundle: ContextBundle, **kwargs) -> str
+    """
+
+    def __init__(self):
+        # Dictionary mapping placeholder strings to handler functions
+        self.handlers: Dict[str, HandlerProtocol] = {}
+
+    def register(self, placeholder: str, handler: HandlerProtocol):
+        """
+        Register a handler function for a specific placeholder.
+
+        Args:
+            placeholder: The placeholder string to replace (e.g., "{{JSON_FILE_NAME_PLACEHOLDER}}")
+            handler: A callable that takes (code, context_bundle, **kwargs) and returns modified code
+        """
+        self.handlers[placeholder] = handler
+
+    def replace(
+        self, template_code: str, context_bundle: ContextBundle, **kwargs: Any
+    ) -> str:
+        """
+        Replace all registered placeholders in the template code.
+
+        Args:
+            template_code: The template code containing placeholders
+            context_bundle: Context information about the kernel
+            **kwargs: Additional keyword arguments passed to handler functions
+
+        Returns:
+            The code with all placeholders replaced
+        """
+        code = template_code
+        for placeholder, handler in self.handlers.items():
+            code = handler(code, context_bundle, **kwargs)
+        return code
+
+
+class DefaultPlaceholderReplacer(PlaceholderReplacer):
+    """
+    Default implementation of PlaceholderReplacer.
+
+    Handles the following placeholders:
+    - {{JSON_FILE_NAME_PLACEHOLDER}}: Replaced with the JSON file name
+    - # {{KERNEL_SYSPATH_PLACEHOLDER}}: Replaced with sys.path setup code
+    - # {{KERNEL_IMPORT_PLACEHOLDER}}: Replaced with kernel import statement
+    - # {{KERNEL_INVOCATION_PLACEHOLDER}}: Replaced with kernel invocation code
+    """
+
+    def __init__(self):
+        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(
+            "# {{KERNEL_INVOCATION_PLACEHOLDER}}", self._replace_kernel_invocation
+        )
+
+    def _replace_json_filename(
+        self, code: str, context_bundle: ContextBundle, **kwargs
+    ) -> str:
+        """Replace the JSON file name placeholder."""
+        temp_json_path = kwargs.get("temp_json_path")
+        if temp_json_path is None:
+            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."""
+        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."""
+        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
+    ) -> str:
+        """Replace the kernel invocation placeholder."""
+        source_code = context_bundle.kernel_info.source_code
+        pos_args, kw_args = _parse_kernel_signature(source_code)
+        invocation_snippet = _generate_invocation_snippet(pos_args, kw_args)
+        return code.replace("# {{KERNEL_INVOCATION_PLACEHOLDER}}", invocation_snippet)

{tritonparse-0.2.4.dev20251013071533 → tritonparse-0.3.0}/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-0.3.0/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"