LiteRT-Compiler-Plug-ins

Wann sollte ich ein Compiler-Plug-in erstellen?

Ein LiteRT-Compiler-Plug-in ist erforderlich, wenn Sie einen bestimmten Hardwarebeschleuniger mit einer Compiler-Abhängigkeit in das LiteRT-Framework einbinden müssen.

Sie sollten ein Compiler-Plug-in erstellen, wenn:

1. Sie haben ein neues Hardware-Backend als Ziel festgelegt, das nicht unterstützt wird.
2. Sie möchten bestimmte Modellvorgänge zur Leistungssteigerung oder zur Verbesserung der Energieeffizienz auf diesen Hardwarebeschleuniger auslagern.
3. Sie benötigen Unterstützung für die AOT-Kompilierung (auf der Workstation) oder die On-Device-Kompilierung.

Das Plug-in fungiert als Brücke, indem es Teile des Machine-Learning-Modells übernimmt und sie in ein Format konvertiert, das von der Zielhardware ausgeführt werden kann. Dazu wird der Compiler des Back-Ends aufgerufen. LiteRT bündelt den benutzerdefinierten Bytecode, der vom Plug-in generiert wird, im .tflite-Modell, sodass er mit der LiteRT-Laufzeit ausgeführt werden kann.

Wie funktionieren Compiler-Plug-ins?

Das LiteRT-Framework verwendet das Compiler-Plug-in während des Ladens des Modells oder der Offline-Vorverarbeitungsphase, um Modell-Untergraphen für die Ausführung auf der Zielhardware zu identifizieren und vorzubereiten.

Der Prozess umfasst zwei Hauptphasen, die vom Framework mithilfe der exportierten Funktionen des Plug-ins orchestriert werden:

1. Partitionierung:Das Plug-in untersucht den gesamten Modellgraphen und identifiziert Teilmengen von Vorgängen, die es unterstützt und die auf der Zielhardware effizient beschleunigt werden können. Diese unterstützten Teilgraphen werden für die Kompilierung „partitioniert“ (markiert) und umrissen.
2. Kompilierung:Das LiteRT-Framework gibt die partitionierten Untergraphen an das Plug-in zurück. Das Plug-in verwendet dann seine interne Logik und möglicherweise externe Toolchains (Compiler), um ein oder mehrere hardwarespezifische Bytecode-Module zu generieren, die die Partitionen implementieren. Dieser Bytecode wird von der Laufzeit (HAL/Treiber) der Zielhardware geladen und ausgeführt.

Das Framework ersetzt die ursprünglichen Teilgraphen durch benutzerdefinierte Vorgänge, die den Hardwaretreiber aufrufen und den vom Plug-in erstellten kompilierten Bytecode übergeben.

LiteRT Dispatch ist das Laufzeit-Analogon für das Compiler-Plug-in. Sie bieten die Möglichkeit, die HAL anhand der Compiler-Ausgabe aufzurufen. Weitere Informationen finden Sie in der Dokumentation zum Dispatch.

AOT im Vergleich zu On-Device

LiteRT kann Compiler-Plug-ins verwenden, um die AOT-Kompilierung über unsere Tools sowie die On-Device-Kompilierung zu unterstützen. Die On-Device-Kompilierung ist flexibler, vollständig in die LiteRT-Laufzeit-APIs integriert und erfordert nur die Verwaltung eines einzelnen Modells. Der AOT-Ablauf kann die Kompilierung entsperren, wenn sie zu ressourcenintensiv ist, um auf dem Gerät ausgeführt zu werden. Das kann bei vielen modernen großen Modellen der Fall sein.

Fallback

LiteRT wurde mit Unterstützung für heterogene Grafiken entwickelt. Alle Vorgänge, die nicht vom Plug-in ausgewählt werden, werden von der CPU ausgeführt oder für die Beschleunigung auf einem anderen Backend zur Verfügung gestellt.

Compiler-Plug-in implementieren

Ein LiteRT-Compiler-Plug-in wird als gemeinsam genutzte Bibliothek implementiert, die eine bestimmte Gruppe von C-Funktionen exportiert, die in der LiteRT C API definiert sind.

Wichtige Schnittstellenfunktionen

Die Kernfunktionen umfassen zwei wichtige Kompilierungsschritte: LiteRtCompilerPluginPartition und LiteRtCompilerPluginCompile.

C-API-Snippets

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

Die Funktionssignatur lautet:

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

