API Trình điều phối hoạt động tuỳ chỉnh LiteRT

Tổng quan

CustomOpDispatcher là API thay thế để xác định các thao tác CPU tuỳ chỉnh và trình phân giải thao tác tuỳ chỉnh trong LiteRT. Nó cung cấp một giao diện rõ ràng hơn để tích hợp các thao tác tuỳ chỉnh vào các mô hình LiteRT.

Tại sao nên dùng CustomOpDispatcher?

Phương pháp truyền thống (Không dùng nữa)

Trước đây, nhà phát triển phải:

  • Tạo cấu trúc TfLiteRegistration theo cách thủ công
  • Triển khai các hàm gọi lại dành riêng cho TFLite (init, prepare, invoke, free)
  • Trực tiếp làm việc với các cấu trúc TFLite cấp thấp (TfLiteContext, TfLiteNode, TfLiteTensor)
  • Sử dụng MutableOpResolver::AddCustom() trực tiếp

Phương pháp CustomOpDispatcher mới

CustomOpDispatcher cung cấp:

  • Lớp trừu tượng rõ ràng trên các thành phần bên trong của TFLite
  • Tích hợp với mô hình đã biên dịch của LiteRT

Flow

┌─────────────────┐
│  User Custom Op │ (Your implementation)
└────────┬────────┘
         │
         ▼
┌─────────────────────────┐
│   LiteRtCustomOpKernel  │ (C API interface)
└────────┬────────────────┘
         │
         ▼
┌─────────────────────────┐
│   CustomOpDispatcher    │ (Bridge layer)
└────────┬────────────────┘
         │
         ▼
┌─────────────────────────┐
│  TfLiteRegistration     │ (TFLite runtime)
└─────────────────────────┘

Các thành phần và tệp chính

Triển khai cốt lõi

Tiêu đề API

Hệ thống tuỳ chọn

  • core/options.h: Cấu trúc lựa chọn cốt lõi với CustomOpOption

Ví dụ về kiểm thử

Tài liệu tham khảo API

Giao diện nhân cốt lõi (C API)

typedef struct {
  LiteRtStatus (*Init)(void* user_data, const void* init_data,
                       size_t init_data_size);
  LiteRtStatus (*GetOutputLayouts)(void* user_data, size_t num_inputs,
                                   const LiteRtLayout* input_layouts,
                                   size_t num_outputs,
                                   LiteRtLayout* output_layouts);
  LiteRtStatus (*Run)(void* user_data, size_t num_inputs,
                      const LiteRtTensorBuffer* inputs, size_t num_outputs,
                      LiteRtTensorBuffer* outputs);
  LiteRtStatus (*Destroy)(void* user_data);
} LiteRtCustomOpKernel;

Lớp cơ sở trừu tượng C++

class CustomOpKernel {
public:
  virtual const std::string& OpName() const = 0;
  virtual int OpVersion() const = 0;
  virtual Expected<void> Init(const void* init_data, size_t init_data_size) = 0;
  virtual Expected<void> GetOutputLayouts(
      const std::vector<Layout>& input_layouts,
      std::vector<Layout>& output_layouts) = 0;
  virtual Expected<void> Run(const std::vector<TensorBuffer>& inputs,
                             std::vector<TensorBuffer>& outputs) = 0;
  virtual Expected<void> Destroy() = 0;
};

Hướng dẫn triển khai

Bước 1: Xác định thao tác tuỳ chỉnh

Triển khai C++

#include "litert/cc/litert_custom_op_kernel.h"

class MyCustomOpKernel : public litert::CustomOpKernel {
public:
  const std::string& OpName() const override {
    return op_name_;
  }

  int OpVersion() const override {
    return 1;
  }

  Expected<void> Init(const void* init_data, size_t init_data_size) override {
    // Initialize any persistent state
    return {};
  }

  Expected<void> GetOutputLayouts(
      const std::vector<Layout>& input_layouts,
      std::vector<Layout>& output_layouts) override {
    // Define output tensor shapes based on inputs
    output_layouts[0] = input_layouts[0];
    return {};
  }

  Expected<void> Run(const std::vector<TensorBuffer>& inputs,
                     std::vector<TensorBuffer>& outputs) override {
    // Lock input buffers for reading
    LITERT_ASSIGN_OR_RETURN(auto input_lock,
        TensorBufferScopedLock::Create<float>(
            inputs[0], TensorBuffer::LockMode::kRead));

    // Lock output buffer for writing
    LITERT_ASSIGN_OR_RETURN(auto output_lock,
        TensorBufferScopedLock::Create<float>(
            outputs[0], TensorBuffer::LockMode::kWrite));

    const float* input_data = input_lock.second;
    float* output_data = output_lock.second;

    // Perform computation
    // ... your custom operation logic ...

    return {};
  }

  Expected<void> Destroy() override {
    // Clean up resources
    return {};
  }

private:
  const std::string op_name_ = "MyCustomOp";
};

Triển khai C

#include "litert/c/litert_custom_op_kernel.h"

