Основная цель интерпретатора 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
# 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
| Категории | Мнемоника | Общий |
|---|---|---|
| 119 | ||
| Поток управления | after_all, случай, если, оптимизация_барьер, в то время как | 5 |
| Перемещение данных | Broadcast_in_dim, объединение, динамический_срез, динамический_обновление_срез, сбор, дополнение, изменение формы, реверс, разброс, срез, сортировка, транспонирование | 12 |
| Распределение | all_gather, all_reduce, all_to_all, коллективное_пермуте, подача, выдача, идентификатор_раздела, получение, уменьшение_разброса, идентификатор_реплики, отправка | 11 |
| Динамизм | динамический_broadcast_in_dim, динамический_конв, динамический_сбор, динамический_йота, динамический_пад, динамический_ресформа, get_dimension_size, реальный_динамический_срез, set_dimension_size | 9 |
| Поэлементно | abs, добавить и, atan2, bitcast_convert, cbrt, ceil, зажим, сравнить, комплексный, преобразовать, косинус, count_leading_zeros, разделить, экспоненциальный, exponential_minus_one, пол, imag, is_finite, log, log_plus_one, логистический, карта, максимум, минимум, умножить, отрицать, не или, popcnt, мощность, реальный, уменьшить_точность, остаток, round_nearest_afz, round_nearest_even, rsqrt, выбрать, сдвиг_влево, сдвиг_вправо_арифметика, сдвиг_вправо_логический, знак, синус, sqrt, вычитание, загар, тан, xor | 48 |
| Расширяемость | custom_call, get_tuple_element, кортеж | 3 |
| Разнообразный | пакетный_норм_град, пакетный_норм_вывод, пакетный_норм_обучение, холеский, константа, БПФ, йота, rng, rng_bit_generator, triangular_solve | 10 |
| Модульность | вызов, функция, модуль, возврат | 4 |
| Не в HLO | широковещательная передача, create_token, сумма перекрестных реплик, точка, einsum, torch_index_select, unary_einsum | 8 |
| Квантование | униформа_деквантизация, униформа_квантизация | 2 |
| Снижение | свертка, dot_general, уменьшить, уменьшить_окно, select_and_scatter | 5 |
Основная цель интерпретатора 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
# 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
| Категории | Мнемоника | Общий |
|---|---|---|
| 119 | ||
| Поток управления | after_all, случай, если, оптимизация_барьер, в то время как | 5 |
| Перемещение данных | Broadcast_in_dim, объединение, динамический_срез, динамический_обновление_срез, сбор, дополнение, изменение формы, реверс, разброс, срез, сортировка, транспонирование | 12 |
| Распределение | all_gather, all_reduce, all_to_all, коллективное_пермуте, подача, выдача, идентификатор_раздела, получение, уменьшение_разброса, идентификатор_реплики, отправка | 11 |
| Динамизм | динамический_broadcast_in_dim, динамический_конв, динамический_сбор, динамический_йота, динамический_пад, динамический_ресформа, get_dimension_size, реальный_динамический_срез, set_dimension_size | 9 |
| Поэлементно | abs, добавить и, atan2, bitcast_convert, cbrt, ceil, зажим, сравнить, комплексный, преобразовать, косинус, count_leading_zeros, разделить, экспоненциальный, exponential_minus_one, пол, imag, is_finite, log, log_plus_one, логистический, карта, максимум, минимум, умножить, отрицать, не или, popcnt, мощность, реальный, уменьшить_точность, остаток, round_nearest_afz, round_nearest_even, rsqrt, выбрать, сдвиг_влево, сдвиг_вправо_арифметика, сдвиг_вправо_логический, знак, синус, sqrt, вычитание, загар, тан, xor | 48 |
| Расширяемость | custom_call, get_tuple_element, кортеж | 3 |
| Разнообразный | пакетная_норма_град, пакетная_норма_вывода, пакетная_норма_тренировка, холески, константа, БПФ, йота, rng, rng_bit_generator, triangular_solve | 10 |
| Модульность | вызов, функция, модуль, возврат | 4 |
| Не в HLO | широковещательная передача, create_token, сумма перекрестных реплик, точка, einsum, torch_index_select, unary_einsum | 8 |
| Квантование | униформа_деквантизация, униформа_квантизация | 2 |
| Снижение | свертка, dot_general, уменьшить, уменьшить_окно, select_and_scatter | 5 |