Привет, Хабр! Меня зовут Владислав Виноградов, я инженер группы исследований и разработки ПО глубокого обучения в YADRO. Моя команда создает и оптимизирует связанное с искусственным интеллектом программное обеспечение. Сегодня я расскажу, как можно разработать тензорный компилятор для процессора на базе открытой архитектуры RISC-V.

Подход сочетает в себе автоматическую кодогенерацию и использование ручных оптимизаций. Это позволяет существенно экономить ресурсы команды для работы над наиболее вычислительно трудоемкими операциями, которые реализуются средствами внешних библиотек. Статья будет полезна, если вы ищете инструменты реализации оптимизирующих компиляторов для эффективного исполнения моделей глубокого обучения или вам интересно посмотреть на пример использования MLIR.

Что такое тензорный компилятор

Тензорным компилятором называют отдельный проект либо компонент в большем фреймворке, который занимается оптимизацией и переводом модели глубокого обучения в исполняемый формат под конкретное устройство. Этим устройством может быть CPU, GPU или специализированные AI-акселераторы.

Тензорный компилятор выполняет ту же работу, что и компиляторы традиционных языков программирования, — комбинирует высокоуровневые оптимизации (планирование исполнения, оптимизации ресурсов) и низкоуровневые (оптимизации уровня машинных инструкций).

Идея исследования

RISC-V — многообещающая архитектура для создания нового поколения процессоров. И она может стать платформой для решения задач AI. При этом эффективное исполнение (инференс) моделей глубокого обучения на любом устройстве требует предварительной оптимизации под конкретную архитектуру. В рамках исследований команды мы решили разработать прототип оптимизирующего тензорного компилятора для RISC-V CPU.

Технологический стек для тензорного компилятора

Основой эксперимента по созданию компилятора стали две технологии — OpenVINO и MLIR.

Почему OpenVINO

OpenVINO — это библиотека с открытым исходным кодом для оптимизации инференса сетей глубокого обучения под различные устройства: CPU, iGPU, GPU, FPGA. Библиотека разрабатывается компанией Intel и в первую очередь ориентирована на оптимизацию именно под их устройства. При этом у нее расширяемая архитектура, которая позволяет писать собственные плагины для сторонних устройств.

OpenVINO предоставляет общее внутреннее представление моделей глубокого обучения и отдельный инструментарий по конвертации из разных фреймворков и форматов в это общее представление. Библиотека значительно упрощает разработку тензорного компилятора, поскольку команда может сфокусироваться на обработке внутреннего представления. В OpenVINO также есть набор утилит по валидации аккуратности и бенчмаркингу, которые экономят ресурсы на начальных этапах разработки.

Почему MLIR

MLIR (Multi-Level Intermediate Representation) — это подпроект более известного фреймворка LLVM, инфраструктура для разработки компиляторов. Она предоставляет все необходимые компоненты для их создания: 

• гибкое и расширяемое внутреннее представление (IR),

• общие алгоритмы и подходы по работе с IR,

• оптимизирующие пассы,

• инструменты анализа и диагностики IR,

• инструменты для тестирования и отладки.

Также в MLIR большое внимание уделено оптимизации работы самого компилятора. В инфраструктуре есть:

• поддержка многопоточной компиляции,

• оптимизированные структуры данных для IR,

• средства профилирования пайплайна компиляции.

OpenVINO поддерживает запуск и оптимизацию моделей под CPU с архитектурой x86 и ARM. Его тензорный компилятор по сути просто планирует операции, которые написаны и оптимизированы вручную. То есть для каждой операции написана своя реализация, оптимизированная под расширения x86/ARM (либо как часть самого плагина OpenVINO, либо как примитив из библиотеки DNNL). Прототипирование такого прямого подхода на RISC-V требует значительных усилий по написанию большой библиотеки оптимизированных операций.

Использование же MLIR позволяет применить гибридный подход. Его идея в том, что разработчики занимаются реализацией для RISC-V трудоемких операций, а код для остальных генерируется автоматически средствами инфраструктуры. Такой подход позволяет команде сфокусироваться на оптимизации наиболее тяжелых bottleneck-операций, таких как свертка или матричные умножения. 

Более того, написанное на MLIR решение можно в теории адаптировать и запустить на разном «железе» и архитектурах: 

• CPU: x86, ARM, RISC-V,

• GPU,

• NPU. 

Есть и еще одна причина выбора MLIR. Процесс компиляции из высокоуровневого представления в низкоуровневый машинный код достаточно сложен из-за cемантического разрыва между ними. Пытаясь сразу прыгнуть из высокоуровневого представления в низкоуровневое, например, в LLVM IR, можно потерять много семантической информации и возможность проводить некоторые высокоуровневые оптимизации. MLIR предлагает постепенную трансформацию и перевод из абстракций высокого уровня в низкоуровневые через разные уровни абстракции. Отсюда Multi-Level в названии технологии.

