# Machine Learning for CUDA+MPI Design Rules

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 5 of its 31 references that are in the map.
- First posted on arXiv 4 Mar 2022.

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

- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [AlphaX: eXploring Neural Architectures with Deep Neural Networks and Monte Carlo Tree Search](https://mlsanity.com/paper/1903.11059/) (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)
- [Customized Monte Carlo Tree Search for LLVM/Polly's Composable Loop Optimization Transformations](https://mlsanity.com/paper/2105.04555/) (2021)

## Abstract (arXiv)

We present a new strategy for automatically exploring the design space of key CUDA+MPI programs and providing design rules that discriminate slow from fast implementations. In such programs, the order of operations (e.g., GPU kernels, MPI communication) and assignment of operations to resources (e.g., GPU streams) makes the space of possible designs enormous. Systems experts have the task of redesigning and reoptimizing these programs to effectively utilize each new platform. This work provides a prototype tool to reduce that burden. In our approach, a directed acyclic graph of CUDA and MPI operations defines the design space for the program. Monte-Carlo tree search discovers regions of the design space that have large impact on the program's performance. A sequence-to-vector transformation defines features for each explored implementation, and each implementation is assigned a class label according to its relative performance. A decision tree is trained on the features and labels to produce design rules for each class; these rules can be used by systems experts to guide their implementations. We demonstrate our strategy using a key kernel from scientific computing -- sparse-matrix vector multiplication -- on a platform with multiple MPI ranks and GPU streams.