Стабильный переводчик HLO

Основная цель интерпретатора StableHLO — предоставить эталонную реализацию семантики опсета StableHLO в соответствии с его спецификацией. Вторичная цель — чтобы реализация точно следовала спецификации, отдавая предпочтение читаемости над производительностью, чтобы обеспечить дополнительную ясность семантике даже самых сложных операций, таких как Convolution , Gather / Scatter и DotGeneral .

На данный момент OpenXLA поддерживает интерпретацию 91 из 96 предусмотренных операций StableHLO. Семантика остальных трех операций ( FftOp , RngOp , RngBitGeneratorOp ) задокументирована в spec.md и завершены первоначальные исследования о том, как двигаться дальше (полный список операций и их последний статус см. в status.md ). Эти окончательные улучшения будут реализованы по мере необходимости сообщества.

Объем

Мы разделили опсет StableHLO на 11 категорий, состоящих в общей сложности из 118 операций (см. Приложение ). Рабочий поток справочной реализации организует работу по реализации интерпретатора для 100 % операций StableHLO, как определено в спецификации StableHLO. Мы планируем завершить всю или почти всю работу в этом рабочем направлении в StableHLO v1.0. Из 96 операций, которые в настоящее время имеют спецификацию, мы можем интерпретировать 91 операцию через OpenXLA (остальные 5 см. в разделе «Особые случаи »).

Спецификация

Основное требование к интерпретатору — соответствие 1:1 спецификации. Спецификация позволяет стандартизировать интерпретатор для аналогичных операций, что приводит к модульной, высококачественной реализации интерпретатора.

Особые случаи

Разнообразный

В этой категории есть разлагаемые операции, будущее которых на данный момент неясно. В этой категории есть три специальные операции, которые интерпретатор на данный момент не поддерживает:

• FftOp
• RngOp
• RngBitGeneratorOp

FftOp относится к категории «Разное», но в отличие от других операций в этой категории, эта операция не имеет прохода расширения, и поддержка этого в StableHLO — это незавершенная работа.

RngOp и RngBitGeneratorOp можно разложить на операции MHLO, но при декомпозиции вводится XlaRngGetAndUpdateStateOp , который является операцией, специфичной для MHLO. Поддерживающая интерпретация этих двух операций является незавершенной работой.

Инструмент для преобразования оставшихся операций в этой категории в операции StableHLO, поддерживаемые интерпретатором, находится в hlo_expand_main.cc .

Нет в ХЛО

Помимо специфицированных операций, эта категория состоит из 8 неспециализированных операций (см. Категории операций StableHLO ), которые планируется вывести из StableHLO. У большинства этих операций уже есть проходы в mhlo для преобразования их в эквивалентные операции StableHLO.

Инструмент для преобразования оставшихся операций в этой категории в эквивалентные операции StableHLO, поддерживаемые интерпретатором, находится в mlir-hlo-opt.cc .

Квантование

Поддержка интерпретатора для операции stablehlo.constant с квантованным типом не поддерживается и отслеживается через #1691 .

Инструкции по использованию

Создание эталонного интерпретатора

Интерпретатор можно собрать и протестировать с помощью Bazel или CMake (предпочтительно). Полные инструкции см. в README.md .

Базель:

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

Для запуска интерпретатора у нас есть инструмент перевода для интерпретации программ StableHLO, написанных на MLIR.

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

Диалект переводчика

Диалект Interpreter содержит различные служебные операции, связанные с интерпретатором. В частности, оператор interpreter.run_parallel (семантику операций и примеры использования см. в InterpreterOps.td ) позволяет интерпретировать операции распространения, а также планирует добавить дополнительные утилиты в зависимости от потребностей сообщества.

Чек-диалект

Диалект Check используется для сравнения значений времени выполнения интерпретатора с ожидаемыми значениями. Выходные данные программы StableHLO можно протестировать с помощью различных операций проверки (семантику операций и примеры использования см. в CheckOps.td ).

Написание тестовых программ

Мы используем инструмент LLVM Lit для запуска и сравнения сгенерированного файла, чтобы сравнить его с выводом интерпретатора (примеры тестов см. в стабильном файле/tests/interpret ).

Тестирование AddOp (пример из 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
}

Операции тестирования в категории «Распространение» требуют запуска их с помощью утилиты interpreter.run_parallel op.

Тестирование AllReduceOp (пример из 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
  }
}

Отладка StableHLO

После выполнения шагов сборки StableHLO двоичные файлы StableHLO для инструментов из stablehlo/tools должны находиться в /build/bin . Для пошагового выполнения кода можно использовать общие инструменты отладки, такие как GDB:

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

Приложение

Преобразование различных операций

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

Преобразование не в операциях HLO

# 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