MLIR предоставляет механизм, называемый диалектами. Они позволяют описывать разные уровни абстракции и домены применения операций, будь то тензорная алгебра, control flow, взаимодействие с потоками или акселератором и так далее. В итоге процесс компиляции представляет собой постепенную трансформацию внутреннего представления из высокоуровневых диалектов в более низкоуровневые.

Трансформация модели глубокого обучения в машинный код

В MLIR между моделью в OpenVINO и машинным кодом лежат 6 уровней абстракций: 

1. Исходное представление модели (input DSL).

2. Внешние библиотеки оптимизированных операций (external libraries).

3. Тензорная алгебра (tensor algebra).

4. Циклы и работа с памятью (loops on buffers).

5. Вызовы функций (function calls).

6. Низкоуровневое внутреннее представление (low-level IR).

У каждого уровня абстракции свой диалект. Дальше мы рассмотрим суть каждого из них.

MLIR предоставляет возможность распечатать внутреннее представление компилятора (IR) в текстовом виде, понятном человеку. Этот формат близок по своей структуре к C-подобному языку. В секциях ниже я буду приводить примеры внутреннего представления именно в этом формате.

Исходное представление модели

Сначала мы переводим внешнюю модель в формате OpenVINO в диалект OV, который написали сами. Он представляет модель один к одному в терминах MLIR и дает отдельные операции, работающие с тензорами.

%0 = ov.convert(%input : tensor<1x3x224x224xui8>) to f32 

%1 = ov.convolution(%0, %weights : tensor<1x3x224x224xf32>,
                    tensor<64x3x5x5xf32>) {
    dilations = array<i64: 2, 2>,
    pads_begin = array<i64: 0, 0>,
    pads_end = array<i64: 0, 0>,
    strides = array<i64: 1, 1>
} : tensor<1x64x216x216xf32> 

%2 = ov.add(%1, %bias : tensor<1x64x216x216xf32>, tensor<64xf32>)

На этом этапе осуществляется приведение внутреннего представления к «каноническому» виду для упрощения последующих этапов. Этот шаг включает в себя:

• Обработку константных путей вычисления.

• Тривиальные оптимизации, например, reshape(reshape(x, shape1), shape2) -> reshape(x, shape2).

• Объединение нескольких операций в одну: convolution+bias.

• Декомпозицию некоторых сложных операций на более простые.

Внешняя библиотека оптимизированных операций

Следующим шагом нужно трансфоровать представление в диалекте OpenVINO в диалект DNNL, который представляет специфику этой внешней библиотеки. Реализацию диалекта мы тоже написали сами. 

Конвертация из OV-диалекта в DNNL делается выборочно для тех операций, что представлены в библиотеке или могут быть выражены через ее API. Например, мы трансформировали связку из конволюции и последующего добавления биаса в операцию, представляющую собой семантику подобной операции из внешней библиотеки:

// ov.convolution + ov.add

%1 = tensor.empty() : tensor<1x64x216x216xf32>
%2 = dnnl.convolution 
  ins(%0 : tensor<1x3x224x224xui8>,
    %weights : tensor<64x3x5x5xf32>,
    %bias : tensor<64xf32>) 
  outs(%1 : tensor<1x64x216x216xf32>) {
    dilations = array<i64: 2, 2>,
    pads_begin = array<i64: 0, 0>,
    pads_end = array<i64: 0, 0>,
    strides = array<i64: 1, 1>
  }

На этом уровне также можно производить оптимизации планирования вызова DNNL/oneDNN. Например: 

• выбрать наиболее подходящую для каждой операции схему расположения данных в памяти — так называемые лейауты (layouts), 

• согласовать лейауты между разными операциями,

• добавить дополнительные операции по изменению лейаута.

// ov.convolution + ov.add 

%1 = tensor.empty() : tensor<1x8x216x216x8xf32>
%2 = dnnl.cast %1 : tensor<1x8x216x216x8xf32>
  to tensor<1x64x216x216xf32, #nChw8c>> // layout

%3 = dnnl.convolution 
  ins(%0 : tensor<1x3x224x224xui8>,
    %weights : tensor<64x3x5x5xf32, #Oihw8o>, // layout
    %bias : tensor<64xf32>)
  outs(%2 : tensor<1x64x216x216xf32, #nChw8c>) { // layout
    dilations = array<i64: 2, 2>,
    pads_begin = array<i64: 0, 0>,
    pads_end = array<i64: 0, 0>,
    strides = array<i64: 1, 1>
  } 
