Deep Learning-Based Approximate Graph-Coloring Algorithm for Register Allocation
TimeThursday, 12 November 202012:50pm - 1:30pm EST
DescriptionGraph-coloring is an NP-hard problem which has a myriad of applications. Register allocation, which is a crucial phase of a good optimizing compiler, relies on graph coloring. Hence, an efficient graph-coloring algorithm is of paramount importance. In this work we try to ‘learn’ a good heuristic for coloring interference graphs that are used in the register allocation phase. We aim to handle moderate-sized interference graphs which have 100 nodes or less. For such graphs we can get the optimal allocation of colors to the nodes. Such optimal coloring is then used to train our Deep Learning (DL) network which is based on several layers of LSTM that output a color for each node of the graph. However, the trained network may allocate the same color to the nodes connected by an edge resulting in an invalid coloring of the interference graph. Since it is difficult to encode constraints in an LSTM to avoid invalid coloring, we augment our deep learning network with a color correction phase that runs after the colors have been allocated by the DL network. Thus, our algorithm is approximate or hybrid in nature consisting of a mix of a DL algorithm followed by a more traditional correction phase. The color correction phase handles the edges with invalid coloring by first trying to reuse a color allocated to other nodes that are not connected to the invalid nodes, failing which it adds a totally new color – thereby breaking the invalid allocation. Our experience with many graphs shows that around 10%-30% edges may get an invalid coloring. We have trained our DL network using several thousand random graphs of varying sparsity(density). On application of our approximate algorithm to various popular graphs found in literature we see that our algorithm does very well when compared to the optimal coloring of these graphs. We have also run our algorithm against LLVM’s popular greedy register allocator (GRA) for several SPEC CPU 2017 benchmarks and notice that the approximate algorithm performs on par or better than GRA for most of these benchmarks.