# Customized Monte Carlo Tree Search for LLVM/Polly's Composable Loop Optimization Transformations

arXiv: https://arxiv.org/abs/2105.04555 · HTML: https://mlsanity.com/paper/2105.04555/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 10 of its 45 references that are in the map.
- First posted on arXiv 10 May 2021.

## Builds on (most cited first, 10 of 10)

- [code2vec: Learning Distributed Representations of Code](https://mlsanity.com/paper/1803.09473/) (2018)
- [Learning to Optimize Tensor Programs](https://mlsanity.com/paper/1805.08166/) (2018)
- [NeuroVectorizer: End-to-End Vectorization with Deep Reinforcement Learning](https://mlsanity.com/paper/1909.13639/) (2019)
- [Autotuning PolyBench Benchmarks with LLVM Clang/Polly Loop Optimization Pragmas Using Bayesian Optimization](https://mlsanity.com/paper/2010.08040/) (2020)
- [Reinforcement Learning and Adaptive Sampling for Optimized DNN Compilation](https://mlsanity.com/paper/1905.12799/) (2019)
- [ProTuner: Tuning Programs with Monte Carlo Tree Search](https://mlsanity.com/paper/2005.13685/) (2020)
- [Autotuning Search Space for Loop Transformations](https://mlsanity.com/paper/2010.06521/) (2020)
- [User-Directed Loop-Transformations in Clang](https://mlsanity.com/paper/1811.00624/) (2018)
- [On the Representation of Partially Specified Implementations and its Application to the Optimization of Linear Algebra Kernels on GPU](https://mlsanity.com/paper/1904.03383/) (2019)
- [Design and Use of Loop-Transformation Pragmas](https://mlsanity.com/paper/1910.02375/) (2019)

## Built on by (most cited first, 4 of 4)

- [ytopt: Autotuning Scientific Applications for Energy Efficiency at Large Scales](https://mlsanity.com/paper/2303.16245/) (2023)
- [BaCO: A Fast and Portable Bayesian Compiler Optimization Framework](https://mlsanity.com/paper/2212.11142/) (2022)
- [Performance Optimization using Multimodal Modeling and Heterogeneous GNN](https://mlsanity.com/paper/2304.12568/) (2023)
- [Machine Learning for CUDA+MPI Design Rules](https://mlsanity.com/paper/2203.02530/) (2022)

## Abstract (arXiv)

Polly is the LLVM project's polyhedral loop nest optimizer. Recently, user-directed loop transformation pragmas were proposed based on LLVM/Clang and Polly. The search space exposed by the transformation pragmas is a tree, wherein each node represents a specific combination of loop transformations that can be applied to the code resulting from the parent node's loop transformations. We have developed a search algorithm based on Monte Carlo tree search (MCTS) to find the best combination of loop transformations. Our algorithm consists of two phases: exploring loop transformations at different depths of the tree to identify promising regions in the tree search space and exploiting those regions by performing a local search. Moreover, a restart mechanism is used to avoid the MCTS getting trapped in a local solution. The best and worst solutions are transferred from the previous phases of the restarts to leverage the search history. We compare our approach with random, greedy, and breadth-first search methods on PolyBench kernels and ECP proxy applications. Experimental results show that our MCTS algorithm finds pragma combinations with a speedup of 2.3x over Polly's heuristic optimizations on average.