Trình bổ trợ trình biên dịch LiteRT

Khi nào tôi nên tạo một trình bổ trợ trình biên dịch?

Bạn cần có Trình bổ trợ trình biên dịch LiteRT khi cần tích hợp một bộ tăng tốc phần cứng cụ thể với một phần phụ thuộc trình biên dịch vào khung LiteRT.

Bạn nên tạo một trình bổ trợ trình biên dịch nếu:

  1. Bạn đang nhắm đến một phần cứng phụ trợ mới không được hỗ trợ.
  2. Bạn muốn giảm tải các thao tác mô hình cụ thể cho bộ tăng tốc phần cứng đó để cải thiện hiệu suất hoặc hiệu quả sử dụng năng lượng.
  3. Bạn cần hỗ trợ biên dịch AOT (trên máy trạm) hoặc biên dịch trên thiết bị.

Trình bổ trợ này đóng vai trò là cầu nối, lấy các phần của mô hình học máy và chuyển đổi chúng thành một định dạng mà phần cứng mục tiêu của bạn có thể thực thi, bằng cách sử dụng lệnh gọi đến trình biên dịch của phần phụ trợ. LiteRT sẽ gói mã byte tuỳ chỉnh do trình bổ trợ tạo vào mô hình .tflite, giúp mô hình này có thể thực thi bằng thời gian chạy LiteRT.

Trình bổ trợ trình biên dịch hoạt động như thế nào?

Khung LiteRT sử dụng trình bổ trợ trình biên dịch trong giai đoạn tải mô hình hoặc xử lý trước ngoại tuyến để xác định và chuẩn bị các đồ thị con của mô hình để thực thi trên phần cứng mục tiêu.

Quy trình này bao gồm 2 giai đoạn chính do khung điều phối bằng cách sử dụng các hàm đã xuất của trình bổ trợ:

  1. Phân vùng: Trình bổ trợ này kiểm tra toàn bộ biểu đồ mô hình và xác định các tập hợp con của những thao tác mà trình bổ trợ hỗ trợ và có thể tăng tốc hiệu quả trên phần cứng mục tiêu. Các đồ thị con được hỗ trợ này được "phân vùng" (đánh dấu) để biên dịch và phác thảo.
  2. Biên dịch: Khung LiteRT chuyển các đồ thị con được phân vùng trở lại trình bổ trợ. Sau đó, trình bổ trợ sẽ sử dụng logic nội bộ và có thể là các chuỗi công cụ (trình biên dịch) bên ngoài để tạo một hoặc nhiều mô-đun mã byte dành riêng cho phần cứng triển khai các phân vùng. Bytecode này là những gì mà thời gian chạy (HAL/trình điều khiển) của phần cứng đích sẽ tải và thực thi.

Khung này sẽ thay thế các đồ thị con ban đầu bằng các thao tác tuỳ chỉnh gọi trình điều khiển phần cứng, truyền cùng với mã byte đã biên dịch do trình bổ trợ tạo.

LiteRT Dispatch là phiên bản tương tự thời gian chạy cho trình bổ trợ trình biên dịch. Chúng cung cấp phương tiện gọi vào HAL dựa trên đầu ra của trình biên dịch. Để biết thêm thông tin chi tiết, hãy tham khảo tài liệu về việc gửi.

AOT so với Trên thiết bị

LiteRT có thể sử dụng các trình bổ trợ trình biên dịch để hỗ trợ quá trình biên dịch AOT thông qua công cụ của chúng tôi, cũng như quá trình biên dịch trên thiết bị. Quá trình biên dịch trên thiết bị linh hoạt hơn, được nội bộ hoá hoàn toàn trong API thời gian chạy LiteRT và chỉ yêu cầu quản lý một mô hình duy nhất. Quy trình AOT có thể bỏ chặn quá trình biên dịch khi quá tốn nhiều tài nguyên để chạy trên thiết bị. Đây có thể là trường hợp của nhiều mô hình lớn hiện đại.

Dự phòng

LiteRT được xây dựng với khả năng hỗ trợ các biểu đồ không đồng nhất. Mọi thao tác không được chọn bởi trình bổ trợ sẽ được chuyển sang CPU hoặc được cung cấp để tăng tốc trên một chương trình phụ trợ khác.

Triển khai trình bổ trợ biên dịch

Trình bổ trợ trình biên dịch LiteRT được triển khai dưới dạng một thư viện dùng chung xuất một tập hợp cụ thể các hàm C được xác định trong API C của LiteRT.

Các hàm giao diện thiết yếu

Chức năng cốt lõi xoay quanh 2 bước biên dịch chính: LiteRtCompilerPluginPartition và LiteRtCompilerPluginCompile.

Chức năng Mục đích
LiteRtCompilerPluginPartition Chọn và đánh dấu tất cả các thao tác được hỗ trợ trong một đồ thị con mô hình nhất định (bước Phân vùng).
LiteRtCompilerPluginCompile$ Tạo mã byte dành riêng cho phần cứng cho các phân vùng được chọn trước (bước Biên dịch).

