StabilerHLO-Dolmetscher

Das Hauptziel des StableHLO-Interpreters ist es, Implementierung entsprechend der Semantik des StableHLO-Opsets Spezifikation zu ändern. Sekundäres Ziel ist es, bei der Implementierung genau zu beobachten, die Lesbarkeit gegenüber der Leistung bevorzugt, um für mehr Klarheit zu sorgen. mit der Semantik selbst der verbundensten Vorgänge wie Convolution, Gather/Scatter und DotGeneral.

Derzeit unterstützt OpenXLA die Interpretation von 91 von 96 Spezifikationen StableHLO-Vorgänge Für die verbleibenden drei Vorgänge (FftOp, RngOp, RngBitGeneratorOp) wie ihre Semantik in spec.md und haben erste Untersuchungen zum weiteren Vorgehen durchgeführt (siehe status.md finden Sie eine vollständige Liste der Vorgänge mit dem aktuellen Status. Diese finalen werden bei Bedarf in der Community implementiert.

Umfang

Wir haben das StableHLO-Opset in 11 Kategorien mit 118 Operationen in (siehe Anhang). Referenzimplementierung In Workstream wird die Arbeit an der Implementierung eines Dolmetschers organisiert. für 100% der StableHLO-Vorgänge gemäß Definition in der StableHLO-Spezifikation. Wir sind Sie planen, alle oder fast alle Arbeiten in diesem Workstream in StableHLO abzuschließen. Version 1.0. Von den 96 Vorgängen, die derzeit eine Spezifikation haben, können wir 91 Vorgänge durch OpenXLA (siehe Sonderfälle für die übrigen fünf Fälle).

Spezifikation

Die wichtigste Voraussetzung für den Dolmetscher ist eine direkte Korrespondenz mit dem Spezifikation. Die Spezifikation ermöglicht die Standardisierung des Interpreters für ähnliche Vorgänge, zu einer modularen, hochwertigen Implementierung des Dolmetschers führen.

Sonderfälle

Sonstige

Diese Kategorie enthält zerlegbare Operationen, deren Zukunft noch unklar ist. Es sind drei vordefinierte Operationen in dieser Kategorie, die vom Interpreter nicht unterstützt werden. für den Moment:

• FftOp
• RngOp
• RngBitGeneratorOp

FftOp ist als „Sonstiges“ kategorisiert, aber im Gegensatz zu anderen Vorgängen in dieser Kategorie Dieser Vorgang hat keinen Erweiterungs-Pass. Die Unterstützung in StableHLO ist ein In Arbeit.

RngOp und RngBitGeneratorOp können in MHLO-Vorgänge zerlegt werden, aber der der Zerlegung führt zu einem XlaRngGetAndUpdateStateOp, das ein MHLO-spezifisches Op. Eine unterstützende Interpretation dieser beiden Vorgänge ist WIP.

Das Tool zum Umwandeln der verbleibenden Operationen in dieser Kategorie in StableHLO-Operationen, die Der Interpreter befindet sich in hlo_expand_main.cc.

Nicht in HLO

Abgesehen von den spezifizierten Operationen besteht diese Kategorie aus acht nicht spezifizierten Operationen (siehe StableHLO Ops Categories), die in Zukunft aus StableHLO verschoben. Für die meisten dieser Vorgänge sind Karten/Tickets in Mhlo nach in StableHLO-Äquivalente umwandeln.

Das Tool zum Umwandeln der verbleibenden Vorgänge in dieser Kategorie in entsprechende StableHLO-Vorgänge die vom Interpreter unterstützt werden, befindet sich in mlir-hlo-opt.cc.

Quantisierung

Interpreter-Unterstützung für stablehlo.constant-Vorgänge mit quantisiertem Typ ist nicht unterstützt und erfasst über #1691.

Verwendungsanleitung

Referenzinterpreter erstellen

Der Interpreter kann über Bazel oder CMake (bevorzugt) erstellt und getestet werden. Vollständige finden Sie unter README.md.

Das ist z. B. der Fall.

bazel build //...

CMake:

mkdir -p build && cd build

cmake .. -GNinja \
  -DLLVM_ENABLE_LLD="$LLVM_ENABLE_LLD" \
  -DCMAKE_BUILD_TYPE=Release \
  -DLLVM_ENABLE_ASSERTIONS=On \
  -DMLIR_DIR=${PWD}/../llvm-build/lib/cmake/mlir

Zum Ausführen des Interpreters gibt es ein Übersetzungstool zum Interpretieren von StableHLO-Programmen. die in MLIR geschrieben sind.

stablehlo-translate --interpret <path/to/program>

Der Dolmetscherdialekt

Der Dialekt Interpreter enthält verschiedene Dienstprogrammvorgänge im Zusammenhang mit dem Dolmetscher. Konkret enthält die interpreter.run_parallel (siehe InterpreterOps.td für Vorgangssemantik und Beispielverwendung) op ermöglicht die Interpretation von Verteilungsvorgängen und mehr werden je nach Bedarf ergänzt.

Der Prüfdialekt

Der Dialekt Check wird verwendet, um die Interpreter-Laufzeitwerte mit den erwarteten Werten zu vergleichen Werte. StableHLO-Programmausgaben können über verschiedene Prüfvorgänge getestet werden (siehe CheckOps.td für die Semantik von Operationen und die Beispielverwendung).

Testprogramme schreiben

Wir verwenden das lit-Tool von LLVM, Mit der generierten Datei vergleichen, um mit der Ausgabe des Interpreters zu vergleichen (siehe stablehlo/tests/interpret) z. B. Tests).

