Dolmetscherdesign

Datenmodell

StableHLO-Programme sind Berechnungen über Tensoren (n-dimensionale Arrays), die im aktuellen Modell mit der Klasse Tensor implementiert werden. In der zugrunde liegenden Speicherklasse detail::Buffer des Tensor-Objekts wird der mlir::ShapedType des Tensors zusammen mit einem mlir::HeapAsmResourceBlob-Objekt gespeichert, das ein änderbares Blob von Tensordaten darstellt und als zusammenhängendes Byte-Array in Haupt-zu-Klein-Reihenfolge angelegt ist. detail::Buffer-Objekte werden gezählt, um die Speicherverwaltung zu vereinfachen.

Einzelne Elemente eines Tensors werden mithilfe der Klasse Element dargestellt, die eine diskriminierte Union zum Speichern verwendet, die entweder APInt, APFloat oder pair<APFloat,APFloat> enthält. Der letzte wird zum Speichern von Elementen mit komplexen Typen verwendet.

Tensor hat die folgenden APIs, um mit den einzelnen Elementen zu interagieren:

  • Element Tensor::get(llvm::ArrayRef<int64_t> index): Zum Extrahieren eines einzelnen Tensorelements bei einem mehrdimensionalen Index index als Element-Objekt.
  • void Tensor::set(llvm::ArrayRef<int64_t> index, Element element);: Zum Aktualisieren eines Element-Objekts element in einen Tensor bei dem mehrdimensionalen Index index.

So funktioniert der Dolmetscher

Die Eingabefunktion für den Interpreter ist

SmallVector<Tensor> eval(func::FuncOp func, ArrayRef<Tensor> args);

die folgende Aufgaben ausführt:

  1. Verfolgt die SSA-Argumente von func und die zugehörigen Tensor-Laufzeitwerte, die in args bereitgestellt werden, mithilfe der Symboltabellenzuordnung M.
  2. Für jeden Vorgang innerhalb von func in SSACFG-Reihenfolge:
    • Löst eval im Vorgang aus. Extrahieren Sie für jeden SSA-Operanden des Vorgangs dessen Laufzeitwert aus M, der als Argument für den eval-Aufruf bereitgestellt wird.
    • Verfolgt die SSA-Ergebnisse des Vorgangs und des ausgewerteten Werts in M.

Der in (2) erwähnte eval auf Vorgangsebene ist für die Implementierung der Ausführungssemantik des Vorgangs verantwortlich. Es folgt ein Beispiel für stablehlo::AddOp. Im Beispiel werden einzelne Elemente der Tensoren lhs und rhs paarweise als Element-Objekte extrahiert, die dann hinzugefügt werden. Das Ergebnis der Addition, ein Element-Objekt, wird im endgültigen result-Tensor gespeichert.

Tensor eval(AddOp op, const Tensor &lhs, const Tensor &rhs) {
  Tensor result(op.getType());

  for (auto it = result.index_begin(); it != result.index_end(); ++it)
    result.set(*it, lhs.get(*it) + rhs.get(*it));

  return result;
}

Insgesamt ist das Design des Interpreters für die Lesbarkeit von Implementierungen von eval-Funktionen für einzelne Vorgänge optimiert, da er als Referenzimplementierung für StableHLO dient. Anstatt eval als Vorlagenfunktion zu definieren und mit Elementtypen zu parametrisieren, werden beispielsweise Details zur Verarbeitung verschiedener Elementtypen in Element::operator+ usw. zusammengefasst, um die Implementierung von eval zu vereinfachen.

Interpreter für konstante Faltung verwenden

Wir können den Interpretermechanismus verwenden, um Operationen mit konstanten Operandenwerten zu falten. Das folgende Code-Snippet zeigt eine Idee für die Implementierung des Faltvorgangs von stablehlo::AddOp mit Gleitkommazahlenwerten:

OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
  auto attrs = adaptor.getOperands();
  DenseElementsAttr lhsData = attrs[0].dyn_cast<DenseElementsAttr>();
  DenseElementsAttr rhsData = attrs[1].dyn_cast<DenseElementsAttr>();
  if (!lhsData || !rhsData) return {};

  auto lhs = Tensor(lhsData);
  auto rhs = Tensor(rhsData);
  auto result = eval(*this, lhs, rhs);

  SmallVector<APFloat> values;
  for (auto i = 0; i < result.getNumElements(); ++i) {
    Element element = result.get(i);
    values.push_back(element.getValue().cast<FloatAttr>().getValue());
  }

  return DenseElementsAttr::get(result.getType(), values);
}

