CoolFace
Datasetpublic

Tay-eb/mlir-rl-dataset

MLIR-RL Benchmark Datasets Code: https://github.com/Modern-Compilers-Lab/MLIR-RL Paper: Bendib et al. 2024, A Reinforcement Learning Environment for Automatic Code Optimization in the MLIR Compiler — https://arxiv.org/abs/2409.11068 What is MLIR-RL? MLIR-RL is a reinforcement-learning environment for automatic loop-nest optimization in MLIR. An agent observes linalg operations (matmul, conv, …) and learns sequences of transformations — tiling, interchange… See the full description on the dataset page: https://huggingface.co/datasets/Tay-eb/mlir-rl-dataset.

sourceHugging Faceotherupdated 1d agoView on Hugging Face
0likes61downloads
Dataset Card

MLIR-RL Benchmark Datasets

Code: https://github.com/Modern-Compilers-Lab/MLIR-RL Paper: Bendib et al. 2024, A Reinforcement Learning Environment for Automatic Code Optimization in the MLIR Compiler — https://arxiv.org/abs/2409.11068

What is MLIR-RL?

MLIR-RL is a reinforcement-learning environment for automatic loop-nest optimization in MLIR. An agent observes linalg operations (matmul, conv, …) and learns sequences of transformations — tiling, interchange, vectorization, fusion — that maximize speedup over the unoptimized baseline.

What is this dataset?

MLIR loop-nest benchmarks for training RL autoschedulers. Single-operation kernels, multi-op blocks extracted from neural networks, and full source models — all in MLIR with timing wrappers for baseline measurement. This Hub repo is the distributable copy (data/ is gitignored in the source repo).

Dataset Structure

data/
├── ops_and_blocks/       # Primary dataset — ~8,093 flat .mlir files
├── mlir_rl_v1_paper/     # Paper reproduction — ~1,357 flat .mlir files
└── full_models/          # Source ONNX/MLIR models before extraction (~29GB)
DatasetFilesContent
ops_and_blocks/~8,093Primary training set. Single ops + multi-op blocks from 18 NN models: BERT, ALBERT, DistilBERT, GPT-2 (medium), BART, T5, ViT-B/16, ResNet-50, ResNeXt-50, MobileNetV3-Small, EfficientNet-B0, ConvNeXt-Tiny, VGG-16, GAT, GIN, Whisper-base (encoder), YOLOv8m (backbone), Llama 3.2 1B. See docs/data/NEW_DATASET.md in the GitHub repo for per-model rationale.
mlir_rl_v1_paper/~1,357Paper reproduction set. NN ops (add, matmul, conv, relu, pooling_nchw_max) plus Lattice QCD kernels (baryon_*, dibaryon_*, hexaquark_*), plus 3 full models (model_mobile_net_v2, model_res_net, model_vgg).
full_models/~70Source models before extraction (.onnx, .onnx.data, _linalg.mlir, _torch.mlir). ~29GB — skip unless you re-run extraction.

Usage

For train/eval splits, file format, and training pipeline instructions, see the official GitHub repo: https://github.com/Modern-Compilers-Lab/MLIR-RL

Citation

If you use this dataset, please cite the original paper:

bibtex
@misc{bendib2024mlirrl,
  title  = {A Reinforcement Learning Environment for Automatic Code
            Optimization in the MLIR Compiler},
  author = {Bendib, Nazim and Aouadj, Iheb Nassim and Baghdadi, Riyadh},
  year   = {2024},
  eprint = {2409.11068},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url    = {https://arxiv.org/abs/2409.11068}
}