Skip to content

04. Developer Guide

FindHao edited this page Oct 5, 2025 · 8 revisions

This guide is for developers who want to contribute to TritonParse, understand its architecture, or extend its functionality.

πŸ—οΈ Architecture Overview

High-Level Architecture

TritonParse consists of three main components:

β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”    β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”    β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
β”‚   Python Backend     β”‚    β”‚   Processing         β”‚    β”‚   Frontend UI        β”‚
β”‚                      β”‚    β”‚                      β”‚    β”‚                      β”‚
β”‚ β€’ Structured Logging │──▢│ β€’ Log Parsing        │──▢│ β€’ React Interface    β”‚
β”‚ β€’ Triton Hooks       β”‚    β”‚ β€’ Source Mapping     β”‚    β”‚ β€’ IR Visualization   β”‚
β”‚ β€’ Trace Generation   β”‚    β”‚ β€’ Data Compression   β”‚    β”‚ β€’ Code Comparison    β”‚
β”‚                      β”‚    β”‚ β€’ Launch Analysis    β”‚    β”‚ β€’ File Diff View     β”‚
β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜    β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜    β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
         β”‚                           β”‚                            
         β”‚                           β–Ό                            
         β”‚                  β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”             
         β”‚                  β”‚   Reproducer         β”‚             
         β”‚                  β”‚                      β”‚             
         └─────────────────▢│ β€’ Script Generation  β”‚             
                            β”‚ β€’ Tensor Rebuilding  β”‚             
                            β”‚ β€’ Template System    β”‚             
                            β”‚ β€’ Context Bundling   β”‚             
                            β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜             

Component Roles

Python Backend - Captures Triton compilation events and generates structured logs. Core logic for IR parsing, trace processing, and source mapping.

Processing Pipeline - Transforms raw NDJSON logs into structured, compressed data with source mappings and launch analysis.

Frontend UI - React/TypeScript web application for interactive visualization using Monaco Editor for code display.

Reproducer System - Generates standalone Python scripts to reproduce kernel executions from trace files with full context.

πŸ’‘ See Project Structure below for detailed file organization.

πŸ“ Project Structure

