Popcorn Kernels

A comprehensive toolkit for generating PyTorch-Triton code pairs through GitHub repository analysis and synthetic program generation, designed for machine learning model training on code compilation tasks.

Overview

Popcorn Kernels combines two powerful approaches to create diverse datasets of PyTorch neural network modules paired with their corresponding Triton GPU kernel implementations:

1. Real-world Code Mining: Scrapes and analyzes thousands of PyTorch repositories from GitHub
2. Synthetic Program Generation: Uses language models to create diverse, valid PyTorch programs

The toolkit generates datasets where PyTorch eager execution code is paired with optimized Triton kernels produced by PyTorch's Inductor compiler, enabling research into code generation, compilation optimization, and neural architecture understanding.

Projects

🔍 GitHub PyTorch Index

Scrapes PyTorch repositories from GitHub, extracts neural network modules, and evaluates them with torch.compile to generate PyTorch-Triton code pairs.

Key Features:

• Parallel repository downloading and processing
• Automated nn.Module extraction and validation
• torch.compile evaluation with Inductor backend
• Sharded processing for large-scale data generation
• Comprehensive error handling and logging

🎯 Synthetic Torch Modules

Generates synthetic PyTorch programs using language models, creating diverse neural architectures that complement real-world scraped code.

Key Features:

• LLM-driven program synthesis from operator specifications
• Automated validation and testing pipeline
• Configurable operator complexity and diversity
• KernelBench contamination filtering
• Parallel generation with yield optimization

Quick Start

1. Install dependencies:

   uv pip install -r requirements.txt
   

2. [Optional] Generate synthetic data:

   cd synthetic_torch_modules/
   python3 generate_synth_torch.py .parallel num_total_samples=1000
   

3. Run the full pipeline (GitHub scraping + evaluation):

   cd github_pytorch_index/
   ./scripts/run_full_pipline.sh --jobs=8 --run-dir=runs/experiment1
   <!-- If you did step 2 then run -->
   cd github_pytorch_index/
   ./scripts/run_full_pipline.sh --jobs=8 --run-dir=runs/combined \
     --synthetic-data-dir=../synthetic_torch_modules/generated_programs/
   

Output

The pipeline generates parquet datasets containing:

• scrape_dataset.parquet: Real GitHub repository data
• synthetic_dataset.parquet: Language model generated programs
• dataset.parquet: Combined dataset

Each entry contains PyTorch source code paired with corresponding Triton kernel implementations generated by torch.compile's Inductor backend.

Requirements

• Python 3.12.9+
• PyTorch 2.7.1+ with CUDA support
• 10+ GB storage for repository downloads
• GPU recommended for torch.compile evaluation

Note: To reproduce results from the KernelBook dataset, you must use PyTorch 2.5.0. The requirements.txt specifies PyTorch 2.7.1 for security and policy reasons, but you can downgrade if needed for result reproduction:

pip install torch==2.5.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

Contributing

See individual project READMEs for detailed development instructions and architecture documentation.

License

This project builds upon pytorch-jit-paritybench and follows its licensing terms.