Afinal, o que é um kernel?

Dois números em uma GPU

Você já sabe fazer c = a + b em Python. Já fez isso mil vezes. Aqui vamos falar sobre isso em tensores do PyTorch. Um tensor é apenas um array de números. Colocá-lo em uma GPU significa que esse array vive na memória da GPU em vez da memória comum da CPU. Quando a e b são dois tensores em uma GPU, essa linha é executada rápido o suficiente para que você nunca pense nisso.

Agora reduza. Digamos que a e b sejam dois floats únicos, ambos na GPU. Mesma linha. c = a + b. O que realmente roda no chip?

A resposta é um kernel. Neste contexto, um kernel é um pequeno programa que a GPU executa em algum dado. Não é o kernel do sistema operacional que seu laptop usa para inicializar. Não é o kernel matemático de um livro de álgebra linear. A palavra é reutilizada com frequência e isso não é culpa sua. No mundo das GPUs, um kernel significa apenas: uma pequena função que a GPU é instruída a executar, agora mesmo, em paralelo, nos dados que você forneceu a ela.

Ao final disso, você será capaz de olhar para um trecho de código PyTorch e contar quantos kernels a GPU executará. Isso parece um truque simples, e é, mas também é o primeiro passo para se livrar da sensação de "GPU como caixa preta", aquela em que seu modelo está lento e você não sabe por quê. Toda pergunta que você pode fazer sobre desempenho de GPU eventualmente se resume a "quais kernels rodaram e o que eles estavam fazendo". Então é por aqui que começamos.

Seu primeiro kernel

Vamos tornar a e b um pouco maiores. Tensores de comprimento 8 desta vez. Ainda uma linha PyTorch: c = a + b.

Quando você executa isso, sua CPU (a máquina que está rodando o Python) diz à GPU: ei, execute este programa nestes dados. Essa instrução é chamada de lançamento. O que é lançado é um kernel: um programa, pronto para rodar. Os lançamentos em si são baratos, microssegundos cada. O que está ao redor do lançamento (os dados indo para a GPU, os resultados voltando) é onde o custo real reside, e é isso que vamos contar ao longo do caminho.

Dentro do kernel, o trabalho real é feito por pequenos trabalhadores chamados threads. Uma GPU tem milhares deles disponíveis. Para nossa soma de comprimento 8, 8 threads pegam o trabalho: thread 0 lida com o elemento 0, thread 1 lida com o elemento 1, e assim por diante até a thread 7. Cada thread executa o mesmo pequeno programa: ler um elemento de a, ler o elemento correspondente de b, somá-los, escrever o resultado em c.

(Na prática, a GPU lança threads em grupos de tamanho fixo chamados warps, sempre 32 threads em placas NVIDIA, e mascara os extras quando seu array não divide igualmente. Pode ignorar por enquanto.)

Então temos uma linha PyTorch, um lançamento, um kernel, 8 threads fazendo 8 somas. Agora vamos contar o que realmente atravessou o chip. Para fazer a soma, cada thread precisa de seu elemento de a e seu elemento de b. Isso são 8 leituras de a e 8 leituras de b. Então cada thread escreve seu resultado em c. Isso são 8 escritas.

Essas leituras e escritas vão para a grande memória localizada ao lado do chip da GPU. Em placas de datacenter (A100, H100), essa memória é chamada de HBM (high-bandwidth memory). Em placas de consumo (RTX 4090) e nas T4s estilo Colab que as pessoas provavelmente vão tentar usar, é chamada de GDDR. De qualquer forma, é uma memória rápida ao lado do chip, e vamos chamá-la apenas de memória da GPU. É rápida, mas não é gratuita, e cada viagem até ela tem um custo.

Um kernel = um lançamento = uma passagem pelos dados. Seja o que for que o kernel faça internamente, as leituras e escritas em suas bordas (as viagens até a memória da GPU para buscar entradas, a viagem de volta para escrever saídas) são a parte que custa. Essa é a forma completa.

Nada disso muda quando os tensores ficam maiores. Mesma linha PyTorch, mesmo kernel, apenas mais threads. Se a e b têm um milhão de elementos cada, a GPU lança o mesmo kernel com um enxame maior de threads. A matemática escala, os bytes movidos escalam, o modelo mental não. Uma linha, um kernel.

O que acontece entre duas operações

1c = (a + b).relu()

Você conhece Python bem o suficiente para saber que isso são duas operações, uma soma e depois uma relu, encadeadas. Em um interpretador, isso são duas chamadas de função. Em uma GPU, no PyTorch eager, isso são dois lançamentos de kernel: um para a soma, um para a relu. Até aqui, nada surpreendente.

O que é realmente interessante é o que acontece entre os dois kernels.

