Есть вопрос? Присоединяйтесь к сообществу на форуме TensorFlow. Посетите форум.

Архитектура XLA

Зачем мы построили XLA?

У нас было несколько целей для работы XLA с TensorFlow:

  • Увеличьте скорость выполнения. Компилируйте подграфы, чтобы сократить время выполнения недолговечных операций, чтобы устранить накладные расходы из среды выполнения TensorFlow, объедините конвейерные операции, чтобы уменьшить накладные расходы на память, и специализироваться на известных тензорных формах, чтобы обеспечить более агрессивное распространение констант.

  • Улучшение использования памяти. Анализируйте и планируйте использование памяти, в принципе устраняя множество промежуточных буферов хранения.

  • Уменьшите зависимость от пользовательских операций. Устранение необходимости во множестве настраиваемых операций за счет повышения производительности автоматически объединяемых операций низкого уровня, чтобы соответствовать производительности пользовательских операций, которые были объединены вручную.

  • Уменьшите площадь, занимаемую мобильными устройствами. Исключите среду выполнения TensorFlow, заблаговременно скомпилировав подграф и выпустив пару объект / файл заголовка, которую можно связать напрямую с другим приложением. Результаты могут на несколько порядков сократить занимаемое мобильным устройством место.

  • Улучшение портативности. Упростите написание новой серверной части для нового оборудования, после чего большая часть программ TensorFlow будет работать на этом оборудовании без изменений. Это контрастирует с подходом к специализации отдельных монолитных операций для нового оборудования, который требует переписывания программ TensorFlow для использования этих операций.

Как работает XLA?

Язык ввода для XLA называется «HLO IR» или просто HLO (высокоуровневые операции). Семантика HLO описана на странице Семантика операций . Удобнее всего рассматривать HLO как компилятор IR .

XLA берет графы («вычисления»), определенные в HLO, и компилирует их в машинные инструкции для различных архитектур. XLA является модульным в том смысле, что его легко вставить в альтернативный бэкэнд для нацеливания на какую-то новую архитектуру аппаратного обеспечения . Бэкэнд ЦП для x64 и ARM64, а также бэкэнд NVIDIA GPU находятся в дереве исходного кода TensorFlow.

На следующей диаграмме показан процесс компиляции в XLA:

XLA поставляется с несколькими оптимизациями и проходами анализа, которые не зависят от цели, такими как CSE , слияние операций, не зависящих от цели, и анализ буфера для выделения памяти времени выполнения для вычислений.

После шага, не зависящего от цели, XLA отправляет вычисление HLO в бэкэнд. Серверная часть может выполнять дальнейшую оптимизацию на уровне HLO, на этот раз с учетом конкретной информации и потребностей. Например, серверная часть графического процессора XLA может выполнять объединение операций, выгодное специально для модели программирования графического процессора, и определять, как разделить вычисления на потоки. На этом этапе серверные модули могут также сопоставить определенные операции или их комбинации с оптимизированными вызовами библиотеки по шаблону.

Следующим шагом является генерация кода для конкретной цели. Бэкэнды CPU и GPU, включенные в XLA, используют LLVM для низкоуровневого IR, оптимизации и генерации кода. Эти бэкенды выдают LLVM IR, необходимый для эффективного представления вычислений XLA HLO, а затем вызывают LLVM для выдачи собственного кода из этого LLVM IR.

Бэкэнд GPU в настоящее время поддерживает GPU NVIDIA через бэкэнд LLVM NVPTX; бэкэнд ЦП поддерживает несколько ISA ЦП.