Tłumacz rozmowy w StableHLO

Głównym celem tłumaczenia rozmowy StableHLO jest dostarczenie zgodnie z semantyką ustawienia StableHLO. specyfikacji. Dodatkowym celem jest ścisłe przestrzeganie specyfikację, nadając priorytet czytelności nad wydajnością, aby zwiększyć przejrzystość nawet w najważniejszych operacjach, takich jak Convolution, Gather/Scatter i DotGeneral.

Obecnie OpenXLA obsługuje interpretację 91 z 96 zastosowanych Operacje StableHLO. Pozostałe 3 operacje (FftOp, RngOp, RngBitGeneratorOp) wykazały ich semantykę opisaną w spec.md i mają przeprowadzili wstępną analizę na temat dalszych działań (patrz status.md , gdzie znajdziesz pełną listę operacji i ich najnowszy stan). Te ostatnie ulepszenia będą wdrażane w razie potrzeby w ramach społeczności.

Zakres

Podzieliliśmy grupę StableHLO na 11 kategorii zawierających 118 operacji (patrz Dodatek). Implementacja referencyjna przepływ pracy porządkuje pracę nad wdrożeniem tłumacza dla 100% operacji StableHLO zgodnie z definicją w specyfikacji StableHLO. Jesteśmy planujecie wykonanie wszystkich lub prawie wszystkich zadań w tym strumieniu pracy w StableHLO Wersja 1.0. Spośród 96 operacji mających obecnie specyfikację możemy interpretować 91 OpenXLA (pozostałe 5 znajdziesz w przypadkach specjalnych).

Specyfikacja

Głównym wymaganiem od tłumaczy jest korespondencja 1:1 z specyfikacja. Specyfikacja pozwala ustandaryzować tłumacza w podobnych operacjach, prowadzi do modułowej implementacji tłumacza o wysokiej jakości.

Przypadki szczególne

Różne

W tej kategorii znajdują się operacje rozkładane, a przyszłość nie jest w tej chwili nieznana. OK to 3 określone operacje w tej kategorii, których tłumacz nie obsługuje ten moment:

• FftOp
• RngOp
• RngBitGeneratorOp

Działanie FftOp jest sklasyfikowane jako Różne, ale w przeciwieństwie do innych działań w tej kategorii to działanie nie ma karnetu rozwijania, a obsługa go w StableHLO jest W trakcie opracowywania.

RngOp i RngBitGeneratorOp można podzielić na operacje MHLO, ale dekompozycja wprowadza element XlaRngGetAndUpdateStateOp, który jest specyficzny dla MHLO op. Interpretacja tych 2 operacji jest w trakcie opracowywania.

Narzędzie do konwertowania pozostałych operacji w tej kategorii na operacje StableHLO, które są tłumacz znajduje się w pliku hlo_expand_main.cc.

Brak w HLO

Oprócz działań wskazanych w tej kategorii składa się z 8 nieoczekiwanych operacji (patrz kategorii operacji StableHLO), które planujemy przeniesione ze StableHLO. Większość z tych operacji ma już karty mhlo na i przekonwertować je na operacje równoważne StableHLO.

Narzędzie do konwertowania pozostałych operacji w tej kategorii na odpowiadające im operacje StableHLO w pliku mlir-hlo-opt.cc.

Kwantyfikacja

Obsługa tłumaczenia rozmowy w przypadku operacji stablehlo.constant z typem kwantyzowanym to nieobsługiwane i śledzone przez 1691.

Instrukcje

Tworzenie tłumaczenia rozmowy

Tłumacz można utworzyć i przetestować w Bazel lub CMake (preferowane). Pełny instrukcji można znaleźć w sekcji README.md.

Bazel:

bazel build //...

Marka:

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

Do uruchomienia tłumaczenia rozmowy mamy narzędzie do tłumaczenia programów StableHLO. napisane w języku MLIR.

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

Dialekt tłumacza

Dialekt Interpreter zawiera różne operacje użytkowe związane z tłumaczem ustnym. Konkretnie chodzi o interpreter.run_parallel (zobacz InterpreterOps.td semantyki operacyjnej i przykładowego użycia) op umożliwia interpretację operacji dystrybucji i nie tylko. plan wykorzystania usług komunalnych w zależności od potrzeb społeczności.

Sprawdzanie dialektu

Dialekt Check służy do porównywania wartości czasu działania tłumaczenia z oczekiwanymi . Dane wyjściowe programu StableHLO można przetestować za pomocą różnych operacji kontrolnych (zobacz CheckOps.td semantyki operacyjnej i przykładowego zastosowania).

Pisanie programów testów

Używamy narzędzia LLVM (lit) do uruchamiania porównanie z wygenerowanym plikiem, aby porównać je z danymi wyjściowymi interpretera (patrz: stablehlo/tests/interpret. np. testy).

Testowanie: AddOp (przykład z 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
}

Testowanie operacji w kategorii Dystrybucja wymaga uruchomienia jej za pomocą interpreter.run_parallel – op.

Testowanie: AllReduceOp (przykład z 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
  }
}

Debugowanie StableHLO

Zgodnie z instrukcjami kompilacji StableHLO pliki binarne StableHLO dla narzędzi w stablehlo/tools to miejsce zamieszkania: /build/bin. Typowe narzędzia do debugowania, takie jak Za pomocą GDB można wykonać kod:

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

Dodatek

Konwertuj różne operacje

# 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>

Konwersja nie należąca do operacji HLO

# 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>

Kategorie operacji StableHLO