Otimização de escalonamento Halide através de aprendizado por reforço

Abstract

Resumo: Escrever programas que usufruam da melhor maneira a capacidade de processamento oferecida por uma determinada arquitetura de hardware não é uma tarefa trivial e tornase cada vez mais trabalhosa à medida em que aumenta a diversidade de equipamentos e plataformas. Em uma linguagem de programação de propósito geral esta tarefa depende muito do compilador e do conhecimento e experiência do programador. Neste trabalho foi avaliado o potencial que uma linguagem de domínio específico chamada Halide oferece para facilitar a implementação de programas de processamento de imagens otimizados para diferentes arquiteturas. A partir de uma característica desta linguagem, chamada escalonamento de execução, é possível empregar mecanismos automatizados para produzir uma implementação com melhor desempenho usufruindo de características específicas do hardware, sem afetar a definição lógica e resultados do programa. Trabalhos anteriores empregaram algoritmos genéticos e diversas técnicas de busca para tentar gerar escalonamentos otimizados, porém não chegaram a uma solução efetiva. No presente trabalho foi desenvolvido uma abordagem de geração do escalonamento de execução empregando aprendizado de máquina, especificamente aprendizado por reforço. O problema foi modelado como um processo de decisão de Markov e incorporado a um ambiente de aprendizado por reforço utilizando um agente que agrega funções de aproximação para espaços de grandes dimensões, baseadas em redes neurais, chamado PPO (Proximal Policy Optimization). Experimentos foram realizados em duas arquiteturas de hardware (CPU e GPU), utilizando programas de processamento de imagens com diferentes níveis de complexidade. Os escalonamentos produzidos pelo agente PPO foram comparados com versões implementadas manualmente por programadores Halide e com soluções geradas pelo recurso de escalonamento automático suportado na linguagem Halide. Os resultados mostram que nos cenários avaliados a abordagem desenvolvida foi capaz de produzir escalonamentos de execução competitivos com as demais soluções comparadas, superando o desempenho em vários casos. Contudo, o modelo desenvolvido ainda não é totalmente automatizado, necessitando da intervenção do programador para mapear quais diretivas de escalonamento serão exploradas pelo agente de aprendizado por reforço. Palavras-chave: Aprendizado por Reforço. Otimização de Código. Linguagem Halide.

Abstract: Writing programs that best utilize the processing power offered by a particular hardware architecture is not a trivial task and becomes increasingly laborious as the diversity of equipment and platforms increases. In a general-purpose programming language, this task relies heavily on the compiler and the knowledge and experience of the programmer. In this work, we evaluate the potential that a specific domain language called Halide offers to facilitate the implementation of optimized image processing programs for different architectures. From a feature of this language, called execution scheduling, it is possible to employ automated mechanisms to produce a better performing implementation, taking advantage of the specific hardware characteristics, without affecting the logical definition and the program results. Previous works has employed genetic algorithms and several search techniques to try to generate optimized schedules, but they have not reached an effective solution. In the present work was developed an approach of execution scheduling generation using machine learning, specifically reinforcement learning. The problem was modeled as a Markov decision process and incorporated into a reinforcement learning environment using an agent that adds approximation functions for large spaces, based on neural networks, called PPO (Proximal Policy Optimization). Experiments were performed on two hardware architectures (CPU and GPU), using image processing programs with different levels of complexity. The schedules produced by the PPO agent were compared to versions implemented manually by Halide programmers and with solutions generated by the automatic scheduling feature supported in the Halide language. The results show that in the evaluated scenarios the developed approach was able to produce competitive execution schedules compared with the other solutions, surpassing the performance in several cases. However, the developed model is not yet fully automated, requiring the intervention of the programmer to map which scheduling directives will be explored by the reinforcement learning agent. Keywords: Reinforcement Learning. Code Optimization. Halide Language.