Что такое MLIR и как он используется в IREE/TensorRT-LLM?

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

MLIR (Multi-Level Intermediate Representation) — это фреймворк для построения компиляторов, который позволяет представлять машинное обучение (ML) на нескольких уровнях абстракции и выполнять сквозную оптимизацию под разные аппаратные платформы. IREE и TensorRT-LLM — это два фреймворка инференса, которые активно используют MLIR для генерации высокопроизводительного кода: IREE — для универсального исполнения на CPU/GPU/специализированных ускорителях, TensorRT-LLM — для оптимизации больших языковых моделей (LLM) на GPU NVIDIA. MLIR служит общим «клеем», позволяя объединить фронтенды (TensorFlow, PyTorch) с бэкендами (CUDA, Vulkan, Metal, собственные ISA).

1. Термин MLIR (Multi-Level Intermediate Representation)

MLIR — это проект с открытым исходным кодом, стартовавший в Google в 2019 году. Его ключевая идея: вместо одного фиксированного промежуточного представления (как LLVM IR) использовать набор взаимосвязанных диалектов (dialects), каждый из которых описывает определённый уровень абстракции — от высокоуровневых операций ML (свёртки, матричные умножения) до низкоуровневых инструкций конкретного процессора.

Зачем нужен MLIR

• Традиционные компиляторы (LLVM, GCC) плохо справляются с оптимизациями, специфичными для ML (например, слияние операций]], квантование, работа с тензорами переменной формы).
• ML-модели всё чаще запускаются на гетерогенных устройствах (CPU + GPU + TPU + NPU), и для каждого нужен свой бэкенд.
• MLIR предоставляет единую инфраструктуру для написания pass'ов (преобразований), которые работают на разных уровнях, и для генерации кода под множество целей.

Ключевые понятия MLIR

• Dialect — набор операций и типов, объединённых общей семантикой. Примеры: tensor (тензорные операции), linalg (линейная алгебра), arith (арифметика), gpu (GPU-специфичные операции), llvm (LLVM IR).
• Operation — базовая единица вычисления (например, tensor.matmul).
• Pass — проход, который преобразует IR из одного диалекта в другой или оптимизирует внутри одного диалекта.
• Conversion — механизм понижения (lowering) из высокоуровневого диалекта в низкоуровневый.

2. Архитектура MLIR: от модели до исполняемого кода

Типичный пайплайн компиляции ML-модели с использованием MLIR выглядит так:

1. Фронтенд (TensorFlow, PyTorch, ONNX) импортирует модель в высокоуровневый диалект (например, tosa или mhlo).
2. Серия pass'ов выполняет оптимизации: code elimination|удаление мёртвого кода, слияние операций]], квантование, преобразование форм тензоров.
3. Понижение до диалекта linalg (линейная алгебра) или scf (структурный контрольный поток).
4. Дальнейшее понижение до gpu или llvm в зависимости от целевого устройства.
5. Генерация кода (code generation) — создание бинарного кода для конкретного бэкенда (CUDA, SPIR-V, x86).

Пример простого MLIR-кода (диалект tensor):

func.func @main(%arg0: tensor<4x4xf32>, %arg1: tensor<4x4xf32>) -> tensor<4x4xf32> {
  %0 = tensor.matmul %arg0, %arg1 : tensor<4x4xf32> * tensor<4x4xf32> -> tensor<4x4xf32>
  return %0 : tensor<4x4xf32>
}

3. Преимущества MLIR перед традиционными подходами

4. IREE (Intermediate Representation Execution Environment)

IREE — это компилятор и среда выполнения для ML-моделей, разработанный Google. Он использует MLIR как центральный компонент для компиляции моделей под различные бэкенды: CUDA (NVIDIA GPU), Vulkan (кроссплатформенный GPU), Metal (Apple GPU), CPU (через LLVM), а также для специализированных ускорителей (например, Google TPU через PjRt).

Как IREE использует MLIR

• Импортирует модели из TensorFlow, PyTorch, JAX через диалекты stablehlo или mhlo.
• Применяет оптимизации на уровне MLIR: слияние операций, удаление лишних копий, квантование (INT8, FP16).
• Понижает до диалекта iree_linalg_ext (расширенная линейная алгебра) и далее до iree_gpu или iree_llvm.
• Генерирует специализированный код для каждого бэкенда, используя собственные проходы (например, iree-codegen-gpu).
• Результат — сборка (module) в формате VMVX (Virtual Machine eXtreme) или FlatBuffer, которая может быть загружена и выполнена на целевом устройстве.

