Transformações Eficientes no cuDF com Compilação JIT

TL;DR

• A compilação JIT no cuDF permite fusão de kernels compilando kernels fused em tempo de execução, reduzindo a materialização de intermediários e melhorando a localidade de dados.
• A partir do cuDF 25.08, transformações JIT adicionam suporte a operadores não disponíveis no AST, incluindo o operador ternário e funções de string como find e substring.
• Benchmarks mostram ganhos de cerca de 2x–4x em várias traduções de string, com menor contagem de kernels e melhor localidade de cache impulsionando a melhoria.
• Existe um custo inicial de tempo de compilação para JIT (~600 ms por kernel) se nenhum kernel em cache for encontrado; execuções subsequentes são bem mais rápidas (~3 ms para carregar o kernel em cache).

Contexto e antecedentes

O cuDF, parte da suíte RAPIDS, oferece um conjunto amplo de algoritmos ETL para processar dados em GPUs. Para usuários de pandas, o cuDF oferece algoritmos acelerados com zero mudança de código via cudf.pandas. Para desenvolvedores C++, o cuDF expõe um submódulo em C++ que usa visões sem propriedade como entrada e tipos com propriedade como saída, facilitando o raciocínio sobre ciclos de vida dos dados e ampliando a composibilidade das APIs. Uma limitação conhecida desse modelo de várias etapas com materialização de intermediários é o aumento do tráfego de memória na GPU. A fusão de kernels surge como solução poderosa: executar múltiplos cálculos dentro de um único kernel GPU pode melhorar o throughput e reduzir o tráfego de memória. Este post explica como a compilação JIT traz fusão de kernels para o modelo de programação C++ do cuDF, entregando maior throughput e utilização mais eficiente de memória e compute na GPU. Em expressões escalares, um árvore de operandos e operadores é avaliada para produzir uma única coluna de saída. O cuDF oferece três abordagens: precompiled, AST e JIT transform. A abordagem precompiled chama APIs públicas do libcudf para cada operador, recursivamente, computando a árvore. A vantagem é o apoio mais amplo a tipos de dados e operadores, mas a desvantagem é a materialização de intermediários na memória global da GPU entre operações. A abordagem AST usa a API compute_column para percorrer a árvore completa com um kernel especializado, com paralelismo por linha de GPU. O interpretador AST facilita fusão de kernels em cuDF, mas tem limitações de suporte a tipos de dados e operadores. A abordagem JIT transform em cuDF usa NVRTC para compilar um kernel personalizado que realiza transformações arbitrárias, criando kernels fundidos em tempo de execução. NVRTC é uma biblioteca de compilação em tempo de execução para CUDA C++ que cria kernels fundidos durante a execução. A principal vantagem do JIT é que ele utiliza kernels otimizados para a expressão a ser avaliada, permitindo uma alocação eficiente de recursos da GPU e melhor aproveitamento de caches, ao invés de reservar registradores para cenários de pior caso. A partir do cuDF 25.08, o JIT transform adiciona suporte a operadores que o processamento AST não suportava, incluindo o operador ternário para if-else e funções de string como find e substring. O principal inconveniente é o tempo de compilação do kernel (~600 ms) na primeira execução sem kernel em cache, ou a dependência de cache para reduzir esse overhead. O reuso de kernels no cache é detalhado neste post. O repositório rapidsai/cudf no GitHub fornece uma suíte de exemplos string_transforms para demonstrar a manipulação de strings com abordagens precompiled e JIT transform. Os casos de exemplo são UDFs de processamento de strings, como extract_email_jit e extract_email_precompiled, que exploram uma computação que valida o formato de um endereço de e-mail e extrai o provedor, retornando unknown para entradas malformadas. A Figura 2 demonstra o exemplo extract_email_precompiled com a lógica destacada para localizar o ’@’ e o ’.’; a Figura 3 demonstra o exemplo extract_email_jit, que usa um UDF em C++ bruto para simplificar a lógica com if-else e retornos antecipados. O objetivo é ilustrar como o JIT pode tornar o processamento de strings mais eficiente ao reduzir intermediários. O uso do JIT resulta em runtimes mais rápidos quando comparado ao pré-compilado. A principal fonte de melhoria é a menor contagem de kernels. Quando o mesmo trabalho pode ser feito com menos kernels, as localizações de cache são mais eficazes e os registradores da GPU podem manter intermediários, em vez de armazená-los na memória global para kernels subsequentes. A Figura 4 mostra a linha do tempo das execuções de kernels para os três exemplos de transformação de strings, onde barras azuis representam os kernels lançados de forma precompiled e barras verdes representam os kernels lançados pelo JIT transform. A coleta de dados ocorreu com 200M de linhas de entrada (12,5 GB) em hardware NVIDIA GH200 Grace Hopper com recursos assíncronos de memória.

