Плагины компилятора LiteRT

Когда следует создавать плагин компилятора?

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

Вам следует создать плагин компилятора, если:

1. Вы используете новый аппаратный бэкэнд , который не поддерживается.
2. Вы хотите перенести выполнение определенных операций с моделью на этот аппаратный ускоритель для повышения производительности или энергоэффективности.
3. Вам необходима поддержка компиляции AOT (на рабочей станции) или компиляции на устройстве.

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

Как работают плагины компилятора?

В фреймворке LiteRT плагин компилятора используется на этапе загрузки модели или предварительной обработки в автономном режиме для идентификации и подготовки подграфов модели к выполнению на целевом оборудовании.

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

1. Разделение на подмножества: Плагин анализирует весь граф модели и определяет подмножества операций, которые он поддерживает и может эффективно ускорять на целевом оборудовании. Эти поддерживаемые подграфы «разделяются» (помечаются) для компиляции и описываются.
2. Компиляция: Фреймворк LiteRT передает разделенные подграфы обратно плагину. Затем плагин использует свою внутреннюю логику и, возможно, внешние наборы инструментов (компиляторы) для генерации одного или нескольких аппаратно-специфических модулей байт-кода, реализующих эти разделы. Именно этот байт-код в конечном итоге загрузит и выполнит среда выполнения целевого оборудования (HAL/драйвер).

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

LiteRT Dispatch — это аналог плагина компилятора в среде выполнения. Он предоставляет средства для вызова функций HAL на основе выходных данных компилятора. Для получения более подробной информации обратитесь к документации по Dispatch .

AOT против On-Device

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

Отступать

LiteRT разработан с поддержкой гетерогенных графов. Любая операция, не выбранная плагином, будет либо передана на выполнение процессору, либо доступна для ускорения на другом бэкэнде.

Реализация плагина компилятора

Плагин компилятора LiteRT реализован в виде разделяемой библиотеки, которая экспортирует определенный набор функций C, определенных в API LiteRT C.

Основные функции интерфейса

Основная функциональность основана на двух ключевых этапах компиляции: LiteRtCompilerPluginPartition и LiteRtCompilerPluginCompile .

Фрагменты кода API на языке C

// Name associated with the manufacturer this plugin relates to.
LITERT_CAPI_EXPORT const char* LiteRtGetCompilerPluginSocManufacturer();

// Create and initialize the plugin instance.
LITERT_CAPI_EXPORT LiteRtStatus
LiteRtCreateCompilerPlugin(LiteRtCompilerPlugin* compiler_plugin,
                           LiteRtEnvironmentOptions env, LiteRtOptions options);

// Choose ops for compilation.
// This is the PARTITION step.
LITERT_CAPI_EXPORT LiteRtStatus LiteRtCompilerPluginPartition(
    LiteRtCompilerPlugin compiler_plugin, const char* soc_model,
    LiteRtSubgraph subgraph, LiteRtOpList selected_ops);

// Prepare result to pass to the runtime for given model containing partitioned
// subgraphs. This is the COMPILE step.
LITERT_CAPI_EXPORT LiteRtStatus LiteRtCompilerPluginCompile(
    LiteRtCompilerPlugin compiler_plugin, const char* soc_model,
    LiteRtModel partitions, LiteRtCompiledResult* compiled_result);

1. Функция разделения

Сигнатура функции следующая:

LITERT_CAPI_EXPORT LiteRtStatus LiteRtCompilerPluginPartition(
    LiteRtCompilerPlugin compiler_plugin, const char* soc_model,
    LiteRtSubgraph subgraph, LiteRtOpList selected_ops);

Что делает функция partition : Это этап выбора . Плагин перебирает операции во входном LiteRtSubgraph . Для каждой операции, которую поддерживает целевое оборудование и которую оно может ускорить, плагин добавляет эту операцию в список LiteRtOpList$, предоставленный в параметре selected_ops . Фреймворк LiteRt использует этот список для определения границ разделов, которые будут отправлены на заключительный этап компиляции.

По умолчанию LiteRT группирует все выбранные операции в максимально возможные подграфы. Для более точного разбиения на подграфы при выборе операций можно использовать индекс, который дополнительно разделяет эти подграфы.

2. Функция компиляции

Сигнатура функции следующая:

LITERT_CAPI_EXPORT LiteRtStatus LiteRtCompilerPluginCompile(
    LiteRtCompilerPlugin compiler_plugin, const char* soc_model,
    LiteRtModel partitions, LiteRtCompiledResult* compiled_result);

Что делает функция compile : Это этап генерации . Входные partitions представляют собой модель, в которой все выбранные подграфы изолированы. Плагин обрабатывает эти разделы, вызывая свой собственный набор инструментов для генерации байт-кода для целевого оборудования. Ожидается, что выходные данные плагина будут содержать точку входа для каждого подграфа, переданного на компиляцию. В большинстве случаев это либо отдельные модули байт-кода для каждого входного подграфа, либо один модуль байт-кода с несколькими точками входа.

Тип данных, возвращаемых функцией compile : Функция LiteRtCompilerPluginCompile возвращает свой результат, используя выходной параметр LiteRtCompiledResult .

Объект LiteRtCompiledResult представляет собой непрозрачный (по отношению к LiteRT) дескриптор структуры, управляемой плагином. Он отображает результат компиляции и содержит две основные части информации:

1. Модули байт-кода: один или несколько буферов в оперативной памяти, содержащих исполняемый байт-код, специфичный для конкретного оборудования (т.е. скомпилированные инструкции).
2. Информация о вызове: Метаданные для каждого раздела. Они обеспечивают сопоставление i -го входного подграфа с модулем байт-кода результата и идентификатором точки входа в этот модуль.

Пример реализации