tritonparse/
β”œβ”€β”€ tritonparse/                 # Python package
β”‚   β”œβ”€β”€ __init__.py
β”‚   β”œβ”€β”€ __main__.py              # CLI entry point
β”‚   β”œβ”€β”€ run.py                   # Main CLI with argparse subcommands (parse, reproduce)
β”‚   β”œβ”€β”€ structured_logging.py    # Core logging infrastructure to capture events
β”‚   β”œβ”€β”€ trace_processor.py       # Processes raw trace files, groups events, and generates diffs
β”‚   β”œβ”€β”€ ir_parser.py             # Extracts source location information from various IRs
β”‚   β”œβ”€β”€ mapper.py                # Creates bidirectional mappings between different IRs and Python source code
β”‚   β”œβ”€β”€ event_diff.py            # Logic for comparing kernel launch events
β”‚   β”œβ”€β”€ utils.py                 # Main CLI entrypoint (`unified_parse`) and other parsing utilities
β”‚   β”œβ”€β”€ source_type.py           # Source type definitions (e.g., TTIR, PTX)
β”‚   β”œβ”€β”€ sourcemap_utils.py       # Helper functions for source mapping
β”‚   β”œβ”€β”€ common.py                # Common utilities and helper functions
β”‚   β”œβ”€β”€ shared_vars.py           # Shared state and variables for the package
β”‚   β”œβ”€β”€ tp_logger.py             # Logger configuration
β”‚   └── reproducer/              # Reproducer system for generating standalone scripts
β”‚       β”œβ”€β”€ __init__.py
β”‚       β”œβ”€β”€ orchestrator.py      # Main reproducer entry point
β”‚       β”œβ”€β”€ cli.py               # CLI argument handling for reproduce command
β”‚       β”œβ”€β”€ placeholder_replacer.py  # Template placeholder substitution
β”‚       β”œβ”€β”€ utils.py             # Tensor creation and path utilities
β”‚       β”œβ”€β”€ ingestion/           # Trace file parsing
β”‚       β”‚   β”œβ”€β”€ __init__.py
β”‚       β”‚   └── ndjson.py        # NDJSON trace parsing and context extraction
β”‚       β”œβ”€β”€ templates/           # Code generation templates
β”‚       β”‚   β”œβ”€β”€ __init__.py
β”‚       β”‚   β”œβ”€β”€ loader.py        # Template loading utilities
β”‚       β”‚   └── example.py       # Default reproducer template
β”‚       └── providers/           # Provider-specific integrations (if any)
β”œβ”€β”€ website/                     # React web application for visualization
β”‚   β”œβ”€β”€ src/
β”‚   β”‚   β”œβ”€β”€ components/          # Reusable React components
β”‚   β”‚   β”‚   β”œβ”€β”€ ArgumentViewer.tsx   # Displays kernel arguments
β”‚   β”‚   β”‚   β”œβ”€β”€ CodeViewer.tsx       # Displays IR code with Monaco editor
β”‚   β”‚   β”‚   β”œβ”€β”€ DiffViewer.tsx       # Side-by-side diff view for text
β”‚   β”‚   β”‚   β”œβ”€β”€ DiffComparisonView.tsx  # File diff comparison component
β”‚   β”‚   β”‚   β”œβ”€β”€ CodeComparisonView.tsx # Compares two IRs with line mappings
β”‚   β”‚   β”‚   └── AutotuneAnalysis.tsx # Displays autotuning results
β”‚   β”‚   β”œβ”€β”€ pages/               # Main application pages
β”‚   β”‚   β”‚   β”œβ”€β”€ KernelOverview.tsx # Main analysis view for a kernel, combining all components
β”‚   β”‚   β”‚   β”œβ”€β”€ CodeView.tsx     # A focused view for a single IR file
β”‚   β”‚   β”‚   └── FileDiffView.tsx # File diff page for cross-trace comparison
β”‚   β”‚   β”œβ”€β”€ context/             # React context providers
β”‚   β”‚   β”‚   └── FileDiffSession.tsx  # File diff state management
β”‚   β”‚   β”œβ”€β”€ utils/               # Utility functions
β”‚   β”‚   β”‚   └── dataLoader.ts    # Data loading and processing from parsed logs
β”‚   β”‚   β”œβ”€β”€ App.tsx              # Main application component routing to pages
β”‚   β”‚   └── main.tsx             # Application entry point
β”‚   β”œβ”€β”€ public/                  # Static assets (e.g., images, sample data)
β”‚   β”œβ”€β”€ package.json             # Frontend dependencies and scripts
β”‚   └── vite.config.ts           # Vite build configuration
β”œβ”€β”€ tests/                       # Test suite for the Python package
β”œβ”€β”€ docs/                        # Project documentation
β”œβ”€β”€ .github/                     # GitHub Actions workflows
β”œβ”€β”€ .ci/                         # CI scripts
β”œβ”€β”€ pyproject.toml               # Python project configuration
β”œβ”€β”€ Makefile                     # Development commands
└── README.md                    # Project overview

πŸ”§ Development Environment Setup

Prerequisites

  • Python >= 3.10
  • Node.js >= 22.0.0
  • Triton >= 3.4.0 (latest version recommended)
  • Git for version control

Installation Steps

# 1. Clone repository
git clone https://github.com/meta-pytorch/tritonparse.git
cd tritonparse

# 2. Install Python dependencies
make install-dev

# 3. Install website dependencies
cd website
npm install

Verify Installation

# Python: Check formatting and run tests
make format-check
make lint-check
python -m unittest tests.test_tritonparse.TestTritonparseCPU -v

# Website: Start dev server
cd website
npm run dev

πŸ› οΈ Development Workflow

Code Style and Formatting

We use a comprehensive formatting pipeline:

Tool Purpose Configuration
Black Code formatting pyproject.toml
usort Import sorting pyproject.toml
Ruff Linting Built-in rules

Essential Commands

# Format code
make format

# Check formatting
make format-check

# Run linting
make lint-check

# Run tests
python -m unittest tests.test_tritonparse -v

# Website development
cd website && npm run dev

Development Quality Checks

Before committing, ensure:

  1. Code is formatted: make format
  2. Linting passes: make lint-check
  3. Tests pass: python -m unittest tests.test_tritonparse -v
  4. Website builds: cd website && npm run build

πŸ—οΈ Backend Development

Core Components

Structured Logging (structured_logging.py) - Captures Triton compilation and launch events. Main functions: init() and init_with_env() for initialization.

