Операторы управления потоком

Зачем использовать операторы управления потоком?

PyTorch позволяет писать модели на обычном Python, в том числе использовать такие конструкции управления потоком, как операторы if, циклы for и циклы while. Это обеспечивает большую гибкость, но создаёт трудности при компиляции.

Рассмотрим простой пример:

if mod.static_config == 0:
    return f(x)
return g(x)

Ветви могут содержать совершенно разные операции. Если бы мы пытались компилировать обе ветви каждый раз при выполнении оператора if, количество путей выполнения росло бы экспоненциально и очень быстро становилось бы неуправляемым.

Для решения этой проблемы torch.compile использует специализацию и защитные проверки. При трассировке модели компилятор выбирает путь выполнения на основе текущего значения предиката (специализация), а затем добавляет защитную проверку, чтобы проверить это предположение во время выполнения. Если проверка не проходит, компилятор выполняет повторную компиляцию.

То же самое относится к циклам: компилятор разворачивает их и добавляет защитную проверку количества итераций. В результате получается линейный вычислительный граф, который легко оптимизировать.

Этот подход хорошо работает для статического управления потоком, но не подходит в нескольких случаях:

• Управление потоком, зависящее от данных: если предикат зависит от значения тензора, компилятор не может выбрать ветвь во время компиляции, поскольку значение ещё неизвестно. Аналогично, он не может развернуть цикл while, если количество итераций зависит от значений тензора. Компилятор решает эту проблему, разбивая граф и возвращаясь к Python, что также делает невозможным запуск модели без среды выполнения Python (например, на периферийных устройствах).
• Управление потоком, зависящее от динамической формы: если количество итераций цикла или предикат ветвления зависят от динамического размера тензора, специализация приводит к тому, что скомпилированный код работает только для этого конкретного размера. При каждом изменении размера компилятору приходится выполнять повторную компиляцию.
• Большие вычислительные графы: даже при статическом количестве итераций разворачивание длинного цикла приводит к тому, что граф растёт линейно с числом итераций, хотя каждая итерация выполняет одни и те же действия. Это увеличивает время компиляции и расход памяти.

Операторы управления потоком решают эти проблемы, представляя управление потоком в виде явных операторов, которые понимает компилятор. Вместо того чтобы устранять управление потоком посредством специализации, эти операторы сохраняют его в скомпилированном графе.

Доступные операторы

• Управление потоком — Cond
• Управление потоком — Switch
• Управление потоком — While Loop
• Управление потоком — Scan
• Управление потоком — Associative Scan
• Управление потоком — Map