API LiteRT Custom Op Dispatcher

Présentation

CustomOpDispatcher est l'API de remplacement pour définir des opérations de processeur personnalisées et des résolveurs d'opérations personnalisés dans LiteRT. Elle fournit une interface plus propre pour intégrer des opérations personnalisées dans les modèles LiteRT.

Pourquoi utiliser CustomOpDispatcher ?

Approche traditionnelle (obsolète)

Auparavant, les développeurs devaient :

  • Créer manuellement des structures TfLiteRegistration
  • Implémenter des fonctions de rappel spécifiques à TFLite (init, prepare, invoke, free)
  • Travailler directement avec des structures TFLite de bas niveau (TfLiteContext, TfLiteNode, TfLiteTensor)
  • Utiliser directement MutableOpResolver::AddCustom()

Nouvelle approche CustomOpDispatcher

CustomOpDispatcher fournit les éléments suivants :

  • Couche d'abstraction propre sur les composants internes de TFLite
  • Intégration au modèle compilé de LiteRT

Flow

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

Composants et fichiers clés

Implémentation de base

En-têtes d'API

Système d'options

  • core/options.h : structure des options principales avec CustomOpOption

Exemples de tests

Documentation de référence de l'API

Interface du noyau principal (API C)

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;

Classe de base abstraite 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;
};

Guide d'implémentation

Étape 1 : Définissez votre opération personnalisée

Implémentation 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";
};

Implémentation 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;
}

Étape 2 : Enregistrez l'opération personnalisée

Inscription à 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));

Inscription 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);

Étape 3 : Exécuter le modèle

Exécution 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;

Exécution 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]);