افزونه‌های کامپایلر LiteRT

چه زمانی باید یک افزونه کامپایلر ایجاد کنم؟

افزونه کامپایلر LiteRT زمانی ضروری است که نیاز به ادغام یک شتاب‌دهنده سخت‌افزاری خاص با وابستگی به کامپایلر در چارچوب LiteRT داشته باشید.

شما باید یک افزونه کامپایلر ایجاد کنید اگر:

1. شما در حال هدف قرار دادن یک سخت‌افزار جدید هستید که پشتیبانی نمی‌شود.
2. شما می‌خواهید عملیات مدل خاصی را برای بهبود عملکرد یا بهره‌وری انرژی به آن شتاب‌دهنده سخت‌افزاری واگذار کنید .
3. شما به پشتیبانی از کامپایل AOT (روی ایستگاه کاری) یا کامپایل روی دستگاه نیاز دارید.

این افزونه به عنوان یک پل عمل می‌کند، بخش‌هایی از مدل یادگیری ماشین را می‌گیرد و آنها را با استفاده از فراخوانی کامپایلر backend به فرمتی تبدیل می‌کند که سخت‌افزار هدف شما بتواند آن را اجرا کند. LiteRT بایت‌کد سفارشی تولید شده توسط افزونه را در مدل .tflite قرار می‌دهد و آن را با استفاده از زمان اجرای LiteRT قابل اجرا می‌کند.

افزونه‌های کامپایلر چگونه کار می‌کنند؟

چارچوب LiteRT از افزونه کامپایلر در طول بارگذاری مدل یا مرحله پیش‌پردازش آفلاین برای شناسایی و آماده‌سازی زیرگراف‌های مدل برای اجرا روی سخت‌افزار هدف استفاده می‌کند.

این فرآیند شامل دو مرحله اصلی است که توسط چارچوب با استفاده از توابع خروجی افزونه تنظیم می‌شوند:

1. پارتیشن‌بندی: این افزونه کل گراف مدل را بررسی می‌کند و زیرمجموعه‌هایی از عملیات را که پشتیبانی می‌کند و می‌تواند به طور موثر روی سخت‌افزار هدف تسریع کند، شناسایی می‌کند. این زیرگراف‌های پشتیبانی‌شده برای کامپایل «پارتیشن‌بندی» (علامت‌گذاری) شده و طرح کلی آنها مشخص می‌شود.
2. کامپایل: چارچوب LiteRT زیرگراف‌های پارتیشن‌بندی شده را به افزونه بازمی‌گرداند. سپس افزونه از منطق داخلی خود و احتمالاً زنجیره ابزارهای خارجی (کامپایلر) برای تولید یک یا چند ماژول بایت‌کد مخصوص سخت‌افزار که پارتیشن‌ها را پیاده‌سازی می‌کنند، استفاده می‌کند. این بایت‌کد همان چیزی است که در نهایت توسط زمان اجرای سخت‌افزار هدف (HAL/درایور) بارگذاری و اجرا خواهد شد.

این چارچوب، زیرگراف‌های اصلی را با عملیات سفارشی جایگزین می‌کند که درایور سخت‌افزار را فراخوانی می‌کنند و بایت‌کد کامپایل‌شده‌ی ایجاد شده توسط افزونه را ارسال می‌کنند.

LiteRT Dispatch معادل زمان اجرا برای افزونه کامپایلر است. آنها امکان فراخوانی خروجی کامپایلر داده شده به HAL را فراهم می‌کنند. برای جزئیات بیشتر، به مستندات dispatch مراجعه کنید.

AOT در مقابل On-Device

LiteRT می‌تواند از افزونه‌های کامپایلر برای پشتیبانی از کامپایل AOT از طریق ابزارهای ما و همچنین کامپایل روی دستگاه استفاده کند. کامپایل روی دستگاه انعطاف‌پذیرتر است، کاملاً در APIهای زمان اجرای LiteRT داخلی‌سازی شده است و فقط به مدیریت یک مدل واحد نیاز دارد. جریان AOT می‌تواند زمانی که اجرای آن روی دستگاه بسیار فشرده از نظر منابع است، که ممکن است در مورد بسیاری از مدل‌های بزرگ امروزی صدق کند، کامپایل را از حالت انسداد خارج کند.

بازگشت به عقب

LiteRT با پشتیبانی از گراف‌های ناهمگن ساخته شده است. هر عملیاتی که توسط افزونه انتخاب نشود، به CPU واگذار می‌شود یا برای شتاب‌دهی در یک backend دیگر در دسترس قرار می‌گیرد.

پیاده‌سازی یک افزونه کامپایلر

یک افزونه کامپایلر LiteRT به عنوان یک کتابخانه مشترک پیاده‌سازی شده است که مجموعه‌ای خاص از توابع C تعریف شده در LiteRT C API را صادر می‌کند.

توابع ضروری رابط کاربری

عملکرد اصلی حول دو مرحله کلیدی کامپایل می‌چرخد: LiteRtCompilerPluginPartition و LiteRtCompilerPluginCompile .

قطعه کدهای 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);

۱. تابع پارتیشن

امضای تابع عبارت است از:

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