Log Processing (utils.py) - Transforms raw logs into analyzable format. Entry point: unified_parse() for parsing NDJSON logs, extracting source mappings, and compressing data.

Source Mapping (extract_source_mappings.py) - Correlates lines between different IR stages (TTIR, TTGIR, LLIR, PTX, AMDGCN).

πŸ’‘ See inline code documentation for detailed function signatures and parameters.

Adding New Features

  1. Define the new data: Determine what new information needs to be captured.
  2. Update structured_logging.py: Add logic to capture the new data within the appropriate hooks (e.g., pre-compilation, post-compilation).
  3. Modify trace_processor.py: If the new data requires special processing or aggregation (like the launch analysis), add the logic here.
  4. Update unified_parse(): Ensure the new data is handled correctly during the main parsing routine.
  5. Write tests: Add unit and integration tests to tests/ to validate the new feature.

CLI Architecture

The CLI uses argparse with two main subcommands: parse and reproduce.

# Parse subcommand
tritonparse parse ./logs/ --out ./parsed_output

# Reproduce subcommand
tritonparse reproduce trace.ndjson --line 1 --out-dir repro_output

Implementation Files:

  • run.py - Main CLI with argparse subparsers
  • __main__.py - Entry point that calls run.main()
  • utils.py - _add_parse_args() function
  • reproducer/cli.py - _add_reproducer_args() function

πŸ’‘ See Usage Guide for complete CLI reference and parameters.

Reproducer System

Generates standalone Python scripts from trace files.

Core Components:

  • Orchestrator - Main reproduce() function, loads events and generates scripts
  • Ingestion - Extracts kernel context, arguments, and metadata from NDJSON
  • Templates - Customizable code templates (default: example.py)
  • Placeholder Replacement - Substitutes template placeholders for imports, invocations, paths
  • Utils - Rebuilds kernel arguments from JSON, handles tensor creation

Custom Templates Example:

# my_template.py
import torch
# {{KERNEL_IMPORT_PLACEHOLDER}}

if __name__ == "__main__":
    # {{KERNEL_INVOCATION_PLACEHOLDER}}
    print("Custom execution!")

Use with: tritonparse reproduce trace.ndjson --line 1 --template my_template.py

πŸ’‘ See Testing section for running tests.

🎨 Frontend Development

Technology Stack

  • React 19 - UI framework
  • TypeScript - Type safety
  • Vite - Build tool and dev server
  • Tailwind CSS - Styling
  • Monaco Editor - IDE-quality code display

Key Components

Data Loading (utils/dataLoader.ts) - Loads and processes trace files from URLs or local files.

Code Viewer (components/CodeViewer.tsx) - Displays IR code with syntax highlighting, line numbers, and source mapping.

IR Code View (components/CodeComparisonView.tsx) - Side-by-side IR viewing with synchronized scrolling and interactive line mapping.

File Diff View (pages/FileDiffView.tsx) - Cross-trace kernel comparison with customizable diff options.

Adding Frontend Features

  1. Update data loader - Modify dataLoader.ts for new data fields
  2. Create components - Add React components in website/src/components/
  3. Integrate - Add components to pages (e.g., KernelOverview.tsx)
  4. Style - Use Tailwind CSS for consistent styling
  5. Test - Verify with npm run build and manual testing

Testing Frontend Changes

cd website

# Development server
npm run dev

# Type checking
npm run build

# Linting
npm run lint

# Test with sample data
# Load ./public/f0_fc0_a0_cai-.ndjson in browser

πŸ“Š Data Flow

End-to-End Data Flow

Python Code
     β”‚
     β–Ό
Triton Compilation
(triggers Hook Events)
     β”‚
     β–Ό
Structured Logging
     β”‚
     β–Ό
Raw NDJSON Logs
     β”‚
     β–Ό
Log Processing
  - Source Mapping
  - Launch Analysis
     β”‚
     β–Ό
Compressed Data (.gz)
     β”‚
     β–Ό
Web Interface
     β”‚
     β–Ό
Interactive Visualization

πŸ’‘ Data Format: Events are logged as NDJSON, processed into structured format with IR files and source mappings, then compressed as .ndjson.gz for the web interface.

πŸ” Debugging and Development Tools

Debug Logging

# Enable debug logging
export TRITONPARSE_DEBUG=1

