Abstract:We propose a novel solution for the Register Allocation problem, leveraging multi-agent hierarchical Reinforcement Learning. We formalize the constraints that precisely define the problem for a given instruction-set architecture, while ensuring that the generated code preserves semantic correctness. We also develop a gRPC based framework providing a modular and efficient compiler interface for training and inference. Experimental results match or outperform the LLVM register allocators, targeting Intel x86 and ARM AArch64.
Abstract:We propose IR2Vec, a Concise and Scalable encoding infrastructure to represent programs as a distributed embedding in continuous space. This distributed embedding is obtained by combining representation learning methods with data and control flow information to capture the syntax as well as the semantics of the input programs. Our embeddings are obtained from the Intermediate Representation (IR) of the source code, and are both language as well as machine independent. The entities of the IR are modelled as relationships, and their representations are learned to form a seed embedding vocabulary. This vocabulary is used along with the flow analyses information to form a hierarchy of encodings based on various levels of program abstractions. We show the effectiveness of our methodology on a software engineering task (program classification) as well as optimization tasks (Heterogeneous device mapping and Thread coarsening). The embeddings generated by IR2Vec outperform the existing methods in all the three tasks even when using simple machine learning models. As we follow an agglomerative method of forming encodings at various levels using seed embedding vocabulary, our encoding is naturally more scalable and not data-hungry when compared to the other methods.