TensorFlow accélère l'apprentissage machine avec MLIR

Le sous-projet MLIR révélé par les ingénieurs travaillant sur le framework d'apprentissage machine TensorFlow de Google est destiné à servir de langage intermédiaire commun pour les frameworks d'apprentissage machine. MLIR, abréviation de Multi-Level Intermediate Representation, permettra de compiler les projets utilisant TensorFlow et d'autres bibliothèques d'apprentissage machine dans un code plus efficace capable de tirer le meilleur parti du matériel sous-jacent. De plus, à terme, MLIR pourrait être utilisé par tous les compilateurs et faire profiter d’autres projets de ses avantages d'optimisation, en dehors des seuls projets d'apprentissage machine.

Le langage MLIR ne ressemble pas au C++ ou à Python. Il introduit une étape de compilation intermédiaire entre ces langages de niveau supérieur et le code machine. Le framework de compilation LLVM utilise une représentation intermédiaire, ou IR, qui lui est propre. L'un des fondateurs de LLVM, Chris Lattner, est co-créateur du MLIR. Faire du MLIR un co-projet LLVM pourrait favoriser une adoption plus générale.

La compilation parallélisée, point fort de MLIR

Dans un diaporama présenté à la conférence EuroLLLVM organisée les 8 et 9 avril à Bruxelles, Chris Lattner et sa collègue Tatiana Shpeisman ont expliqué comment TensorFlow générait déjà plusieurs RI en interne, mais que ces RI disparates ne tiraient pas profit les uns des autres. MLIR fournit un IR standard unique pour tous ces sous-systèmes TensorFlow. TensorFlow migre actuellement pour utiliser le MLIR en interne. Autre avantage du MLIR : la compilation parallélisée. Avec le MLIR, un compilateur peut travailler sur différents segments de code en parallèle, ce qui permet aux modèles d'apprentissage machine - et à d'autres types d'applications - une mise en production plus rapide.

Le MLIR pourrait apporter d'autres avantages aux langages et aux frameworks en dehors de l'apprentissage machine. Par exemple, les langages basés sur LLVM comme Swift et Rust ont dû développer leurs propres IR internes, parce que de nombreuses optimisations utilisées dans ces langages ne peuvent être exprimées en LLVM. Le MLIR pourrait devenir un moyen d’expression standard de ces optimisations, avec la possibilité de les réutiliser à leur tour pour d'autres langages. Le projet MLIR est open source. Une spécification officielle est disponible pour les développeurs souhaitant générer du MLIR.