Derzeit arbeiten wir nicht aktiv daran, den Interpreter in das kontinuierliche Zusammenklappen einzubinden, da wir nicht planen, Ordner für StableHLO zu implementieren. Wir planen jedoch, den Interpreter für die konstante Faltung in MHLO in Zukunft zu nutzen. Dann wird die Ergonomie des Code-Snippets oben verbessert (z. B. könnte eine Hilfsfunktion eingerichtet werden, die konstante Operanden in Tensor-Objekte packt und Tensor-Ergebnisse in OpFoldResult entpackt).

StableHLO-Interpreter testen

Der Interpreter nimmt als Eingaben (A) ein StableHLO-Programm und (B) Datenwerte an, die in das Programm eingespeist werden, und generiert Ausgabedatenwerte, die mit den vom Nutzer bereitgestellten erwarteten Datenwerten abgeglichen werden. Die Datenwerte (B) sind im Programm selbst mithilfe von stablehlo.constant-Vorgängen hartcodiert. Der Interpreter wertet das Eingabeprogramm aus. Die Ausgaben des zu testenden Vorgangs werden wie unten gezeigt mithilfe von Prüfungen (z.B. check.expect_eq, check.expect_almost_eq) geprüft. check.expect_eq und check.expect_eq_const prüfen die bitweise Gleichheit für jeden unterstützten Typ. check.expect_almost_eq und check.expect_almost_eq_const prüfen auf eine nahezu Gleichheit innerhalb einer Toleranz, wie in der Testrichtlinie (G6) für Gleitkomma- und komplexe Typen erläutert.

// CHECK-LABEL: Evaluated results of function: add_op_test_ui4
func.func @add_op_test_ui4() {
  %0 = stablehlo.constant dense<[0, 2]> : tensor<2xui4>
  %1 = stablehlo.constant dense<[15, 3]> : tensor<2xui4>
  %2 = stablehlo.add %0, %1 : tensor<2xui4>
  check.expect_eq_const %2, [15, 5] : tensor<2xui4>
  func.return
}

Ein Test-Dienstprogramm stablehlo-translate --interpret (Code) ist für das Parsen des Programms und die Interpretation jeder Funktion verantwortlich, einschließlich der Vorgänge, aus denen die Funktion besteht. Wir haben eine dedizierte Testsuite, die aus mehreren Tests besteht, die für jeden StableHLO-Vorgang verschiedenes Laufzeitverhalten ausführen. Die Tests finden Sie hier (z.B. interpret_*.mlir).

Testrichtlinien

(G1) Müssen wir für jeden Vorgang alle unterstützten Typen testen?

Wir können eine Kombination der folgenden Regeln verwenden, um zu entscheiden:

  1. Wenn bei der Implementierung eines Vorgangs Code in der entsprechenden eval-Funktion vorhanden ist, um einen bestimmten Typ zu verarbeiten, müssen Tests vorhanden sein, um diesen Typ abzudecken. Beispielsweise gibt es für den add-Vorgang exklusiven Code zur Verarbeitung von Ganzzahl-, booleschen, Gleitkomma- und komplexen Typen. Daher benötigen wir einen Test für jede Typkategorie.

  2. Wenn ein Satz von Typen in der entsprechenden eval-Funktion einheitlich verarbeitet wird, sollte ein einzelner Test für alle diese Typen ausreichen. Für den add-Vorgang werden beispielsweise alle Varianten von Ganzzahltypen (si4, u4, si8, u8 usw.) mit llvm::APInt-APIs gleich behandelt. Daher können wir das Hinzufügen von Tests für jede dieser Varianten überspringen und stattdessen einen einzigen repräsentativen Test hinzufügen. Um Unklarheiten bei der Auswahl des Vertreters zu vermeiden, sollten wir uns an die folgenden Richtlinien halten:

    • Wenn alle einheitlich verarbeiteten Typen denselben primitiven Typ haben (d.h., wenn es sich alle um Ganzzahl-, Gleitkomma- oder komplexe Typen handelt), wählen Sie den Typ mit der maximalen Bitbreite aus.
    • Wenn alle einheitlich verarbeiteten Typen eine Mischung aus primitiven Typen haben, wählen Sie den Typ mit dem folgenden primitiven Typ in absteigender Reihenfolge der Präferenz aus: Ganzzahl, Gleitkommazahl, boolescher Wert, komplex.

(G2) Wie entscheiden wir über die Anzahl der Tests, die erforderlich sind, um das Verhalten einer Operation abzudecken?

Das Ziel besteht darin, die Logik des Interpreters für den Vorgang (d.h. alle Eckfälle der Implementierung) mit einer minimalen Anzahl von Tests umfassend abzudecken. Für eine gute Verwaltbarkeit ist es wichtig, die Anzahl der Tests möglichst gering zu halten. Je weniger Tests wir haben, desto einfacher ist es, diese zu prüfen und dafür zu sorgen, dass sie den Vorgang umfassend abdecken. Daher erwarten wir, dass die meisten einfacheren Vorgänge am Ende nur einen Test haben. Wenn eine umfassende Abdeckung aus gutem Grund nicht praktikabel ist, können Sie bei >= 90 % aufhören. Dies wird von Fall zu Fall im Rahmen der Überprüfung von Pull-Anfragen entschieden.