Приведенные ниже фрагменты кода иллюстрируют, как базовый плагин может реализовывать основные функции. Этот пример взят из полностью функционального примера в litert/vendors/examples/

Идентификация и настройка плагинов

Эти функции предоставляют фреймворку основную информацию о плагине и оборудовании.

// Define the plugin's internal state structure
struct LiteRtCompilerPluginT {};

// Identify the manufacturer
const char* LiteRtGetCompilerPluginSocManufacturer() {
  return "AcmeCorp"; // Example manufacturer name
}

// Specify the supported hardware (in this example, it supports kLiteRtHwAcceleratorNpu)
LiteRtStatus LiteRtGetCompilerPluginSupportedHardware(
    LiteRtCompilerPlugin compiler_plugin,
    LiteRtHwAccelerators* supported_hardware) {
  // ... argument checking ...
  *supported_hardware = kLiteRtHwAcceleratorNpu;
  return kLiteRtStatusOk;
}

Логика разделения пространства ( LiteRtCompilerPluginPartition )

В этом примере показано, как плагин выбирает ограниченный набор операций ( mul , sub и конкретную составную операцию) только в том случае, если все входные и выходные данные являются 32-битными числами с плавающей запятой. Обычно определение того, следует ли выбирать ту или иную операцию, включает вызов обработчика проверки в цепочке инструментов компилятора бэкенда.

LiteRtStatus LiteRtCompilerPluginPartition(LiteRtCompilerPlugin compiler_plugin,
                                          const char* soc_model,
                                          LiteRtSubgraph subgraph,
                                          LiteRtOpList selected_ops) {

  // Iterate over ops and check criteria for selection
  // (using a C++ wrapper namespace '::litert' for convenience).
  // `subgraph` is a single subgraph from the original model, as such
  // this function will be called for each subgraph in the original model.

  ::litert::Subgraph main_subgraph(subgraph);
  for (const auto& op : main_subgraph.Ops()) {
    // 1. Check a constraint: require all tensors to be Float32
    bool only_f32 = true;
    // ... logic to check input/output types ...
    if (!only_f32) {
      continue;
    }

    // 2. Check op codes and push to selected_ops list
    if (op.Code() == kLiteRtOpCodeTflMul) {
      LITERT_RETURN_IF_ERROR(LiteRtPushOp(selected_ops, op.Get(), 0));
    } else if (op.Code() == kLiteRtOpCodeTflSub) {
      LITERT_RETURN_IF_ERROR(LiteRtPushOp(selected_ops, op.Get(), 0));
    } else if (op.Code() == kLiteRtOpCodeShloComposite) {
      // Example of checking composite op options
      // ... logic to check for "odml.rms_norm" name ...
      LITERT_RETURN_IF_ERROR(LiteRtPushOp(selected_ops, op.Get(), 0));
    }
  }
  return kLiteRtStatusOk;
}

Перед запуском компиляции LiteRT проверит и «выделит» все выбранные операции в новые подграфы в новой промежуточной модели. Именно эта промежуточная модель передается в процесс компиляции.

Логика компиляции ( LiteRtCompilerPluginCompile )

Эта функция принимает разделенные подграфы и генерирует пользовательский объект LiteRtCompiledResult . В этом примере для каждого раздела генерируется отдельный модуль байт-кода, который необходимо скомпилировать. В реальных случаях это обычно включает преобразование операций LiteRT в типы для библиотеки компилятора бэкэнда. Функция «компиляция» в функциональном примере плагина создает удобочитаемую строку, которая кодирует граф.

// Internal structure defining the compiled output
struct LiteRtCompiledResultT {
  std::vector<std::string> byte_code;   // The hardware bytecode buffers
  std::vector<std::string> per_op_data; // Per-call metadata (CallInfo)
};

LiteRtStatus LiteRtCompilerPluginCompile(
    LiteRtCompilerPlugin compiler_plugin, const char* soc_model,
    LiteRtModel partitions, LiteRtCompiledResult* compiled_result) {

  // 1. Create the internal result structure
  auto model = litert::Model::CreateFromNonOwnedHandle(partitions);
  const auto num_partitions = model.NumSubgraphs();
  auto result = std::make_unique<LiteRtCompiledResultT>();
  result->byte_code.resize(num_partitions);
  result->per_op_data.resize(num_partitions);

  // 2. Iterate and compile each partition
  for (auto i = 0; i < num_partitions; ++i) {
    // CompileSinglePartition is an internal helper that converts the subgraph
    // into the target hardware's format and stores it in result->byte_code.
    // In the case of the example this is just a stringification of the graph.

    // ... internal call to CompileSinglePartition ...
    // Example: result.byte_code[i] = generated_hw_code;
    // Example: result.per_op_data[i] = absl::StrFormat("Partition_%d", i);

    // The "per_op_data" is a unique identifier associated to the `ith` partition.
    // This is analogous to the name of a function in a library.
    // This is only meaningful when the plugin is preparing single modules with multiple entry points.
  }

  // 3. Pass ownership of the result back to the framework
  *compiled_result = result.release();

  return kLiteRtStatusOk;
}

// Functions to expose the compiled result data to the framework
LiteRtStatus LiteRtGetCompiledResultByteCode(
    LiteRtCompiledResult compiled_result, LiteRtParamIndex byte_code_idx,
    const void** byte_code, size_t* byte_code_size) {
  // ... implementation reads from compiled_result->byte_code ...
}
// ... other LiteRtGetCompiledResult* functions ...

Использование и проверка

LiteRT предоставляет различные инструменты для применения плагинов компилятора к файлам моделей, выполнения результатов и проверки/бенчмаркинга. См. документацию по набору тестов ускорителя и документацию по бенчмаркингу и профилированию .