Wtyczki kompilatora LiteRT

Kiedy należy utworzyć wtyczkę kompilatora?

Wtyczka kompilatora LiteRT jest niezbędna, gdy chcesz zintegrować konkretny akcelerator sprzętowy z zależnością kompilatora z platformą LiteRT.

Wtyczkę kompilatora należy utworzyć, jeśli:

1. Kierujesz reklamy na nowy backend sprzętowy, który nie jest obsługiwany.
2. Chcesz przenieść określone operacje modelu na akcelerator sprzętowy, aby zwiększyć wydajność lub efektywność energetyczną.
3. Wymagasz obsługi kompilacji AOT (na stacji roboczej) lub kompilacji na urządzeniu.

Wtyczka działa jak pomost, pobierając części modelu uczenia maszynowego i konwertując je na format, który może być wykonywany przez docelowy sprzęt, za pomocą wywołania kompilatora backendu. LiteRT łączy niestandardowy kod bajtowy wygenerowany przez wtyczkę z modelem .tflite, dzięki czemu można go wykonywać za pomocą środowiska wykonawczego LiteRT.

Jak działają wtyczki kompilatora?

Framework LiteRT używa wtyczki kompilatora podczas wczytywania modelu lub w fazie wstępnego przetwarzania offline, aby identyfikować i przygotowywać podgrafy modelu do wykonania na docelowym sprzęcie.

Proces ten obejmuje 2 główne etapy koordynowane przez platformę za pomocą funkcji eksportowanych przez wtyczkę:

1. Partycjonowanie: wtyczka analizuje cały wykres modelu i identyfikuje podzbiory operacji, które obsługuje i które może skutecznie przyspieszyć na docelowym sprzęcie. Te obsługiwane podgrafy są „dzielone” (oznaczane) na potrzeby kompilacji i wyświetlane w postaci konturów.
2. Kompilacja: platforma LiteRT przekazuje podzielone podgrafy z powrotem do wtyczki. Wtyczka używa następnie swojej logiki wewnętrznej i ewentualnie zewnętrznych łańcuchów narzędzi (kompilatorów) do wygenerowania co najmniej jednego modułu kodu bajtowego specyficznego dla sprzętu, który implementuje partycje. Ten kod bajtowy zostanie ostatecznie wczytany i wykonany przez środowisko wykonawcze (HAL/sterownik) docelowego sprzętu.

Platforma zastępuje oryginalne podgrafy niestandardowymi operacjami, które wywołują sterownik sprzętu, przekazując skompilowany kod bajtowy utworzony przez wtyczkę.

LiteRT Dispatch to odpowiednik w czasie działania wtyczki kompilatora. Umożliwiają one wywoływanie HAL na podstawie danych wyjściowych kompilatora. Więcej informacji znajdziesz w dokumentacji dotyczącej wysyłania.

AOT a kompilacja na urządzeniu

LiteRT może używać wtyczek kompilatora do obsługi kompilacji AOT za pomocą naszych narzędzi, a także kompilacji na urządzeniu. Kompilacja na urządzeniu jest bardziej elastyczna, w pełni zintegrowana z interfejsami API środowiska wykonawczego LiteRT i wymaga tylko zarządzania jednym modelem. Proces AOT może odblokować kompilację, gdy jest ona zbyt zasobochłonna, aby można ją było uruchomić na urządzeniu. Może to mieć miejsce w przypadku wielu współczesnych dużych modeli.

Wycofanie

LiteRT obsługuje wykresy heterogeniczne. Każda operacja, która nie została wybrana przez wtyczkę, zostanie przekazana do procesora lub udostępniona do przyspieszenia na innym backendzie.

Implementowanie wtyczki kompilatora

Wtyczka kompilatora LiteRT jest implementowana jako biblioteka udostępniona, która eksportuje określony zestaw funkcji C zdefiniowanych w interfejsie LiteRT C API.

Podstawowe funkcje interfejsu

Główna funkcjonalność opiera się na 2 kluczowych etapach kompilacji: LiteRtCompilerPluginPartition i LiteRtCompilerPluginCompile.

Fragmenty kodu w C API

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

Podpis funkcji to:

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

