Houve muita discussão após a apresentação do kernel MLA Jalapeño no HotChips e os comentários subsequentes do SemiAnalysis. Como equipe de hardware da OpenAI, mal tocamos nessa pequena pepita de ouro: o fato de que a IA está escrevendo nossos kernels e que, quando o faz, não precisamos realmente entender o que o kernel faz linha por linha. Deixamos de lado, de forma evidente: como algo assim é possível? Qual é a maneira correta de pensar sobre isso, em comparação com um método mais tradicional de geração de código? O kernel otimizado é tão sólido quanto o não otimizado?
Quanto à minha formação, trabalho com compiladores para aceleradores há mais de uma década. Comecei o XLA, que é uma excelente infraestrutura de compilador com uma equipe e esforço excepcionais entre empresas. Nos últimos 2 anos ou mais na OpenAI, tenho tentado reconceituar como os compiladores devem funcionar na era da IA. Novas formulações de compiladores aproveitarão os pontos fortes existentes, mas é impossível negar que há uma ferramenta nova e poderosa para alavancar no kit de ferramentas.
Esta será uma jornada um tanto longa, mas espero iluminar como a IA está sendo usada para a automação da melhoria de programas de computador; ou seja, compilação otimizadora. Acredito que, por meio da IA, podemos experimentar algo que consideramos "compiladores 2.0". A IA é fundamentalmente menos limitada no que pode propor, e o que ela propõe é resultado do treinamento e do contexto do modelo, o que me leva a classificá-la como um "otimizador estocástico" – isso pode apresentar desafios, mas, como veremos, também é uma fonte de grandes pontos fortes…
Muita pesquisa acadêmica e aplicação industrial já estão caminhando nessa direção, descobrindo rapidamente o potencial do envolvimento da IA no reino dos compiladores otimizadores, mas estamos em um ponto em que uma explicação geral se faz necessária.
Contexto
Compiladores recebem programas e geram versões traduzidas ou melhoradas desses programas.
Programas, tanto no lado de entrada quanto no de saída, têm semântica que nos diz o que os programas significam, o que eles podem fazer e como raciocinar sobre essas coisas que podem fazer.
Aqueles de nós que trabalham com compiladores os veem como funções puras – eles recebem uma estrutura de dados e geram uma estrutura de dados que deve ter semântica correspondente.
Às vezes, nossos compiladores focam em "redução" ou "tradução". Por exemplo, eles podem receber C e gerar assembly x86-64, que muitas vezes consideramos de "nível mais baixo". Mas frequentemente eles fazem mais do que apenas tradução como uma subparte desse processo…
Nossos compiladores, na prática, focam em "otimização". Eles podem receber uma estrutura de dados que representa o programa – em nossa linguagem, uma "Representação Intermediária" (IR) – e tentam produzir uma versão melhor desse programa. Às vezes, "melhor" significa que leva menos ciclos para executar, às vezes significa que terá menos código desnecessário, às vezes significa especializar para coisas que podemos provar que "devem ser verdadeiras" sobre o programa (avaliação parcial).
Agora, considere brevemente que os LLMs foram originalmente criados para traduzir texto humano de um idioma para outro. Claramente, a tradução está em sua alçada. E podemos ver, através do uso de LLMs em tarefas diárias, que eles também podem escrever novas soluções e melhorar soluções existentes. Muitos de nós, programadores, também temos experiência em pedir a um LLM "otimize este trecho de código" e eles notavelmente conseguem. (No entanto, precisamos saber se eles otimizaram o código corretamente, ao que chegaremos!) Isso serve apenas para destacar que os LLMs têm as capacidades que buscamos em um compilador otimizador.
Otimização e Otimalidade
Compiladores otimizadores estão, como era de se esperar, tentando aumentar a otimalidade do programa em que estão trabalhando, por algum objetivo (geralmente tempo de execução). Isso é tão difícil de fazer no caso geral, para um programa arbitrário, que existe um teorema chamado teorema do pleno emprego para engenheiros de compiladores. (Só descobri isso depois que escolhi ser engenheiro de compiladores, mas ainda assim me trouxe conforto!)
"Superotimizadores" são um subcampo incrível dos compiladores otimizadores. Imagine que existe um determinado programa, e podemos dizer o que ele faz através da semântica. Qual é o programa mais otimizado que tem essa mesma semântica? É isso que os superotimizadores tentam resolver, e é efetivamente um problema de busca…
Imagine que estou tentando encontrar o programa mais curto que tem a mesma semântica, e tenho uma maneira de perguntar se um programa candidato tem a mesma semântica. Eu poderia, hipoteticamente, enumerar todos os programas em ordem objetiva e escolher o menor que tivesse a mesma semântica.
No entanto, enumerar todos os programas em ordem objetiva parece bastante intratável. Um dos meus artigos acadêmicos favoritos, feito em 2013, intitulado "STOKE" (Superotimização Estocástica), perguntou: "bem, e se apenas ajustarmos os programas aleatoriamente repetidamente, eventualmente observaremos o melhor programa?" Eles propuseram que, através de uma caminhada aleatória (e com nosso amigo OG de aprendizado de máquina, Markov Chain Monte Carlo / Metropolis-Hastings), eventualmente você veria aquele programa ótimo.
O ajuste de Monte Carlo é tipicamente burro (você escolhe um ajuste aleatoriamente), mas também rápido. Os LLMs são muito inteligentes (muitos tokens de raciocínio), mas comparativamente lentos.
E se, em vez do ajuste burro/rápido de Monte Carlo, tivéssemos LLMs descobrindo as direções para as quais levar os programas? Teríamos um otimizador estocástico que era muito inteligente, guiando nosso programa pelo espaço de programas otimizados.
Intuições para Otimização
Vamos dar um passo atrás. Considere a pessoa que você conhece que melhor personifica "otimiza ao máximo trechos de código". Para encurtar, vamos chamá-la de "Ollie, o otimizador". Ollie provavelmente tem uma intuição inata sobre quais tipos de ajustes de código podem dar frutos. Ollie provavelmente tenta algumas coisas para ver se funcionam e, se não derem certo, desfaz e tenta outra coisa. Mas ele tem alguma intuição sobre quais tipos de coisas são possíveis e como podem ser capazes de superar o compilador.
Essas intuições que Ollie tem geralmente estão além do que os compiladores fazem. Embora os compiladores otimizadores modernos sejam bastante impressionantes em seus resultados, eles são baseados em regras e heurísticas bastante simples. Em jargão técnico, eles são baseados na ideia de uma transformação local de fluxo de dados que é executada até o ponto fixo. Também ordenamos as considerações em fases; ou seja, construímos pipelines de compilador para considerar A e depois B, mas não o problema composto AB. Escalonadores e alocadores de registradores são um exemplo notório disso, muitos PhDs foram tentados no escalonador-alocador de registradores composto (para obter os benefícios de colapsar a ordenação de fases), mas eles têm sido desafiadores para fazer funcionar na prática.
É por isso que a experiência de Ollie é valiosa. Frequentemente, Ollie sabe como equilibrar vários problemas NP-completos com heurísticas que são feitas sob medida para a situação. Portanto, há mais consciência e sensibilidade de contexto personalizadas. Ollie também é capaz de empregar técnicas que os compiladores otimizadores podem não aplicar de forma lucrativa, especialmente em combinação, coisas como desmembramento ou criação de ABIs personalizados ou transformações para permitir a vetorização, ou a infinidade de outras coisas que nos fazem resmungar "eu gostaria que o compilador tivesse uma maneira de fazer isso…"
Agora considere que a IA, através de quaisquer recursos de raciocínio que tenha, pode ser capaz de atuar como um mini Ollie. Pode não ter a intuição correspondente em termos do que dará certo, mas tem uma noção do que pode ser lucrativo, e pode dar muitas, muitas tentativas.
Com essa abordagem, ao contrário do artigo STOKE, não podemos garantir que, à medida que o tempo vai ao infinito, veremos o programa ótimo, mas como a IA tem recursos de raciocínio "mais humanos", ela pode realmente obter tração significativa semelhante à humana por unidade de tempo.
Relacionando de Volta: Kernel MLA
Deixe-me começar dizendo: eu não sei qual código de baixo nível a IA gerou para nosso kernel MLA Jalapeño, mas sei como digitar o numpy para o MLA.
No compilador XLA em que trabalhei anteriormente, fundiríamos essas operações numpy em aglomerados e, em seguida, usaríamos um metaprograma chamado "emissor" para reduzi-lo a loops, instruções e primitivas de nível inferior.
Quando o compilador / programa emissor XLA fazia isso, eu não precisava me importar com qual assembly saía. Para nosso otimizador estocástico, a IA conceitualmente assume o lugar do meta-programa emissor – ela tanto reduz quanto otimiza, e podemos pedir que ela otimize cada vez mais em direção ao roofline.
Espero que isso deixe claro onde a IA se encaixa e como ela é análoga a um componente em um sistema de compilador otimizador existente. Também é útil pensar: qual camada consideramos como "código assembly" agora está subindo. Quando você digita C++ normal e o compila com -O3 (o nível de otimização típico mais alto), você não espera entender o assembly que sai, mesmo que entendesse o C++ que digitou. Estamos fazendo a coisa análoga aqui, mas com uma especificação de entrada de nível mais alto e mais matemática.
Agora, uma questão fundamental é como verificamos se o programa que obtemos da IA é de fato equivalente à descrição de nível superior / numpy. Esse mecanismo de verificação estabelece a solidez do processo de otimização estocástica da IA. Espero que uma postagem futura do blog entre em mais detalhes sobre isso, mas, por enquanto, basta dizer que testar a equivalência semântica é possível e nós o fazemos. Programas de aceleradores são particularmente adequados para contratos fortes e completos que podemos verificar "são exatamente o que o programa otimizado por IA faz", pois são bastante matemáticos e orientados a fluxo de dados em seu contexto amplo.
Observe que muitas técnicas relevantes nesta área foram pioneiras por esforços no subcampo da síntese de programas. Enquanto os compiladores otimizadores dizem: "aqui está um programa com semântica, melhore-o, mas com semântica equivalente!", a síntese de programas diz: "acredito que existe um programa com essa semântica, tente encontrar o melhor que puder". A síntese de programas é um problema mais difícil do que a compilação otimizadora, mas também é fundamentalmente menos limitada. É efetivamente o que humanos como Ollie fazem quando superam o compilador otimizador, e é algo que a IA agora pode nos ajudar a automatizar. A IA pode tirar "inspiração" do programa original, mas não precisa apenas realizar pequenas transformações locais nele. Compiladores otimizadores clássicos não verão "ah, você escreveu um bubble sort" e, ao entender o contrato, trocá-lo por um quick-sort, mas tanto Ollie quanto a IA são capazes de fazer isso. Isso é o que nos coloca mais no regime de síntese de programas com otimização estocástica do que no regime clássico de compilador otimizador.
Tudo isso se junta no fato de que você pode começar com algo que "não está muito longe do numpy", esperar 48 horas e ter um kernel otimizado com a mesma semântica, como mostramos em nossa palestra no HotChips:
Como o slide também observa, em nossa máquina, muitas vezes somos capazes de observar a IA superando o desempenho de nossos especialistas humanos mesmo nos kernels que considerávamos bastante ajustados. Muitas vezes, ainda há uma porcentagem alcançável decente devido apenas às muitas variedades de combinações / permutações que podem precisar ser exploradas. Estas são frequentemente intratavelmente tediosas para um engenheiro de desempenho humano.
Recapitulação e Conclusão
Um compilador, no final das contas, é apenas uma função. Damos nosso programa a essa função e recebemos de volta uma versão melhor do nosso programa. Esperamos que o programa que obtemos e o programa que inserimos tenham a mesma semântica.
Compiladores otimizadores tradicionais melhoram os programas por meio de regras de fluxo de dados e heurísticas. Elas são totalmente compreensíveis em sua proveniência, mas também podem ser mais limitadas nos movimentos que podem fazer.
Em contraste, a IA, como um otimizador estocástico, só precisa "pensar muito" e gerar algo. Seus movimentos não são tão fundamentalmente limitados, tornando-os mais análogos ao nosso otimizador humano especialista. Precisamos de maneiras de verificar se os programas que ela gera são sólidos e implementam a mesma semântica que inserimos, e nós as temos em vigor. E esse tipo de otimização por IA é particularmente adequado para operações matemáticas que têm contratos muito fortes. Os contratos evitam a necessidade de entender o que o kernel faz linha por linha.
É assim que obtivemos o kernel MLA gerado por IA!
