Application custom DECODE operator support (part 1)

tensorflow/lite/micro/kernels/decode.cc

@@ -13,6 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <utility>
+
 #include "tensorflow/lite/micro/c/common.h"
 #include "tensorflow/lite/micro/kernels/decode_state.h"
 #include "tensorflow/lite/micro/kernels/internal/compatibility.h"
@@ -49,6 +51,27 @@ TfLiteStatus SetOutputTensorData(TfLiteContext* context, const TfLiteNode* node,
   return kTfLiteOk;
 }
 
+DecodeState* GetDecodeStateFromCustomRegistration(const TfLiteContext* context,
+                                                  uint8_t type) {
+  const MicroContext* mc = GetMicroContext(context);
+  const MicroContext::CustomDecodeRegistration* registrations;
+  size_t registrations_count;
+  std::tie(registrations, registrations_count) =
+      mc->GetCustomDecodeRegistrations();
+  if (registrations == nullptr) {
+    return nullptr;
+  }
+
+  for (size_t i = 0; i < registrations_count; i++) {
+    auto& reg = registrations[i];
+    if (reg.type == type && reg.create_state != nullptr) {
+      return reg.create_state(reg, *context, mc->GetAlternateProfiler());
+    }
+  }
+
+  return nullptr;
+}
+
 TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   const size_t num_inputs = NumInputs(node);
   const size_t num_outputs = NumOutputs(node);
@@ -113,21 +136,22 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
         dsp = DecodeState::CreateDecodeStateHuffman(
             context, micro_context->GetAlternateProfiler());
         break;
-      case DecodeState::kDcmTypeCustom:
-        MicroPrintf("Custom decode type not yet supported");
-        break;
       default:
-        MicroPrintf("unsupported decode type %u",
-                    DecodeState::Type(*ancillary));
+        uint32_t type = DecodeState::Type(*ancillary);
+        if (type >= DecodeState::kDcmTypeCustomFirst &&
+            type <= DecodeState::kDcmTypeCustomLast) {
+          dsp = GetDecodeStateFromCustomRegistration(context, type);
+        } else {
+          MicroPrintf("unsupported decode type %u", type);
+        }
         break;
     }
 