(G3) Wie wäre es, Tests für die Interpreter-Infrastruktur hinzuzufügen?

Die Infrastruktur des Dolmetschers ist größtenteils einfach und kann unserer Vertrauensbasis hinzugefügt werden. Der einzige nicht triviale Teil ist, wie verschiedene Typen in den zugrunde liegenden Interpreterspeicher gepackt und daraus entpackt werden. Wie in G1 erläutert, werden nur die unterschiedlichen Vorgangstypen getestet. Dadurch kann es vorkommen, dass der Verpackungs-/Unverpackungscode, der verschiedenen Varianten von Ganzzahl-/Gleitkommatypen entspricht, während des Tests nicht vollständig abgedeckt wird. Damit eine vollständige Abdeckung gewährleistet ist, können wir eine Operation wie constant auswählen, die alle StableHLO-Elementtypen unterstützt, und umfassende Tests schreiben.

(G4) Wenn die Implementierung eines Vorgangs von anderen Vorgängen abhängt, sollten wir dann Tests für Letzteren schreiben?

Nein. Die Implementierung von batch_norm_grad kann beispielsweise auf divide, subtract, multiply und anderen basieren. Wir sollten die letzteren Vorgänge nicht testen.

(G5) Sollten wir Tests schreiben, um das durch die Implementierung definierte / nicht definierte Verhalten auszuführen?

Wir sollten keine Tests schreiben, die das von der Implementierung definierte oder undefinierte Verhalten des Vorgangs ausführen. Tests mit implementierungsdefinierten Verhaltensweisen zeigen ein lokales Verhalten des Interpreters, das nicht verallgemeinert werden sollte. Tests mit undefiniertem Verhalten tragen nicht zum Verständnis des Verhaltens des Vorgangs bei.

(G6) Mit welcher Genauigkeit muss das erwartete Ergebnis beim Schreiben von Tests für Gleitkommatypen in Prüfungen angegeben werden?

Bei Elementaroperationen (Addition, Subtraktion, Multiplikation, Division und Quadrat) wird erwartet, dass eine Implementierung gemäß der IEEE-Spezifikation ein gerundetes Ergebnis innerhalb von 0,5 ULP des mathematisch genauen Ergebnisses liefert. Wir können uns aber sicher vorstellen, dass das erwartete Ergebnis, das aus diesen Vorgängen kommt, höchstens eine ULP voneinander entfernt sein sollte. Dies funktioniert jedoch möglicherweise nicht bei transzendentalen Funktionen (sine, cosine usw.), für die die Genauigkeitsgarantien eine Implementierung definiert haben (Begründung).

Die aktuelle Implementierung verwendet einen "One-Size-fits-all"-Toleranzwert von 0,0001. Das folgende Beispiel veranschaulicht die obige Toleranz in Aktion.

func.func @check_tolerance() {
  %0 = stablehlo.constant dense<0.2> : tensor<f32>

  // The following check succeeds as %0 is almost equal to the provided
  // constant modulo the tolerance, mentioned above.
  check.expect_almost_eq_const %0, dense<0.19999> : tensor<f32>

  // The following check fails as %0 is not bitwise equal to the provided
  // constant.
  check.expect_eq_const %0, dense<0.19999> : tensor<f32>

  func.return
}

Dies ist nur der erste Schritt zum Testen der numerischen Genauigkeit von StableHLO-Vorgängen. Derzeit ist dies ein Bereich, der in der StableHLO-Spezifikation zu wenig angegeben ist. Wir arbeiten kontinuierlich daran, Nr. 1156 zu ermitteln, da er auf der Grundlage unserer Erfahrungen bei der Verwendung von StableHLO in der Praxis und auf Feedback von Stakeholdern basiert. Im weiteren Verlauf dieses Verfahrens werden wir die Infrastruktur entsprechend aktualisieren.

(G7) Gab es bei den Tests etwas zum Codierungsstil?

  1. Achten Sie darauf, den tatsächlichen Namen der Ein-/Ausgaben zu verwenden, anstatt standardmäßig SSA-Werte (z. B. %0, %1 usw.) zu verwenden.
  2. Stellen Sie sicher, dass für die Tests ein sauber gedrucktes Format verwendet wird, sofern vorhanden.

(G8) Sollten wir das bereits in der Spezifikation enthaltene Beispiel verwenden? Ja (aus Gründen der Vollständigkeit der Tests).