Compiladores e Runtimes

Compiladores e Runtimes

Visão de alto nível de XLA/TensorRT e por que a compilação importa.

O que “Compiladores & Runtimes” Significam no ML Moderno

No software clássico, um compilador (compiler) traduz código-fonte (ex.: C++) em código de máquina, e um tempo de execução (runtime) fornece serviços enquanto o programa roda (alocação de memória, escalonamento, E/S, etc.). Em aprendizado de máquina (machine learning), as mesmas ideias se aplicam — mas o “programa” geralmente é uma computação com tensores (tensor computation) (multiplicações de matrizes, convoluções, atenção, normalização) executada em hardware especializado como GPUs e TPUs.

Dois ecossistemas de compilação de ML amplamente usados ilustram isso bem:

  • XLA (Accelerated Linear Algebra): um compilador usado pelo TensorFlow e pelo JAX (e tempos de execução relacionados) para otimizar programas tensoriais para GPUs/TPUs/CPUs.
  • TensorRT: o compilador/tempo de execução de inferência da NVIDIA que otimiza redes neurais treinadas para GPUs da NVIDIA.

Mesmo quando você nunca escreve CUDA, ainda assim está dependendo de uma pilha de tecnologia de compilador e tempo de execução para transformar um modelo como uma Arquitetura Transformer em kernels de alta vazão que saturam o hardware.

Este artigo oferece uma visão de alto nível de por que a compilação importa, o que compiladores/tempos de execução de fato fazem e como XLA e TensorRT se encaixam em fluxos reais de treinamento e inferência.

Por que a Compilação Importa para o Desempenho e o Custo em ML

Cargas de trabalho de aprendizado profundo (deep learning) frequentemente são limitadas não por “quantos FLOPs” um chip teoricamente consegue fazer, mas por quão eficientemente conseguimos:

  • mover tensores pelas hierarquias de memória,
  • fundir operações para reduzir gravações intermediárias,
  • escolher a melhor implementação de kernel para uma dada forma (shape) e tipo de dado (datatype),
  • paralelizar trabalho entre núcleos/SMs e entre dispositivos.

Esses são problemas de compilador.

1) Grafos de ML têm estrutura que compiladores podem explorar

Um passo típico de modelo não é código arbitrário; ele é uma composição majoritariamente regular de ops tensoriais. Por exemplo, uma sequência comum em inferência é:

  1. multiplicação de matrizes / convolução
  2. adicionar viés (bias)
  3. aplicar ativação (ReLU/GELU)
  4. normalização / dropout (dropout) (treinamento)
  5. atenção (Transformers)

Se executado de forma ingênua, cada op grava sua saída na memória global, e então a próxima op lê de volta. Esse tráfego de memória é caro — muitas vezes o custo dominante (ver Memória e Largura de Banda).

Um compilador pode fundir essas ops em menos kernels e reduzir movimentação de memória.

2) Lançamentos de kernel e sobrecarga de Python são gargalos reais

Em frameworks em modo eager (eager-mode), você pode lançar milhares de pequenos kernels de GPU por passo. Cada lançamento de kernel tem sobrecarga; em Python, a sobrecarga de despacho (dispatch overhead) adiciona ainda mais. Compiladores reduzem o número de lançamentos e podem gerar kernels maiores e mais eficientes.

3) O hardware é complexo e está evoluindo

GPUs diferentes (e TPUs) têm diferentes estratégias ideais de tileamento (tiling), uso de Tensor Cores (tensor cores), uso de memória compartilhada (shared memory) e suporte a tipos de dado. Compiladores codificam conhecimento específico de hardware e podem autoajustar (autotuning) escolhas.

É por isso que a portabilidade de desempenho (performance portability) depende cada vez mais de compilação, em vez de kernels escritos à mão por toda parte. (Ainda assim, bibliotecas do fornecedor (vendor libraries) como cuBLAS/cuDNN continuam sendo blocos de construção cruciais; compiladores frequentemente orquestram e selecionam chamadas de biblioteca em vez de gerar tudo do zero.)

4) A compilação viabiliza decisões de precisão + layout

Para inferência, você frequentemente quer FP16/BF16/INT8/INT4 para vazão e menor custo (ver Quantização). Um compilador pode:

  • inserir conversões de tipo (casts) com segurança,
  • escolher kernels que usam Tensor Cores,
  • manter acumulações em maior precisão quando necessário,
  • escolher layouts de memória que melhorem coalescência e localidade de cache.

