A Multi-Level Superoptimizer for Tensor Programs
<|reference_start|>A Multi-Level Superoptimizer for Tensor Programs: We introduce Mirage, the first multi-level superoptimizer for tensor programs. A key idea in Mirage is $\mu$Graphs, a uniform representation of tensor programs at the kernel, thread block, and thread levels of the GPU compute hierarchy. $\mu$Graphs enable Mirage to discover novel optimizations that combine algebraic transformations, schedule transformations, and generation of new custom kernels. To navigate the large search space, Mirage introduces a pruning technique based on abstraction that significantly reduces the search space and provides a certain optimality guarantee. To ensure that the optimized $\mu$Graph is equivalent to the input program, Mirage introduces a probabilistic equivalence verification procedure with strong theoretical guarantees. Our evaluation shows that Mirage outperforms existing approaches by up to 3.5$\times$ even for DNNs that are widely used and heavily optimized. Mirage is publicly available at https://github.com/mirage-project/mirage.<|reference_end|>
@article{wu2024a,
    title={A Multi-Level Superoptimizer for Tensor Programs},
    author={Mengdi Wu and Xinhao Cheng and Oded Padon and Zhihao Jia},
    journal={arXiv preprint arXiv:2405.05751},
    year={2024},
    archivePrefix={arXiv},
    eprint={2405.05751},
    primaryClass={cs.LG cs.AI cs.PL}
}