# Run with debug output
python your_script.py

Development Utilities

# Check log file contents
head -n 10 ./logs/*.ndjson

# Inspect compressed data
zcat ./parsed_output/*.gz | head -n 20

# Test parsing pipeline
python -c "
import tritonparse.utils
tritonparse.utils.unified_parse('./logs/', './test_output/', verbose=True)
"

Browser Developer Tools

// Enable frontend debug logging
localStorage.setItem('tritonparse-debug', 'true');

// Inspect loaded data
console.log(window.tritonparseData);

// Test data processing
import { processKernelData } from './utils/dataLoader';
console.log(processKernelData(rawData));

πŸ§ͺ Testing

Test Structure

tests/
β”œβ”€β”€ test_tritonparse.py         # Main test suite
β”œβ”€β”€ test_add.py                 # Manual test example
└── example_output/             # Sample data

Running Tests

# All tests
python -m unittest tests.test_tritonparse -v

# CPU-only tests
python -m unittest tests.test_tritonparse.TestTritonparseCPU -v

# GPU tests (requires CUDA)
python -m unittest tests.test_tritonparse.TestTritonparseCUDA -v

# Manual test
cd tests
TORCHINDUCTOR_FX_GRAPH_CACHE=0 python test_add.py

Writing Tests

Follow this workflow when adding end-to-end tests in TestTritonparseCUDA:

  1. Define test method - Create method starting with test_
  2. Define Triton kernel - Write kernel to test feature
  3. Setup environment - Create temp directories with tempfile.mkdtemp()
  4. Initialize logging - Call tritonparse.structured_logging.init()
  5. Run kernel - Execute to generate compilation/launch events
  6. Parse logs - Call tritonparse.utils.unified_parse()
  7. Assert results - Verify output files and contents
  8. Cleanup - Use try...finally to remove temp directories

Example Structure:

@unittest.skipUnless(torch.cuda.is_available(), "CUDA not available")
def test_new_feature(self):
    temp_dir = tempfile.mkdtemp()
    tritonparse.structured_logging.init(temp_dir + "/logs", enable_trace_launch=True)
    try:
        # Run kernel and generate logs
        kernel[(grid,)](args, BLOCK_SIZE=128)
        
        # Parse and verify
        tritonparse.utils.unified_parse(source=log_path, out=parsed_path)
        self.assertGreater(len(os.listdir(parsed_path)), 0)
    finally:
        shutil.rmtree(temp_dir)
        tritonparse.structured_logging.clear_logging_config()

πŸ’‘ See existing tests in tests/test_tritonparse.py for complete examples.

πŸ“¦ Release Process

Version Management

Versions are managed in:

  • pyproject.toml - Python package version
  • website/package.json - Frontend version

Release Steps

  1. Update version numbers
  2. Update CHANGELOG.md
  3. Run full test suite
  4. Build and test website
  5. Create GitHub release
  6. Deploy to GitHub Pages

GitHub Actions

CI/CD pipeline includes:

  • Format checking - Code style validation
  • Linting - Code quality checks
  • Testing - Python and frontend tests
  • Website deployment - Automatic GitHub Pages deployment

🀝 Contributing Guidelines

Pull Request Process

  1. Fork the repository
  2. Create feature branch: git checkout -b feature-name
  3. Make changes following coding standards
  4. Add tests for new functionality
  5. Run formatting: make format
  6. Run tests: make lint-check && python -m unittest tests.test_tritonparse -v
  7. Submit pull request

Code Review Process

  • All PRs require review by core maintainers
  • CI checks must pass before merge
  • Documentation updates required for new features
  • Tests required for new functionality

Issue Reporting

When reporting issues:

  1. Use issue templates provided
  2. Include system information
  3. Provide reproduction steps
  4. Include error messages and logs

πŸ“š Additional Resources

Documentation

Community

External Resources

πŸ”— Next Steps

For new developers:

  1. Complete the Installation Guide
  2. Read the Usage Guide to understand the tool
  3. Explore the codebase starting with simple components
  4. Run the test suite to verify your setup
  5. Join GitHub Discussions for community support

For experienced contributors:

  1. Check GitHub Issues for open tasks
  2. Review the Architecture Deep Dive for advanced topics
  3. Contribute to documentation improvements
  4. Propose new features through GitHub Discussions
Clone this wiki locally