Um Modelo Mental: o Pipeline de Compilação em ML

Embora as implementações diferenciem, a maioria dos compiladores de ML segue um fluxo semelhante:

1) Capturar um grafo (graph) (ou representação intermediária (IR)) a partir do Python

Código de framework costuma ser dinâmico (fluxo de controle em Python, formas dinâmicas), então o compilador precisa de uma representação estável:

  • Rastreamento (tracing): executar uma vez, registrar ops (comum em jit do JAX, grafos do TensorFlow, rastreamento/compilação do PyTorch).
  • Construção de grafo (graph building): construir um grafo estático diretamente (historicamente, modo grafo do TensorFlow).
  • Híbrido: permitir dinamismo limitado com guardas (guards) (comum no torch.compile do PyTorch).

2) Otimizar em alto nível

Nesta etapa, o compilador pode aplicar transformações algébricas e no grafo:

  • Dobrar constantes (constant folding): computar subgrafos constantes em tempo de compilação.
  • Eliminação de subexpressões comuns (common subexpression elimination): reutilizar computações repetidas.
  • Fusão de operadores (operator fusion): combinar cadeias elementwise e ops produtor-consumidor.
  • Propagação de layout (layout propagation): escolher layouts do tipo NCHW vs NHWC ou layouts em blocos para Tensor Cores.
  • Eliminar transposições/reestruturações (reshapes) redundantes.

3) Rebaixar (lowering) para kernels específicos de hardware

Isso inclui:

  • selecionar implementações (ex.: qual algoritmo de convolução),
  • tileamento e vetorização (vectorization),
  • mapear para threads/warps/blocks (GPU) ou matrizes sistólicas (systolic arrays) de TPU,
  • gerar código ou chamar bibliotecas do fornecedor.

Frequentemente, a parte de “geração de código (codegen)” é uma combinação de:

  • kernels gerados sob medida para padrões fundidos, mais
  • chamadas para cuBLAS/cuDNN (GPU) ou bibliotecas específicas de TPU para GEMM/conv.

4) Execução em tempo de execução

Um tempo de execução então lida com:

  • alocação e reutilização de memória,
  • lançar kernels em streams (streams),
  • escalonamento, sincronização e eventos (events),
  • cache de artefatos compilados,
  • lidar com formas dinâmicas (se suportado),
  • coletivas multi-dispositivo (collectives) (frequentemente via NCCL em GPUs).

Entender o tempo de execução importa porque “otimizações de tempo de compilação” só ajudam se o tempo de execução conseguir executá-las com eficiência.

Otimizações-Chave de Compilador (O Que Você Geralmente Ganha “De Graça”)

Fusão de operadores (a grande)

Considere um bloco simplificado de inferência:

y = relu(x @ W + b)

Sem fusão:

  1. o kernel de multiplicação geral de matrizes (GEMM) grava x @ W na memória
  2. o kernel de soma lê, grava o resultado
  3. o kernel de ReLU lê, grava o resultado

Com fusão (conceitualmente):

  • a GEMM produz tiles, e bias+ReLU são aplicados antes de gravar a saída final.

Isso reduz pressão sobre a largura de banda de memória e lançamentos de kernel.

Seleção de kernel + autoajuste

Para a mesma convolução, pode haver múltiplos algoritmos (GEMM implícita, baseado em FFT, Winograd, etc.). Compiladores/tempos de execução fazem benchmark de candidatos (“táticas”) e escolhem o mais rápido para uma dada forma e GPU.

Isso é uma parte importante da história de desempenho do TensorRT, e também aparece no XLA e em backends de biblioteca.

Planejamento de memória

Compiladores podem reutilizar buffers cujas janelas de vida não se sobrepõem. Um bom plano de memória:

  • reduz o pico de memória,
  • melhora o comportamento de cache,
  • pode reduzir sobrecarga do alocador.

Isso afeta diretamente a viabilidade do tamanho de lote e a vazão (novamente, intimamente ligado a Memória e Largura de Banda).

Propagação de layout e tipo de dado

Um compilador pode manter tensores em layouts mais amigáveis a Tensor Cores e reduzir transposições. Ele também pode propagar FP16/BF16/INT8 onde for seguro.

Este é um dos motivos pelos quais “apenas fazer cast para FP16” às vezes tem desempenho pior: desempenho ótimo frequentemente precisa de escolhas coordenadas de layout + kernel, não apenas mudanças de tipo de dado.