-    status = SetOutputTensorData(context, node, i / 2, output);
-    if (status != kTfLiteOk) {
-      break;
-    }
-
     if (dsp != nullptr) {
+      status = SetOutputTensorData(context, node, i / 2, output);
+      if (status != kTfLiteOk) {
+        break;
+      }
       status = dsp->Setup(*input, *ancillary, *output);
       if (status != kTfLiteOk) {
         break;

tensorflow/lite/micro/kernels/decode_state.h

@@ -26,6 +26,7 @@ limitations under the License.
 
 namespace tflite {
 namespace micro {
+
 class DecodeState {
  public:
   DecodeState() = delete;
@@ -71,7 +72,8 @@ class DecodeState {
   static constexpr uint8_t kDcmTypeLUT = 0;
   static constexpr uint8_t kDcmTypeHuffman = 1;
   static constexpr uint8_t kDcmTypePrune = 2;
-  static constexpr uint8_t kDcmTypeCustom = 127;
+  static constexpr uint8_t kDcmTypeCustomFirst = 128;
+  static constexpr uint8_t kDcmTypeCustomLast = 255;
 
   static constexpr size_t kDcmSizeInBytes = 16;
 

tensorflow/lite/micro/kernels/decode_state_huffman.h

@@ -24,6 +24,7 @@ limitations under the License.
 
 namespace tflite {
 namespace micro {
+
 class DecodeStateHuffman : public DecodeState {
  public:
   DecodeStateHuffman() = delete;

tensorflow/lite/micro/kernels/decode_state_huffman_test.cc

@@ -269,7 +269,7 @@ TEST(DecodeStateHuffmanTest, DecodeHuffmanTable16BitsInt16Fail) {
   tflite::testing::TestDecode<encodes.size() + ancillaries.size(),
                               outputs.size()>(
       encodes, ancillaries, outputs, expected, tflite::Register_DECODE(),
-      nullptr, kTfLiteError);
+      nullptr, nullptr, kTfLiteError);
 }
 
 TEST(DecodeStateHuffmanTest, DecodeHuffmanTable32BitsInt8) {

tensorflow/lite/micro/kernels/decode_state_lut.h

@@ -23,6 +23,7 @@ limitations under the License.
 
 namespace tflite {
 namespace micro {
+
 class DecodeStateLut : public DecodeState {
  public:
   DecodeStateLut() = delete;

tensorflow/lite/micro/kernels/decode_state_prune.h

@@ -23,6 +23,7 @@ limitations under the License.
 
 namespace tflite {
 namespace micro {
+
 class DecodeStatePrune : public DecodeState {
  public:
   DecodeStatePrune() = delete;

tensorflow/lite/micro/kernels/decode_state_prune_test.cc

@@ -573,7 +573,7 @@ TEST(DecodeStatePruneTest, DecodePruneQuantizedInvalidZeroPointInt16) {
   tflite::testing::TestDecode<kEncodes.size() + kAncillaries.size(),
                               kOutputs.size()>(
       kEncodes, kAncillaries, kOutputs, kExpected, tflite::Register_DECODE(),
-      nullptr, kTfLiteError);
+      nullptr, nullptr, kTfLiteError);
 }
 
 TF_LITE_MICRO_TESTS_MAIN

tensorflow/lite/micro/kernels/decode_test.cc

@@ -66,6 +66,76 @@ constexpr int kEncodedShapeLUT[] = {1, sizeof(kEncodedLUT)};
 constexpr int8_t kExpectLUT0[] = {1, 2, 3, 4, 4, 3, 2, 1};
 constexpr int16_t kExpectLUT1[] = {5, 6, 7, 8, 8, 7, 6, 5};
 
+//
+// Custom DECODE test data
+//
+constexpr int kDecodeTypeCustom = 200;
+
+constexpr int8_t kAncillaryDataCustom[] = {0x42};
+
+constexpr uint8_t kDcmCustom[tflite::DecodeState::kDcmSizeInBytes] = {
+    kDecodeTypeCustom,  // type: custom
+    1,                  // DCM version: 1
+};
+
+// Align the tensor data the same as a Buffer in the TfLite schema
+alignas(16) const uint8_t kEncodedCustom[] = {0x42, 0x43, 0x40, 0x46,
+                                              0x4A, 0x52, 0x62, 0x02};
+
+// Tensor shapes as TfLiteIntArray
+constexpr int kOutputShapeCustom[] = {1, 8};
+constexpr int kEncodedShapeCustom[] = {1, sizeof(kEncodedCustom)};
+
+constexpr int8_t kExpectCustom[] = {0x00, 0x01, 0x02, 0x04,
+                                    0x08, 0x10, 0x20, 0x40};
+
+class DecodeStateCustom : public tflite::DecodeState {
+ public:
+  DecodeStateCustom() = delete;
+
+  DecodeStateCustom(const TfLiteContext* context,
+                    tflite::MicroProfilerInterface* profiler)
+      : DecodeState(context, profiler) {}
+
+  TfLiteStatus Setup(const TfLiteTensor& input, const TfLiteTensor& ancillary,
+                     const TfLiteTensor& output) override {
+    return kTfLiteOk;
+  }
+
+  TfLiteStatus Decode(const TfLiteEvalTensor& input,
+                      const TfLiteEvalTensor& ancillary,
+                      const TfLiteEvalTensor& output) override {
+    const uint8_t* inp = tflite::micro::GetTensorData<uint8_t>(&input);
+    TF_LITE_ENSURE(const_cast<TfLiteContext*>(context_), inp != nullptr);
+    uint8_t* outp = tflite::micro::GetTensorData<uint8_t>(
+        const_cast<TfLiteEvalTensor*>(&output));
+    TF_LITE_ENSURE(const_cast<TfLiteContext*>(context_), outp != nullptr);
+    const uint8_t* vp = tflite::micro::GetTensorData<uint8_t>(&ancillary);
+    TF_LITE_ENSURE(const_cast<TfLiteContext*>(context_), vp != nullptr);
+    vp += kDcmSizeInBytes;
+
+    // simple XOR de-obfuscation
+    std::transform(inp, inp + input.dims->data[0], outp,
+                   [vp](uint8_t i) { return i ^ *vp; });
+
+    return kTfLiteOk;
+  }
+
+  static DecodeState* CreateDecodeStateCustom(
+      const tflite::MicroContext::CustomDecodeRegistration& registration,
+      const TfLiteContext& context, tflite::MicroProfilerInterface* profiler) {
+    alignas(DecodeStateCustom) static uint8_t buffer[sizeof(DecodeStateCustom)];
+    DecodeState* instance = new (buffer) DecodeStateCustom(&context, profiler);
+    return instance;
+  }
+
+ protected:
+  ~DecodeStateCustom() override = default;
+
+ private:
+  TF_LITE_REMOVE_VIRTUAL_DELETE
+};
+
 }  // namespace
 
 using tflite::testing::AncillaryData;
@@ -244,4 +314,112 @@ TEST(DecodeTest, DecodeWithAltDecompressionMemory) {
       encodes, ancillaries, outputs, expected, tflite::Register_DECODE(), &amr);
 }
 
+TEST(DecodeTest, DecodeWithCustomRegistration) {
+  // Align the tensor data the same as a Buffer in the TfLite schema
+  alignas(16) int8_t output_data[std::size(kExpectCustom)] = {};
+  alignas(16) const AncillaryData<int8_t, std::size(kAncillaryDataCustom)>
+      kAncillaryData = {{kDcmCustom}, {kAncillaryDataCustom}};
+
+  constexpr int kAncillaryShapeCustom[] = {1, sizeof(kAncillaryData)};
+
+  const TfLiteIntArray* const encoded_dims =
+      tflite::testing::IntArrayFromInts(kEncodedShapeCustom);
+  static const TensorInDatum tid_encode = {
+      kEncodedCustom,
+      *encoded_dims,
+  };
+  static constexpr std::initializer_list<const TensorInDatum*> encodes = {
+      &tid_encode,
+  };
+
+  const TfLiteIntArray* const ancillary_dims =
+      tflite::testing::IntArrayFromInts(kAncillaryShapeCustom);
+  static const TensorInDatum tid_ancillary = {
+      &kAncillaryData,
+      *ancillary_dims,
+  };
+  static constexpr std::initializer_list<const TensorInDatum*> ancillaries = {
+      &tid_ancillary};
+
+  const TfLiteIntArray* const output_dims =
+      tflite::testing::IntArrayFromInts(kOutputShapeCustom);
+  constexpr int kOutputZeroPointsData[] = {0};
+  const TfLiteIntArray* const kOutputZeroPoints =
+      tflite::testing::IntArrayFromInts(kOutputZeroPointsData);
+  const TfLiteFloatArray kOutputScales = {kOutputZeroPoints->size};
+  static const TensorOutDatum tod = {
+      output_data, *output_dims, kTfLiteInt8, kOutputScales, *kOutputZeroPoints,
+      0,           {},
+  };
+  static constexpr std::initializer_list<const TensorOutDatum*> outputs = {
+      &tod};
+
+  const std::initializer_list<const void*> expected = {kExpectCustom};
+
+  const std::initializer_list<tflite::MicroContext::CustomDecodeRegistration>
+      cdr = {{
+          DecodeStateCustom::CreateDecodeStateCustom,
+          kDecodeTypeCustom,
+      }};
+
+  tflite::testing::TestDecode<encodes.size() + ancillaries.size(),
+                              outputs.size()>(
+      encodes, ancillaries, outputs, expected, tflite::Register_DECODE(),
+      nullptr, &cdr);
+}
+
+TEST(DecodeTest, DecodeWithCustomMismatchedRegistration) {
+  // Align the tensor data the same as a Buffer in the TfLite schema
+  alignas(16) int8_t output_data[std::size(kExpectCustom)] = {};
+  alignas(16) const AncillaryData<int8_t, std::size(kAncillaryDataCustom)>
+      kAncillaryData = {{kDcmCustom}, {kAncillaryDataCustom}};
+
+  constexpr int kAncillaryShapeCustom[] = {1, sizeof(kAncillaryData)};
+
+  const TfLiteIntArray* const encoded_dims =
+      tflite::testing::IntArrayFromInts(kEncodedShapeCustom);
+  static const TensorInDatum tid_encode = {
+      kEncodedCustom,
+      *encoded_dims,
+  };
+  static constexpr std::initializer_list<const TensorInDatum*> encodes = {
+      &tid_encode,
+  };
+
+  const TfLiteIntArray* const ancillary_dims =
+      tflite::testing::IntArrayFromInts(kAncillaryShapeCustom);
+  static const TensorInDatum tid_ancillary = {
+      &kAncillaryData,
+      *ancillary_dims,
+  };
+  static constexpr std::initializer_list<const TensorInDatum*> ancillaries = {
+      &tid_ancillary};
+
+  const TfLiteIntArray* const output_dims =
+      tflite::testing::IntArrayFromInts(kOutputShapeCustom);
+  constexpr int kOutputZeroPointsData[] = {0};
+  const TfLiteIntArray* const kOutputZeroPoints =
+      tflite::testing::IntArrayFromInts(kOutputZeroPointsData);
+  const TfLiteFloatArray kOutputScales = {kOutputZeroPoints->size};
+  static const TensorOutDatum tod = {
+      output_data, *output_dims, kTfLiteInt8, kOutputScales, *kOutputZeroPoints,
+      0,           {},
+  };
+  static constexpr std::initializer_list<const TensorOutDatum*> outputs = {
+      &tod};
+
+  const std::initializer_list<const void*> expected = {kExpectCustom};
+
+  const std::initializer_list<tflite::MicroContext::CustomDecodeRegistration>
+      cdr = {{
+          DecodeStateCustom::CreateDecodeStateCustom,
+          kDecodeTypeCustom + 1,
+      }};
+
+  tflite::testing::TestDecode<encodes.size() + ancillaries.size(),
+                              outputs.size()>(
+      encodes, ancillaries, outputs, expected, tflite::Register_DECODE(),
+      nullptr, &cdr, kTfLiteError);
+}
+
 TF_LITE_MICRO_TESTS_MAIN

tensorflow/lite/micro/kernels/decode_test_helpers.h

@@ -81,6 +81,8 @@ void ExecuteDecodeTest(
     const std::initializer_list<const void*>& expected,
     TfLiteStatus expected_status,
     const std::initializer_list<MicroContext::AlternateMemoryRegion>* amr =
+        nullptr,
+    const std::initializer_list<MicroContext::CustomDecodeRegistration>* cdr =
         nullptr) {
   int kInputArrayData[kNumInputs + 1] = {kNumInputs};
   for (size_t i = 0; i < kNumInputs; i++) {
@@ -102,6 +104,11 @@ void ExecuteDecodeTest(
                                                          amr->size());
   }
 
+  if (cdr != nullptr) {
+    runner.GetFakeMicroContext()->SetCustomDecodeRegistrations(cdr->begin(),
+                                                               cdr->size());
+  }
+
   TfLiteStatus status = runner.InitAndPrepare();
   if (status == kTfLiteOk) {
     status = runner.Invoke();
@@ -152,6 +159,8 @@ void TestDecode(
     const TFLMRegistration& registration,
     const std::initializer_list<MicroContext::AlternateMemoryRegion>* amr =
         nullptr,
+    const std::initializer_list<MicroContext::CustomDecodeRegistration>* cdr =
+        nullptr,
     const TfLiteStatus expected_status = kTfLiteOk) {
   TfLiteTensor tensors[kNumInputs + kNumOutputs] = {};
 
@@ -185,7 +194,7 @@ void TestDecode(
   }
 
   ExecuteDecodeTest<kNumInputs, kNumOutputs>(tensors, registration, expected,
-                                             expected_status, amr);
+                                             expected_status, amr, cdr);
 }
 
 }  // namespace testing

tensorflow/lite/micro/micro_context.cc

@@ -174,4 +174,14 @@ void MicroContext::ResetDecompressionMemoryAllocations() {
   std::fill_n(decompress_regions_allocations_, decompress_regions_size_, 0);
 }
 
+TfLiteStatus MicroContext::SetCustomDecodeRegistrations(
+    const CustomDecodeRegistration* registrations, size_t count) {
+  if (custom_decode_registrations_ != nullptr) {
+    return kTfLiteError;
+  }
+  custom_decode_registrations_ = registrations;
+  custom_decode_registrations_size_ = count;
+  return kTfLiteOk;
+}
+
 }  // namespace tflite