Huawei ouvre le code source d’AscendNPU IR pour Triton et Ascend 950

Un architecte de compilateur est monté sur la scène de la rencontre technique d’HyperAI le 1er août avec un problème qui reste généralement invisible pour les utilisateurs : comment transformer quelques lignes de code tensoriel en tâches réparties entre la mémoire, les cœurs et les pipelines intégrés à la puce. Hai Lijuan, architecte d’AscendNPU IR chez Huawei, a présenté la couche de compilation ouverte qui relie Triton et d’autres langages au matériel Ascend, notamment Ascend 950.

Triton abaisse le seuil d’entrée pour écrire des opérateurs d’IA grâce à une syntaxe proche de Python et à une programmation par tuiles au niveau des blocs. Les développeurs décrivent la manière dont les données sont divisées et traitées ; le compilateur peut automatiser la coalescence de la mémoire et la gestion de la mémoire partagée. AscendNPU IR transpose ce modèle vers la puce. Fondé sur MLIR, il relie les langages de haut niveau en amont à LLVM IR en aval, puis compile le résultat en binaires Ascend.

La conception sépare ce qui peut fonctionner sur différents matériels de ce qui doit suivre les spécificités d’une puce donnée. HFusion fournit une couche indépendante du matériel pour l’optimisation de fusion multidimensionnelle. HIVM associe ensuite les tenseurs logiques à la mémoire physique, affecte les opérations cube et vectorielles à leurs cœurs respectifs et active la vectorisation pour les instructions hautes performances. La pile couvre les architectures Ascend allant de A2/A3 à Ascend 950.

Ascend 950 introduit un changement plus profond au sein de cette couche de mappage. HIVM s’étend du SIMD fondé sur la mémoire au SIMD fondé sur les registres et au SIMT, tandis que les cœurs cube et vectoriels peuvent communiquer via une liaison étroitement couplée sur la puce, au lieu d’échanger des données par l’intermédiaire de la mémoire globale. De nouvelles fonctions du compilateur prennent en charge les flux de contrôle complexes, les boucles imbriquées, le pipeline de type flash attention et un mode dans lequel deux cœurs vectoriels calculent chacun la moitié des données avec un cœur cube.

Concrètement, qu’est-ce que cela change ? Les développeurs peuvent réutiliser davantage d’approches de programmation familières tout en ciblant les puces Ascend, au lieu de reconstruire chaque opérateur autour des détails de bas niveau liés à la mémoire et aux instructions. Le code est désormais ouvert à la co-construction via Triton-Ascend et AscendNPU IR. Pour les équipes qui ne disposent pas d’environnement Ascend, la communauté de Huawei propose HiDevLab, avec 100 heures gratuites de temps de calcul. La publication du code source ne prouve pas à elle seule les performances sur des charges de travail indépendantes, mais elle offre aux développeurs externes un espace pour tester, adapter et améliorer la chaîne de compilation.