tritonparse 0.1.1__tar.gz

tritonparse-0.1.1/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2019, pytorch
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

tritonparse-0.1.1/PKG-INFO

@@ -0,0 +1,10 @@
+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/README.md

@@ -0,0 +1,124 @@
+# 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.1.1/pyproject.toml

@@ -0,0 +1,34 @@
+[build-system]
+requires = ["setuptools>=40.8.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "tritonparse"
+version = "0.1.1"
+dependencies = [
+    "triton>3.3.1",
+]
+requires-python = ">=3.10"
+
+[project.optional-dependencies]
+test = [
+    "coverage>=7.0.0",
+]
+
+
+[tool.setuptools.packages.find]
+include = ["tritonparse*"]
+
+[tool.black]
+line-length = 88
+target-version = ["py310"]
+
+[tool.ufmt]
+formatter = "black"
+sorter = "usort"
+
+[tool.usort]
+first_party_detection = false
+
+[project.urls]
+"Homepage" = "https://github.com/meta-pytorch/tritonparse"

tritonparse-0.1.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

tritonparse-0.1.1/tests/test_add.py

@@ -0,0 +1,81 @@
+"""
+Simple Triton kernel for tensor addition. This file is not included in unit tests.
+
+Test Plan:
+```
+TORCHINDUCTOR_FX_GRAPH_CACHE=0 TRITONPARSE_DEBUG=1 python tests/test_add.py
+```
+"""
+
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+import tritonparse.structured_logging
+import tritonparse.utils
+
+log_path = "./logs"
+tritonparse.structured_logging.init(log_path, enable_trace_launch=True)
+
+os.environ["TORCHINDUCTOR_FX_GRAPH_CACHE"] = "0"
+
+
+@triton.jit
+def add_kernel(
+    a_ptr,
+    b_ptr,
+    c_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    a = tl.load(a_ptr + offsets, mask=mask)
+    b = tl.load(b_ptr + offsets, mask=mask)
+    c = a + b
+    tl.store(c_ptr + offsets, c, mask=mask)
+
+
+def tensor_add(a, b):
+    n_elements = a.numel()
+    c = torch.empty_like(a)
+    BLOCK_SIZE = 1024
+    grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
+    add_kernel[grid](a, b, c, n_elements, BLOCK_SIZE)
+    return c
+
+
+def simple_add(a, b):
+    return a + b
+
+
+def test_tensor_add():
+    torch.manual_seed(0)
+    size = (1024, 1024)
+    a = torch.randn(size, device="cuda", dtype=torch.float32)
+    b = torch.randn(size, device="cuda", dtype=torch.float32)
+
+    # Test Triton kernel
+    c_triton = tensor_add(a, b)
+    c_triton.sum()
+    tensor_add(a, b)
+    print("Triton kernel executed successfully")
+
+    # Test torch.compile
+    compiled_add = torch.compile(simple_add)
+    c_compiled = compiled_add(a, b)
+    c_compiled.sum()
+    print("Torch compiled function executed successfully")
+
+
+if __name__ == "__main__":
+    test_tensor_add()
+    # Use improved unified_parse with explicit output directory
+    tritonparse.utils.unified_parse(
+        source=log_path, out="./parsed_output", overwrite=True
+    )

tritonparse-0.1.1/tests/test_add2.py

@@ -0,0 +1,81 @@
+"""
+Simple Triton kernel for tensor addition. This file is not included in unit tests.
+
+Test Plan:
+```
+TORCHINDUCTOR_FX_GRAPH_CACHE=0 TRITONPARSE_DEBUG=1 python tests/test_add.py
+```
+"""
+
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+import tritonparse.structured_logging
+import tritonparse.utils
+
+log_path = "./logs"
+tritonparse.structured_logging.init(log_path, enable_trace_launch=True)
+
+os.environ["TORCHINDUCTOR_FX_GRAPH_CACHE"] = "0"
+
+
+@triton.jit
+def add_kernel(
+    a_ptr,
+    b_ptr,
+    c_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    a = tl.load(a_ptr + offsets, mask=mask)
+    b = tl.load(b_ptr + offsets, mask=mask)
+    c = a + b
+    tl.store(c_ptr + offsets, c, mask=mask)
+
+
+def tensor_add(a, b):
+    n_elements = a.numel()
+    c = torch.empty_like(a)
+    BLOCK_SIZE = 1024
+    grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
+    add_kernel[grid](a, b, c, n_elements, BLOCK_SIZE)
+    return c
+
+
+def simple_add(a, b):
+    return a + b
+
+
+def test_tensor_add():
+    torch.manual_seed(0)
+    size = (1024, 1024)
+    a = torch.randn(size, device="cuda", dtype=torch.float32)
+    b = torch.randn(size, device="cuda", dtype=torch.float32)
+
+    # # Test Triton kernel
+    # c_triton = tensor_add(a, b)
+    # c_triton.sum()
+    # tensor_add(a, b)
+    # print("Triton kernel executed successfully")
+
+    # Test torch.compile
+    compiled_add = torch.compile(simple_add)
+    c_compiled = compiled_add(a, b)
+    c_compiled.sum()
+    print("Torch compiled function executed successfully")
+
+
+if __name__ == "__main__":
+    test_tensor_add()
+    # Use improved unified_parse with explicit output directory
+    tritonparse.utils.unified_parse(
+        source=log_path, out="./parsed_output", overwrite=True
+    )

tritonparse-0.1.1/tests/test_reproducer.py

@@ -0,0 +1,128 @@
+import json
+from pathlib import Path
+
+from tritonparse.reproducer import (
+    generate_allocation_snippet,
+    generate_from_ndjson,
+    generate_kwargs_dict,
+)
+
+
+class DummyProvider:
+    def generate_code(
+        self,
+        system_prompt: str,
+        user_prompt: str,
+        *,
+        temperature: float = 0.2,
+        max_tokens: int = 8192,
+        stop=None,
+        extra=None,
+    ) -> str:
+        # Return a tiny script that always runs successfully
+        return """
+print("dummy ok")
+"""
+
+
+def _write_minimal_ndjson(path: Path) -> None:
+    comp_event = {
+        "event_type": "compilation",
+        "payload": {
+            "metadata": {
+                "hash": "h1",
+                "num_warps": 4,
+                "num_stages": 2,
+                "arch": "sm_90",
+                "backend_name": "ptx",
+                "triton_version": "3.0.0",
+            },
+            "python_source": {
+                "code": """
+import triton
+import triton.language as tl
+
+@triton.jit
+def kernel(X_ptr):
+    pass
+""",
+            },
+        },
+    }
+    launch_event = {
+        "event_type": "launch",
+        "grid": [1],
+        "compilation_metadata": {
+            "hash": "h1",
+            "num_warps": 4,
+            "num_stages": 2,
+        },
+        "extracted_args": {
+            "X": {
+                "type": "tensor",
+                "shape": [8, 8],
+                "dtype": "float32",
+                "device": "cuda:0",
+                "stride": [8, 1],
+                "is_contiguous": True,
+                "numel": 64,
+            },
+            "BLOCK": {"type": "constexpr", "value": 128},
+        },
+    }
+    with path.open("w", encoding="utf-8") as f:
+        f.write(json.dumps(comp_event) + "\n")
+        f.write(json.dumps(launch_event) + "\n")
+
+
+def test_param_generator_snippet_basic(tmp_path):
+    # Make a pseudo bundle and verify snippet contains expected constructs
+    bundle = {
+        "tensor_args": {
+            "X": {
+                "type": "tensor",
+                "shape": [4, 4],
+                "dtype": "float32",
+                "device": "cuda:0",
+                "stride": [4, 1],
+                "is_contiguous": True,
+            }
+        },
+        "args": {
+            "X": {
+                "type": "tensor",
+                "shape": [4, 4],
+                "dtype": "float32",
+                "device": "cuda:0",
+                "stride": [4, 1],
+                "is_contiguous": True,
+            },
+            "N": {"type": "constexpr", "value": 4},
+        },
+    }
+    snippet = generate_allocation_snippet(bundle)
+    assert "torch.empty" in snippet
+    assert "device = 'cuda:0'" in snippet
+    kwargs = generate_kwargs_dict({"launch": {"kwargs": {"N": 4}}})
+    assert kwargs == {"N": 4}
+
+
+def test_orchestrator_with_dummy_provider(tmp_path):
+    ndjson = tmp_path / "trace.ndjson"
+    _write_minimal_ndjson(ndjson)
+
+    out_py = tmp_path / "repro.py"
+    res = generate_from_ndjson(
+        str(ndjson),
+        provider=DummyProvider(),
+        launch_index=0,
+        out_py=str(out_py),
+        execute=True,
+        retries=0,
+        temperature=0.0,
+        max_tokens=256,
+    )
+    assert out_py.exists()
+    assert res.get("returncode", 0) == 0
+    # stdout should contain our dummy output
+    assert "dummy ok" in (res.get("stdout") or "")