Décorateurs JIT Python : trois stratégies d’implémentation

Cet article explore le modèle populaire de décorateur JIT en Python, notamment son utilisation dans les bibliothèques JAX et Triton. L’auteur implémente trois décorateurs JIT à partir de zéro à l’aide d’un exemple simplifié : basé sur l’AST, basé sur le bytecode et basé sur le traçage. L’approche basée sur l’AST manipule directement l’arbre de syntaxe abstraite ; l’approche basée sur le bytecode utilise l’interpréteur de bytecode Python ; et l’approche basée sur le traçage construit un IR d’expression en traçant l’exécution de la fonction au moment de l’exécution. L’article détaille les avantages et les inconvénients de chaque approche et utilise JAX et Numba comme exemples pour illustrer leurs stratégies dans les applications du monde réel.