В этом документе описываются проходы оптимизации и преобразования HLO в компиляторе XLA .
Введение
Один проход HLO может состоять из одной или нескольких оптимизаций и преобразований компилятора, и XLA обеспечивает несколько сотен таких проходов. HLO фокусируется только на форме (например, матрице 3x4) и семантике операций над массивами, что упрощает оптимизацию или преобразование.
Например:
AlgebraicSimplifier: проход, выполняющий ряд упрощений и оптимизаций, в основном арифметических. В том числе:- При делении на константу выполняется оптимизация, преобразующая операцию в умножение на инверсию константы.
HloRematerialization: проход, который пересчитывает выбранные выражения в вычислении, чтобы уменьшить нагрузку на память, вызванную долгоживущими диапазонами значений в форме массива.
Информация о разработчике
Базовый класс для пропусков HLO можно найти в xla/hlo/pass/hlo_pass_interface.h . Пропуск HLO не должен расширять этот класс напрямую, а должен расширять HloModulePass .
См. также XLA HLO Pass Framework .
Инструменты и испытания
XLA поставляется с несколькими инструментами командной строки, включая hlo-opt. Этот инструмент позволяет выполнять отдельные проходы независимо от этапов компиляции на заданной платформе. Подробнее см. в разделе «Инструменты» .
Информацию о написании модульных тестов для проходов HLO см. в разделе Тестирование проходов HLO .
Примеры аппаратно-независимых проходов HLO
В этом разделе описаны несколько примеров проходов, общих для бэкендов XLA. Некоторые проходы могут быть специализированными для конкретных бэкендов, но их высокоуровневая функциональность схожа.
Общие проходы или аппаратно-независимые проходы можно найти в xla/hlo/transforms .
Рематериализация
См. также HloRematerialization .
Выборочно пересчитывает выражения в графе HLO для экономии памяти. Обеспечивает экономию памяти за счёт увеличения вычислительных мощностей. Может сократить использование памяти на десятки процентов и требуется для запуска многих крупных моделей.
Алгебраический упроститель
См. также AlgebraicSimplifier .
Набор упрощений, оптимизаций и канонизаций. Аналог прохода instcombine в LLVM .
Постоянное сворачивание
См. также HloConstantFolding .
Заменяет выражения, которые могут быть вычислены во время компиляции, их константными эквивалентами.
Устранение мертвого кода
См. также HloDCE .
Удаляет операции с неиспользованными результатами (быстрая реализация).
Сглаживание графика вызовов
См. также FlattenCallGraph .
Проход легализации, который преобразует граф вызовов HLO в дерево путём клонирования вычислений. Требуется, поскольку память статически выделяется операциям HLO, а не зависит от динамического контекста вызова.
Reshape Mover
См. также ReshapeMover .
Переформирование и транспонирование могут быть затратными, особенно в TPU. Этот проход перемещает, переформировывает и транспонирует элементы поэлементно, позволяя объединять или исключать эти операции.
Устранение HLO нулевого размера
См. также ZeroSizedHloElimination .
HLO поддерживает массивы нулевого размера (одно или несколько измерений имеют нулевую границу). Этот проход упрощает граф, заменяя операции нулевого размера константами нулевого размера.
Примеры проходов HLO для TPU
Проходы, специфичные для бэкэнда TPU.
Параллелизм модели
Разделение программы XLA на несколько ядер выполняется на уровне HLO, а конвейер TPU HLO включает ряд проходов для поддержки многоядерного выполнения.
Пространственное разбиение
См. также ShardingPropagation .
Перейти к поддержке разделения операций между устройствами по непакетным измерениям.
Обработка bfloat16
См. также BFloat16ConversionFolding , BFloat16MixedPrecisionRemoval и BFloat16Propagation .
TPU поддерживают bfloat16 как более компактное представление чисел с плавающей точкой с более низкой точностью, чем 32-битные числа с плавающей точкой. Использование bfloat16 сокращает объём занимаемой памяти и её пропускную способность. Конвейер HLO TPU включает различные проходы для замены чисел с плавающей точкой на bfloat16 в программе и распространения точности по всему графу.
Пропуска по легализации
См. также GatherExpander и BatchNormExpander .
Проходы, которые преобразуют неподдерживаемый HLO в форму, которую бэкэнд может испускать или для которой бэкэнд производит более эффективное понижение.
Пример прохода HLO для конкретного графического процессора
Проходы, специфичные для бэкэнда GPU, находятся в xla/service/gpu . Эти проходы можно идентифицировать как классы, определённые в namespace gpu .
cuDNN Rewriter
См. также CudnnFusedConvRewriter и CudnnNormRewriter .
Переписывает операции комбинированной свертки и нормирования в соответствующие им библиотечные вызовы в cuDNN.
Примеры проходов HLO для конкретного ЦП
Проходы, специфичные для бэкэнда ЦП, находятся в xla/service/cpu . Эти проходы можно идентифицировать как классы, определённые в namespace cpu .
Канонизация свертки
См. также ConvCanonicalization .
Канонизирует свертки, чтобы их можно было быстро реализовать в Eigen.
Операция «Распараллеливание»
См. также ParallelTaskAssigner .
Разделяет HLO на задачи для выполнения в отдельных потоках.
Анализ пройден
Проходы анализа не считаются проходами HLO, поскольку они не преобразуют HLO и не могут расширять HloModulePass . Общие анализы находятся в xla/hlo/analysis .
Примеры прохождения анализа
Анализ потока данных
См. также HloDataflowAnalysis .
Определяет все значения HLO на графике и их использование.
Анализ псевдонимов
См. также HloAliasAnalysis .
Определяет обязательные связи между значениями в программе.
Анализ стоимости вычислений
См. также HloCostAnalysis .
Вычисляет количество FLOP и использование памяти для всех операций в программе.
Проверка HLO
См. также HloVerifier .
Проверяет различные инварианты графа HLO.