Tempos de Execução: A Outra Metade da História

Um compilador produz um plano executável, mas o tempo de execução determina quão suavemente ele roda em produção.

Principais responsabilidades do tempo de execução incluem:

  • Orquestração de execução: despachar kernels, gerenciar streams CUDA, sobrepor computação e cópias de memória.
  • Alocador de memória: pooling, cache, controle de fragmentação.
  • Tratamento de formas: lidar com formas dinâmicas ou múltiplos “perfis”.
  • Comunicação multi-dispositivo: coletivas (all-reduce, all-gather) para treinamento distribuído.
  • Cache de compilação: evitar compilação repetida entre execuções/réplicas.

Se você já observou “a primeira requisição é lenta, as seguintes são rápidas”, você vivenciou efeitos de compilação e cache no tempo de execução.

XLA na Prática (JAX e TensorFlow)

O que é o XLA

XLA é um compilador para álgebra linear e computações tensoriais. Ele recebe grafos/IR de frameworks (notavelmente JAX e TensorFlow) e os rebaixa para executáveis otimizados para CPU, GPU ou TPU.

Uma representação intermediária chave no XLA é frequentemente chamada de HLO (High Level Optimizer) IR, que é bem adequada à otimização do grafo inteiro (whole-graph optimization) (fusão, layout, simplificação algébrica).

Por que as pessoas usam XLA

  • Menos lançamentos de kernel e mais fusão.
  • Otimização do programa inteiro (whole-program optimization): XLA pode otimizar através de fronteiras de ops que um tempo de execução por-op não consegue.
  • Suporte a TPU: XLA é central na pilha de software da TPU.
  • Particionamento SPMD (SPMD partitioning): XLA pode transformar computações para execução multi-dispositivo (especialmente comum em TPU; também relevante em pilhas de GPU com o tempo de execução correto).

Exemplo: `jit` do JAX

No JAX, jit compila uma função Python em um executável do XLA:

import jax
import jax.numpy as jnp

@jax.jit
def mlp_step(x, w1, b1, w2, b2):
    h = jnp.tanh(x @ w1 + b1)
    y = h @ w2 + b2
    return y

x  = jnp.ones((1024, 4096), dtype=jnp.float16)
w1 = jnp.ones((4096, 4096), dtype=jnp.float16)
b1 = jnp.zeros((4096,), dtype=jnp.float16)
w2 = jnp.ones((4096, 1024), dtype=jnp.float16)
b2 = jnp.zeros((1024,), dtype=jnp.float16)

y = mlp_step(x, w1, b1, w2, b2)

O que acontece conceitualmente:

  • Na primeira chamada com um novo conjunto de formas/tipos, o JAX rastreia a função e a compila com XLA.
  • Chamadas subsequentes reutilizam o executável compilado (caminho rápido).
  • O XLA pode fundir ops elementwise (como + b e tanh) ao redor de GEMMs quando possível, e ele escolherá implementações eficientes de GEMM.

Pegadinha prática: a compilação é especializada por forma

Se suas entradas mudam de forma com frequência (ex.: comprimentos de sequência variáveis), você pode disparar recompilações e ver picos de latência. Soluções incluem:

  • preencher (padding)/bucketizar (bucketing) formas,
  • usar recursos de polimorfismo de forma (shape polymorphism) (quando disponíveis),
  • controlar o comportamento de formas dinâmicas.

TensorFlow e XLA

O TensorFlow pode usar XLA de múltiplas formas, incluindo “compilação XLA” para partes de um grafo. Na prática, os mesmos temas se aplicam: captura de grafo, otimização/fusão e então execução via um tempo de execução.

O XLA é mais benéfico quando o grafo tem:

  • muitas ops pequenas que podem ser fundidas,
  • formas estáveis,
  • sensibilidade de desempenho à sobrecarga de lançamento de kernel.

TensorRT na Prática (Inferência em GPU NVIDIA)

O que é o TensorRT

TensorRT é o otimizador e tempo de execução de inferência da NVIDIA para GPUs da NVIDIA. Ele pega uma rede treinada (frequentemente via ONNX ou integrações com frameworks) e produz um motor (engine) otimizado para uma GPU e configuração alvo.

O TensorRT é tipicamente usado para implantação (deployment)/inferência, não para treinamento. Seu foco de design é:

  • baixa latência (latency),
  • alta vazão (throughput),
  • precisão mista (mixed precision) e quantização,
  • seleção e fusão agressivas de kernels.

