Plug-in del compilatore LiteRT

Quando dovrei creare un plug-in del compilatore?

Un plug-in del compilatore LiteRT è necessario quando devi integrare un acceleratore hardware specifico con una dipendenza del compilatore nel framework LiteRT.

Devi creare un plug-in del compilatore se:

1. Stai scegliendo come target un nuovo backend hardware non supportato.
2. Vuoi eseguire l'offload di operazioni specifiche del modello su questo acceleratore hardware per migliorare le prestazioni o l'efficienza energetica.
3. Hai bisogno del supporto per la compilazione AOT (sulla workstation) o per la compilazione sul dispositivo.

Il plug-in funge da ponte, prendendo parti del modello di machine learning e convertendole in un formato che l'hardware di destinazione può eseguire, utilizzando una chiamata al compilatore del backend. LiteRT raggruppa il bytecode personalizzato generato dal plug-in nel modello .tflite, rendendolo eseguibile utilizzando il runtime LiteRT.

Come funzionano i plug-in del compilatore?

Il framework LiteRT utilizza il plug-in del compilatore durante il caricamento del modello o la fase di pre-elaborazione offline per identificare e preparare i sottografi del modello per l'esecuzione sull'hardware di destinazione.

Il processo prevede due fasi principali orchestrate dal framework utilizzando le funzioni esportate del plug-in:

1. Partizionamento:il plug-in esamina l'intero grafico del modello e identifica sottoinsiemi di operazioni che supporta e che può accelerare in modo efficiente sull'hardware di destinazione. Questi sottografi supportati sono "partizionati" (contrassegnati) per la compilazione e delineati.
2. Compilazione:il framework LiteRT restituisce i sottografi partizionati al plug-in. Il plug-in utilizza quindi la sua logica interna ed eventualmente toolchain (compilatori) esterni per generare uno o più moduli bytecode specifici per l'hardware che implementano le partizioni. Questo bytecode è ciò che il runtime (HAL/driver) dell'hardware di destinazione alla fine caricherà ed eseguirà.

Il framework sostituisce i sottografi originali con operazioni personalizzate che richiamano il driver hardware, passando il bytecode compilato creato dal plug-in.

LiteRT Dispatch è l'analogo del runtime per il plug-in del compilatore. Forniscono il mezzo per chiamare l'HAL dato l'output del compilatore. Per saperne di più, consulta la documentazione sulla distribuzione.

AOT e On-Device

LiteRT può utilizzare i plug-in del compilatore per supportare la compilazione AOT tramite i nostri strumenti, nonché la compilazione on-device. La compilazione on-device è più flessibile, completamente internalizzata all'interno delle API di runtime di LiteRT e richiede solo la gestione di un singolo modello. Il flusso AOT può sbloccare la compilazione quando è troppo intensiva di risorse per essere eseguita sul dispositivo, come può accadere con molti modelli di grandi dimensioni contemporanei.

Fallback

LiteRT è stato creato con il supporto per grafici eterogenei. Qualsiasi operazione non selezionata dal plug-in verrà lasciata alla CPU o resa disponibile per l'accelerazione su un altro backend.

Implementazione di un plug-in del compilatore

Un plug-in del compilatore LiteRT viene implementato come libreria condivisa che esporta un insieme specifico di funzioni C definite nell'API C di LiteRT.

Funzioni essenziali dell'interfaccia

La funzionalità principale si basa su due passaggi di compilazione chiave: LiteRtCompilerPluginPartition e LiteRtCompilerPluginCompile.

Snippet 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. Il ruolo di partizione

La firma della funzione è:

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

Funzione partition: questa è la fase di selezione. Il plug-in esegue l'iterazione sulle operazioni nell'LiteRtSubgraph di input. Per ogni operazione supportata e accelerabile dall'hardware di destinazione, il plug-in aggiunge l'operazione a LiteRtOpList$ fornito nel parametro selected_ops. Il framework LiteRt utilizza questo elenco per definire i limiti delle partizioni che verranno inviate per il passaggio di compilazione finale.

Per impostazione predefinita, LiteRT raggruppa tutte le operazioni selezionate nei sottodag più grandi possibili. Per un partizionamento più granulare, è possibile associare un indice quando si selezionano le operazioni che servono a suddividere ulteriormente questi sottografi.

2. La funzione Compile

La firma della funzione è:

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

Cosa fa la funzione compile: questa è la fase di generazione. L'input partitions rappresenta un modello in cui tutti i sottografi selezionati sono stati isolati. Il plug-in elabora queste partizioni richiamando la toolchain specifica per generare il bytecode per l'hardware di destinazione. È previsto che l'output del plug-in fornisca un punto di ingresso per ogni sottografo passato per la compilazione. Nella maggior parte dei casi, si tratta di singoli moduli di bytecode per ogni sottografo di input o di un singolo modulo di bytecode con più punti di accesso.

Tipo di dati restituiti da compile: la funzione LiteRtCompilerPluginCompile restituisce l'output utilizzando il parametro out LiteRtCompiledResult.

LiteRtCompiledResult è un handle opaco (rispetto a LiteRT) di una struttura gestita dal plug-in. Rappresenta l'output della compilazione e contiene due informazioni principali:

1. Moduli di bytecode:uno o più buffer di memoria non elaborati contenenti il bytecode eseguibile specifico dell'hardware (ovvero le istruzioni compilate).
2. Informazioni sulle chiamate:metadati per ogni partizione. Questo fornisce la mappatura dal i° sottografo di input a un modulo di bytecode dei risultati e all'identificatore del punto di ingresso in quel modulo.

Esempio di implementazione

I seguenti snippet illustrano come un plug-in di base potrebbe implementare le funzioni principali. Questo esempio è tratto da un esempio completamente funzionale in litert/vendors/examples/

Identificazione e configurazione dei plug-in

Queste funzioni forniscono al framework informazioni di base sul plug-in e sull'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;
}

Logica di partizionamento (LiteRtCompilerPluginPartition)

Questo esempio mostra il plug-in che seleziona un insieme limitato di operazioni (mul, sub e un'operazione composita specifica) solo se tutti gli input e gli output sono float a 32 bit. In genere, per determinare se un'operazione deve essere selezionata viene chiamata una hook di convalida nella toolchain del compilatore 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;
}

Prima della compilazione, LiteRT convalida e "delinea" tutte le operazioni selezionate in nuovi sottografi in un nuovo modello intermedio. Questo modello intermedio viene passato alla compilazione.

Logica di compilazione (LiteRtCompilerPluginCompile)

Questa funzione prende i sottografi partizionati e genera un LiteRtCompiledResult personalizzato. Questo esempio genera un modulo bytecode autonomo per ogni partizione da compilare. In casi reali, di solito ciò comporta la conversione delle operazioni LiteRT in tipi nella libreria del compilatore backend. La "compilazione " del plug-in dell'esempio funzionale crea una stringa leggibile che codifica il grafico.

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

Utilizzo e convalida

LiteRT fornisce vari strumenti per applicare i plug-in del compilatore ai file modello, eseguire il risultato e convalidare/benchmark. Consulta la documentazione della suite di test dell'acceleratore e la documentazione su benchmarking e profilazione.