Transformations efficaces dans cuDF via compilation JIT

TL;DR

• cuDF utilise la compilation JIT (NVRTC) pour créer des kernels fusionnés au runtime pour des transformations arbitraires, réduisant le nombre de kernels et améliorant le débit.
• Les transformations JIT peuvent supplanter les chemins précompilés en diminuant les transferts de mémoire GPU et en améliorant la localité du cache, notamment pour les UDFs complexes.
• À partir de cuDF 25.08, JIT supporte l’opérateur ternaire (if-else) et les fonctions de chaînes telles que find et substring.
• La compilation initiale du JIT peut prendre environ 600 ms par kernel, avec un chargement en cache d’environ 3 ms. Des caches réchauffés donnent des appels ultérieurs encore plus rapides.
• Le JIT apporte des bénéfices pratiques sur des volumes de données importants, avec des gains autour de 2x à 4x pour certaines transformations de chaînes et la possibilité de traiter environ 30% de données en plus avant d’atteindre la limite mémoire GPU sur le Grace Hopper (GH200).

Contexte et arrière-plan

RAPIDS cuDF propose un ensemble étendu d’algorithmes ETL pour le traitement des données sur les GPUs. Pour les utilisateurs de pandas, cuDF offre des algorithmes accélérés sans changement de code via cudf.pandas. Pour les développeurs C++, cuDF expose un sous-module C++ qui accepte des vues non propriétaires en entrée et renvoie des types propriétaires en sortie, facilitant la gestion de la durée de vie et de la propriété des données GPU et améliorant la composition des API. Un inconvénient majeur des approches qui matérialisent des intermédiaires est le coût des transferts mémoire. La fusion de kernels permet d’effectuer plusieurs calculs dans un seul kernel, réduisant les transferts et améliorant la localité des caches. Ce billet explique comment la compilation JIT introduit la fusion de kernels dans le modèle de programmation C++ de cuDF, apportant un débit plus élevé et une utilisation plus efficace du compute et de la mémoire GPU. Dans le cadre du traitement, les expressions peuvent être représentées comme un arbre d’opérandes et d’opérateurs, chaque feuille étant une colonne ou un scalaire et chaque intersection un opérateur. Pour les expressions arithmétiques, trois options d’évaluation existent dans cuDF : précompilée, AST et transformation JIT. L’approche précompilée appelle l’API libcudf pour chaque opérateur, mais peut matérialiser des intermédiaires dans la mémoire globale. L’AST utilise compute_column et une kernel spécialisée pour parcourir et évaluer l’arbre, avec un modèle de parallélisme par ligne. LInterpreter kernel AST est utile pour la fusion de kernels, mais présente des limites en support de types et d’opérateurs. L’approche JIT utilise NVRTC pour compiler dynamiquement un kernel pour effectuer les transformations. L’approche JIT présente des avantages lorsque la travail peut être effectué à l’aide d’un seul kernel fusionné, optimisant les ressources GPU et évitant des écritures multiples en mémoire. En date du cuDF 25.08, JIT ajoute le support de certains opérateurs non pris en charge par l’exécution AST, notamment l’opérateur ternaire et des fonctions de chaînes comme find et substring. L’inconvénient principal est le temps de compilation initial du kernel (~600 ms), qui peut être amorti par la mise en cache JIT; lors des exécutions ultérieures dans le même processus, le chargement est rapide (~3 ms). Les exemples string_transforms dans le dépôt rapidsai/cudf démontrent des transformations sur chaînes en utilisant à la fois les approches précompilées et JIT, en se concentrant sur des UDF de traitement de chaînes. Les cas extract_email_jit et extract_email_precompiled illustrent le calcul qui vérifie un format d’adresse et extrait le fournisseur (provider) d’un e-mail, en gérant les entrées malformées comme unknown. Ces démonstrations montrent comment le JIT peut produire des kernels plus efficaces et des temps d’exécution plus courts grâce à des transformations fusionnées et à une utilisation optimisée des registres et de la mémoire. Les résultats montrent une réduction du nombre total de kernels et une meilleure locality, avec des gains de performance qui s’amplifient lorsque les données augmentent.

Quoi de neuf

Le chemin de transformation JIT dans cuDF permet de générer des kernels au runtime via NVRTC, autorisant la fusion de kernels pour des expressions et UDFs diverses. Les points clés incluent :