LiteRtStatus MyOp_Init(void* user_data, const void* init_data,
                       size_t init_data_size) {
  // Initialize state
  return kLiteRtStatusOk;
}

LiteRtStatus MyOp_GetOutputLayouts(void* user_data, size_t num_inputs,
                                   const LiteRtLayout* input_layouts,
                                   size_t num_outputs,
                                   LiteRtLayout* output_layouts) {
  // Set output shape to match first input
  output_layouts[0] = input_layouts[0];
  return kLiteRtStatusOk;
}

LiteRtStatus MyOp_Run(void* user_data, size_t num_inputs,
                     const LiteRtTensorBuffer* inputs, size_t num_outputs,
                     LiteRtTensorBuffer* outputs) {
  // Lock buffers
  void* input_addr;
  LITERT_RETURN_IF_ERROR(LiteRtLockTensorBuffer(
      inputs[0], &input_addr, kLiteRtTensorBufferLockModeRead));

  void* output_addr;
  LITERT_RETURN_IF_ERROR(LiteRtLockTensorBuffer(
      outputs[0], &output_addr, kLiteRtTensorBufferLockModeWrite));

  // Perform computation
  float* in = (float*)input_addr;
  float* out = (float*)output_addr;
  // ... your custom operation logic ...

  // Unlock buffers
  LITERT_RETURN_IF_ERROR(LiteRtUnlockTensorBuffer(inputs[0]));
  LITERT_RETURN_IF_ERROR(LiteRtUnlockTensorBuffer(outputs[0]));

  return kLiteRtStatusOk;
}

LiteRtStatus MyOp_Destroy(void* user_data) {
  // Clean up
  return kLiteRtStatusOk;
}

Bước 2: Đăng ký Thao tác tuỳ chỉnh

Đăng ký C++

#include "litert/cc/litert_environment.h"
#include "litert/cc/litert_compiled_model.h"
#include "litert/cc/litert_options.h"

// Create environment
LITERT_ASSERT_OK_AND_ASSIGN(Environment env, Environment::Create({}));

// Load model
Model model = /* load your model */;

// Create options and register custom op
LITERT_ASSERT_OK_AND_ASSIGN(Options options, Options::Create());
options.SetHardwareAccelerators(kLiteRtHwAcceleratorCpu);

// Register custom op kernel
MyCustomOpKernel my_custom_op;
ASSERT_TRUE(options.AddCustomOpKernel(my_custom_op));

// Create compiled model with custom op
LITERT_ASSERT_OK_AND_ASSIGN(CompiledModel compiled_model,
                            CompiledModel::Create(env, model, options));

Đăng ký C

#include "litert/c/litert_environment.h"
#include "litert/c/litert_compiled_model.h"
#include "litert/c/litert_options.h"

// Create options
LiteRtOptions options;
LiteRtCreateOptions(&options);
LiteRtSetOptionsHardwareAccelerators(options, kLiteRtHwAcceleratorCpu);

// Define kernel
LiteRtCustomOpKernel kernel = {
    .Init = MyOp_Init,
    .GetOutputLayouts = MyOp_GetOutputLayouts,
    .Run = MyOp_Run,
    .Destroy = MyOp_Destroy,
};

// Register custom op
LiteRtAddCustomOpKernelOption(options, "MyCustomOp", 1, &kernel, NULL);

// Create environment
LiteRtEnvironment env;
LiteRtCreateEnvironment(0, NULL, &env);

// Create compiled model
LiteRtCompiledModel compiled_model;
LiteRtCreateCompiledModel(env, model, options, &compiled_model);

Bước 3: Thực thi Mô hình

Thực thi C++

// Create buffers
LITERT_ASSERT_OK_AND_ASSIGN(auto input_buffers,
                            compiled_model.CreateInputBuffers());
LITERT_ASSERT_OK_AND_ASSIGN(auto output_buffers,
                            compiled_model.CreateOutputBuffers());

// Fill input data
input_buffers[0].Write<float>(your_input_data);

// Run inference
compiled_model.Run(input_buffers, output_buffers);

// Read output
LITERT_ASSERT_OK_AND_ASSIGN(auto lock,
    TensorBufferScopedLock::Create<const float>(
        output_buffers[0], TensorBuffer::LockMode::kRead));
const float* results = lock.second;

Thực thi C

// Create buffers (see test files for complete buffer creation)
LiteRtTensorBuffer input_buffers[num_inputs];
LiteRtTensorBuffer output_buffers[num_outputs];
// ... buffer creation code ...

// Write input data
void* input_addr;
LiteRtLockTensorBuffer(input_buffers[0], &input_addr,
                      kLiteRtTensorBufferLockModeWrite);
memcpy(input_addr, your_data, data_size);
LiteRtUnlockTensorBuffer(input_buffers[0]);

// Run inference
LiteRtRunCompiledModel(compiled_model, 0, num_inputs, input_buffers,
                       num_outputs, output_buffers);

// Read output
void* output_addr;
LiteRtLockTensorBuffer(output_buffers[0], &output_addr,
                      kLiteRtTensorBufferLockModeRead);
// Process output_addr
LiteRtUnlockTensorBuffer(output_buffers[0]);