Typinferenz

StableHLO wurde ursprünglich vom MHLO-Dialekt per Bootstrapping übertragen und die MHLO-Implementierung der Typinferenz übernommen. Der Implementierungsfortschritt wird in status.md verfolgt.

Die folgenden vorgeschlagenen Richtlinien sollen die Implementierung qualitativ hochwertiger Verifier und Formfunktionen für StableHLO-Vorgänge gewährleisten.

Vorschlag

Diese Vorschläge gelten sowohl für die Überprüfung vorhandener Implementierungen als auch für das Erreichen neuer Vorgänge, bis eine umfassende Abdeckung erreicht ist.

(P1) StableHLO-Spezifikation als zentrale Datenquelle verwenden

Die spec ist die zentrale Datenquelle für alle Prüfer und Formfunktionen der StableHLO-Vorgänge. Die vorhandenen Verifizierungs- und Formfunktionen jeder Operation müssen noch einmal überprüft werden, um vollständig auf die Spezifikation abgestimmt zu sein. Beachten Sie, dass das Spezifikationsdokument ständig weiterentwickelt wird. Wenn die Spezifikation für einen Vorgang nicht verfügbar ist, sollte stattdessen die XLA-Implementierung verwendet werden, einschließlich xla/service/shape_inference.cc und xla/service/hlo_verifier.cc. Die XLA-Implementierung deckt keine unbegrenzte Dynamik ab. Bei unbegrenzter Dynamik wenden wir daher gesunden Menschenverstand an, bis der Dynamik-RFC verfügbar ist.

(P2) ODS optimal nutzen

In ODS-Dateien (z. B. StablehloOps.td) werden Vorgänge mit Merkmalen und Typen für alle Operanden/Attribute/Ergebnisse definiert und die Überprüfungen durchgeführt. Daher ist in den Verifizierer- oder Formfunktionen für die Eigenschaften, die bereits von der ODS garantiert werden, KEIN Bestätigungscode erforderlich. Entfernen Sie den Bestätigungscode, wenn er mit ODS dupliziert wird, da er nie ausgelöst wird.

Müssen wir Tests für die Einschränkungen aus der ODS hinzufügen? Weitere Informationen findest du unter Testrichtlinien festlegen.

(P3) Bestätigungscode in Prüfern und Formfunktionen verwalten

Beides:

  • verifiers: implementiert durch Op::verify() und
  • Formfunktionen: implementiert durch InferTypeOpInterfaces wie Op::inferReturnTypes() oder Op::inferReturnTypeComponents

kann einen Bestätigungscode zur Überprüfung von Operanden/Attributen/Ergebnissen enthalten. Eine anfängliche Aufteilung könnte so aussehen: Lassen Sie die Prüfer die Operanden/Attribute prüfen, dann lassen die Formfunktionen nur abgeleitete Ergebnistypen berechnen und die Kompatibilität mit den echten Ergebnistypen vergleichen. In Wirklichkeit bringt diese Aufteilung jedoch einige Probleme mit sich:

  • Die Formfunktion kann von den automatisch generierten build()-Funktionen aufgerufen werden, ohne dass zuerst der Prüfer aufgerufen wird. Die zugehörigen Eingaben müssen also auch in der Formfunktion verifiziert werden.
  • Duplizierter Code: In Verifikationsanbietern verarbeiten wir beispielsweise die Operanden und prüfen dann einige Zwischenergebnisse. In Formfunktionen sind diese Zwischenergebnisse nützlich, um die Endergebnisse abzuleiten. Diese Zwischenergebnisse müssen zweimal berechnet werden.
  • Verwaltungsaufwand: Überprüfungen eines Vorgangs werden in zwei verschiedenen Methoden durchgeführt.

Die Lösung sieht so aus:

  1. Für die meisten Vorgänge ohne Regionen (z. B. PadOp): Versuchen Sie, den gesamten Bestätigungscode in die Formfunktionen einzufügen, und verwerfen Sie Verifikationsanbieter komplett. Wenn dies nicht möglich ist, da Rückgabetypen nicht abgeleitet werden können (z. B. mit ReshapeOp oder BroadcastInDimOp), erstellen Sie einen Prüfer, der die erforderliche Verifizierungslogik enthält. In der Regel benötigen ableitende Vorgänge wie AddOp möglicherweise weiterhin einen Verifizierer, um zusätzliche Prüfungen durchzuführen, da Einschränkungen eines bereitgestellten Rückgabetyps geprüft werden, auf den in den Typ-/Forminferenzmethoden nicht zugegriffen werden kann.

  2. Vorgänge mit Regionen (z. B. ReduceOp/IfOp; eine vollständige Liste finden Sie hier): Die automatisch generierten Builder verwenden keine Regionen als Parameter. Wenn diese Builder also eine Typinferenz beinhalten, wird die Formfunktion mit leeren Regionen aufgerufen (siehe dieses Beispiel).

    1. Wenn die Regionen nicht für die Typinferenz erforderlich sind (wie ReduceOp), geben Sie die regionsbezogene Verifizierungslogik in Verifier statt in die Formfunktionen ein. Duplizieren Sie Code, wenn sich das nicht vermeiden lässt.

    2. Wenn die Regionen für die Typinferenz erforderlich sind (IfOp/CaseOp/MapOp), muss die Formfunktion zusätzlich prüfen, ob die Regionen explizit leer sind, auch wenn die ODS möglicherweise bereits das Vorhandensein in der Op-Definition garantiert.

(P4) Testrichtlinien festlegen

Müssen Tests für Überprüfungen, die von ODS abgedeckt sind, hinzugefügt/gepflegt werden?

Das tun wir nicht. Die Tests sollten sich auf die Verifizierungs- und Formfunktionen konzentrieren, während bei Änderungen an ODS eine erneute Überprüfung dieses Vorgangs erforderlich ist.

Behalten Sie jedoch die fehlenden Elemente im Auge: Wenn der Vorgang beispielsweise das Merkmal SameOperandsAndResultShape enthält, das nur Formen, aber nicht den Elementtyp prüft, sind für die Überprüfung der Elementtypen von Operanden/Ergebnissen weiterhin Tests erforderlich.

Wo werden Verifikationsanbieter und Typinferenzen getestet?

ops_stablehlo.mlir enthält die positiven Fälle von Vorgängen und (mindestens) einen negativen Test für jeden Verifizierungsfehler. Außerdem kann geprüft werden, ob der abgeleitete Rückgabetyp mit dem tatsächlichen Ergebnistyp kompatibel ist (nicht mit diesem identisch).

infer_stablehlo.mlir prüft anhand der Zeile hlo_test_infer.get_return_type_components"(%x):..., ob die Formfunktion eines Vorgangs vorhanden ist, und prüft, ob der abgeleitete Typ genau wie erwartet übereinstimmt. Im Allgemeinen ein positiver Test pro Operation.

Deine Aufgabe:

Wenn Sie die Verifizierungs- und/oder Shape-Funktion einer Operation implementieren oder überprüfen, sollten Sie Folgendes beachten:

  1. Tragen Sie alle positiven und negativen Fälle in ops_stablehlo.mlir ein.

  2. Fügen Sie infer_stablehlo.mlir einen einzelnen positiven Test hinzu, um die Schnittstelle zu testen.

  3. (Optional) Wenn eine Operation kompliziert ist und viele Tests enthalten kann, sollten Sie im selben Ordner eine separate Testdatei mit dem Namen verify_<op_name>.mlir oder verify_<your_topic>.mlir hinzufügen.