Что такое operator fusion в компиляторах и какие паттерны fusion существуют?

Краткий тезис

Operator fusion (слияние операторов]]) — это техника оптимизации в компиляторах машинного обучения, при которой несколько последовательных или параллельных операций (например, матричное умножение, активация, нормализация) объединяются в один kernel (вычислительное ядро). Это снижает накладные расходы на запуск ядер, уменьшает количество чтений/записей в глобальную память и повышает эффективность использования кэша. Основные паттерны: pointwise fusion (поэлементные операции), reduction fusion (свёртки/редукции) и horizontal fusion (независимые операции, выполняемые на одном ядре).

1. Термин: Operator fusion (слияние операторов)

Operator fusion — это этап компиляции вычислительного графа нейронной сети, при котором несколько узлов графа (операторов) объединяются в один fused kernel. Цель — минимизировать число вызовов ядер на GPU/TPU и оптимизировать использование памяти.

Зачем нужно слияние

• Kernel launch overhead: каждый вызов ядра на GPU требует передачи параметров, синхронизации — это микросекунды, но при тысячах вызовов накапливается.
• Memory bandwidth bottleneck: если два оператора последовательно читают и пишут промежуточные тензоры в глобальную память, это создаёт узкое место. Fusion позволяет хранить промежуточные результаты в регистрах или shared memory.
• Улучшение локальности данных: объединённый kernel может переиспользовать данные из кэша, не выгружая их в DRAM.

Пример без fusion

A = matmul(X, W)   # запись в глобальную память
B = relu(A)        # чтение A, запись B
C = add(B, bias)   # чтение B, запись C

Три ядра, три чтения/записи больших тензоров.

С fusion

fused_kernel(X, W, bias) -> C
# внутри: matmul -> relu -> add, всё в регистрах

Одно ядро, один проход по данным.

2. Паттерны operator fusion

Существует несколько основных паттернов, которые компиляторы (TVM, XLA, TensorRT, MLIR) распознают и применяют автоматически.

2.1 Pointwise fusion (поэлементное слияние)

Объединяет несколько поэлементных операций (ReLU, Sigmoid, Add, Mul, Clamp) в одно ядро. Эти операции не меняют форму тензора и независимы по элементам.

Пример: relu(add(x, bias)) → один kernel, который для каждого элемента вычисляет max(0, x_i + bias_i).

Когда применимо почти всегда, если между операциями нет изменения формы (reshape, transpose) или редукции.

2.2 Reduction fusion (слияние с редукцией)

Объединяет редукционную операцию (sum, mean, max, softmax, layer normalization) с предшествующими поэлементными операциями. Редукции требуют агрегации по оси, поэтому fusion сложнее: нужно аккуратно управлять частичными суммами.

Пример: layer_norm(relu(matmul(X, W))) — компилятор может объединить matmul, relu и layer norm в один kernel, вычисляя нормализацию на лету.

Паттерны

• Fuse into reduction: поэлементная операция перед редукцией (например, sum(relu(x))).
• Fuse after reduction: редукция, затем поэлементная (например, add(softmax(x), bias)).

2.3 Horizontal fusion (горизонтальное слияние)

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

Пример: два параллельных свёрточных слоя с разными ядрами, применяемые к одному входу. Вместо двух отдельных kernel launch — один fused kernel, который вычисляет обе свёртки одновременно, переиспользуя загрузку входного тензора.

Когда применимо в архитектурах с несколькими параллельными ветвями (Inception, ResNeXt, некоторые трансформеры).

2.4 Composite fusion (композитное слияние)

Комбинация pointwise и reduction fusion в цепочке. Например, add(bias) -> relu -> layer_norm — все три этапа могут быть слиты.

3. Как fusion реализуется в компиляторах ML

3.1 Анализ графа

Компилятор строит вычислительный граф (IR — intermediate representation). Затем применяет pattern matching (поиск шаблонов) для обнаружения последовательностей, подходящих под fusion.

Пример в TVM (Apache TVM):

# TVM relay graph
x = relay.var("x")
w = relay.var("w")
b = relay.var("b")
matmul = relay.nn.dense(x, w)
add = relay.add(matmul, b)
relu = relay.nn.relu(add)
# TVM автоматически сливает add и relu в один kernel

3.2 Генерация fused kernel

После обнаружения шаблона компилятор генерирует код одного ядра (CUDA, OpenCL, Triton). Внутри ядра:

• Загружаются входные данные (например, блок матрицы).
• Выполняется вся цепочка операций в регистрах / shared memory.
• Записывается результат.

3.3 Ограничения

• Сложность редукций: редукции требуют синхронизации между потоками, что усложняет слияние.
• Изменение формы: reshape, transpose, split — часто блокируют fusion, так как нарушают поэлементное соответствие.
• Динамические формы: если размер тензора неизвестен на этапе компиляции, fusion может быть невозможен.

4. Примеры фреймворков и компиляторов, использующих fusion

5. Влияние fusion на производительность

До fusion

• 3 kernel launch → ~10–20 мкс накладных расходов.
• 3 чтения/записи больших тензоров (например, 1 ГБ) → ~3× bandwidth.

После fusion

• 1 kernel launch.
• 1 чтение входа, 1 запись выхода.
• Ускорение может достигать 2–5× на операциях с интенсивным вводом-выводом (memory-bound).

Когда fusion не помогает

• Если операции уже compute-bound (например, большие matmul) — накладные расходы на запуск незначительны.
• Если fusion приводит к увеличению использования регистров (register pressure) и спарсинг (spilling) в локальную память.

6. Связь operator fusion с Agentic RAG

Хотя вопрос напрямую про компиляторы, в контексте Agentic RAG fusion важен для:

• Оптимизации inference агентов: агенты часто используют несколько вызовов LLM (планирование, поиск, генерация). Fusion в компиляторе ускоряет каждый вызов модели.
• Компиляции графов инструментов: если агент использует нейросетевые инструменты (например, эмбеддеры, ранжировщики), fusion ускоряет их выполнение.
• Сборки пайплайнов: в Agentic RAG может быть цепочка: эмбеддинг запроса → retrieval → re-ranking → генерация. Операции внутри каждого этапа (например, attention + softmax) могут быть слиты.

7. Пет-проект для закрепления

Задача Реализовать простой компилятор, который выполняет pointwise fusion для двух операций (ReLU и Add) на CPU с использованием Numpy.

Инструменты Python, Numpy, (опционально) Numba для JIT-компиляции.

Шаги:

1. Напишите две функции: add(x, bias) и relu(x).
2. Измерьте время выполнения последовательного вызова (два прохода по массиву).
3. Напишите fused-функцию fused_add_relu(x, bias), которая за один проход делает max(0, x + bias).
4. Сравните время выполнения на большом массиве (10^8 элементов). Зафиксируйте ускорение.
5. (Дополнительно) Используйте Numba @njit для генерации fused kernel и сравните с ручным циклом.

Ожидаемый результат Fused-версия будет в ~1.5–2 раза быстрее за счёт уменьшения количества проходов по памяти и снижения накладных расходов на вызов функций.

8. Связь с другими вопросами

9. Навигация

• Предыдущий: 321
• Следующий: 323
• Индекс: 00. Индекс разборов

Навигация

• Предыдущий: 321
• Следующий: 323
• Индекс: 00. Индекс разборов