Funktionsweise der partition-Funktion:Dies ist die Auswahlphase. Das Plug-in durchläuft die Vorgänge in der Eingabe LiteRtSubgraph. Für jeden Vorgang, den die Zielhardware unterstützt und beschleunigen kann, fügt das Plug-in diesen Vorgang der im Parameter selected_ops bereitgestellten LiteRtOpList$hinzu. Das LiteRt-Framework verwendet diese Liste, um die Grenzen der Partitionen zu definieren, die für den letzten Kompilierungsschritt gesendet werden.

Standardmäßig gruppiert LiteRT alle ausgewählten Vorgänge in den größtmöglichen untergeordneten DAGs. Für eine detailliertere Partitionierung kann beim Auswählen von Vorgängen ein Index zugeordnet werden, um diese Untergraphen weiter aufzuteilen.

2. Die Funktion „Kompilieren“

Die Funktionssignatur lautet:

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

Funktionsweise der compile-Funktion:Dies ist die Generierungsphase. Die Eingabe partitions stellt ein Modell dar, in dem alle ausgewählten Teilgraphen isoliert wurden. Das Plug-in verarbeitet diese Partitionen und ruft die zugehörige Toolchain auf, um den Bytecode für die Zielhardware zu generieren. Es wird erwartet, dass die Ausgabe des Plug-ins einen Einstiegspunkt für jeden Untergraphen bietet, der für die Kompilierung übergeben wird. In den meisten Fällen sind das entweder einzelne Bytecode-Module für jeden Eingabe-Subgraph oder ein einzelnes Bytecode-Modul mit mehreren Einstiegspunkten.

Typ der von compile zurückgegebenen Daten:Die Funktion LiteRtCompilerPluginCompile gibt ihre Ausgabe über den Out-Parameter LiteRtCompiledResult zurück.

LiteRtCompiledResult ist ein (in Bezug auf LiteRT) undurchsichtiges Handle für eine vom Plug-in verwaltete Struktur. Sie stellt die Ausgabe der Kompilierung dar und enthält zwei wichtige Informationen:

1. Bytecode-Module:Ein oder mehrere Rohspeicherpuffer, die den hardwarespezifischen ausführbaren Bytecode (d.h. kompilierte Anweisungen) enthalten.
2. Anrufinformationen:Metadaten für jede Partition. Dies stellt die Zuordnung vom i-ten Eingabe-Subgraph zu einem Ergebnis-Bytecode-Modul und einem Einstiegspunkt-Identifier in diesem Modul bereit.

Beispielimplementierung

Die folgenden Snippets veranschaulichen, wie ein einfaches Plug-in die Kernfunktionen implementieren kann. Dieses Beispiel stammt aus einem voll funktionsfähigen Beispiel in litert/vendors/examples/.

Plug‑in-Identifizierung und ‑Einrichtung

Diese Funktionen liefern dem Framework grundlegende Informationen zum Plug-in und zur 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;
}

Partitionierungslogik (LiteRtCompilerPluginPartition)

In diesem Beispiel wählt das Plug-in nur dann eine begrenzte Anzahl von Vorgängen (mul, sub und einen bestimmten zusammengesetzten Vorgang) aus, wenn alle Ein- und Ausgaben 32-Bit-Gleitkommazahlen sind. Normalerweise wird bei der Entscheidung, ob ein Vorgang ausgewählt werden soll, ein Validierungshook in der Compiler-Toolchain des Backends aufgerufen.

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

Vor der Kompilierung ruft LiteRT alle ausgewählten Vorgänge ab und fasst sie in einem neuen Zwischenmodell in neuen Teilgraphen zusammen. Dieses Zwischenmodell wird an die Kompilierung übergeben.

Zusammenstellungslogik (LiteRtCompilerPluginCompile)

Diese Funktion verwendet die partitionierten Untergraphen und generiert einen benutzerdefinierten LiteRtCompiledResult. In diesem Beispiel wird für jede zu kompilierende Partition ein eigenständiges Bytecode-Modul generiert. In der Praxis geht es dabei in der Regel darum, LiteRT-Vorgänge in Typen für die Backend-Compiler-Bibliothek zu konvertieren. Bei der „Kompilierung“ des funktionalen Beispiel-Plug-ins wird ein für Menschen lesbarer String erstellt, der den Graphen codiert.

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

Nutzung und Validierung

LiteRT bietet verschiedene Tools zum Anwenden von Compiler-Plug-ins auf Modelldateien, zum Ausführen des Ergebnisses und zum Validieren/Benchmarking. Weitere Informationen finden Sie in der Dokumentation zur Accelerator-Testsuite und in der Dokumentation zu Benchmarking und Profilerstellung.