AddOp wird getestet (Beispiel aus interpret_add.mlir):

// RUN: stablehlo-translate --interpret %s

func.func @add_op_scalar() {
  %0 = stablehlo.constant dense<2> : tensor<i4>
  %1 = stablehlo.constant dense<3> : tensor<i4>
  %2 = stablehlo.add %0, %1 : tensor<i4>
  check.expect_eq_const %2, dense<5> : tensor<i4>
  func.return
}

Für Testvorgänge in der Kategorie "Verteilung" muss die Ausführung über den interpreter.run_parallel-Dienstprogrammvorgang

AllReduceOp wird getestet (Beispiel aus all_reduce.mlir):

// RUN: stablehlo-translate --interpret %s

module @cross_replica {
  func.func public @all_reduce(%operand : tensor<4xi64>) -> tensor<4xi64> {
    %result = "stablehlo.all_reduce"(%operand) ({
      ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
        %0 = stablehlo.add %arg0, %arg1 : tensor<i64>
        stablehlo.return %0 : tensor<i64>
    }) {
      replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
      channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
    } : (tensor<4xi64>) -> tensor<4xi64>
    return %result : tensor<4xi64>
  }
  func.func public @main() {
    %inputs0 = stablehlo.constant dense<[1, 2, 3, 4]> : tensor<4xi64>
    %inputs1 = stablehlo.constant dense<[5, 6, 7, 8]> : tensor<4xi64>
    %results:2 = "interpreter.run_parallel"(%inputs0, %inputs1) {
      programs=[[@all_reduce], [@all_reduce]]
    } : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>, tensor<4xi64>)
    check.expect_eq_const %results#0, dense<[6, 8, 10, 12]> : tensor<4xi64>
    check.expect_eq_const %results#1, dense<[6, 8, 10, 12]> : tensor<4xi64>
    func.return
  }
}

Fehlerbehebung bei StableHLO

Nach den StableHLO-Build-Schritten werden die StableHLO-Binärdateien für Tools in stablehlo/tools sollte sich in /build/bin befinden. Gängige Debugging-Tools wie Mit GDB kann der Code schrittweise durchlaufen werden:

gdb --args ./build/bin/stablehlo-translate -allow-unregistered-dialect --interpret ./stablehlo/tests/interpret/<test>.mlir

Anhang

Verschiedene Vorgänge konvertieren

# batch_norm_grad
hlo-expand --batch_norm_grad_expander <path/to/hlo_module>

# batch_norm_inference
hlo-expand --batch_norm_inference_expander <path/to/hlo_module>

# batch_norm_training
hlo-expand --batch_norm_training_expander <path/to/hlo_module>

# cholesky
hlo-expand --cholesky_expander <path/to/hlo_module>

# constant
# Supported in StableHLO interpreter.

# fft
# TBD

# iota
# Supported in StableHLO interpreter.

# rng
# TBD

# rng_bit_generator
# TBD

# triangular_solve
hlo-expand --triangular_solve_expander <path/to/hlo_module>

Konvertierung nicht in HLO-Vorgänge

# broadcast
mlir-hlo-opt -mhlo-legalize-broadcast-to-broadcast-in-dim <path/to/input>

# create_token
mlir-hlo-opt -mhlo-legalize-create-token-to-after-all <path/to/input>

# cross-replica-sum
mlir-hlo-opt -mhlo-legalize-cross-replica-sum-to-all-reduce <path/to/input>

# dot
mlir-hlo-opt -mhlo-legalize-dot-to-dot-general <path/to/input>

# einsum
mlir-hlo-opt -mhlo-legalize-einsum-to-dot-general <path/to/input>

# torch_index_select
mlir-hlo-opt -mhlo-legalize-torch-index-select-to-gather <path/to/input>

# unary_einsum
mlir-hlo-opt --canonicalize -mhlo-legalize-einsum-to-dot-general <path/to/input>

StableHLO Ops-Kategorien