Inférence de type

StableHLO a été initialisé à l'origine à partir du dialecte MHLO et a hérité de l'implémentation MHLO de l'inférence de type. La progression de la mise en œuvre est suivie dans status.md.

Les consignes proposées ci-dessous visent à garantir l'implémentation de vérificateurs et de fonctions de forme de haute qualité pour les opérations StableHLO.

Proposition

Ces propositions s'appliquent à la fois au réexamen de mises en œuvre existantes et à la réalisation de nouvelles opérations jusqu'à une couverture complète.

(P1) Utiliser la spécification StableHLO comme source fiable

La spec est la source de référence pour tous les vérificateurs et toutes les fonctions de forme des opérations StableHLO. Les vérificateurs et les fonctions de forme existantes de chaque opération doivent être réexaminés afin d'être entièrement conformes à la spécification. Notez que le document de spécification n'arrête pas d'évoluer. Si la spécification d'une opération n'est pas disponible, utilisez plutôt l'implémentation XLA comme source de référence, y compris xla/service/shape_inference.cc et xla/service/hlo_verifier.cc. L'implémentation XLA ne couvre pas le dynamisme illimité. Par conséquent, pour un dynamisme illimité, nous utiliserons le bon sens jusqu'à ce que le RFC du dynamisme soit disponible.

(P2) Exploitez tout le potentiel d'ODS

Les fichiers ODS (comme StablehloOps.td) définissent les opérations avec des caractéristiques et des types pour chaque opérande/attribut/résultat, et effectuent les vérifications. Ainsi, AUCUN code de validation n'est nécessaire dans les vérificateurs ou les fonctions de forme pour les propriétés déjà garanties par l'ODS. Supprimez le code de validation s'il est dupliqué avec ODS, car ils ne seront jamais déclenchés.

Devons-nous ajouter des tests pour les contraintes de l'ODS ? Veuillez consulter Établir les consignes de test.

(P3) Gérer le code de validation dans les vérificateurs et les fonctions de forme

Les deux:

  • Vérificateurs: implémentés par Op::verify() et
  • Fonctions de forme: implémentées par des InferTypeOpInterface comme Op::inferReturnTypes() ou Op::inferReturnTypeComponents

peut disposer d'un code de validation pour vérifier les opérandes/attributs/résultats. La division initiale peut se présenter comme suit: laissez les vérificateurs vérifier les opérandes/attributs, puis laissez les fonctions de forme calculer uniquement les types de résultats déduits et vérifier la compatibilité avec les types de résultats réels. Cependant, cette division présente en réalité quelques problèmes:

  • La fonction de forme peut être appelée par les fonctions build() générées automatiquement, sans appeler d'abord l'outil de vérification. Les entrées associées doivent donc être vérifiées également dans la fonction de forme.
  • Code en double: par exemple, dans les vérificateurs, nous traitons les opérandes, puis nous vérifions des résultats intermédiaires. Dans les fonctions de forme, ces résultats intermédiaires sont ensuite utiles pour déduire les résultats finaux. Ces résultats intermédiaires doivent être calculés deux fois.
  • Charge de maintenance: les vérifications d'une opération se trouvent dans deux méthodes différentes.

La solution est la suivante:

  1. Pour la plupart des opérations sans région (comme PadOp): intégrez tout le code de validation dans les fonctions de forme et supprimez complètement les vérificateurs.

  2. Pour les opérations avec des régions (comme ReduceOp/IfOp ; la liste complète est disponible ici): les compilateurs générés automatiquement n'utilisent pas les régions comme paramètres. Par conséquent, si ces compilateurs impliquent une inférence de type, la fonction de forme est appelée avec des régions vides (voir cet exemple).

    1. Si les régions ne sont pas nécessaires pour l'inférence de type (comme ReduceOp), placez la logique de vérification liée à la région dans des vérificateurs au lieu des fonctions de forme. Dupliquez une partie du code si elle est inévitable.

    2. Si les régions sont nécessaires pour l'inférence de type (IfOp/CaseOp/MapOp), la fonction de forme doit également vérifier que les régions ne sont pas explicitement vides, même si l'ODS peut déjà garantir leur existence dans la définition de l'opération.

(P4) Établir des consignes de test

Dois-je ajouter/maintenir des tests pour les vérifications couvertes par ODS ?

ce n'est pas le cas. Les tests doivent se concentrer sur les vérificateurs et les fonctions de forme, tandis que les modifications apportées à ODS nécessitent une réévaluation de cette opération.

Mais faites attention aux éléments manquants: par exemple, si l'opération contient la caractéristique SameOperandsAndResultShape, qui ne vérifie que les formes, mais pas le type d'élément, la vérification des types d'éléments des opérandes/résultats nécessite toujours des tests.

Où placer les tests des vérificateurs et l'inférence de type ?

ops_stablehlo.mlir contient les cas positifs d'opérations et (au moins) un test négatif pour chaque erreur de vérification. Elle permet également de vérifier que le type renvoyé inféré est compatible avec (différent du type) le type de résultat réel.

infer_stablehlo.mlir vérifie l'existence de la fonction de forme d'une opération par ligne avec hlo_test_infer.get_return_type_components"(%x):... et vérifie que le type déduit correspond exactement à vos attentes. En général, un test positif par opération.

Que faire ?

Lorsque vous implémentez ou réexaminez la fonction de vérification et/ou de forme d'une opération:

  1. Placez tous les cas positifs et négatifs dans ops_stablehlo.mlir.

  2. Ajoutez un test positif unique dans infer_stablehlo.mlir pour tester l'interface.

  3. (Facultatif) Si une opération est complexe et peut contenir de nombreux tests, envisagez d'ajouter un fichier de test distinct nommé verify_<op_name>.mlir ou verify_<your_topic>.mlir dans le même dossier.