• Fusion de kernels : un kernel JIT peut exécuter plusieurs opérations d’une même expression, réduisant la materialisation des intermediaries.
• Support étendu d’opérateurs : cuDF 25.08 ajoute le support du ternaire et des fonctions de chaînes comme find et substring dans JIT, comblant des lacunes par rapport à l’AST.
• Gains de performance : les benchmarks montrent des gains de 2x à 4x pour certaines transformations de chaînes (ex.: localize_phone) et de 1x à 2x pour des cas plus simples (ex.: extract_email, format_phone).
• Échelle et mémoire : le JIT permet de traiter environ 30% de données en plus avant d’atteindre la limite mémoire de 100 Go sur Grace Hopper, grâce à moins d’intermédiaires.
• Démarrage et cache : à la première exécution, le kernel est compilé JIT et stocké dans un cache (taille typique ~130 Ko). Si le kernel n’est pas dans le cache, la compilation prend ~600 ms; s’il est présent, le chargement prend ~3 ms. Les exécutions ultérieures dans le même processus n’impliquent pas de compilation.
• Déploiement pratique : les exemples string_transforms, comme extract_email_jit vs extract_email_precompiled, illustrent comment une transformation peut être exprimée par un UDF en C++ avec un flux de contrôle clair. Les démonstrations de string_transforms indiquent que le JIT peut réduire le nombre de kernels et améliorer la locality, ce qui se traduit par des performances accrues sur des workloads plus lourds. Les graphiques temporels des launches de kernels montrent une réduction des kernels par rapport à l’approche précompilée, avec des gains plus marqués à mesure que la taille des données augmente.

Pourquoi cela compte pour les développeurs et les entreprises

• Performance et efficacité : la fusion de kernels via JIT réduit l’ensemble des passes et des données déplacées, augmentant le débit des transformations complexes impliquant des UDFs.
• Efficacité des ressources : moins de kernels signifie une meilleure utilisation des registres et moins d’écritures en mémoire globale.
• Support d’expressions complexes : les opérateurs avancés, désormais disponibles en JIT, élargissent le champ des transformations ETL sans dépendre exclusivement de pipelines précompilés.
• Considérations opérationnelles : le coût initial de compilation et la gestion du cache JIT sont des facteurs à prendre en compte; les caches préremplis permettent de bénéficier des premières lignes sans attendre la compilation.
• Écosystème et adoption : cuDF continue d’enrichir ses exemples et capacités de transformations sur chaînes, renforçant l’intégration CUDA dans l’écosystème RAPIDS.

Détails techniques ou Implémentation

• Fusion et compilation JIT : NVRTC génère des kernels spécialisés pour une expression donnée, autorisant une exécution fusionnée et évitant des intermediaries multiples.
• Support des opérateurs et évolution : depuis cuDF 25.08, ternaire et fonctions chaînes (find, substrings) sont pris en charge par JIT.
• Caractéristiques de performance : l’avantage principale vient de la réduction du nombre total de kernels et d’une meilleure locality du cache et usage des registres. Les tests sur 200 millions de lignes (≈12,5 Go) sur Grace Hopper démontrent des bénéfices clairs pour JIT par rapport au précompilé dans la suite string_transforms.
• Coûts d’initialisation et cache : compilation JIT ~600 ms par kernel sans cache; si le kernel est dans le cache, le chargement est ~3 ms. Le chemin du cache est contrôlé par LIBCUDF_KERNEL_CACHE_PATH.
• Taille des données et mémoire : le JIT permet de traiter ~30% de données en plus avant d’atteindre les 100 Go de mémoire GPU.
• Exemples et reproductibilité : les exemples string_transforms dans le dépôt rapidsai/cudf montrent des cas précompilés et JIT pour des transformations de chaînes.

Points-clés

• JIT offre la fusion de kernels pour les transformations arbitraires.
• Le support des opérateurs augmente avec la version 25.08 (ternaire, find, substring).
• Le coût initial peut être amorti via le cache JIT; les gains augmentent avec la taille des jeux de données.
• Les gains rapportés vont typiquement de 2x à 4x pour certaines transformations de chaînes.

Références

• https://developer.nvidia.com/blog/efficient-transforms-in-cudf-using-jit-compilation