Using q-learning to select the best among functionally equivalent implementations
Meggie van den Oever, Lauren E. Grimley, Richard Veras · 2022
High performance code generation for computationally intensive kernels is a persistent challenge for developers. Given a target architecture and a specific operation, the developer must tune that operation to the lowest-level details of the architecture. This problem is exacerbated by the fact that different architectural targets necessitate different implementations, and even the slightest adjustment to the operation may require large changes in the implementation in order to achieve performance. For performance critical applications this generation is typically performed by hand. However, this level of programming is difficult in terms of the domain knowledge required, and yields coded implementations that increase that challenge of reasoning about the correctness of the problem. Automatic code generation would address these issues. At the very least, by automating the application of the various code transformations needed for performance, this should reduce the issue of correctness, as long as these transformations only lead to correct implementations in the search space. In this paper, we look at a subset of correct implementations of an operation, all valid static schedules of instructions of one particular mix of instructions. We then explore the use of Reinforcement Learning in order to search for the optimal implementation in this subset for the target operation. This work is the first step in automating the exploration of correct implementations using Reinforcement Learning for automatic code generation.