O que há de novo

A rota JIT transform do cuDF permite gerar kernels fundidos em tempo de execução para a transformação ou UDF específica. Isso é feito pela compilação de um kernel personalizado com NVRTC, permitindo execução especializada e fundida sem a necessidade de materializar intermediários entre etapas. Principais capacidades novas incluem:

• Fusão baseada em JIT de expressões, incluindo operações escalares e de string que não eram suportadas pelo AST anteriormente.
• Suporte a operadores como o ternário e funções de string como find e substring na trajetória JIT (desde cuDF 25.08).
• Benefícios de performance demonstrados em tarefas de transformação de strings; o caminho JIT reduz a contagem de kernels, melhora a localidade de cache e mantém mais dados nos registradores da GPU.
• Medições com o NVIDIA GH200, usando 200M de linhas e 12,5 GB de dados, evidenciam ganhos do JIT com a fusão de kernels.
• Um mecanismo de cache de kernel é utilizado para amortizar o custo de compilação. A primeira execução de um kernel JIT pode levar cerca de 600 ms para compilar, enquanto carregá-lo a partir do cache leva cerca de 3 ms. Subsequentes execuções no mesmo processo não têm overhead adicional. Na prática, os exemplos string_transforms demonstram como o uso de JIT pode simplificar a lógica de processamento de strings, por exemplo, com o extract_email_jit: ele localiza os caracteres e separa o provedor em uma única etapa, reduzindo a necessidade de construção de várias colunas intermediárias que ocorrem na abordagem precompiled.

Por que isso importa (impacto para desenvolvedores/empresas)

Para desenvolvedores que constroem pipelines de dados em GPUs, transformar o modo como as transformações são executadas pode levar a maior throughput e menor uso de memória. Fusão de kernels reduz overhead de lançamento e tráfego de memória entre operações, o que é particularmente relevante para workloads de strings e transformações complexas. Do ponto de vista corporativo, a capacidade de acelerar UDFs sem reescrever código em C++ ou adotar novas abstrações facilita ciclos de desenvolvimento mais curtos. A ampliação do suporte de cuDF 25.08 para operadores ternários e funções de string amplia as situações em que a fusão JIT pode ser aplicada, promovendo maior escalabilidade para análises em GPUs modernas. Em termos práticos, a JIT também ajuda a processar tamanhos de dados maiores antes de atingir limites de memória da GPU devido à menor materialização de intermediários.

Detalhes técnicos ou Implementação

• Abordagens de avaliação em cuDF: precompiled, AST e JIT transform. Precompiled usa APIs libcudf por operador, com amplo suporte, mas pode exigir intermediários. AST percorre a árvore de expressão com um kernel dedicado, com fusão potencial, porém com limitações de suporte a operadores/tipos. JIT transform utiliza NVRTC para gerar um kernel fundido em tempo de execução para a expressão ou UDF.
• Vantagens do JIT: kernels fundidos específicos para a transformação, melhor alocação de recursos e localidade de cache, reduzindo o número total de kernels.
• Novos operadores no JIT desde cuDF 25.08: o operador ternário e funções de string como find e substring.
• Cache de kernels e aquecimento: a primeira execução exige o tempo de compilação (~600 ms por kernel) na ausência de kernel em cache; carregar do cache leva ~3 ms. Execuções subsequentes no mesmo processo não têm overhead de JIT.
• Medições envolvem 200M de linhas, 12,5 GB, e hardware GH200 com memória assíncrona, com comparação entre kernels precompiled (azuis) e JIT (verdes).

Notas de implementação na prática

• Para facilitar testes e implantação, a RAPIDS fornece containers Docker com releases e builds noturnos para cuDF, incluindo binários pré-compilados e o submódulo libcudf.
• Os usuários podem instalar binários via canal Conda rapidsai-nightly, e acessar os exemplos string_transforms e testes associados.

Principais conclusões

• Transformações JIT permitem kernels fundidos em tempo de execução, reduzindo tráfego de memória e aumentando o throughput para transformações de strings e escalares.
• O cuDF 25.08 expande o suporte JIT para incluir operadores como ternário e funções de string como find e substring.
• Observa-se, em cenários práticos, ganhos de 2x–4x para algumas cargas de trabalho de transformação de strings, impulsionados por menos kernels e melhor localidade de cache; dados maiores tendem a apresentar benefícios mais perceptíveis.
• O custo inicial de compilação do JIT (~600 ms por kernel) pode ser amortizado com caches apropriados; execuções subsequentes tendem a ser rápidas (~3 ms para carregar o kernel cache).
• O uso efetivo do JIT pode depender do pré-carregamento de kernels no cache para obter vantagens já nas primeiras leituras de dados.

Referências

• NVIDIA blog: Efficient Transforms in cuDF Using JIT Compilation