کاری که تابع partition انجام می‌دهد: این مرحله انتخاب است. افزونه روی عملیات موجود در LiteRtSubgraph ورودی تکرار می‌کند. برای هر عملیاتی که سخت‌افزار هدف پشتیبانی می‌کند و می‌تواند آن را تسریع کند، افزونه آن عملیات را به LiteRtOpList$ ارائه شده در پارامتر selected_ops اضافه می‌کند. چارچوب LiteRt از این لیست برای تعریف مرزهای پارتیشن‌هایی که برای مرحله کامپایل نهایی ارسال می‌شوند، استفاده می‌کند.

به طور پیش‌فرض، LiteRT تمام عملیات‌های انتخاب شده را در بزرگترین زیرگراف‌های ممکن گروه‌بندی می‌کند. برای تقسیم‌بندی دقیق‌تر، می‌توان هنگام انتخاب عملیات‌ها، یک شاخص مرتبط کرد که به تجزیه بیشتر این زیرگراف‌ها کمک می‌کند.

۲. تابع کامپایل

امضای تابع عبارت است از:

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

کاری که تابع compile انجام می‌دهد: این مرحله تولید است. partitions ورودی مدلی را نشان می‌دهند که در آن تمام زیرگراف‌های انتخاب‌شده ایزوله شده‌اند. افزونه این پارتیشن‌ها را پردازش می‌کند و بایت‌کد مخصوص خود را برای سخت‌افزار هدف تولید می‌کند. انتظار می‌رود خروجی افزونه یک نقطه ورودی برای هر زیرگراف ارسالی برای کامپایل فراهم کند. در بیشتر موارد، این یا ماژول‌های کد بایت جداگانه برای هر زیرگراف ورودی است، یا یک ماژول کد بایت واحد با چندین نقطه ورودی.

نوع داده‌ای که توسط compile برگردانده می‌شود: تابع LiteRtCompilerPluginCompile خروجی خود را با استفاده از پارامتر خروجی LiteRtCompiledResult برمی‌گرداند.

LiteRtCompiledResult یک هندل مات (نسبت به LiteRT) برای ساختاری است که توسط افزونه مدیریت می‌شود. این هندل خروجی کامپایل را نشان می‌دهد و شامل دو بخش اصلی اطلاعات است:

1. ماژول‌های کد بایت: یک یا چند بافر حافظه خام که حاوی کد بایت اجرایی مخصوص سخت‌افزار (یعنی دستورالعمل‌های کامپایل شده) هستند.
2. اطلاعات فراخوانی: فراداده برای هر پارتیشن. این، نگاشت از زیرگراف ورودی i ام به یک ماژول کد بایت نتیجه و شناسه نقطه ورود به آن ماژول را فراهم می‌کند.

پیاده‌سازی مثال

قطعه کدهای زیر نشان می‌دهند که چگونه یک افزونه‌ی ساده می‌تواند توابع اصلی را پیاده‌سازی کند. این مثال از یک مثال کاملاً کاربردی در litert/vendors/examples/ گرفته شده است.

شناسایی و تنظیم افزونه

این توابع اطلاعات اولیه‌ای در مورد افزونه و سخت‌افزار در اختیار فریم‌ورک قرار می‌دهند.

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

منطق پارتیشن‌بندی ( LiteRtCompilerPluginPartition )

این مثال نشان می‌دهد که افزونه فقط در صورتی که همه ورودی‌ها و خروجی‌ها اعشاری ۳۲ بیتی باشند، مجموعه محدودی از عملیات ( mul ، sub و یک عملیات مرکب خاص) را انتخاب می‌کند. معمولاً تعیین اینکه آیا یک عملیات باید انتخاب شود یا خیر، شامل فراخوانی یک قلاب اعتبارسنجی در زنجیره ابزار کامپایلر 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;
}

قبل از فراخوانی کامپایل، LiteRT تمام عملیات انتخاب شده را در یک مدل میانی جدید، اعتبارسنجی و در زیرگراف‌های جدید "طرح کلی" می‌کند. این مدل میانی همان چیزی است که به کامپایل منتقل می‌شود.

منطق کامپایل ( LiteRtCompilerPluginCompile )

این تابع زیرگراف‌های پارتیشن‌بندی شده را می‌گیرد و یک LiteRtCompiledResult سفارشی تولید می‌کند. این مثال برای هر پارتیشنی که قرار است کامپایل شود، یک ماژول بایت‌کد مستقل تولید می‌کند. در موارد واقعی، این معمولاً شامل تبدیل عملیات LiteRT به انواعی برای کتابخانه کامپایلر backend است. "کامپایل" افزونه مثال تابعی، یک رشته قابل خواندن توسط انسان ایجاد می‌کند که گراف را کدگذاری می‌کند.

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

استفاده و اعتبارسنجی

LiteRT ابزارهای مختلفی را برای اعمال افزونه‌های کامپایلر به فایل‌های مدل، اجرای نتیجه و اعتبارسنجی/محاسبه ارائه می‌دهد. به مستندات مجموعه تست شتاب‌دهنده و مستندات محک‌زنی و پروفایل‌سازی مراجعه کنید.