Компания Amazon представила компилятор NNVM, предназначенный для компиляции предоставляемых системами машинного обучения высокоуровневых графов вычислений в оптимизированный набор машинных кодов. Код проекта распространяется под лицензией Apache 2.0.

Система поддерживает компиляцию моделей в форматах OpenML (поддерживаются фреймворки Keras и Caffe), Apache MXNet и недавно представленного FaceBook и Microsoft открытого формата ONNX (Open Neural Network Exchange), при помощи которого могут передавать модели обучения из фреймворков Caffe2, PyTorch и CNTK (Cognitive Toolkit). На выходе может генерироваться код для различных бэкендов, включая ядра для вычислений на стороне GPU при помощи CUDA, OpenCL и Metal, а также промежуточный код LLVM, на основе которого могут формироваться машинные инструкции для архитектур x86 и ARM или представление WebAssembly.

Компиляция включает в себя несколько стадий:

• Формирование промежуточного представления графа вычислений на основе данных, полученных от различных фреймворков машинного обучения.
• Оптимизация полученного графа вычислений и выделение операторов с подпрограммами обработки данных в графе.
• Компиляция операторов в исполняемые модули и развёртывание для различных аппаратных бэкендов с минимальными зависимостями. Для построения кода на основе промежуточного представления используется стек TVM, предоставляющий предметно ориентированный язык для обеспечения работы используемых в модели операторов при помощи набора базовых реализаций, оптимизированных для различных целевых аппаратных платформ, а также вычислительных примитивов, таких как многопоточность, кэширование и разбиение циклов на блоки (tiling).

Архитектура NNVM позволяет легко добавлять новые фронтэнды, операторы и оптимизации графа, без необходимости изменения базовых интерфейсов. Полученные после компиляции модули могут быть сгенерированы в виде кода на c++, python, javascript, java, objective-c для запуска на серверах, мобильных устройствах, встраиваемых системах и в web-браузере.

Проведение тестирования генерации кода для CPU ARM и GPU NVIDIA показало, что формируемый компилятором итоговый код превосходит по производительности фреймворк MXNet. Размер формируемого модуля зависит в основном от размера runtime TVM, который при сборке для Raspberry Pi и мобильных устройств занимает около 300 Кб.