TensorCompile

TensorCompile — это AOT тензорный компилятор, созданный с нуля. Проект транслирует вычислительные графы нейронных сетей из формата ONNX в высокооптимизированный машинный код (ассемблер x86-64) с использованием инфраструктуры MLIR и LLVM.

Проект успешно компилирует и исполняет реальные архитектуры компьютерного зрения, такие как SqueezeNet.

---

Ключевые особенности

• Парсинг ONNX: Поддержка операций Conv (включая padding и strides), Relu, MaxPool, GlobalAveragePool, Concat, Flatten, MatMul и Linear.
• MLIR Lowering Pipeline: Математическое представление графов через диалекты tensor и linalg, спуск до циклов (scf), управление выделением памяти (bufferization -> memref) и преобразование в диалект llvm.
• C-ABI Runner: Механизм нативного запуска сгенерированного ассемблера из C++ кода.
• Математическая верификация: покрытие базовых операций модульными тестами (GTest) и сквозное сравнение вывода (inferece) с эталонными результатами PyTorch.
• Визуализация графов: Автоматическая генерация .dot файлов (Graphviz) для визуального анализа вычислительного графа сети.

---

Структура проекта

TensorCompile/
├── src/                # Исходный код компилятора (Frontend и MLIR Codegen)
├── include/            # Заголовочные файлы
├── tests/              # Unit-тесты на базе GTest
├── data/               # Директория для ONNX моделей и бинарных дампов тензоров
├── runner.cpp          # Механизм запуска сгенерированного машинного кода
├── export_squeezenet.py # Скрипт подготовки модели и эталонных данных (PyTorch)
└── CMakeLists.txt      # Конфигурация сборки

---

Зависимости

• C++17 совместимый компилятор (GCC / Clang)
• CMake (версии 3.15+)
• LLVM 17 и MLIR (включая mlir_c_runner_utils)
• Python 3 с библиотеками torch, torchvision, numpy, onnx
• Graphviz (опционально, для визуализации графов)

---

Сборка и запуск (на примере SqueezeNet)

Пайплайн работы разделен на 3 этапа: подготовка модели, кодогенерация (AOT) и нативное исполнение.

1. Подготовка модели и эталонных данных

Скрипт загрузит архитектуру SqueezeNet, отключит biases, задаст константные веса для предотвращения переполнения float32 и сохранит чистый входной тензор.

python3 export_squeezenet.py

(В папке data/ появятся squeezenet.onnx, input_224.bin и output_ref.bin)

2. Сборка компилятора и генерация машинного кода

Собираем сам проект компилятора через CMake:

cmake -B build
cmake --build build

Прогоняем модель через компилятор для получения ассемблера (output_asm.s):

./build/etc_compiler data/squeezenet.onnx

3. Линковка и запуск нативного Runner'а

Компилируем C++ runner вместе со сгенерированным ассемблером и прилинковываем библиотеку утилит MLIR:

g++ runner.cpp output_asm.s -o runner_exec -no-pie \\
    -L/usr/lib/llvm-17/lib -lmlir_c_runner_utils \\
    -Wl,-rpath=/usr/lib/llvm-17/lib

Запускаем инференс сети:

./runner_exec

(Ожидаемый результат: AOT компилятор и эталон PyTorch должны выдать идентичные значения с минимальной погрешностью).

---

Тестирование

В проекте настроены Unit-тесты для верификации парсера и конвейера генерации MLIR. Запуск тестов:

cd build
ctest --output-on-failure

---

Визуализация

При каждом запуске компилятора генерируется файл graph.dot.