Supports External Context in TfLite Micro

[deqiangc]

The processor that TFLM runs on may support an additional hardware accelerator. For an optimized kernel to use this hardware accelerator, the application may need to pass down additional information to the optimized kernel. This can achieved by supporting the TfLiteExternalContext concept. The TfLiteExternalContext is already supported by TfLite (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/c/common.h#L72). This issue is to add this support with micro's contain in TfLite Micro.

The motivating PR for this issue in PR#477.

The next comment will outline a design proposal.

[deqiangc]

This describes TFLite Micro’s implementation choice of TFLite’s external context concept. The motivation request originally comes from PR#477, in which external context is used to pass Cadence’s accelerator information to a specific optimized kernel for a specific OP.

The key to this design is to decide on a label that represents what the expected external data structure actually is so that the client and kernel can be sure everything's set up as expected. The following describes multiple options to represent that label in TFLM.

Option 1(NOT ADOPTED) (Accelerator Specific Enum in TfLiteExternalContextType)

TfLite has the TfLiteExternalContextType enum to represent what specific target that an external context is referring to. So one natural choice is to add new values to TfLiteExternalContextType enum similar as below

typedef enum TfLiteExternalContextType {
  …
  kTfLiteMicroCadenceContext = 4,
  kTfLiteMaxExternalContexts = 5
} TfLiteExternalContextType;

However, there can be multiple hardware vendors with different accelerators in the diverse TFLM embedded environments (Potential examples include Cadence, Qualcomm, ARM). This path will lead us to adding more and more values in the TfLiteExternalContextType, which increases TfLite’s memory footprint unnecessarily (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/interpreter.h#L814). Additionally, it contaminates common.h with hardware vendor information unnecessarily.

Option 2 (NOT ADOPTED) (Standalone Enum or Int)

TFLM can choose to define its own standalone enum or even just integer as a TFLM replacement of the TfLiteExternalContextType enum. But then comes the problem of how to set and how to get this external context. The set is relatively straightforward compared with the get; so let us focus on get next. Currently, TFLM still uses TfLiteContext which has the following API to get external context.

TfLiteExternalContext* (*GetExternalContext)(struct TfLiteContext*,TfLiteExternalContextType);

As you can see, the new type is incompatible with this api.

One way to get around this is to force typecast TfLiteExternalContextType to the new type inside TFLM Get(Set)ExternalContext. However, this is messy and actually could have unpredictable outcomes because type casting between enums with different value ranges may lead to undefined value (https://riptutorial.com/cplusplus/example/18751/enum-conversions).

Other approaches include adding a new api into TfLiteContext or defining a new TfLiteMicroContext, both of which seem extreme for the purpose.

Option 3 (NOT ADOPTED) (Single Micro External Context Enum in TfLiteExternalContextType And Per Vendor Per Chipset Enum Specialization inside Micro Repo)

We define one single micro external context in TfLiteExternalContextType.

typedef enum TfLiteExternalContextType {
  …
  kTfLiteMicroAcceleratorContext = 4,   // include micro_external_context.h to use.
                                       // Micro supports one accelator context
                                       // per target. The interpretation of the
                                       // acccelator context varies by targets.
  kTfLiteMaxExternalContexts = 5
} TfLiteExternalContextType;

When TfLiteExternalContextType == kTfLiteMicroAcceleratorContext, the TfLiteExternalContext must be specialized to

typedef enum TfLiteMicroExternalContextSubType {
  kTfLiteCadenceExternalContext = 0, // .
  kTfLiteMicroMaxExternalContextSubType = 1,
} TfLiteMicroExternalContextSubType;

struct TfLiteMicroExternalContext: struct TfLiteExternalContext {
  // When type in TfLiteExternalContext is kTfLiteMicroAcceleratorContext,
  // the pass in TfLiteExternalContext must be of TfLiteMicroExternalContext;
  TfLiteMicroExternalContextSubType subtype_;
  // 
  void * context_data_; 
}

As you can see, we can add more accelerators into TFLM external context subtype with this approach and the cost to TFLite is fixed (stay constant with more accelerators in TFLM) minimal (just one more entry in its external context type).

[deqiangc]

Option 4 (ADOPTED) (Function in kernel_util.h)

Even Option 3 requires changes to TfLite code base and also requires an OP to do ugly typecast due to the context ->GetExternalContext is not exactly returning the type that we really care about.

An alternative design pattern is to design interface directly in kernel_util and abstract the unfortunate typecast inside the framework and let the OP directly getting an void * pointer to
the payload.

This design pattern can be shown to have a clean API for the point of view of an OP. The following outlines how the API looks like using external context and micro graph as an example.

From the Kernel’s perspective, they use api in micro/kernels/kernel_utils.h as below

// Returns a blob of payload data. The payload is subjected to interpretation by the OP.
// This is the recommended API for an OP to get an external context. OP should use this instead of directly calling GetExternalContext function in context.
// Example usage:
//
// An application can set external context through interpreter as below
//     interpreter->SetMicroExternalContext(pointer_to_your_payload);
//
//  Inside an OP that needs this payload, it get the payload pointer by:
//    Prepare(TfliteContext * context) {
//       ...
//       payload_ptr = reinterp_cast<your_data_type>(GetMicroExternalContext(context))
//       ...
//    }
//
void * GetMicroExternalContext(TfLiteContext* context);

For the application, they use the below Interpreter API to set:

void SetMicroExternalContext(void* external_context);

An implementation sketch is available at #648

Note that:
In this approach, TFLM does not provide a central registry for different vendors’ accelerators etc.
This relieves TFLM from the duty of ensuring the validity of payload which TFLM cannot do a good job anyway b/c there are so many combinations of vendors, chipsets, and even chipset versions.
If a vendor would like to do error checking in the validity payload, then we would recommend putting a magic hash word as the first word of the payload structure and verify based on that.

[github-actions]

"This issue is being marked as stale due to inactivity. Remove label or comment to prevent closure in 5 days."

[github-actions]

"This issue is being closed because it has been marked as
stale for 5 days with no further activity."