Plug-ins do compilador LiteRT

Quando devo criar um plug-in do compilador?

Um plug-in do compilador LiteRT é necessário quando você precisa integrar um acelerador de hardware específico com uma dependência do compilador na estrutura LiteRT.

Crie um plug-in do compilador se:

1. Você está segmentando um novo back-end de hardware que não é compatível.
2. Você quer descarregar operações específicas do modelo para esse acelerador de hardware para melhorar o desempenho ou a eficiência energética.
3. Você precisa de suporte para compilação AOT (na estação de trabalho) ou no dispositivo.

O plug-in atua como uma ponte, pegando partes do modelo de aprendizado de máquina e convertendo-as em um formato que o hardware de destino pode executar, usando uma chamada ao compilador do back-end. O LiteRT agrupa o bytecode personalizado gerado pelo plug-in no modelo .tflite, tornando-o executável usando o tempo de execução do LiteRT.

Como os plug-ins do compilador funcionam?

O framework LiteRT usa o plug-in do compilador durante o carregamento do modelo ou a fase de pré-processamento off-line para identificar e preparar subgrafos de modelo para execução no hardware de destino.

O processo envolve duas fases principais organizadas pelo framework usando as funções exportadas do plug-in:

1. Particionamento:o plug-in inspeciona todo o gráfico do modelo e identifica subconjuntos de operações que ele oferece suporte e pode acelerar com eficiência no hardware de destino. Esses subgrafos compatíveis são "particionados" (marcados) para compilação e delineados.
2. Compilação:o framework LiteRT transmite os subgrafos particionados de volta ao plug-in. Em seguida, o plug-in usa a lógica interna e possivelmente cadeias de ferramentas (compiladores) externas para gerar um ou mais módulos de bytecode específicos do hardware que implementam as partições. Esse bytecode é o que o tempo de execução (HAL/driver) do hardware de destino vai carregar e executar.

O framework substitui os subgrafos originais por operações personalizadas que invocam o driver de hardware, transmitindo o bytecode compilado criado pelo plug-in.

O LiteRT Dispatch é o análogo de tempo de execução para o plug-in do compilador. Eles fornecem os meios de chamar a HAL com base na saída do compilador. Para mais detalhes, consulte a documentação de envio.

AOT x no dispositivo

O LiteRT pode usar plug-ins de compilador para oferecer suporte à compilação AOT com nossas ferramentas, além da compilação no dispositivo. A compilação no dispositivo é mais flexível, totalmente internalizada nas APIs de tempo de execução do LiteRT e exige apenas o gerenciamento de um único modelo. O fluxo AOT pode desbloquear a compilação quando ela é muito intensiva em recursos para ser executada no dispositivo, o que pode acontecer com muitos modelos grandes contemporâneos.

Fallback

O LiteRT foi criado com suporte para gráficos heterogêneos. Qualquer operação não selecionada pelo plug-in será deixada para a CPU ou disponibilizada para aceleração em outro back-end.

Implementar um plug-in do compilador

Um plug-in do compilador LiteRT é implementado como uma biblioteca compartilhada que exporta um conjunto específico de funções C definidas na API C do LiteRT.

Funções essenciais da interface

A funcionalidade principal gira em torno de duas etapas de compilação principais: LiteRtCompilerPluginPartition e LiteRtCompilerPluginCompile.

Snippets da 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. A função de partição

A assinatura da função é:

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

O que a função partition faz:esta é a fase de seleção. O plug-in itera as operações na entrada LiteRtSubgraph. Para cada operação que o hardware de destino aceita e pode acelerar, o plug-in adiciona essa operação ao LiteRtOpList$ fornecido no parâmetro selected_ops. O framework LiteRt usa essa lista para definir os limites das partições que serão enviadas para a etapa de compilação final.

Por padrão, o LiteRT agrupa todas as operações selecionadas nos maiores sub-DAGs possíveis. Para um particionamento mais refinado, um índice pode ser associado ao selecionar operações que servem para dividir ainda mais esses subgrafos.

2. A função "compile"

A assinatura da função é:

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

O que a função compile faz:essa é a fase de geração. A entrada partitions representa um modelo em que todos os subgrafos selecionados foram isolados. O plug-in processa essas partições, invocando a cadeia de ferramentas específica para gerar o bytecode para o hardware de destino. Espera-se que a saída do plug-in forneça um ponto de entrada para cada subgrafo transmitido para compilação. Na maioria dos casos, são módulos de bytecode individuais para cada subgrafo de entrada ou um único módulo de bytecode com vários pontos de entrada.

Tipo de dados retornados por compile:a função LiteRtCompilerPluginCompile retorna a saída usando o parâmetro out LiteRtCompiledResult.

O LiteRtCompiledResult é um identificador opaco (em relação ao LiteRT) para uma estrutura gerenciada pelo plug-in. Ele representa a saída da compilação e contém duas informações principais:

1. Módulos de bytecode:um ou mais buffers de memória bruta que contêm o bytecode executável específico do hardware (ou seja, instruções compiladas).
2. Informações da chamada:metadados de cada partição. Isso fornece o mapeamento do iº subgrafo de entrada para um módulo de bytecode de resultado e um identificador de ponto de entrada nesse módulo.

Implementação de exemplo

Os snippets a seguir ilustram como um plug-in básico pode implementar as funções principais. Este exemplo foi extraído de um exemplo totalmente funcional em litert/vendors/examples/

Identificação e configuração de plug-ins

Essas funções fornecem ao framework informações básicas sobre o plug-in e o hardware.

// 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;
}

Lógica de particionamento (LiteRtCompilerPluginPartition)

Este exemplo mostra o plug-in selecionando um conjunto limitado de operações (mul, sub e uma operação composta específica) somente se todas as entradas e saídas forem números de ponto flutuante de 32 bits. Normalmente, para determinar se uma operação deve ser selecionada, é necessário incluir uma chamada para um hook de validação no conjunto de ferramentas do compilador do back-end.

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;
}

Antes de chamar a compilação, o LiteRT vai validar e "delinear" todas as operações selecionadas em novos subgrafos em um novo modelo intermediário. Esse modelo intermediário é o que é transmitido para a compilação.

Lógica de compilação (LiteRtCompilerPluginCompile)

Essa função usa os subgrafos particionados e gera um LiteRtCompiledResult personalizado. Este exemplo gera um módulo de bytecode independente para cada partição a ser compilada. Em casos reais, isso geralmente envolve a conversão de operações do LiteRT em tipos para a biblioteca do compilador de back-end. O plug-in de exemplo funcional "compilação" cria uma string legível que codifica o gráfico.

// 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 ...

Uso e validação

O LiteRT oferece várias ferramentas para aplicar plug-ins de compilador a arquivos de modelo, executar o resultado e validar/fazer comparativos de mercado. Consulte a documentação do conjunto de testes de aceleradores e a documentação de comparativos de mercado e criação de perfis.