Plug-ins du compilateur LiteRT

Quand dois-je créer un plug-in de compilation ?

Un plug-in de compilation LiteRT est nécessaire lorsque vous devez intégrer un accélérateur matériel spécifique avec une dépendance du compilateur dans le framework LiteRT.

Vous devez créer un plug-in de compilateur si :

1. Vous ciblez un nouveau backend matériel qui n'est pas compatible.
2. Vous souhaitez décharger des opérations de modèle spécifiques sur cet accélérateur matériel pour améliorer les performances ou l'efficacité énergétique.
3. Vous avez besoin d'une assistance pour la compilation AOT (sur la station de travail) ou la compilation sur l'appareil.

Le plug-in sert de pont. Il prend des parties du modèle de machine learning et les convertit dans un format que votre matériel cible peut exécuter, à l'aide d'un appel au compilateur du backend. LiteRT regroupe le bytecode personnalisé généré par le plug-in dans le modèle .tflite, ce qui le rend exécutable à l'aide du runtime LiteRT.

Comment fonctionnent les plug-ins de compilation ?

Le framework LiteRT utilise le plug-in du compilateur lors de la phase de chargement du modèle ou de prétraitement hors connexion pour identifier et préparer les sous-graphiques du modèle à exécuter sur le matériel cible.

Le processus comporte deux phases principales orchestrées par le framework à l'aide des fonctions exportées du plug-in :

1. Partitionnement : le plug-in inspecte l'ensemble du graphique du modèle et identifie les sous-ensembles d'opérations qu'il prend en charge et qu'il peut accélérer efficacement sur le matériel cible. Ces sous-graphes compatibles sont "partitionnés" (marqués) pour la compilation et décrits.
2. Compilation : le framework LiteRT renvoie les sous-graphiques partitionnés au plug-in. Le plug-in utilise ensuite sa logique interne et éventuellement des chaînes d'outils (compilateurs) externes pour générer un ou plusieurs modules de bytecode spécifiques au matériel implémentant les partitions. Ce bytecode est ce que le runtime (HAL/pilote) du matériel cible chargera et exécutera en fin de compte.

Le framework remplace les sous-graphes d'origine par des opérations personnalisées qui appellent le pilote matériel, en transmettant le bytecode compilé créé par le plug-in.

LiteRT Dispatch est l'équivalent de l'environnement d'exécution pour le plug-in du compilateur. Ils permettent d'appeler la HAL en fonction de la sortie du compilateur. Pour en savoir plus, consultez la documentation sur Dispatch.

AOT ou sur l'appareil

LiteRT peut utiliser des plug-ins de compilation pour prendre en charge la compilation AOT via nos outils, ainsi que la compilation sur l'appareil. La compilation sur l'appareil est plus flexible, entièrement internalisée dans les API d'exécution LiteRT et ne nécessite que la gestion d'un seul modèle. Le flux AOT peut débloquer la compilation lorsqu'elle est trop gourmande en ressources pour s'exécuter sur l'appareil, ce qui peut être le cas avec de nombreux grands modèles contemporains.

Action de remplacement

LiteRT est conçu pour prendre en charge les graphiques hétérogènes. Toute opération non sélectionnée par le plug-in sera laissée au processeur ou mise à disposition pour l'accélération sur un autre backend.

Implémenter un plug-in de compilation

Un plug-in de compilateur LiteRT est implémenté en tant que bibliothèque partagée qui exporte un ensemble spécifique de fonctions C définies dans l'API C LiteRT.

Fonctions d'interface essentielles

La fonctionnalité principale repose sur deux étapes de compilation clés : LiteRtCompilerPluginPartition et LiteRtCompilerPluginCompile.

Extraits de code de l'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. Fonction de partition

La signature de la fonction est la suivante :

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

Fonctionnement de la fonction partition : il s'agit de la phase de sélection. Le plug-in itère sur les opérations dans l'entrée LiteRtSubgraph. Pour chaque opération que le matériel cible prend en charge et peut accélérer, le plug-in ajoute cette opération à la LiteRtOpList$ fournie dans le paramètre selected_ops. Le framework LiteRt utilise cette liste pour définir les limites des partitions qui seront envoyées pour l'étape de compilation finale.

Par défaut, LiteRT regroupe toutes les opérations sélectionnées dans les sous-DAG les plus grands possibles. Pour un partitionnement plus précis, un index peut être associé lors de la sélection des opérations, ce qui permet de décomposer davantage ces sous-graphes.

2. La fonction "compile"

La signature de la fonction est la suivante :

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

Fonction de compile : il s'agit de la phase de génération. L'entrée partitions représente un modèle dans lequel tous les sous-graphes sélectionnés ont été isolés. Le plug-in traite ces partitions en invoquant sa chaîne d'outils spécifique pour générer le bytecode pour le matériel cible. La sortie du plug-in doit fournir un point d'entrée pour chaque sous-graphe transmis pour la compilation. Dans la plupart des cas, il s'agit de modules de bytecode individuels pour chaque sous-graphe d'entrée ou d'un module de bytecode unique avec plusieurs points d'entrée.

Type de données renvoyé par compile : la fonction LiteRtCompilerPluginCompile renvoie sa sortie à l'aide du paramètre out LiteRtCompiledResult.

LiteRtCompiledResult est un handle opaque (par rapport à LiteRT) vers une structure gérée par le plug-in. Il représente le résultat de la compilation et contient deux informations principales :

1. Modules de code octet : un ou plusieurs tampons de mémoire brute contenant le bytecode exécutable spécifique au matériel (c'est-à-dire les instructions compilées).
2. Informations sur l'appel : métadonnées pour chaque partition. Cela fournit le mappage du sous-graphe d'entrée i à un module de code octet de résultat et à un identifiant de point d'entrée dans ce module.

Exemple de configuration

Les extraits suivants illustrent comment un plug-in de base peut implémenter les fonctions principales. Cet exemple est tiré d'un exemple entièrement fonctionnel dans litert/vendors/examples/.

Identification et configuration des plug-ins

Ces fonctions fournissent au framework des informations de base sur le plug-in et le matériel.

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

Logique de partitionnement (LiteRtCompilerPluginPartition)

Cet exemple montre que le plug-in ne sélectionne qu'un ensemble limité d'opérations (mul, sub et une opération composite spécifique) uniquement si toutes les entrées et sorties sont des floats 32 bits. En général, pour déterminer si une opération doit être sélectionnée ou non, un appel à un crochet de validation est inclus dans la chaîne d'outils du compilateur du 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;
}

Avant d'appeler la compilation, LiteRT validera et "décrira" toutes les opérations sélectionnées dans de nouveaux sous-graphiques d'un nouveau modèle intermédiaire. Ce modèle intermédiaire est celui qui est transmis à la compilation.

Logique de compilation (LiteRtCompilerPluginCompile)

Cette fonction prend les sous-graphiques partitionnés et génère un LiteRtCompiledResult personnalisé. Cet exemple génère un module de bytecode autonome pour chaque partition à compiler. Dans les cas réels, cela implique généralement de convertir les opérations LiteRT en types dans la bibliothèque du compilateur backend. La "compilation " du plug-in d'exemple fonctionnel crée une chaîne lisible qui encode le graphique.

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

Utilisation et validation

LiteRT fournit divers outils pour appliquer des plug-ins de compilateur aux fichiers de modèle, exécuter le résultat et effectuer la validation/l'évaluation. Consultez la documentation de la suite de tests des accélérateurs et la documentation sur les benchmarks et le profilage.