Complementos del compilador de LiteRT

¿Cuándo debo crear un complemento del compilador?

Un complemento del compilador de LiteRT es necesario cuando necesitas integrar un acelerador de hardware específico con una dependencia del compilador en el framework de LiteRT.

Debes crear un complemento del compilador en los siguientes casos:

1. Estás segmentando anuncios para un nuevo backend de hardware que no se admite.
2. Deseas descargar operaciones específicas del modelo en ese acelerador de hardware para mejorar el rendimiento o la eficiencia energética.
3. Necesitas compatibilidad con la compilación AOT (en la estación de trabajo) o la compilación en el dispositivo.

El complemento actúa como un puente, ya que toma partes del modelo de aprendizaje automático y las convierte en un formato que el hardware de destino puede ejecutar, mediante una llamada al compilador del backend. LiteRT agrupa el bytecode personalizado que genera el complemento en el modelo .tflite, lo que permite ejecutarlo con el entorno de ejecución de LiteRT.

¿Cómo funcionan los complementos del compilador?

El framework de LiteRT usa el complemento del compilador durante la fase de carga del modelo o de preprocesamiento sin conexión para identificar y preparar los subgrafos del modelo para su ejecución en el hardware de destino.

El proceso consta de dos fases principales coordinadas por el framework a través de las funciones exportadas del complemento:

1. Particionamiento: El complemento inspecciona todo el grafo del modelo y, luego, identifica subconjuntos de operaciones que admite y que puede acelerar de manera eficiente en el hardware de destino. Estos subgrafos admitidos se "particionan" (marcan) para la compilación y se describen.
2. Compilación: El framework de LiteRT pasa los subgrafos particionados de vuelta al complemento. Luego, el complemento usa su lógica interna y, posiblemente, cadenas de herramientas externas (compiladores) para generar uno o más módulos de bytecode específicos del hardware que implementan las particiones. Este bytecode es lo que el tiempo de ejecución (HAL/controlador) del hardware de destino cargará y ejecutará.

El framework reemplaza los subgrafos originales por operaciones personalizadas que invocan el controlador de hardware y pasan el código de bytes compilado creado por el complemento.

LiteRT Dispatch es el análogo del tiempo de ejecución para el complemento del compilador. Proporcionan los medios para llamar a la HAL a partir de la salida del compilador. Para obtener más detalles, consulta la documentación de envío.

Comparación entre AOT y On-Device

LiteRT puede usar complementos del compilador para admitir la compilación AOT a través de nuestras herramientas, así como la compilación en el dispositivo. La compilación en el dispositivo es más flexible, está completamente internalizada dentro de las APIs de tiempo de ejecución de LiteRT y solo requiere la administración de un solo modelo. El flujo de AOT puede desbloquear la compilación cuando es demasiado intensivo en recursos para ejecutarse en el dispositivo, lo que puede ocurrir con muchos modelos grandes contemporáneos.

Resguardo

LiteRT se creó con compatibilidad para gráficos heterogéneos. Cualquier operación que no seleccione el complemento se dejará para la CPU o estará disponible para la aceleración en otro backend.

Implementa un complemento del compilador

Un complemento del compilador de LiteRT se implementa como una biblioteca compartida que exporta un conjunto específico de funciones en C definidas en la API de LiteRT en C.

Funciones de interfaz esenciales

La funcionalidad principal gira en torno a dos pasos clave de compilación: LiteRtCompilerPluginPartition y LiteRtCompilerPluginCompile.

Fragmentos de la API de 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. La función de partición

La firma de la función es la siguiente:

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

Qué hace la función partition: Esta es la fase de selección. El complemento itera sobre las operaciones en el LiteRtSubgraph de entrada. Para cada operación que el hardware de destino admite y puede acelerar, el complemento agrega esa operación a LiteRtOpList$ proporcionado en el parámetro selected_ops. El framework de LiteRt usa esta lista para definir los límites de las particiones que se enviarán para el paso de compilación final.

De forma predeterminada, LiteRT agrupará todas las operaciones seleccionadas en los sub-DAG más grandes posibles. Para una partición más detallada, se puede asociar un índice cuando se seleccionan operaciones que sirven para dividir aún más estos subgrafos.

2. La función Compile

La firma de la función es la siguiente:

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

Qué hace la función compile: Esta es la fase de generación. La entrada partitions representa un modelo en el que se aislaron todos los subgrafos seleccionados. El complemento procesa estas particiones y, luego, invoca su cadena de herramientas específica para generar el bytecode para el hardware de destino. Se espera que el resultado del complemento proporcione un punto de entrada para cada subgrafo que se pase para la compilación. En la mayoría de los casos, se trata de módulos de código de bytes individuales para cada subgrafo de entrada o de un solo módulo de código de bytes con varios puntos de entrada.

Tipo de datos que devuelve compile: La función LiteRtCompilerPluginCompile devuelve su resultado con el parámetro de salida LiteRtCompiledResult.

El LiteRtCompiledResult es un identificador opaco (con respecto a LiteRT) para una estructura administrada por el complemento. Representa el resultado de la compilación y contiene dos datos principales:

1. Módulos de código de bytes: Uno o más búferes de memoria sin procesar que contienen el código de bytes ejecutable específico del hardware (es decir, instrucciones compiladas).
2. Información de la llamada: Metadatos de cada partición Esto proporciona la asignación del subgrafo de entrada número i a un módulo de código de bytes de resultado y un identificador de punto de entrada en ese módulo.

Ejemplo de implementación

En los siguientes fragmentos, se ilustra cómo un complemento básico podría implementar las funciones principales. Este ejemplo se tomó de un ejemplo completamente funcional en litert/vendors/examples/.

Identificación y configuración del complemento

Estas funciones proporcionan al framework información básica sobre el complemento y el 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 partición (LiteRtCompilerPluginPartition)

En este ejemplo, se muestra el complemento que selecciona un conjunto limitado de operaciones (mul, sub y una operación compuesta específica) solo si todas las entradas y salidas son números de punto flotante de 32 bits. Por lo general, determinar si se debe seleccionar una operación incluirá una llamada a un hook de validación en la cadena de herramientas del compilador del backend.

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 llamar a la compilación, LiteRT validará y "esbozará" todas las operaciones seleccionadas en nuevos subgrafos en un nuevo modelo intermedio. Este modelo intermedio es el que se pasa a la compilación.

Lógica de compilación (LiteRtCompilerPluginCompile)

Esta función toma los subgrafos particionados y genera un LiteRtCompiledResult personalizado. En este ejemplo, se genera un módulo de código de bytes independiente para cada partición que se compilará. En casos reales, esto suele implicar la conversión de operaciones de LiteRT en tipos para la biblioteca del compilador de backend. La "compilación" del complemento de ejemplo funcional crea una cadena legible que codifica el 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 y validación

LiteRT proporciona varias herramientas para aplicar complementos del compilador a archivos de modelos, ejecutar el resultado y realizar validaciones o comparativas. Consulta la documentación del paquete de pruebas del acelerador y la documentación de comparativas y generación de perfiles.