Conceitos-chave: motores, táticas e perfis

  • Motor: o artefato compilado que o TensorRT constrói para uma arquitetura de GPU e rede específicas.
  • Táticas (tactics): implementações alternativas de kernel para camadas (especialmente convoluções/GEMMs). O TensorRT faz benchmark das táticas para escolher a mais rápida.
  • Perfis de otimização (optimization profiles): para formas dinâmicas, você especifica formas mín/ótima/máx; o TensorRT constrói planos que funcionam em todo esse intervalo.

Exemplo: construindo um motor TensorRT a partir de ONNX (ilustrativo)

Em sistemas reais, você adicionaria tratamento de erro e configuraria memória/workspace com cuidado, mas o fluxo de trabalho central se parece com:

import tensorrt as trt

logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)
parser = trt.OnnxParser(network, logger)

with open("model.onnx", "rb") as f:
    parser.parse(f.read())

config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 2 << 30)  # 2GB

# Enable FP16 if supported
if builder.platform_has_fast_fp16:
    config.set_flag(trt.BuilderFlag.FP16)

serialized_engine = builder.build_serialized_network(network, config)
with open("model.trt", "wb") as f:
    f.write(serialized_engine)

Em tempo de implantação, você carrega o motor e roda inferência por meio do tempo de execução do TensorRT (frequentemente integrado em uma pilha de serving (serving stack)).

Por que o TensorRT tende a ser rápido

  • Ele realiza fusão de camadas (layer fusion) e autoajuste de kernel especificamente para inferência.
  • Ele suporta caminhos FP16/BF16/INT8 fortemente integrados ao hardware da NVIDIA.
  • Ele otimiza reutilização de memória e pode reduzir sobrecarga no caminho quente de inferência.

Pegadinha prática: formas dinâmicas exigem planejamento

Se você atende tamanhos de lote ou comprimentos de sequência variáveis, geralmente precisa de:

  • perfis de otimização explícitos,
  • benchmark cuidadoso nas formas ótimas que importam para você,
  • às vezes múltiplos motores para diferentes regimes (ex.: sequências curtas vs longas).

XLA vs TensorRT: Como Pensar a Diferença

Eles se sobrepõem em “compilar programas tensoriais”, mas miram pontos diferentes do ciclo de vida.

XLA (compilador mais geral, frequentemente treinamento + pesquisa)

  • Integrado aos modelos de execução de JAX/TF.
  • Ótimo para fusão de função inteira e redução de sobrecarga de Python.
  • Frequentemente usado em loops de treinamento e pesquisa.
  • Suporta TPU como backend de primeira classe.
  • A compilação pode acontecer em estilo compilação JIT (just-in-time compilation) durante o desenvolvimento; cache é importante.

TensorRT (compilador/tempo de execução especializado em inferência para NVIDIA)

  • Focado em implantação e inferência de produção em GPUs da NVIDIA.
  • Forte em seleção de táticas, otimização de kernel de baixo nível e precisão mista.
  • Tipicamente consome um modelo exportado (ex.: ONNX).
  • O tempo de construção do motor pode ser não trivial, mas o desempenho em tempo de execução é excelente.

Um padrão comum na indústria é:

  • treinar em PyTorch/JAX/TF,
  • exportar para ONNX (ou usar integrações diretas),
  • otimizar e servir com TensorRT em GPUs da NVIDIA.

Por que a Compilação Afeta Latência, Vazão e Custo de Formas Diferentes

Latência

A compilação ajuda a latência ao:

  • reduzir lançamentos de kernel,
  • fundir sequências limitadas por memória,
  • escolher kernels de menor latência.

Mas também pode piorar a latência p99 se:

  • você compila no caminho da requisição (JIT em tempo de execução),
  • variabilidade de formas causa compilação repetida,
  • você tem faltas de cache (novas versões de modelo, novas formas).

Vazão

A vazão melhora quando a compilação:

  • aumenta a intensidade aritmética (arithmetic intensity) (mais computação por byte movido),
  • melhora ocupação (occupancy) e uso de Tensor Cores,
  • reduz pontos de sincronização.

Isso importa especialmente para grandes Redes Neurais e Transformers, onde blocos de atenção e MLP podem ser otimizados intensamente.

Custo