// re-layout
%4 = tensor.empty() : tensor<1x64x216x216xf32>
%5 = dnnl.reorder ins(%3 : tensor<1x64x216x216xf32, #nChw8c>)
          outs(%4 : tensor<1x64x216x216xf32>)

Тензорная алгебра

Далее мы переходим в ветку автоматической кодогенерации. Здесь используется диалект linalg, который доступен в самом MLIR. На этом уровне мы сводим разные операции к более общему представлению тензорной алгебры. Эту конвертацию делаем после того, как оптимизировали планировщик операций на предыдущем уровне. В результате получаем возможность автоматически генерировать код для добавляемых операций, например, смены лейаута.

// ov.convert

%0 = tensor.empty() : tensor<1x3x224x224xf32>
%1 = linalg.map { arith.uitofp } // generic op
    ins(%input : tensor<1x3x224x224xi8>) 
    outs(%0: tensor<1x3x224x224xf32>)
      
// ...
      
// dnnl.reorder
      
%4 = dnnl.cast %1 : tensor<1x64x216x216xf32, #nChw8c>
    to tensor<1x8x216x216x8xf32>>
%5 = tensor.empty() : tensor<1x64x216x216xf32>
%6 = tensor.unpack %3 inner_dims_pos = [1] inner_tiles = [8] into %5 
    : tensor<1x8x216x216x8xf32>> -> tensor<1x64x216x216xf32> // generic op

Циклы и работа с памятью

Четвертый шаг понижения семантики — это диалект affine и переход от операций алгебры над тензорами к обычным циклам по работе с памятью.  

На этом этапе мы можем решать задачи выделения памяти — например, для того, чтобы выделять ее заранее, а не при каждом вызове инференса. Представление похоже по семантике на операции с памятью в языке С: прочитать из памяти, сохранить в нее данные, произвести операции над элементами и так далее.

// ov.convert

%1 = memref.alloc : memref<123456xi8> // memory planning
  
%2 = memref.view %1[0] : memref<123456xi8> to 
              memref<1x3x224x224xf32> 

%3 = memref.collapse_shape %input [[0, 1, 2, 3]] :
    memref<1x3x224x224xi8> into memref<150528xi8>
%4 = memref.collapse_shape %2 [[0, 1, 2, 3]] :
    memref<1x3x224x224xf32> into memref<150528xf32>

affine.for %ind = 0 to 150528 { // loops on buffers
  %i0 = affine.load %3[%ind] : memref<150528xi8>
  %i1 = arith.uitofp %i0: i8 to f32
  affine.store %i1, %4[%ind] : memref<150528xf32> 
}

Вызов функций внешних библиотек

Следующий шаг — это трансформация из DNNL-диалекта непосредственно в вызовы функций из внешних библиотек. Для этого используется диалект func.

Архитектура на базе MLIR-диалектов не ограничивает количество подключаемых оптимизированных внешних библиотек, которые отвечают за инференс слоев нейросети. В этом примере мы рассматриваем реализацию на примере интеграции нашего  MLIR-компилятора с библиотекой DNNL/oneDNN, оптимизированной под RISC-V.

На этом этапе мы также можем разделить для некоторых сложных операций этап инициализации примитива из внешней библиотеки и этап непосредственно вызова этого примитива средствами MLIR.

// dnnl.convolution

func.func @initialize(%ctx: !llvm.ptr) { // initialization step
  %prim = call @get_primitive(%ctx, %c0)
  call @add_primitive_arg(%prim, %c4, %src_dims, %src_strides)
  call @create_convolution(%prim, %conv_strides, %conv_pads)
}

func.func @infer(%ctx: !llvm.ptr, %input: memref<1x3x224x224xi8>, 
                         %out: memref<1x64x216x216xf32>) {
  %prim = call @get_primitive(%ctx, %c0)
  %ptr = memref.extract_aligned_pointer_as_index %4
  call @set_primitive_arg(%prim, %c0, %ptr)
  call @execute_primitive(%prim) // execution step
}

Низкоуровневое внутреннее представление

На последнем шаге компиляции к единому низкоуровневому представлению сводятся как представление вызова функций внешней библиотеки, так и автоматически сгенерированный код. Таким представлением выступает LLVM-диалект, который является зеркалом LLVM IR.

...
^bb6:  // pred: ^bb5
 %11 = llvm.mul %5, %0  : i64
 %12 = llvm.mul %7, %4  : i64
 %13 = llvm.add %11, %12  : i64
 %14 = llvm.add %13, %9  : i64
 %15 = llvm.getelementptr %arg1[%14] : (!llvm.ptr<f32>, i64) -> !llvm.ptr<f32>
 %16 = llvm.load %15 : !llvm.ptr<f32>
 %17 = llvm.getelementptr %arg12[%14] : (!llvm.ptr<f32>, i64) -> !llvm.ptr<f32>
 %18 = llvm.load %17 : !llvm.ptr<f32>
 %19 = llvm.fadd %16, %18  : f32
 %20 = llvm.getelementptr %arg23[%14] : (!llvm.ptr<f32>, i64) -> !llvm.ptr<f32>
 llvm.store %19, %20 : !llvm.ptr<f32>
 %21 = llvm.add %9, %3  : i64
 llvm.br ^bb5(%21 : i64)
…

Генерация машинного кода

Финальный этап работы тензорного компилятора — это генерация машинного кода средствами LLVM и его бэкэнда для RISC-V.

Итого компиляция представляет собой пошаговую оптимизацию и перевод внутреннего представления во все более низкоуровневые диалекты.

При разработке экспериментального тензорного компилятора мы использовали гибридный подход. Вручную мы реализуем свертки, матричные умножения и data-dependent операции (например, DetectionOutput из SSD-моделей детектирования). Все остальное делает кодогенерация MLIR. Получились следующие результаты: 

• 17 операций — ручная реализация библиотеки ядер,

• 38 операций — автоматическая кодогенерация.

Запуск на «железе» и выводы

В качестве тестового стенда для разработанного плагина OpenVINO с интегрированным MLIR-компилятором мы использовали плату MangoPI с одноядерным RISC-V процессором CPU RV64GCV @ 408 МГц и пиковой производительностью 1.63 GFlops @ 408 МГц. Данная плата поддерживает векторное расширение RISC-V версии 0.7.1 (RVV 0.7.1), но при этом кодогенерация в бэкенде RISC-V LLVM работает только с версией 1.0. 

По этим причинам в кодогенерации слоев мы использовали скалярный, а не векторизованный код, хотя слои, написанные во внешней библиотеке, использовали возможности RVV версии 0.7.1. Плату с векторным расширением 1.0 для тестов нам пока не удалось найти, из-за этого мы еще не смогли проверить векторизацию в кодогенерации на реальном железе, только на эмуляторе.

Мы сравнили вызовы всех операций через внешнюю библиотеку и гибридный режим, когда вызовы во внешнюю библиотеку отправляются только для самых тяжелых слоев, а все остальные генерируются автоматически. Для представленных моделей команда оптимизаторов также реализовала все необходимые слои во внешней библиотеке. Вот результаты запуска на реальном «железе»:

В итоге производительность гибридного режима сопоставима с производительностью использования внешней библиотеки. В отдельных случаях результаты гибридного режима даже лучше. Это объясняется тем, что в библиотеке оптимизированных операций код написан для общего случая (разных размеров тензоров, разных раскладок по памяти и т.п.), а при автоматической генерации компилятор генерирует код под конкретные параметры, известные на этапе компиляции. Это позволяет использовать дополнительные оптимизации, такие как loop unrolling и более оптимальные паттерны доступа в память. При этом для некоторых моделей общее число слоев с автогенерацией достаточно существенное — 20-30%. 

Автоматическая генерация забрала на себя немалое число слоев, но время их работы существенно не влияет на общее время работы нейронной сети. Зато освобождаются ресурсы команды, которая может сфокусироваться на оптимизации более тяжеловесных слоев.

Вывод, который мы сделали при реализации прототипа тензорного компилятора, — гибридный подход имеет право на жизнь. 

Мы также отметили несколько минусов в подобном подходе. Первый — это достаточно высокий порог входа в MLIR для людей, которые не сталкивались со спецификой написания компиляторов и LLVM-фреймворком. И второй минус — отсутствие поддержки RVV 0.7.1 в LLVM.

Дальнейшие шаги

В дальнейших исследованиях мы хотим поработать над тремя направлениями: 

1. Развитие оптимизации для автоматической кодогенерации.

2. Бо́льшая кооперация с библиотекой ядер.

3. Увеличение покрытия поддерживаемых моделей. 

Во-первых, мы хотим заняться дальнейшим улучшением оптимизации для автоматической кодогенерации. Рассмотреть более сложные паттерны оптимизации, векторизации для слоев, и увеличить покрытие слоев, для которых производится кодогенерация. 

Второе направление, которое мы планируем изучить, — это бо́льшая кооперация компилятора с внешней библиотекой ядер. Например, когда некоторые слои конволюции на стадии кодогенерации разбиваются на более мелкие примитивы или микроядра, небольшие матричные умножения. Их реализация будет во внешней библиотеке ядер. Тогда кодогенерация будет заниматься планированием микроядер и их оптимизацией по типу фьюзинга с другими операциями, например, активациями. Это избавит команду оптимизаторов от необходимости реализовывать обвязки вокруг микроядер и поддержку фьюзинга.

И третье — мы хотим увеличивать покрытие поддерживаемых моделей. Например, добавить поддержку моделей с динамическими размерами тензоров и моделей для обработки естественного языка (RNN, transformers).

Если вам будет интересно прочитать туториал про MLIR и процесс написания своих диалектов, пишите в комментариях. Мы с командой можем развить эту тему в отдельном тексте.