GPUMODE/KernelBook

1---2license: other3language:4- en5---6## Overview7 8`dataset_permissive{.json/.parquet}` is a curated collection of pairs of pytorch programs and equivalent triton code (generated by torch inductor) which can be used to train models to translate pytorch code to triton code.9The triton code was generated using **PyTorch 2.5.0** so for best results during evaluation / running the triton code we recommend using that version of pytorch.10 11## Dataset Creation12 13The dataset was created through the following process:14 151. **Repository Collection**: PyTorch repositories were collected from GitHub using repositories (and associated hashes) from the [Stack v1](https://huggingface.co/datasets/bigcode/the-stack).162. **PyTorch Module Extraction**: We extracted the pytorch code from the repositories, and seperated them into individual `torch.nn` modules with appropriate dependencies.173. **Creating Unit Tests**: We created unit tests for each module to ensure that the code was working as expected. Code in which could not create unit tests for was removed.184. **Extracting Triton Code**: We used torch.compile in order to produce triton code from the pytorch code.195. **Transorming Triton Code**: We transformed the triton code into one which resembled the format seen in [KernelBench](https://github.com/ScalingIntelligence/KernelBench).205. **Metadata Enrichment**: Each repository entry was enriched with metadata such as license information, star count, and commit SHA.21 22The scripts to do this yourself can be found [here](https://github.com/pytorch-labs/popcorn-kernels/tree/main/github_pytorch_index)23 24## Data Structure25 26Each entry in the dataset contains the following fields:27 28| Field | Description |29|-------|-------------|30| `repo_name` | The name of the repository in the format `username/repository` |31| `licenses` | List of licenses associated with the repository |32| `stars` | Number of GitHub stars the repository has |33| `sha` | The commit SHA hash used for version reference |34| `repo_link` | Direct link to the repository at the specific commit (GitHub URL) |35| *Additional fields* | The dataset may contain other repository-specific information |36 37## File Formats38 39The dataset is available in two formats:40 411. **JSON**: `dataset_permissive.json` - A human-readable format that can be easily parsed by most programming languages.422. **Parquet**: `dataset_permissive.parquet` - A columnar storage format optimized for analytics and big data processing.43 44## Usage Examples45 46### Loading the Dataset in Python47 48#### Using JSON:49```python50import json51 52# Load the JSON version53with open('dataset_permissive.json', 'r') as f:54    repos = json.load(f)55 56# Example: Print the first 5 repository names57for repo in repos[:5]:58    print(repo['repo_name'])59```60 61#### Using Parquet:62```python63import pandas as pd64 65# Load the Parquet version66df = pd.read_parquet('dataset_permissive.parquet')67 68# Example: Get repositories with more than 1000 stars69popular_repos = df[df['stars'] > 1000]70print(f"Number of popular repositories: {len(popular_repos)}")71```72 73## License Information74 75This dataset is licensed under the [June 9 Researcher Reciprocity License](LICENSE).76 77You are free to use, share, analyze, transform, and redistribute the material for research, education, benchmarking, publication, commercial analysis, and other lawful purposes, provided you give appropriate credit.78 79This license adapts the Open RAIL-D structure and adds one specific use restriction: training, fine-tuning, distillation, synthetic-data generation for training, embedding for training, or otherwise using this dataset to improve an AI model or AI service requires Researcher Reciprocity.80 81> If you train on it, you let us generate.82 83Covered AI model and service providers may not use this dataset while imposing terms that prevent GPU Mode, dataset contributors, or authorized researchers from generating outputs, evaluating models, benchmarking, publishing research, or exploring their own research ideas on materially equal terms to ordinary users.84 85The `dataset_permissive` contains code derived from repositories with permissive licenses, including but not limited to MIT, Apache-2.0, BSD, MPL, Unlicense, and zlib. Users should still verify the license of individual repositories before using their code in production or commercial settings.86 87## Citation88 89```90@software{kernelbook2025,91    title={KernelBook},92    author={Paliskara, Sahan and Saroufim, Mark},93    year={2025},94    month={5},95    url={https://huggingface.co/datasets/GPUMODE/KernelBook},96}97```