XLA (Accelerated Linear Algebra) — это компилятор машинного обучения (ML), который оптимизирует линейную алгебру, обеспечивая улучшение скорости выполнения и использования памяти. На этой странице представлен краткий обзор целей и архитектуры компилятора XLA.
Цели
Сегодня XLA поддерживает несколько интерфейсов платформы машинного обучения (включая PyTorch, TensorFlow и JAX) и является частью проекта OpenXLA — экосистемы технологий компилятора с открытым исходным кодом для машинного обучения, разработанной совместно ведущими организациями, занимающимися аппаратным и программным обеспечением машинного обучения. До создания проекта OpenXLA XLA разрабатывался внутри проекта TensorFlow, но основные цели остались прежними:
Улучшите скорость выполнения. Компилируйте подграфы, чтобы сократить время выполнения кратковременных операций и устранить накладные расходы во время выполнения, объединить конвейерные операции для уменьшения накладных расходов на память и специализировать известные формы тензоров, чтобы обеспечить более агрессивное распространение констант.
Улучшение использования памяти. Анализируйте и планируйте использование памяти, устраняя множество промежуточных буферов хранения.
Уменьшите зависимость от пользовательских операций. Устраните необходимость во множестве пользовательских операций, улучшив производительность автоматически объединенных низкоуровневых операций, чтобы она соответствовала производительности пользовательских операций, которые изначально были объединены вручную.
Улучшите портативность. Упростите написание новой серверной части для нового оборудования, чтобы большая часть моделей машинного обучения могла работать на этом оборудовании без изменений. Это контрастирует с подходом специализации отдельных монолитных операций для нового оборудования, который требует переписывания моделей для использования этих операций.
Как это работает
Компилятор XLA берет графы моделей из платформ машинного обучения, определенных в StableHLO , и компилирует их в машинные инструкции для различных архитектур. StableHLO определяет набор операций с поддержкой версий (HLO = операции высокого уровня), который обеспечивает уровень переносимости между платформами ML и компилятором.
В общем, процесс компиляции, который преобразует граф модели в целевой оптимизированный исполняемый файл, включает в себя следующие шаги:
XLA выполняет несколько встроенных проходов оптимизации и анализа графа StableHLO, которые не зависят от цели, например CSE , объединение операций, не зависящих от цели, и анализ буфера для выделения памяти времени выполнения для вычислений. На этом этапе оптимизации XLA также преобразует диалект StableHLO во внутренний диалект HLO.
XLA отправляет вычисления HLO на серверную часть для дальнейшей оптимизации уровня HLO, на этот раз с учетом целевой информации и потребностей. Например, серверная часть графического процессора может выполнять объединение операций, что полезно конкретно для модели программирования графического процессора, и определять, как разделить вычисления на потоки. На этом этапе серверные части также могут сопоставлять определенные операции или их комбинации с оптимизированными вызовами библиотеки.
Затем серверная часть выполняет генерацию кода для конкретной цели. Серверные части ЦП и графического процессора, включенные в XLA, используют LLVM для низкоуровневого IR, оптимизации и генерации кода. Эти серверные части создают LLVM IR, необходимый для эффективного представления вычислений HLO, а затем вызывают LLVM для создания собственного кода из этого LLVM IR.
В рамках этого процесса компилятор XLA является модульным в том смысле, что в него можно легко вставить альтернативный бэкэнд для работы с какой-либо новой аппаратной архитектурой . Серверная часть графического процессора в настоящее время поддерживает графические процессоры NVIDIA через серверную часть LLVM NVPTX. Серверная часть ЦП поддерживает несколько ЦП ISA.