Компиляторы
MLIR: Multi-Level IR
PyTorch запускает transformer и компилирует его на лету для GPU. TensorFlow оптимизирует conv2d для TPU. Apple Core ML адаптирует модель под Neural Engine. Все три системы столкнулись с одной задачей: нужно несколько уровней абстракции - от 'матрица умножений' до 'конкретные CUDA intrinsics'. MLIR - это инфраструктура которая сделала создание таких multi-level компиляторов инженерной задачей, а не исследовательской.
- **PyTorch 2.0 torch.compile**: использует MLIR через TorchDynamo для JIT компиляции Python/PyTorch кода - достигает 2-4x speedup на GPU inference без изменения user кода
- **Google XLA/OpenXLA**: StableHLO dialect (на основе MLIR) стал стандартом для portability ML моделей между фреймворками - одна модель компилируется для TPU, GPU и CPU через один IR
- **Apple ANE compiler**: Core ML использует MLIR для компиляции ML моделей на Apple Neural Engine - нейросетевые операции lowering до аппаратных инструкций ANE через dialect pipeline
MLIR Dialects
MLIR (Multi-Level Intermediate Representation) - инфраструктура для создания компиляторов с несколькими уровнями абстракции. Ключевая концепция - dialect: набор типов и операций для конкретного домена. `linalg` dialect - линейная алгебра, `affine` dialect - полиэдральные преобразования, `gpu` dialect - GPU операции, `scf` dialect - control flow. Разные dialects сосуществуют в одном IR.
Google перевёл TensorFlow XLA на MLIR в 2019-2020 годах: StableHLO dialect стал стандартом для обмена ML моделями между фреймворками. PyTorch 2.0 torch.compile использует MLIR через TorchDynamo -> TorchMLIR pipeline. Apple Core ML compiler и OpenXLA (Google) - оба построены на MLIR.
Зачем MLIR позволяет смешивать разные dialects в одном IR файле?
Progressive Lowering
Progressive lowering - пошаговый спуск от высокоуровневых dialects к низкоуровневым. Каждый шаг - это набор ConversionPattern, которые заменяют операции одного dialect операциями другого. Цепочка для ML: `linalg -> affine -> scf -> cf -> llvm dialect -> LLVM IR -> машинный код`.
XLA (Accelerated Linear Algebra) использует MLIR для оптимизации ML моделей: fusion циклов (объединение нескольких операций в один kernel), тайлинг для cache efficiency, SIMD векторизация. На TPU: XLA -> MLIR -> TPU-specific dialect -> XLA runtime. Speedup от fusion: 2-10x меньше memory bandwidth для transformer-like операций.
Почему progressive lowering предпочтительнее одного большого шага компиляции?
MLIR Operations и Types
MLIR Operation - унифицированный узел IR: имеет operands, results, attributes, regions (вложенный IR), successor blocks, location info. Regions позволяют вкладывать полноценный IR внутрь операций: функции - это регион в FuncOp, тело if-else - регион в scf.if. Это мощнее чем LLVM IR где функции - верхнеуровневые сущности.
MLIR используется в production: Google TPU compiler, Apple ANE (Apple Neural Engine) compiler, Intel OpenVINO, AMD ROCm (через MLIR GPU dialect), ARM Compute Library. Стандартные dialects: `arith`, `func`, `memref`, `tensor`, `linalg`, `affine`, `scf`, `cf`, `vector`, `gpu`, `spirv`, `nvgpu`. Кастомный dialect создаётся за ~500 строк TableGen + C++.
Что такое Region в MLIR Operation и зачем он нужен?
Rewrite Patterns и Dialect Conversion
RewritePattern - правило замены одной операции IR другими. PatternRewriter применяет паттерны с cost model (benefit). DialectConversion - специализация для перевода между dialects: определяется target (какие операции легальны) и паттерны для illegal -> legal. Greedy pattern application находит fixpoint автоматически.
MLIR Linalg-on-tensors pipeline (2023) показывает 1.5-2x speedup для transformer inference по сравнению с XLA без MLIR. Это достигается через better fusion: вместо материализации каждого tensor в памяти, fusion объединяет операции и держит данные в регистрах/L1 кеше. AMD ROCm использует MLIR для GPU kernel генерации начиная с ROCm 5.x.
MLIR - это просто новая версия LLVM IR с другим синтаксисом
MLIR - это инфраструктура для создания компиляторных IR с нуля через dialects; LLVM IR - это один конкретный IR, MLIR позволяет создавать десятки специализированных IR для разных доменов
LLVM IR фиксирован: типы, инструкции, семантика заданы раз и навсегда. MLIR позволяет определить любые типы и операции через TableGen. GPU-специфичные операции, ML tensor типы, hardware-specific ISA - всё это MLIR dialects которые невозможно выразить в LLVM IR
Что делает `applyPatternsAndFoldGreedily` в MLIR?
Итоги
- Dialect - набор типов и операций для конкретного домена; разные dialects сосуществуют в одном IR и поэтапно опускаются до LLVM IR
- Progressive lowering: linalg -> affine -> scf -> cf -> llvm; каждый уровень позволяет применять оптимизации специфичные для этой абстракции
- RewritePattern + greedy fixpoint: паттерн-матчинг на IR операциях с автоматическим поиском fixpoint; базис для fusion, tiling, vectorization
Связанные темы
MLIR расширяет концепции LLVM для ML и DSL компиляции:
- LLVM инфраструктура — MLIR - это надстройка над LLVM идеями; финальный lowering идёт в LLVM dialect -> LLVM IR
- LLVM Passes — MLIR passes используют тот же принцип что LLVM passes: анализ + трансформация IR
- DSL компиляторы — SQL, GLSL shader, конфигурационные языки используют схожие multi-level IR принципы
Вопросы для размышления
- MLIR позволяет смешивать dialects в одном IR. Как компилятор гарантирует корректность на границах dialects при mixed-dialect программах?
- Affine dialect предполагает полиэдральную модель циклов (статически известные bounds, affine access patterns). Как компилятор обрабатывает dynamic-size операции которые не вписываются в affine model?
- PyTorch torch.compile JIT-компилирует Python на основе трассировки. Какие проблемы возникают при компиляции Python кода с data-dependent control flow?