Custo frequentemente é “horas de GPU” ou “dólares de GPU por milhão de tokens”. A compilação reduz custo ao:

  • aumentar tokens/s ou exemplos/s por GPU,
  • permitir que GPUs menores ou mais baratas atinjam SLOs (service level objectives),
  • habilitar inferência em menor precisão com segurança.

Em produção, melhorias de custo frequentemente vêm de combinar compilação com:

Orientação Prática: Quando e Como Usar Compilação

Quando a compilação mais ajuda

A compilação tipicamente traz grandes ganhos quando você tem:

  • Muitas ops pequenas (cadeias elementwise, layer norm, bias-adds, funções de ativação)
  • Formas estáveis (tamanho de lote e comprimento de sequência fixos, ou um pequeno número de buckets)
  • Execução repetida (passos de treinamento, inferência em lote, serviços de longa duração)
  • Subutilização de GPU no perfilamento (profiling) (muitos kernels pequenos, baixa ocupação de SM (SM occupancy))

Quando os ganhos podem ser limitados

  • O modelo é dominado por poucas GEMMs/convoluções grandes já bem tratadas por bibliotecas do fornecedor.
  • As formas mudam o tempo todo (causando recompilações).
  • Você está limitado por pipelines de entrada (input pipelines), pré-processamento de CPU (CPU preprocessing) ou E/S de rede (network I/O), e não por computação no dispositivo.

O que medir

Antes e depois de habilitar um compilador, meça:

  • latência ponta a ponta (end-to-end latency) (p50/p95/p99),
  • vazão em tamanhos de lote realistas,
  • utilização de GPU e largura de banda de memória,
  • tempo de compilação (compile time) e taxa de acerto de cache (cache hit rate),
  • deriva de acurácia (accuracy drift) ao habilitar menor precisão.

As ferramentas variam por pilha (Nsight Systems para timelines de GPU, profilers do framework, etc.), mas o objetivo é sempre o mesmo: verificar que a compilação reduz tempo gasto em tráfego de memória e na sobrecarga de lançamento de kernel.

Um padrão amigável para produção: compilar offline, rodar rápido online

  • Para TensorRT: construir motores antecipadamente (CI/CD ou pipeline de build), enviar artefatos, evitar construir no caminho de serving.
  • Para XLA/JAX: aquecer o serviço executando entradas representativas na inicialização, para que a compilação aconteça antes de atender tráfego real; garantir bucketização de formas para maximizar reutilização de cache.

Como Compiladores Interagem com o Resto do Sistema

Compiladores não existem isoladamente; eles fazem parte de uma pilha de “sistemas de computação”:

  • CUDA e kernels de GPU: entender o básico como streams e lançamentos de kernel ajuda a interpretar por que fusão importa (Noções Básicas de CUDA).
  • Limites e largura de banda de memória: fusão e reutilização de buffers atacam diretamente gargalos de largura de banda (Memória e Largura de Banda).
  • Quantização: compiladores/tempos de execução implementam e aceleram caminhos de baixa precisão (Quantização).
  • Escalonamento e utilização do cluster: inferência mais rápida pode mudar estratégias de batching e eficiência de empacotamento de GPU (ver Escalonamento de GPU e Filas de Cluster).

Na prática, os melhores resultados vêm de tratar compilação como uma alavanca entre várias: arquitetura do modelo, precisão, batching e orquestração em nível de sistema interagem entre si.

Resumo

  • Compilação importa em ML porque reduz tráfego de memória, funde ops, seleciona kernels eficientes e adapta a execução ao hardware.
  • Tempos de execução importam porque gerenciam memória, escalonamento, cache e execução multi-dispositivo — frequentemente determinando se ganhos do compilador se traduzem em ganhos reais.
  • XLA é um compilador geral de ML usado intensamente em JAX/TF para otimização de programa inteiro (comumente treinamento e pesquisa, também inferência).
  • TensorRT é um compilador/tempo de execução de inferência altamente otimizado e especializado em GPUs da NVIDIA, enfatizando seleção de táticas, fusão e precisão mista.
  • Os maiores ganhos práticos vêm de formas estáveis, execução repetida e medição cuidadosa de latência/vazão — enquanto se gerencia a sobrecarga de compilação via cache e aquecimento.

Se você já entende ideias em nível de modelo como Retropropagação e noções básicas de hardware como Introdução ao Hardware, compiladores e tempos de execução são a próxima camada que explica por que o mesmo modelo pode rodar 2–10× mais rápido (ou mais lento) dependendo de como ele é rebaixado para o dispositivo.