Quando a soma termina, ela precisa colocar seu resultado em algum lugar. Esse lugar é a memória da GPU. A soma escreve um array intermediário completo (vamos chamá-lo de tmp) na memória. Então, um momento depois, a relu é lançada, e seu primeiro trabalho é ler esse mesmo array tmp de volta da memória. Ela lê o array inteiro, aplica relu a cada elemento e escreve o resultado em c.

Conte o tráfego de memória para esses dois kernels:

• A soma: lê a, lê b, escreve tmp. Três transferências do tamanho do array.
• A relu: lê tmp, escreve c. Mais duas.

Cinco transferências do tamanho do array no total. Compare isso com a soma de comprimento 8 sozinha da seção anterior, que foram três. Adicionar .relu() à cadeia não custou apenas o processamento da relu. Custou uma viagem de ida e volta completa do array através da memória da GPU, porque tmp teve que ser escrito apenas para que o próximo kernel pudesse lê-lo de volta.

Nada foi armazenado em cache. tmp não pôde ficar em um registrador ou em um cache local rápido. Ele foi para a memória da GPU (do tipo lento e distante) e voltou. Os dois kernels são estranhos um para o outro. Eles têm que passar dados através do único meio que ambos sabem usar: a memória da GPU.

Por que o PyTorch faz isso dessa forma? Porque no modo eager, quando você escreve a + b, o PyTorch o executa agora. Ele não sabe que você estava prestes a chamar .relu() em seguida. Cada operação é despachada no momento em que sua linha Python é executada. Não há plano, não há antecipação. Cada operação é independente, produz um array real e o entrega ao que vem a seguir através da memória.

Este é o padrão para se lembrar. Cada valor intermediário em seu código PyTorch é fisicamente escrito na memória da GPU e lido de volta pela próxima operação. Cada um. É isso que a "contagem de kernels" realmente mede. Cada kernel extra é outra viagem de ida e volta que seus dados têm que fazer através da memória da GPU.

Fusão: duas operações, um kernel

Imagine um kernel que faz tudo isso de uma vez: lê seu elemento de a, lê seu elemento de b, soma-os, aplica relu ao resultado (tudo dentro do kernel, em um pequeno espaço de trabalho por thread que nunca sai do chip), e só então escreve o valor final em c. O intermediário (a + b) ainda existe, mas apenas dentro do kernel, no espaço de trabalho privado de cada thread. Ele nunca é escrito na memória da GPU. tmp, como um array real, não existe.

Conte as transferências agora. Leituras de a: 1 por elemento. Leituras de b: 1 por elemento. Escritas de c: 1 por elemento. Três transferências do tamanho do array. A mesma matemática da versão de dois kernels, mas com duas viagens de ida e volta a menos.

No comprimento 8, isso é um erro de arredondamento. Ninguém se importa. No comprimento de 1 milhão, ou 100 milhões, essas viagens extras se tornam uma grande parte do tempo de execução, e o relógio de parede reflete isso. Por que o tráfego de memória acaba dominando dessa forma é o assunto do Artigo 2 desta série, então vou deixar o "porquê" de lado aqui. O ponto agora é apenas: mesma matemática, menos viagens, mais rápido na prática.

Esse truque, combinar operações que seriam kernels separados em um único kernel para que o intermediário nunca precise visitar a memória da GPU, tem um nome. É chamado de fusão. Essa é a palavra toda. Essa é a ideia toda.

Agora a parte complicada. Escrever esse kernel combinado manualmente parece fácil para soma + relu. Duas operações. Uma linha de "processamento" no meio. Mas o código PyTorch real tem dezenas de operações encadeadas, cada uma com suas próprias formas, tipos de dados e regras de broadcasting. Escrever um kernel fundido que lida com tudo isso corretamente é um trabalho de engenharia real. Geralmente você não escreveria esses kernels elementares de rotina manualmente.

Boas notícias: o PyTorch vem com uma ferramenta que faz a reescrita para você, exatamente para esse tipo de caso, automaticamente. É chamada de torch.compile.

Você quase certamente já viu torch.compile usado assim:

1model = torch.compile(model)

Uma linha. Alguém na internet te disse que isso acelera as coisas. Aqui está o que realmente faz, em português claro: em vez de executar suas operações uma de cada vez como o modo eager faz, o torch.compile captura as operações de tensor que sua função realiza, procura oportunidades para combiná-las e gera código otimizado. Chamadas posteriores que correspondam às mesmas premissas podem reutilizar esse código.

A fusão que fizemos no papel acima (soma e relu, compartilhando um kernel, tmp nunca tocando na memória) é exatamente o tipo de coisa que o torch.compile fará em seu código automaticamente, desde que as operações sejam simples o suficiente. Quando as pessoas dizem que torch.compile "torna o PyTorch mais rápido", isso é uma grande parte do que querem dizer.

