Résumé

Les charges de travail modernes en intelligence artificielle s'appuient fortement sur des noyaux informatiques optimisés, tant pour l'entraînement que pour l'inférence. Ces noyaux d'IA suivent des schémas de flux de données bien définis, tels que le déplacement de tuiles entre la DRAM et la SRAM, ainsi que l'application d'une séquence de calculs sur ces tuiles. Toutefois, la rédaction de noyaux à haute performance reste complexe, malgré la clarté de ces schémas. Pour atteindre les performances maximales, des optimisations soigneuses et centrées sur l'architecture matérielle sont nécessaires afin d'exploiter pleinement les accélérateurs modernes. Bien que les compilateurs spécifiques à domaine cherchent à réduire la charge liée à l'écriture de noyaux performants, ils peinent souvent à combler les écarts en matière d'utilisabilité et d'expressivité. Dans cet article, nous présentons TileLang, un modèle de programmation généralisé basé sur le découpage en tuiles, visant à améliorer l'efficacité du développement de noyaux d'IA. TileLang sépare l'espace de planification (liaison des threads, disposition, tensorisation et pipeline) du flux de données, et en encapsule les éléments sous forme d'annotations personnalisables et de primitives. Cette approche permet aux utilisateurs de se concentrer sur le flux de données propre au noyau, tout en laissant la majorité des optimisations aux compilateurs. Nous menons des expériences approfondies sur des dispositifs couramment utilisés, à travers de nombreuses évaluations, nos résultats montrent que TileLang atteint des performances de pointe sur des noyaux clés, démontrant ainsi que son paradigme unifié de blocs et de threads, combiné à des capacités de planification transparentes, offre à la fois la puissance et la flexibilité requises par le développement des systèmes d'IA modernes.