Пример команды компиляции IREE

iree-compile model.mlir --iree-hal-target-backends=cuda --o model.vmfb

Преимущества IREE

• Единый пайплайн для множества бэкендов.
• Поддержка stream semantics (асинхронное выполнение).
• Возможность Just-In-Time (JIT) компиляции для динамических форм.

5. TensorRT-LLM

TensorRT-LLM — это библиотека от NVIDIA для оптимизации инференса больших языковых моделей на GPU. Она также использует MLIR, но более узко — для представления и оптимизации графа вычислений LLM.

Как TensorRT-LLM использует MLIR

• Модель (например, LLaMA, GPT) импортируется в собственный диалект TensorRT-LLM (tensorrt_llm), который включает операции, специфичные для трансформеров: attention, rms_norm, silu_mul.
• MLIR-проходы выполняют:
  • Fusion (слияние) операций: например, объединение layernorm + matmul в одну операцию.
  • Квантование (INT4, INT8, FP8) с учётом калибровки.
  • In-flight batching — динамическое объединение запросов.
  • PagedAttention — управление KV-кэшем.
• После оптимизаций MLIR понижается до TensorRT engine (через собственный бэкенд, не через LLVM).
• Результат — высокооптимизированный бинарный файл (plan), который исполняется на GPU.

Пример конфигурации TensorRT-LLM

from tensorrt_llm import Builder
builder = Builder()
builder.set_quant_mode('int4_awq')
builder.build_engine('model.engine')

Преимущества TensorRT-LLM

• Специализированные оптимизации для LLM (слияние attention, квантование).
• Поддержка многократного инференса с динамическим батчингом.
• Интеграция с NVIDIA Triton Inference Server.

6. Сравнение IREE и TensorRT-LLM

7. Пример пайплайна компиляции с MLIR (на примере IREE)

1. Импорт модели (например, PyTorch через torch-mlir):

   torch-mlir-opt model.mlir --convert-torch-to-linalg > model_linalg.mlir
   

2. Оптимизация (слияние операций, удаление лишних reshape):

   iree-opt model_linalg.mlir --iree-optimization-pass-pipeline > model_opt.mlir
   

3. Понижение до целевого диалекта (например, для CUDA):

   iree-opt model_opt.mlir --iree-codegen-gpu-pipeline > model_gpu.mlir
   

4. Генерация кода:

   iree-compile model_gpu.mlir --iree-hal-target-backends=cuda --o model.vmfb
   

5. Запуск:

   import iree.runtime as rt
   vm_module = rt.load_vm_module(rt.VmModule.from_flatbuffer('model.vmfb'))
   result = vm_module.main(input_tensor)
   

8. Роль MLIR в Agentic RAG

В контексте Agentic RAG (агентные системы с retrieval и генерацией) MLIR и фреймворки вроде IREE/TensorRT-LLM решают задачу эффективного развёртывания LLM и эмбеддеров на различных устройствах. Агенты часто требуют низкой задержки (latency) и работы на edge-устройствах (мобильные телефоны, IoT). MLIR позволяет:

• Скомпилировать модель эмбеддингов (например, BERT) под CPU/GPU с минимальным потреблением памяти.
• Оптимизировать LLM для инференса с квантованием и in-flight batching, что критично для многопользовательских RAG-агентов.
• Интегрировать несколько моделей (retriever + генератор) в единый пайплайн, используя общий MLIR-граф.

Таким образом, MLIR является ключевым звеном для создания производственных RAG-систем, работающих на гетерогенном оборудовании.

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

Задача Написать простой MLIR pass, который заменяет операцию tensor.matmul на последовательность arith.mulf + arith.addf (для демонстрации понижения).

Инструменты LLVM/MLIR (сборка из исходников), Python (через mlir.ir), C++.

Шаги:

1. Установить MLIR (например, через pip install mlir или собрать из репозитория llvm-project).
2. Создать файл test.mlir с операцией tensor.matmul.
3. Написать на C++ простой pass, который находит tensor.matmul и заменяет его на циклы с arith операциями (используя scf.for).
4. Зарегистрировать pass в инструменте mlir-opt.
5. Запустить mlir-opt --your-pass test.mlir и проверить результат.

Ожидаемый результат Понимание, как работают преобразования в MLIR, и умение писать собственные оптимизации. Для усложнения можно добавить поддержку разных типов данных (f32, f16) и проверить корректность через запуск на небольшом тензоре.

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

Навигация

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