Para os casos em que o torch.compile não consegue fundir por conta própria (operações personalizadas que não reconhece, reduções incomuns, layouts de memória estranhos), alguém ainda precisa escrever um kernel manualmente. É para isso que servem ferramentas como Triton e CUDA. Um artigo separado.

Veja por si mesmo

Tudo acima foi contar kernels no papel. Hora de contá-los em uma GPU real. Se você tiver alguma máquina com uma GPU CUDA disponível (uma estação de trabalho, um notebook Colab, uma instância na nuvem), pode executar isso em alguns minutos.

A ferramenta é o torch.profiler. Está embutido no PyTorch. Tudo o que faz é registrar o que a GPU realmente fez enquanto seu código rodava e te devolver uma tabela que você pode ler.

Passo 1: a versão eager.

Envolva a linha de duas operações em uma função para ter algo para chamar:

1import torch
2from torch.profiler import profile, ProfilerActivity
3
4def add_relu(a, b):
5    return (a + b).relu()
6
7a = torch.randn(1_000_000, device="cuda")
8b = torch.randn(1_000_000, device="cuda")
9
10with profile(activities=[ProfilerActivity.CUDA]) as prof:
11    add_relu(a, b)
12    torch.cuda.synchronize()
13
14print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

torch.cuda.synchronize() está ali apenas para garantir que a GPU tenha terminado antes de lermos os tempos. O trabalho da GPU é executado de forma assíncrona, e sem a sincronização você às vezes mediria a sobrecarga do lançamento em vez do trabalho real do kernel.

Passo 2: leia a saída.

Sua saída real do profiler terá mais linhas do que você pode esperar. Várias linhas de alocação de memória e contabilidade do PyTorch estarão misturadas. As linhas que nos interessam são os kernels CUDA, as funções reais que a GPU executou. Procure por linhas com kernel no nome. As duas serão mais ou menos assim:

1vectorized_elementwise_kernel<...BinaryFunctor_add...>    1    12us
2vectorized_elementwise_kernel<...threshold_kernel...>     1     9us

Os nomes exatos dos templates mudam entre versões do PyTorch (relu geralmente aparece como threshold porque essa é a operação subjacente, e add às vezes como CUDAFunctor_add). Não tente analisar tudo. Apenas conte as linhas. Duas linhas. Dois kernels. Um para a soma, um para a relu. Exatamente o que dissemos que aconteceria uma seção atrás.

Passo 3: a versão compilada.

Uma mudança de linha. Envolva a função em torch.compile:

1compiled = torch.compile(add_relu)

Agora, aqui está uma pegadinha que vale a pena saber antes de executar. A primeira chamada a uma função torch.compile é lenta, às vezes estranhamente lenta, porque é quando o torch.compile faz seu trabalho: analisar seu código, descobrir o que fundir, gerar o kernel fundido. Se você analisar a primeira chamada, estará medindo a etapa de compilação, não o kernel. Então o padrão é: chame uma vez para aquecer, descarte o resultado, depois analise.

1compiled(a, b)                # aquecimento, descarte
2torch.cuda.synchronize()
3
4with profile(activities=[ProfilerActivity.CUDA]) as prof:
5    compiled(a, b)
6    torch.cuda.synchronize()
7
8print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

Passo 4: leia a saída novamente.

Desta vez:

1triton_poi_fused_add_relu_0    1    14us

Uma linha. Um kernel. O nome até diz o que fez: soma e relu fundidas. Mesma matemática de antes, um lançamento em vez de dois.

Você acabou de fazer o que este artigo vem falando, em uma frase: pediu ao PyTorch para combinar duas operações em um kernel, observou o profiler e confirmou que a contagem caiu de dois para um. Fusão, na prática, na sua máquina.

Se quiser ver de forma mais dramática, tente com alguns tamanhos de tensor diferentes. No comprimento 100, ambas as versões rodam tão rápido que a diferença se perde no ruído. No comprimento de 10 milhões ou 100 milhões, a versão compilada começa a se destacar claramente, porque a viagem de ida e volta que eliminamos é uma parte real do trabalho nessa escala.

Contar kernels não é mais um conselho abstrato. Você tem uma maneira de verificar.

Vamos finalizar

Aqui está tudo de uma vez.

Seu código PyTorch, quando roda em uma GPU, se transforma em uma sequência de kernels. Cada kernel é um lançamento, uma passagem pelos seus dados, uma viagem de ida e volta pela memória da GPU para buscar entradas e escrever saídas. Operações simples se tornam um kernel. Cadeias de operações se tornam um kernel por operação por padrão, com intermediários fazendo viagens de ida e volta pela memória entre eles. O torch.compile pode fundir cadeias simples para você, para que esses intermediários nunca toquem na memória. Menos kernels geralmente significam menos tráfego de memória. E menos tráfego de memória geralmente significa mais rapidez.