Зачем мы построили XLA?
У нас было несколько целей для работы XLA с TensorFlow:
Увеличьте скорость выполнения. Компилируйте подграфы, чтобы сократить время выполнения недолговечных операций, чтобы устранить накладные расходы из среды выполнения TensorFlow, объединяйте конвейерные операции, чтобы уменьшить накладные расходы на память, и специализироваться на известных тензорных формах, чтобы обеспечить более агрессивное распространение констант.
Улучшение использования памяти. Анализируйте и планируйте использование памяти, в принципе устраняя множество промежуточных буферов хранения.
Уменьшите зависимость от нестандартных операций. Устраните необходимость во многих пользовательских операциях, улучшив производительность автоматически слитых низкоуровневых операций, чтобы они соответствовали производительности пользовательских операций, которые были слиты вручную.
Уменьшите площадь, занимаемую мобильными устройствами. Исключите среду выполнения TensorFlow, заблаговременно скомпилировав подграф и выпустив пару объект / файл заголовка, которую можно связать напрямую с другим приложением. Результаты могут на несколько порядков уменьшить занимаемое мобильным устройством место.
Улучшение портативности. Упростите написание новой серверной части для нового оборудования, после чего большая часть программ TensorFlow будет работать на этом оборудовании без изменений. Это контрастирует с подходом к специализации отдельных монолитных операций для нового оборудования, который требует переписывания программ TensorFlow для использования этих операций.
Как работает XLA?
Язык ввода для XLA называется «HLO IR» или просто HLO (высокоуровневые операции). Семантика ПЗ описаны в эксплуатации Семантика странице. Это наиболее удобно думать как ПЗ компилятор IR .
XLA берет графы («вычисления»), определенные в HLO, и компилирует их в машинные инструкции для различных архитектур. XLA имеет модульную конструкцию в том смысле , что это легко слот в альтернативном интерфейсе для целевого некоторой новой архитектуры HW . Бэкэнд ЦП для x64 и ARM64, а также бэкэнд NVIDIA GPU находятся в дереве исходного кода TensorFlow.
На следующей диаграмме показан процесс компиляции в XLA:
XLA поставляется с несколькими проходами оптимизации и анализа , которые являются мишенью-независимой, такими как CSE , целевой независимым от слияния операции, и буфера анализа для выделения памяти во время выполнения для вычисления.
После шага, не зависящего от цели, XLA отправляет вычисление HLO на бэкэнд. Серверная часть может выполнять дальнейшую оптимизацию на уровне HLO, на этот раз с учетом конкретной информации и потребностей. Например, серверная часть графического процессора XLA может выполнять объединение операций, выгодное специально для модели программирования графического процессора, и определять, как разделить вычисление на потоки. На этом этапе серверные модули могут также сопоставить определенные операции или их комбинации с оптимизированными вызовами библиотеки по шаблону.
Следующим шагом является генерация кода для конкретной цели. Центральный процессор и GPU бэкэнды в комплекте с XLA использованием LLVM для ИК низкого уровня, оптимизации и генерации кода. Эти бэкенды выдают LLVM IR, необходимый для эффективного представления вычислений XLA HLO, а затем вызывают LLVM для выдачи собственного кода из этого LLVM IR.
Бэкэнд GPU в настоящее время поддерживает GPU NVIDIA через бэкэнд LLVM NVPTX; бэкэнд ЦП поддерживает несколько ISA ЦП.