Đoạn mã 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. Hàm phân vùng

Chữ ký của hàm là:

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

Chức năng của partition: Đây là giai đoạn lựa chọn. Trình bổ trợ này lặp lại các thao tác trong LiteRtSubgraph đầu vào. Đối với mọi thao tác mà phần cứng mục tiêu hỗ trợ và có thể tăng tốc, trình bổ trợ sẽ thêm thao tác đó vào LiteRtOpList$ được cung cấp trong tham số selected_ops. Khung LiteRt sử dụng danh sách này để xác định ranh giới của các phân vùng sẽ được gửi cho bước biên dịch cuối cùng.

Theo mặc định, LiteRT sẽ nhóm tất cả các hoạt động đã chọn thành các DAG con lớn nhất có thể. Để phân vùng chi tiết hơn, bạn có thể liên kết một chỉ mục khi chọn các thao tác nhằm chia nhỏ thêm các đồ thị con này.

2. Hàm biên dịch

Chữ ký của hàm là:

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

Chức năng của hàm compile: Đây là giai đoạn tạo. Đầu vào partitions biểu thị một mô hình trong đó tất cả các đồ thị con đã chọn đều được tách biệt. Trình bổ trợ này xử lý các phân vùng này, gọi chuỗi công cụ cụ thể để tạo mã byte cho phần cứng mục tiêu. Dự kiến đầu ra của trình bổ trợ sẽ cung cấp một điểm truy cập cho mỗi đồ thị con được truyền để biên dịch. Trong hầu hết các trường hợp, đây là các mô-đun mã byte riêng lẻ cho từng đồ thị con đầu vào hoặc một mô-đun mã byte duy nhất có nhiều điểm truy cập.

Loại dữ liệu do compile trả về: Hàm LiteRtCompilerPluginCompile trả về đầu ra bằng cách sử dụng tham số ngoài LiteRtCompiledResult.

LiteRtCompiledResult là một giá trị nhận dạng không công khai (đối với LiteRT) cho một cấu trúc do trình bổ trợ quản lý. Đây là đầu ra của quá trình biên dịch và chứa 2 phần thông tin chính:

  1. Các mô-đun mã byte: Một hoặc nhiều vùng đệm bộ nhớ thô chứa mã byte thực thi dành riêng cho phần cứng (tức là các chỉ dẫn đã biên dịch).
  2. Thông tin về lệnh gọi: Siêu dữ liệu cho từng phân vùng. Thao tác này cung cấp ánh xạ từ đồ thị con đầu vào thứ i đến một mô-đun mã byte kết quả và mã nhận dạng điểm truy cập vào mô-đun đó.

Ví dụ về cách triển khai

Các đoạn mã sau đây minh hoạ cách một trình bổ trợ cơ bản có thể triển khai các chức năng cốt lõi. Ví dụ này được lấy từ một ví dụ hoạt động đầy đủ trong litert/vendors/examples/

Nhận dạng và thiết lập trình bổ trợ

Các hàm này cung cấp cho khung thông tin cơ bản về trình bổ trợ và phần cứng.

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

Logic phân vùng (LiteRtCompilerPluginPartition)

Ví dụ này cho thấy trình bổ trợ chỉ chọn một số ít thao tác (mul, sub và một thao tác kết hợp cụ thể) nếu tất cả các đầu vào và đầu ra đều là số thực 32 bit. Thông thường, việc xác định xem có nên chọn một thao tác hay không sẽ bao gồm một lệnh gọi đến một hook xác thực trong chuỗi công cụ trình biên dịch của phần phụ trợ.

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

Trước khi gọi quy trình biên dịch, LiteRT sẽ xác thực và "vạch ra" tất cả các hoạt động đã chọn thành các đồ thị con mới trong một mô hình trung gian mới. Mô hình trung gian này là mô hình được truyền đến quá trình biên dịch.

Logic biên dịch (LiteRtCompilerPluginCompile)

Hàm này lấy các đồ thị con được phân vùng và tạo một LiteRtCompiledResult tuỳ chỉnh. Ví dụ này tạo một mô-đun mã byte độc lập cho từng phân vùng cần được biên dịch. Trong trường hợp thực tế, điều này thường liên quan đến việc chuyển đổi các thao tác LiteRT thành các loại cho thư viện trình biên dịch phụ trợ. "Quá trình biên dịch" của trình bổ trợ ví dụ về chức năng sẽ tạo ra một chuỗi ký tự mà con người đọc được để mã hoá biểu đồ.

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

Mức sử dụng và xác thực

LiteRT cung cấp nhiều công cụ để áp dụng các trình bổ trợ trình biên dịch cho tệp mô hình, thực thi kết quả và xác thực/đo điểm chuẩn. Tham khảo tài liệu về bộ kiểm thử trình tăng tốc và tài liệu về hoạt động đo điểm chuẩn và lập hồ sơ.