Działanie funkcji partition: to faza wyboru. Wtyczka iteruje po operacjach w danych wejściowych LiteRtSubgraph. W przypadku każdej operacji, którą obsługuje i może przyspieszyć sprzęt docelowy, wtyczka dodaje tę operację do listy LiteRtOpList$ podanej w parametrze selected_ops. Platforma LiteRT używa tej listy do określania granic partycji, które zostaną wysłane w ostatnim kroku kompilacji.

Domyślnie LiteRT grupuje wszystkie wybrane operacje w największe możliwe podgrafy DAG. Aby uzyskać bardziej szczegółowy podział, podczas wybierania operacji można powiązać indeks, który dodatkowo podzieli te podgrafy.

2. Funkcja kompilacji

Podpis funkcji to:

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

Działanie funkcji compile: to faza generowania. Dane wejściowe partitions reprezentują model, w którym wszystkie wybrane podgrafy zostały odizolowane. Wtyczka przetwarza te partycje, wywołując swój specyficzny łańcuch narzędzi w celu wygenerowania kodu bajtowego dla docelowego sprzętu. Oczekuje się, że dane wyjściowe wtyczki będą zawierać punkt wejścia dla każdego podgrafu przekazanego do kompilacji. W większości przypadków są to pojedyncze moduły kodu bajtowego dla każdego podgrafu wejściowego lub pojedynczy moduł kodu bajtowego z wieloma punktami wejścia.

Typ danych zwracanych przez funkcję compile: funkcja LiteRtCompilerPluginCompile zwraca dane wyjściowe za pomocą parametru wyjściowego LiteRtCompiledResult.

LiteRtCompiledResult to nieprzezroczysty (w odniesieniu do LiteRT) uchwyt do struktury zarządzanej przez wtyczkę. Reprezentuje wynik kompilacji i zawiera 2 główne informacje:

1. Moduły kodu bajtowego: co najmniej jeden surowy bufor pamięci zawierający kod bajtowy wykonywalny na konkretnym sprzęcie (czyli skompilowane instrukcje).
2. Informacje o połączeniach: metadane dotyczące każdej partycji. Zawiera mapowanie od i-tego podgrafu wejściowego do modułu kodu bajtowego wyniku i identyfikatora punktu wejścia w tym module.

Przykładowa implementacja

Poniższe fragmenty kodu pokazują, jak podstawowa wtyczka może implementować podstawowe funkcje. Ten przykład pochodzi z w pełni funkcjonalnego przykładu w litert/vendors/examples/

Identyfikacja i konfiguracja wtyczki

Funkcje te dostarczają platformie podstawowych informacji o wtyczce i sprzęcie.

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

Logika partycjonowania (LiteRtCompilerPluginPartition)

Ten przykład pokazuje, że wtyczka wybiera ograniczony zestaw operacji (mul,sub i określoną operację złożoną) tylko wtedy, gdy wszystkie dane wejściowe i wyjściowe są 32-bitowymi liczbami zmiennoprzecinkowymi. Zwykle określenie, czy operacja powinna zostać wybrana, obejmuje wywołanie haka weryfikacyjnego w łańcuchu narzędzi kompilatora backendu.

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

Przed kompilacją wywołania LiteRT zweryfikuje i „zarysuje” wszystkie wybrane operacje w nowych podgrafach w nowym modelu pośrednim. Ten model pośredni jest przekazywany do kompilacji.

Logika kompilacji (LiteRtCompilerPluginCompile)

Ta funkcja pobiera podzielone podgrafy i generuje niestandardowyLiteRtCompiledResult. Ten przykład generuje samodzielny moduł kodu bajtowego dla każdej partycji do skompilowania. W rzeczywistych przypadkach zwykle wiąże się to z przekształcaniem operacji LiteRT na typy w bibliotece kompilatora backendu. Funkcjonalny przykład wtyczki „kompilacja” tworzy zrozumiały dla człowieka ciąg tekstowy, który koduje wykres.

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

Użycie i weryfikacja

LiteRT udostępnia różne narzędzia do stosowania wtyczek kompilatora do plików modeli, wykonywania wyników oraz weryfikowania i testowania wydajności. Zapoznaj się z dokumentacją pakietu testów akceleratora i dokumentacją dotyczącą testów porównawczych i profilowania.