tritonparse 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

tritonparse/structured_logging.py

@@ -15,6 +15,7 @@ from collections.abc import Mapping
 from dataclasses import asdict, is_dataclass
 from datetime import date, datetime
 from enum import Enum
+from functools import partial
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional, Union
 
@@ -756,7 +757,7 @@ def maybe_trace_triton(
 
     # Add timing information if available
     if times:
-        trace_data["times"] = times
+        trace_data["metadata"]["times"] = times
     # Log the collected information through the tracing system
     trace_structured_triton(
         event_type,
@@ -819,10 +820,18 @@ def extract_arg_info(arg_dict):
     return extracted_args
 
 
-def add_launch_metadata(grid, metadata, arg_dict):
+def add_launch_metadata(grid, metadata, arg_dict, inductor_args=None):
     # Extract detailed argument information
     extracted_args = extract_arg_info(arg_dict)
-    return {"launch_metadata_tritonparse": (grid, metadata._asdict(), extracted_args)}
+    extracted_inductor_args = extract_arg_info(inductor_args) if inductor_args else {}
+    return {
+        "launch_metadata_tritonparse": (
+            grid,
+            metadata._asdict(),
+            extracted_args,
+            extracted_inductor_args,
+        )
+    }
 
 
 class JITHookImpl(JITHook):
@@ -848,6 +857,7 @@ class JITHookImpl(JITHook):
         compile,
         is_manual_warmup: bool,
         already_compiled: bool,
+        inductor_args: Optional[Dict[str, Any]] = None,
     ) -> Optional[bool]:
         """
         Override or set the launch_metadata function for the JIT-compiled kernel.
@@ -875,12 +885,16 @@ class JITHookImpl(JITHook):
                 return True
 
         # Get the current launch_metadata function if it exists
-        current_launch_metadata = getattr(fn.jit_function, "launch_metadata", None)
+        function = getattr(fn, "jit_function", fn)
+
+        current_launch_metadata = getattr(function, "launch_metadata", None)
         if current_launch_metadata is not None:
             log.warning(
                 f"fn {fn} launch_metadata is not None: {current_launch_metadata}. It will be overridden by tritonparse."
             )
-        fn.jit_function.launch_metadata = add_launch_metadata
+        function.launch_metadata = partial(
+            add_launch_metadata, inductor_args=inductor_args
+        )
         return True
 
 
@@ -944,6 +958,7 @@ class LaunchHookImpl(LaunchHook):
             trace_data["extracted_args"] = launch_metadata_tritonparse[
                 2
             ]  # Now contains detailed arg info
+            trace_data["extracted_inductor_args"] = launch_metadata_tritonparse[3]
         trace_structured_triton("launch", metadata_fn=lambda: convert(trace_data))
 
 
@@ -1011,6 +1026,15 @@ def init(trace_folder: Optional[str] = None, enable_trace_launch: bool = False):
     knobs.compilation.listener = maybe_trace_triton
 
 
+def init_with_env():
+    """
+    This function is used to initialize TritonParse with the environment variable TRITON_TRACE_FOLDER and TRITON_TRACE_LAUNCH specifically.
+    It is only supposed to be used in OSS triton's source code.
+    """
+    if triton_trace_folder:
+        init(triton_trace_folder, enable_trace_launch=TRITON_TRACE_LAUNCH)
+
+
 def clear_logging_config():
     """
     Clear all configurations made by init() and init_basic().

tritonparse/tools/readme.md

@@ -0,0 +1 @@
+The tool scripts in this folder are used separately. They are not part of the main tritonparse functionality.

tritonparse-0.2.0.dist-info/METADATA

@@ -0,0 +1,139 @@
+Metadata-Version: 2.4
+Name: tritonparse
+Version: 0.2.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
+Project-URL: Homepage, https://github.com/meta-pytorch/tritonparse
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: triton>3.3.1
+Provides-Extra: test
+Requires-Dist: coverage>=7.0.0; extra == "test"
+Dynamic: license-file
+
+# TritonParse
+
+[![License: BSD-3](https://img.shields.io/badge/License-BSD--3-blue.svg)](https://opensource.org/licenses/BSD-3-Clause)
+[![GitHub Pages](https://img.shields.io/badge/GitHub%20Pages-Deploy-brightgreen)](https://meta-pytorch.org/tritonparse/)
+
+**A comprehensive visualization and analysis tool for Triton IR files** — helping developers analyze, debug, and understand Triton kernel compilation processes.
+
+🌐 **[Try it online →](https://meta-pytorch.org/tritonparse/?json_url=https://meta-pytorch.org/tritonparse/dedicated_log_triton_trace_findhao__mapped.ndjson.gz)**
+
+## ✨ Key Features
+
+- **🚀 Launch Difference Analysis** - Automatically detect and visualize variations in kernel launch parameters, helping you pinpoint performance bottlenecks and debug launch configurations.
+- **🔍 Interactive Visualization** - Explore Triton kernels with detailed metadata and stack traces
+- **📊 Multi-format IR Support** - View TTGIR, TTIR, LLIR, PTX, and AMDGCN in one place
+- **🔄 Side-by-side Comparison** - Compare IR stages with synchronized highlighting
+- **📝 Structured Logging** - Capture detailed compilation and launch events with source mapping
+- **🌐 Ready-to-use Interface** - No installation required, works in your browser
+- **🔒 Privacy-first** - All processing happens locally in your browser, no data uploaded
+
+## 🚀 Quick Start
+
+### 1. Generate Traces
+
+```python
+import tritonparse.structured_logging
+
+# Initialize logging with launch tracing enabled
+tritonparse.structured_logging.init("./logs/", enable_trace_launch=True)
+
+# Your Triton/PyTorch code here
+# ... your kernels ...
+
+# Parse and generate trace files
+import tritonparse.utils
+tritonparse.utils.unified_parse("./logs/")
+```
+The example terminal output is:
+```bash
+tritonparse log file list: /tmp/tmp1gan7zky/log_file_list.json
+INFO:tritonparse:Copying parsed logs from /tmp/tmp1gan7zky to /scratch/findhao/tritonparse/tests/parsed_output
+
+================================================================================
+📁 TRITONPARSE PARSING RESULTS
+================================================================================
+📂 Parsed files directory: /scratch/findhao/tritonparse/tests/parsed_output
+📊 Total files generated: 2
+
+📄 Generated files:
+--------------------------------------------------
+   1. 📝 dedicated_log_triton_trace_findhao__mapped.ndjson.gz (7.2KB)
+   2. 📝 log_file_list.json (181B)
+================================================================================
+✅ Parsing completed successfully!
+================================================================================
+```
+
+### 2. Visualize Results
+
+**Visit [https://meta-pytorch.org/tritonparse/](https://meta-pytorch.org/tritonparse/?json_url=https://meta-pytorch.org/tritonparse/dedicated_log_triton_trace_findhao__mapped.ndjson.gz)** and open your local trace files (.ndjson.gz format).
+
+> **🔒 Privacy Note**: Your trace files are processed entirely in your browser - nothing is uploaded to any server!
+
+## 🛠️ Installation
+
+**For basic usage (trace generation):**
+```bash
+git clone https://github.com/meta-pytorch/tritonparse.git
+cd tritonparse
+pip install -e .
+```
+
+**Prerequisites:** Python ≥ 3.10, Triton ≥ 3.4.0, GPU required (NVIDIA/AMD)
+
+TritonParse relies on new features in Triton. Please install the latest version of Triton:
+```bash
+pip install triton
+```
+
+## 📚 Complete Documentation
+
+| 📖 Guide | Description |
+|----------|-------------|
+| **[🏠 Wiki Home](https://github.com/meta-pytorch/tritonparse/wiki)** | Complete documentation and navigation |
+| **[📦 Installation Guide](https://github.com/meta-pytorch/tritonparse/wiki/01.-Installation)** | Detailed setup for all scenarios |
+| **[📋 Usage Guide](https://github.com/meta-pytorch/tritonparse/wiki/02.-Usage-Guide)** | Complete workflow and examples |
+| **[🌐 Web Interface Guide](https://github.com/meta-pytorch/tritonparse/wiki/03.-Web-Interface-Guide)** | Master the visualization interface |
+| **[🔧 Developer Guide](https://github.com/meta-pytorch/tritonparse/wiki/04.-Developer-Guide)** | Contributing and development setup |
+| **[❓ FAQ](https://github.com/meta-pytorch/tritonparse/wiki/06.-FAQ)** | Frequently asked questions |
+
+## 🛠️ Tech Stack
+
+- **Frontend**: React 19, TypeScript, Vite, Tailwind CSS, Monaco Editor
+- **Backend**: Python with Triton integration, structured logging
+- **Deployment**: GitHub Pages, automatic deployment
+
+## 📊 Understanding Triton Compilation
+
+TritonParse visualizes the complete Triton compilation pipeline:
+
+**Python Source** → **TTIR** → **TTGIR** → **LLIR** → **PTX/AMDGCN**
+
+Each stage can be inspected and compared to understand optimization transformations.
+
+## 🤝 Contributing
+
+We welcome contributions! Please see our **[Developer Guide](https://github.com/meta-pytorch/tritonparse/wiki/04.-Developer-Guide)** for:
+- Development setup
+- Code formatting standards
+- Pull request process
+- Architecture overview
+
+## 📞 Support & Community
+
+- **🐛 Report Issues**: [GitHub Issues](https://github.com/meta-pytorch/tritonparse/issues)
+- **💬 Discussions**: [GitHub Discussions](https://github.com/meta-pytorch/tritonparse/discussions)
+- **📚 Documentation**: [TritonParse Wiki](https://github.com/meta-pytorch/tritonparse/wiki)
+
+## 📄 License
+
+This project is licensed under the BSD-3 License - see the [LICENSE](LICENSE) file for details.
+
+---
+
+**✨ Ready to get started?** Visit our **[Installation Guide](https://github.com/meta-pytorch/tritonparse/wiki/01.-Installation)** or try the **[online tool](https://meta-pytorch.org/tritonparse/)** directly!

tritonparse-0.2.0.dist-info/RECORD

@@ -0,0 +1,24 @@
+tritonparse/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tritonparse/common.py,sha256=VWKFZJG7msIExMH0QNYCO8SHKqUBFtdRuLaCy_K8DFI,13725
+tritonparse/event_diff.py,sha256=yOD6uNxLJroatfx2nEGr-erw24ObOrHU9P6V5pzr8do,4907
+tritonparse/extract_source_mappings.py,sha256=Z6UxFj2cCE5NCWLQTYPKqUpLfbYhqP8xgCl5mvud9KI,1451
+tritonparse/ir_parser.py,sha256=1j1tP9jpUN7wH3e01bKUkUPgTMlNXUdp8LKRCC-WTro,9324
+tritonparse/mapper.py,sha256=prrczfi13P7Aa042OrEBsmRF1HW3jDhwxicANgPkWIM,4150
+tritonparse/shared_vars.py,sha256=-c9CvXJSDm9spYhDOJPEQProeT_xl3PaNmqTEYi_u4s,505
+tritonparse/source_type.py,sha256=nmYEQS8rfkIN9BhNhQbkmEvKnvS-3zAxRGLY4TaZdi8,1676
+tritonparse/sourcemap_utils.py,sha256=qsQmTDuEe9yuUVyxSHRbjTR38gi0hvJEijnPkrJVAV4,2037
+tritonparse/structured_logging.py,sha256=mrApsVigIHZln6De2ElwNTUSZfO61OHXgd08o2X26sM,40430
+tritonparse/tp_logger.py,sha256=vXzY7hMDmVnRBGBhIjFZe3nHZzG5NKKPONGUszJhGgU,242
+tritonparse/trace_processor.py,sha256=QzUOKwnOkBbwTTKBsa5ZMUABPLMJIBFtTcG2SkhO0I8,12771
+tritonparse/utils.py,sha256=wt61tpbkqjGqHh0c7Nr2WlOv7PbQssmjULd6uA6aAko,4475
+tritonparse/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tritonparse/tools/decompress_bin_ndjson.py,sha256=kpt7DM_sSA334F1X45xdkP2OR9LuB27Pc50EkGr6CPM,4144
+tritonparse/tools/format_fix.py,sha256=Ol0Sjui8D7OzHwbamAfGnq8V5Y63uwNaFTKSORN5HkQ,3867
+tritonparse/tools/load_tensor.py,sha256=tfdmNVd9gsZqO6msQBhbXIhOvUzgc83yF64k2GDWPNk,2122
+tritonparse/tools/prettify_ndjson.py,sha256=VOzVWoXpCbaAXYA4i_wBcQIHfh-JhAx7xR4cF_L8yDs,10928
+tritonparse/tools/readme.md,sha256=w6PWYfYnRgoPArLjxG9rVrpcLUkoVMGuRlbpF-o0IQM,110
+tritonparse-0.2.0.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
+tritonparse-0.2.0.dist-info/METADATA,sha256=kHgvrwCWkFfyy2EvE1pRR8OyfP4OamUvaajLhJE7hyo,6151
+tritonparse-0.2.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tritonparse-0.2.0.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
+tritonparse-0.2.0.dist-info/RECORD,,

tritonparse/reproducer/__init__.py

@@ -1,21 +0,0 @@
-"""Reproducer subpackage: generate runnable Triton repro scripts from traces.
-
-Contains:
-- ingestion.ndjson: parse NDJSON and build a context bundle
-- orchestrator: LLM-based code generation with optional execute/repair
-- providers: LLM provider protocol and Gemini provider
-- prompts: simple prompt loader and templates
-- runtime.executor: helper to run generated Python scripts
-- param_generator: synthesize tensor/scalar allocations to reduce LLM burden
-"""
-
-from .ingestion.ndjson import build_context_bundle
-from .orchestrator import generate_from_ndjson
-from .param_generator import generate_allocation_snippet, generate_kwargs_dict
-
-__all__ = [
-    "build_context_bundle",
-    "generate_from_ndjson",
-    "generate_allocation_snippet",
-    "generate_kwargs_dict",
-]

tritonparse/reproducer/__main__.py

@@ -1,81 +0,0 @@
-import argparse
-import sys
-
-
-def main() -> None:
-    p = argparse.ArgumentParser(
-        description=(
-            "Generate a runnable Triton repro script from a tritonparse NDJSON" " trace"
-        )
-    )
-    p.add_argument("--ndjson", required=True, help="Path to NDJSON trace file")
-    p.add_argument(
-        "--launch-index",
-        type=int,
-        default=0,
-        help="Launch index to reproduce",
-    )
-    p.add_argument("--out", default="repro.py", help="Output Python file path")
-    p.add_argument(
-        "--execute",
-        action="store_true",
-        help="Execute the generated script",
-    )
-    p.add_argument(
-        "--retries",
-        type=int,
-        default=0,
-        help="Auto-repair attempts if execution fails",
-    )
-    p.add_argument(
-        "--temperature",
-        type=float,
-        help="Override sampling temperature",
-    )
-    p.add_argument(
-        "--max-tokens",
-        type=int,
-        help="Override max tokens for generation",
-    )
-    args = p.parse_args()
-
-    # Lazy imports to allow `--help` without optional deps installed
-    from .config import load_config
-    from .orchestrator import generate_from_ndjson
-
-    try:
-        from .factory import make_gemini_provider
-    except Exception:  # pragma: no cover
-        print(
-            "Failed to import provider factory. Ensure optional deps are installed (e.g. google-genai).",
-            file=sys.stderr,
-        )
-        raise
-
-    cfg = load_config()
-    try:
-        provider = make_gemini_provider()
-    except ModuleNotFoundError:  # pragma: no cover
-        print(
-            "Gemini provider requires 'google-genai'. Install via: pip install google-genai",
-            file=sys.stderr,
-        )
-        sys.exit(2)
-    temperature = args.temperature if args.temperature is not None else cfg.temperature
-    max_tokens = args.max_tokens if args.max_tokens is not None else cfg.max_tokens
-
-    res = generate_from_ndjson(
-        args.ndjson,
-        provider,
-        launch_index=args.launch_index,
-        out_py=args.out,
-        execute=args.execute,
-        retries=args.retries,
-        temperature=temperature,
-        max_tokens=max_tokens,
-    )
-    print(res)
-
-
-if __name__ == "__main__":  # pragma: no cover
-    main()

tritonparse/reproducer/cli.py

@@ -1,37 +0,0 @@
-import argparse
-
-from .config import load_config
-from .factory import make_gemini_provider
-from .orchestrator import generate_from_ndjson
-
-
-def add_reproducer_subparser(parser: argparse.ArgumentParser) -> None:
-    sub = parser.add_subparsers(dest="subcommand")
-    repro = sub.add_parser(
-        "repro",
-        help="Generate a runnable Triton repro script from NDJSON",
-    )
-    repro.add_argument("--ndjson", required=True)
-    repro.add_argument("--launch-index", type=int, default=0)
-    repro.add_argument("--out", default="repro.py")
-    repro.add_argument("--execute", action="store_true")
-    repro.add_argument("--retries", type=int, default=0)
-
-
-def maybe_handle_reproducer(args: argparse.Namespace) -> bool:
-    if getattr(args, "subcommand", None) != "repro":
-        return False
-    cfg = load_config()
-    provider = make_gemini_provider()
-    res = generate_from_ndjson(
-        args.ndjson,
-        provider,
-        launch_index=args.launch_index,
-        out_py=args.out,
-        execute=args.execute,
-        retries=args.retries,
-        temperature=cfg.temperature,
-        max_tokens=cfg.max_tokens,
-    )
-    print(res)
-    return True

tritonparse/reproducer/config.py

@@ -1,15 +0,0 @@
-import os
-from dataclasses import dataclass
-
-
-@dataclass
-class ReproducerConfig:
-    project: str = os.getenv("GOOGLE_CLOUD_PROJECT", "")
-    location: str = os.getenv("GOOGLE_LOCATION", "us-central1")
-    model: str = os.getenv("TP_REPRO_MODEL", "gemini-2.5-pro")
-    temperature: float = float(os.getenv("TP_REPRO_TEMPERATURE", "0.1"))
-    max_tokens: int = int(os.getenv("TP_REPRO_MAX_TOKENS", "10240"))
-
-
-def load_config() -> ReproducerConfig:
-    return ReproducerConfig()

tritonparse/reproducer/factory.py

@@ -1,16 +0,0 @@
-"""Provider factory for reproducer.
-
-Currently supports Gemini only.
-"""
-
-from .config import load_config
-from .providers.gemini import GeminiProvider
-
-
-def make_gemini_provider() -> GeminiProvider:
-    cfg = load_config()
-    return GeminiProvider(
-        project=cfg.project,
-        location=cfg.location,
-        model=cfg.model,
-    )

tritonparse/reproducer/ingestion/__init__.py

@@ -1,6 +0,0 @@
-"""Ingestion utilities for reproducer.
-
-Currently supports NDJSON trace parsing.
-"""
-
-__all__ = []

tritonparse/reproducer/ingestion/ndjson.py

@@ -1,165 +0,0 @@
-import json
-from typing import Any, Dict, List
-
-
-def _iter_events(path: str):
-    with open(path, "r", encoding="utf-8") as f:
-        for line in f:
-            line = line.strip()
-            if not line:
-                continue
-            try:
-                yield json.loads(line)
-            except json.JSONDecodeError:
-                # skip malformed lines
-                continue
-
-
-def _index_compilations(events: List[Dict[str, Any]]) -> Dict[str, Dict[str, Any]]:
-    idx = {}
-    for e in events:
-        if e.get("event_type") != "compilation":
-            continue
-        payload = e.get("payload") or {}
-        meta = payload.get("metadata") or {}
-        h = meta.get("hash")
-        if h:
-            idx[h] = e
-    return idx
-
-
-def _get_launches(events: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
-    return [e for e in events if e.get("event_type") == "launch"]
-
-
-def _resolve_kernel_source(
-    launch: Dict[str, Any], comp_idx: Dict[str, Dict[str, Any]]
-) -> str:
-    # In new format, launch has top-level compilation_metadata, not payload.*
-    comp_meta = (
-        launch.get("compilation_metadata")
-        or launch.get("payload", {}).get("compilation_metadata")
-        or {}
-    )
-    h = comp_meta.get("hash")
-    if not h:
-        return ""
-    comp = comp_idx.get(h, {})
-    payload = comp.get("payload") or {}
-    py = payload.get("python_source") or {}
-    return py.get("code", "")
-
-
-def _pack_args(args: Dict[str, Any]) -> Dict[str, Any]:
-    packed = {}
-    for k, v in args.items():
-        t = v.get("type") if isinstance(v, dict) else None
-        if t == "tensor":
-            packed[k] = {
-                "type": "tensor",
-                "shape": v.get("shape") if isinstance(v, dict) else None,
-                "dtype": v.get("dtype") if isinstance(v, dict) else None,
-                "device": v.get("device") if isinstance(v, dict) else None,
-                "stride": v.get("stride") if isinstance(v, dict) else None,
-                "is_contiguous": (
-                    v.get("is_contiguous") if isinstance(v, dict) else None
-                ),
-                "numel": v.get("numel") if isinstance(v, dict) else None,
-            }
-        else:
-            # scalar / NoneType etc
-            if isinstance(v, dict):
-                packed[k] = {
-                    "type": v.get("type"),
-                    "value": v.get("value", v.get("repr")),
-                }
-            else:
-                packed[k] = {
-                    "type": None,
-                    "value": v,
-                }
-    return packed
-
-
-# Sentinel and helper to normalize extracted argument values
-_SKIP = object()
-
-
-def _decode_arg(raw: Any):
-    if not isinstance(raw, dict):
-        return raw
-    t = raw.get("type")
-    if t == "tensor":
-        return _SKIP
-    if t == "NoneType":
-        return None
-    return raw.get("value", raw.get("repr"))
-
-
-def build_context_bundle(ndjson_path: str, launch_index: int = 0) -> Dict[str, Any]:
-    events = list(_iter_events(ndjson_path))
-    launches = _get_launches(events)
-    if not launches:
-        raise RuntimeError("No launch events found in NDJSON.")
-    if launch_index < 0 or launch_index >= len(launches):
-        raise IndexError(
-            f"launch_index out of range: {launch_index} (total {len(launches)})"
-        )
-    launch = launches[launch_index]
-    comp_idx = _index_compilations(events)
-    kernel_source = _resolve_kernel_source(launch, comp_idx)
-    # find '@triton.jit' and slice the string
-    jit_marker = "@triton.jit"
-    jit_pos = kernel_source.find(jit_marker)
-    if jit_pos != -1:
-        kernel_source = kernel_source[jit_pos:]
-
-    # flatten launch fields (support both formats)
-    grid = launch.get("grid") or (launch.get("payload", {})).get("grid")
-    comp_meta = (
-        launch.get("compilation_metadata")
-        or (launch.get("payload", {})).get("compilation_metadata")
-        or {}
-    )
-    extracted_args = (
-        launch.get("extracted_args")
-        or (launch.get("payload", {})).get("extracted_args")
-        or {}
-    )
-
-    # compile metadata subset we care about
-    compile_block = {
-        "num_warps": comp_meta.get("num_warps"),
-        "num_stages": comp_meta.get("num_stages"),
-        "arch": comp_meta.get("arch"),
-        "backend": comp_meta.get("backend_name") or comp_meta.get("backend"),
-        "triton_version": comp_meta.get("triton_version"),
-        "hash": comp_meta.get("hash"),
-    }
-
-    # kwargs: include constexpr + explicit scalars used for launch (skip tensor args)
-    kwargs = {}
-    for k, v in extracted_args.items():
-        val = _decode_arg(v)
-        if val is _SKIP:
-            continue
-        kwargs[k] = val
-
-    # tensor args: only tensors
-    tensor_args = {
-        k: v
-        for k, v in extracted_args.items()
-        if isinstance(v, dict) and v.get("type") == "tensor"
-    }
-
-    bundle = {
-        "kernel_source": kernel_source,
-        "compile": compile_block,
-        "launch": {
-            "grid": grid,
-            "kwargs": kwargs,
-        },
-        "args": _pack_args(extracted_args),
-        "tensor_args": _pack_args(tensor_args),
-    }
-    return bundle

tritonparse/reproducer/orchestrator.py

@@ -1,65 +0,0 @@
-from pathlib import Path
-from typing import Any, Dict
-
-from .ingestion.ndjson import build_context_bundle
-from .param_generator import generate_allocation_snippet, generate_kwargs_dict
-from .prompts.loader import render_prompt
-from .providers.base import LLMProvider
-from .runtime.executor import run_python
-
-
-def _excerpt(s: str, n: int = 160):
-    lines = s.splitlines()
-    return "\n".join(lines[:n])
-
-
-def generate_from_ndjson(
-    ndjson_path: str,
-    provider: LLMProvider,
-    *,
-    launch_index=0,
-    out_py="repro.py",
-    execute=False,
-    retries: int = 0,
-    **gen_kwargs,
-) -> Dict[str, Any]:
-    bundle = build_context_bundle(ndjson_path, launch_index=launch_index)
-    # Augment bundle with pre-generated parameter allocation code to reduce LLM burden
-    allocation_snippet = generate_allocation_snippet(bundle)
-    kwargs_dict = generate_kwargs_dict(bundle)
-    context = {
-        **bundle,
-        "allocation_snippet": allocation_snippet,
-        "kwargs_dict": kwargs_dict,
-    }
-    system_prompt = render_prompt("system.txt", context)
-    user_prompt = render_prompt("generate_one_shot.txt", context)
-
-    code = provider.generate_code(system_prompt, user_prompt, **gen_kwargs)
-    Path(out_py).write_text(code, encoding="utf-8")
-
-    if not execute:
-        return {"path": out_py}
-
-    # Execute and optionally repair
-    rc, out, err = run_python(out_py)
-    attempt = 0
-    while rc != 0 and attempt < retries:
-        attempt += 1
-        # Build repair prompt
-        repair_ctx = {
-            "prev_code_excerpt": _excerpt(code, 200),
-            "error_text": err[-4000:] if err else "(no stderr)",
-        }
-        repair_prompt = render_prompt("repair_loop.txt", repair_ctx)
-        code = provider.generate_code(system_prompt, repair_prompt, **gen_kwargs)
-        Path(out_py).write_text(code, encoding="utf-8")
-        rc, out, err = run_python(out_py)
-
-    return {
-        "path": out_py,
-        "returncode": rc,
-        "stdout": out,
-        "stderr": err,
-        "retries_used": attempt,
-    }

tritonparse/reproducer/param_generator.py

@@ -1,142 +0,0 @@
-"""Parameter generator: produce deterministic allocation code from a bundle.
-
-This module reduces LLM burden by emitting Python code that:
-- selects a device
-- seeds RNG
-- allocates tensors with the exact shape/dtype/device/stride
-- prepares scalar/constexpr kwargs
-
-The generated code is intended to be inserted into the final repro script.
-"""
-
-import json
-from typing import Any, Dict, List, Optional
-
-
-def _torch_dtype_expr(dtype: str) -> str:
-    mapping = {
-        "float16": "torch.float16",
-        "bfloat16": "torch.bfloat16",
-        "float32": "torch.float32",
-        "float": "torch.float32",
-        "float64": "torch.float64",
-        "half": "torch.float16",
-        "bf16": "torch.bfloat16",
-        "fp16": "torch.float16",
-        "fp32": "torch.float32",
-        "fp64": "torch.float64",
-        "int8": "torch.int8",
-        "int16": "torch.int16",
-        "int32": "torch.int32",
-        "int64": "torch.int64",
-        "long": "torch.int64",
-        "bool": "torch.bool",
-    }
-    return mapping.get(str(dtype).lower(), "torch.float32")
-
-
-def _compute_storage_numel(shape: List[int], stride: Optional[List[int]]) -> int:
-    if not shape:
-        return 1
-    if not stride:
-        # contiguous default
-        numel = 1
-        for s in shape:
-            numel *= int(s)
-        return numel
-    # minimal storage size (in elements) to support the given logical shape/stride
-    max_index = 0
-    for dim, (sz, st) in enumerate(zip(shape, stride)):
-        if sz <= 0:
-            continue
-        max_index = max(max_index, (int(sz) - 1) * int(st))
-    return int(max_index) + 1
-
-
-def _emit_tensor_alloc(name: str, spec: Dict[str, Any]) -> str:
-    shape = spec.get("shape") or []
-    dtype = _torch_dtype_expr(spec.get("dtype"))
-    device = spec.get("device") or "cuda:0"
-    stride = spec.get("stride")
-
-    # ensure ints
-    shape = [int(s) for s in shape]
-    if stride is not None:
-        stride_list = [int(x) for x in stride]
-    else:
-        stride_list = None
-
-    lines: List[str] = []
-    # allocate backing storage
-    storage_numel = _compute_storage_numel(shape, stride_list)
-    lines.append(
-        f"# {name}: shape={shape}, dtype={dtype}, device={device}, stride={stride_list}"
-    )
-    lines.append(
-        f"_storage_{name} = torch.empty(({storage_numel},), dtype={dtype}, device=device)"
-    )
-    if stride_list:
-        # Create an as_strided view over the 1D storage
-        sizes_expr = str(tuple(shape))
-        strides_expr = str(tuple(stride_list))
-        lines.append(
-            f"{name} = _storage_{name}.as_strided(size={sizes_expr}, stride={strides_expr})"
-        )
-    else:
-        # contiguous allocation
-        size_expr = str(tuple(shape))
-        lines.append(f"{name} = torch.empty({size_expr}, dtype={dtype}, device=device)")
-    return "\n".join(lines)
-
-
-def _emit_scalar(name: str, spec: Dict[str, Any]) -> str:
-    value = spec.get("value")
-    # Preserve JSON-serializable value as-is
-    return f"{name} = {json.dumps(value)}"
-
-
-def generate_allocation_snippet(bundle: Dict[str, Any]) -> str:
-    """Generate a self-contained code snippet that:
-    - imports torch
-    - sets device
-    - seeds RNG
-    - allocates tensors and defines scalars for all args
-    Returns Python source as a string.
-    """
-    tensor_args: Dict[str, Any] = bundle.get("tensor_args", {}) or {}
-    args_all: Dict[str, Any] = bundle.get("args", {}) or {}
-
-    # Pick device from any tensor arg, fallback to cuda:0
-    device = "cuda:0"
-    for spec in tensor_args.values():
-        dev = spec.get("device")
-        if dev:
-            device = str(dev)
-            break
-
-    lines: List[str] = []
-    lines.append("import torch")
-    lines.append(f"device = '{device}'")
-    lines.append("torch.manual_seed(0)")
-    lines.append("if torch.cuda.is_available(): torch.cuda.manual_seed_all(0)")
-    lines.append("")
-
-    # Emit tensors first for names with type==tensor in args_all
-    for name, spec in args_all.items():
-        if isinstance(spec, dict) and spec.get("type") == "tensor":
-            lines.append(_emit_tensor_alloc(name, spec))
-            lines.append("")
-
-    # Emit non-tensor scalars next
-    for name, spec in args_all.items():
-        if not isinstance(spec, dict) or spec.get("type") == "tensor":
-            continue
-        lines.append(_emit_scalar(name, spec))
-    return "\n".join(lines)
-
-
-def generate_kwargs_dict(bundle: Dict[str, Any]) -> Dict[str, Any]:
-    """Return a kwargs dict derived from bundle['launch']['kwargs'] suitable for kernel call."""
-    launch = bundle.get("launch", {}) or {}
-    kwargs = launch.get("kwargs", {}) or {}
-    return kwargs

tritonparse/reproducer/prompts/__init__.py

@@ -1 +0,0 @@
-__all__ = []

tritonparse/reproducer/prompts/loader.py

@@ -1,18 +0,0 @@
-import json
-from pathlib import Path
-from typing import Any, Dict
-
-PROMPTS_DIR = Path(__file__).parent
-
-
-def render_prompt(name: str, context: Dict[str, Any]) -> str:
-    text = (PROMPTS_DIR / name).read_text(encoding="utf-8")
-    # very simple {{key}} replacement for top-level keys; JSON for dicts
-    for k, v in context.items():
-        token = "{{ " + k + " }}"
-        if token in text:
-            if isinstance(v, (dict, list)):
-                text = text.replace(token, json.dumps(v, ensure_ascii=False, indent=2))
-            else:
-                text = text.replace(token, str(v))
-    return text

tritonparse/reproducer/providers/__init__.py

@@ -1 +0,0 @@
-__all__ = []

tritonparse/reproducer/providers/base.py

@@ -1,14 +0,0 @@
-from typing import Any, Dict, List, Optional, Protocol
-
-
-class LLMProvider(Protocol):
-    def generate_code(
-        self,
-        system_prompt: str,
-        user_prompt: str,
-        *,
-        temperature: float = 0.2,
-        max_tokens: int = 8192,
-        stop: Optional[List[str]] = None,
-        extra: Optional[Dict[str, Any]] = None,
-    ) -> str: ...

tritonparse/reproducer/providers/gemini.py

@@ -1,47 +0,0 @@
-import os
-import re
-from typing import Any, Dict, List, Optional
-
-from google.genai import Client
-
-
-def _extract_python_block(s: str) -> str:
-    m = re.search(r"""```python\s+(.*?)```""", s, flags=re.S)
-    return m.group(1).strip() if m else ""
-
-
-class GeminiProvider:
-    def __init__(
-        self, project: str, location: str = "us-central1", model: str = "gemini-2.5-pro"
-    ):
-        # Expect GOOGLE_APPLICATIONS_CREDENTIALS to be set
-        if not os.getenv("GOOGLE_APPLICATION_CREDENTIALS"):
-            raise EnvironmentError("GOOGLE_APPLICATION_CREDENTIALS not set.")
-        self.client = Client(vertexai=True, project=project, location=location)
-        self.model = model
-
-    def generate_code(
-        self,
-        system_prompt: str,
-        user_prompt: str,
-        *,
-        temperature: float = 0.2,
-        max_tokens: int = 8192,
-        stop: Optional[List[str]] = None,
-        extra: Optional[Dict[str, Any]] = None,
-    ) -> str:
-        # Gemini doesn't have a 'system' role in this SDK, prepend system to user
-        full_prompt = f"{system_prompt.strip()}\n\n---\n\n{user_prompt.strip()}"
-        resp = self.client.models.generate_content(
-            model=self.model,
-            contents=full_prompt,
-            config={
-                "temperature": temperature,
-                "max_output_tokens": max_tokens,
-            },
-        )
-        text = getattr(resp, "text", "") or ""
-        code = _extract_python_block(text) or text
-        if not code.strip():
-            raise RuntimeError(f"Empty response from Gemini. Raw: {text[:2000]}")
-        return code

tritonparse/reproducer/runtime/__init__.py

@@ -1 +0,0 @@
-__all__ = []

tritonparse/reproducer/runtime/executor.py

@@ -1,13 +0,0 @@
-import subprocess
-import sys
-
-
-def run_python(path: str, timeout: int = 60):
-    p = subprocess.Popen(
-        [sys.executable, path],
-        stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE,
-        text=True,
-    )
-    out, err = p.communicate(timeout=timeout)
-    return p.returncode, out, err

tritonparse/reproducer/utils/io.py

@@ -1,6 +0,0 @@
-from pathlib import Path
-
-
-def write_text(path: str, content: str, *, encoding="utf-8"):
-    Path(path).parent.mkdir(parents=True, exist_ok=True)
-    Path(path).write_text(content, encoding=encoding)

tritonparse-0.1.1.dist-info/METADATA

@@ -1,10 +0,0 @@
-Metadata-Version: 2.4
-Name: tritonparse
-Version: 0.1.1
-Project-URL: Homepage, https://github.com/meta-pytorch/tritonparse
-Requires-Python: >=3.10
-License-File: LICENSE
-Requires-Dist: triton>3.3.1
-Provides-Extra: test
-Requires-Dist: coverage>=7.0.0; extra == "test"
-Dynamic: license-file

tritonparse-0.1.1.dist-info/RECORD

@@ -1,40 +0,0 @@
-tritonparse/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tritonparse/common.py,sha256=VWKFZJG7msIExMH0QNYCO8SHKqUBFtdRuLaCy_K8DFI,13725
-tritonparse/event_diff.py,sha256=yOD6uNxLJroatfx2nEGr-erw24ObOrHU9P6V5pzr8do,4907
-tritonparse/extract_source_mappings.py,sha256=Z6UxFj2cCE5NCWLQTYPKqUpLfbYhqP8xgCl5mvud9KI,1451
-tritonparse/ir_parser.py,sha256=1j1tP9jpUN7wH3e01bKUkUPgTMlNXUdp8LKRCC-WTro,9324
-tritonparse/mapper.py,sha256=prrczfi13P7Aa042OrEBsmRF1HW3jDhwxicANgPkWIM,4150
-tritonparse/shared_vars.py,sha256=-c9CvXJSDm9spYhDOJPEQProeT_xl3PaNmqTEYi_u4s,505
-tritonparse/source_type.py,sha256=nmYEQS8rfkIN9BhNhQbkmEvKnvS-3zAxRGLY4TaZdi8,1676
-tritonparse/sourcemap_utils.py,sha256=qsQmTDuEe9yuUVyxSHRbjTR38gi0hvJEijnPkrJVAV4,2037
-tritonparse/structured_logging.py,sha256=qWGkAr2oZT8Adxj-EfQqQdX2l-jq9xKi7WhBecYq2bg,39600
-tritonparse/tp_logger.py,sha256=vXzY7hMDmVnRBGBhIjFZe3nHZzG5NKKPONGUszJhGgU,242
-tritonparse/trace_processor.py,sha256=QzUOKwnOkBbwTTKBsa5ZMUABPLMJIBFtTcG2SkhO0I8,12771
-tritonparse/utils.py,sha256=wt61tpbkqjGqHh0c7Nr2WlOv7PbQssmjULd6uA6aAko,4475
-tritonparse/reproducer/__init__.py,sha256=VcCpYVUUmclWotkQmPLlDu5iFOUE4N-4FzcbzXwIow0,773
-tritonparse/reproducer/__main__.py,sha256=ydLMbWx7SFlvAb1erObvKcJ-uxhHNShGRHRrZO6-5ww,2266
-tritonparse/reproducer/cli.py,sha256=nnMgdT4tQzBWjPowOy2a_5QRsnsTMEAA3uegYpLEyRE,1165
-tritonparse/reproducer/config.py,sha256=-hmE5ZqEtYo2WKjXbwMi6k6XzzfQZAbL50UURPvcF3A,478
-tritonparse/reproducer/factory.py,sha256=sFcIjIayfHAqPqMVT8Rnsz9tpMmQXBzoOlKprS1P_1g,341
-tritonparse/reproducer/orchestrator.py,sha256=iD7zZZHE4FU3nNOwNV9SUY2WUcpv_Amg0SvnRxrseEQ,2045
-tritonparse/reproducer/param_generator.py,sha256=m-C_Z1TLd1ZX49EpsWELVfB6tkwOfi-ZHma7wXwz2g4,4654
-tritonparse/reproducer/ingestion/__init__.py,sha256=2AQHxWlUl5JXM4a8F033wzxVnjCVPBEf-4H99kep-OA,99
-tritonparse/reproducer/ingestion/ndjson.py,sha256=_E_dXXjxu438OYomQ1zFFk3jV9Wr1jNoXHiP2gJG7_4,5172
-tritonparse/reproducer/prompts/__init__.py,sha256=da1PTClDMl-IBkrSvq6JC1lnS-K_BASzCvxVhNxN5Ls,13
-tritonparse/reproducer/prompts/loader.py,sha256=n6Of98eEXNz9mI7ZH073X5FihNZD7tI-ehfjN_4yEl0,610
-tritonparse/reproducer/providers/__init__.py,sha256=da1PTClDMl-IBkrSvq6JC1lnS-K_BASzCvxVhNxN5Ls,13
-tritonparse/reproducer/providers/base.py,sha256=DgP_4AdrEf48kstOfBJFvK3pndcHH0vRUGjp6k1bdsY,362
-tritonparse/reproducer/providers/gemini.py,sha256=VlOCdTGRTQdr3c2HMclKFIk-133puGSAjhK_6m6Zj9g,1609
-tritonparse/reproducer/runtime/__init__.py,sha256=da1PTClDMl-IBkrSvq6JC1lnS-K_BASzCvxVhNxN5Ls,13
-tritonparse/reproducer/runtime/executor.py,sha256=AqBFnoEqURoMGDdLC2G3WpHIP3Y4wWGJHEZrjS-NQFM,304
-tritonparse/reproducer/utils/io.py,sha256=95NF9QCGawl-5p5c5yCQHynVBNKS_B_7nIrqnRvAt-E,200
-tritonparse/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tritonparse/tools/decompress_bin_ndjson.py,sha256=kpt7DM_sSA334F1X45xdkP2OR9LuB27Pc50EkGr6CPM,4144
-tritonparse/tools/format_fix.py,sha256=Ol0Sjui8D7OzHwbamAfGnq8V5Y63uwNaFTKSORN5HkQ,3867
-tritonparse/tools/load_tensor.py,sha256=tfdmNVd9gsZqO6msQBhbXIhOvUzgc83yF64k2GDWPNk,2122
-tritonparse/tools/prettify_ndjson.py,sha256=VOzVWoXpCbaAXYA4i_wBcQIHfh-JhAx7xR4cF_L8yDs,10928
-tritonparse-0.1.1.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
-tritonparse-0.1.1.dist-info/METADATA,sha256=qETEInGJRT7fzf-Rl8cAf6QnEa5eJgiYo4rbwBA63yc,287
-tritonparse-0.1.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tritonparse-0.1.1.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
-tritonparse-0.1.1.dist-info/RECORD,,