TFLM compression changes (#2647)
@tensorflow/micro
initial header file changes/additions.
bug=#2646
diff --git a/tensorflow/lite/micro/compression.h b/tensorflow/lite/micro/compression.h
new file mode 100644
index 0000000..43965c2
--- /dev/null
+++ b/tensorflow/lite/micro/compression.h
@@ -0,0 +1,70 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_MICRO_MICRO_COMPRESSION_H_
+#define TENSORFLOW_LITE_MICRO_MICRO_COMPRESSION_H_
+
+#ifdef USE_TFLM_COMPRESSION
+
+#include "tensorflow/lite/c/common.h"
+
+namespace tflite {
+
+//
+// Compressed tensors
+//
+
+static constexpr const char* kCompressionMetadataString = "TFLM_COMPRESSION";
+
+enum class CompressionScheme : uint8_t {
+ kBinQuant,
+};
+
+// TODO(ddavis-2015): pack struct
+struct LookupTableData {
+ static constexpr size_t kMaxBitWidth = 7;
+ static constexpr size_t kMaxValueTableChannelStride = 128;
+
+ const void* value_table; // Pointer into FlatBuffer Values.
+ uint8_t value_table_channel_stride; // elements per channel
+ uint8_t compressed_bit_width : 3; // 1 to 7 bits
+ bool is_per_channel_quantized : 1; // tensor is per-channel quantized
+ bool use_alternate_axis : 1; // shape default channel:
+ // 0 = first, 1 = last
+ uint8_t reserved : 3;
+};
+
+union CompressionData {
+ LookupTableData* lut_data;
+};
+
+// TODO(ddavis-2015): pack struct
+struct CompressionTensorData {
+ CompressionScheme scheme;
+ CompressionData data;
+};
+
+// TODO(ddavis-2015): pack struct
+struct CompressedTensorList {
+ // Sparsely populated array with the same number of elements as there are
+ // tensors in the Subgraph. An alternative would include a tensor index in
+ // the struct for each and walk the list on look up. This could be slow.
+ CompressionTensorData** tensors;
+};
+
+} // namespace tflite
+
+#endif // USE_TFLM_COMPRESSION
+#endif // TENSORFLOW_LITE_MICRO_MICRO_COMPRESSION_H_
diff --git a/tensorflow/lite/micro/fake_micro_context.h b/tensorflow/lite/micro/fake_micro_context.h
index 46d8a9b..7cf9c68 100644
--- a/tensorflow/lite/micro/fake_micro_context.h
+++ b/tensorflow/lite/micro/fake_micro_context.h
@@ -1,4 +1,4 @@
-/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -30,7 +30,12 @@
~FakeMicroContext() = default;
FakeMicroContext(TfLiteTensor* tensors, SingleArenaBufferAllocator* allocator,
- MicroGraph* micro_graph);
+ MicroGraph* micro_graph
+#ifdef USE_TFLM_COMPRESSION
+ ,
+ const CompressedTensorList* compressed_tensors = nullptr
+#endif // USE_TFLM_COMPRESSION
+ );
void* AllocatePersistentBuffer(size_t bytes) override;
TfLiteStatus RequestScratchBufferInArena(size_t bytes,
@@ -50,6 +55,24 @@
void* external_context() override;
MicroGraph& graph() override;
+#ifdef USE_TFLM_COMPRESSION
+
+ // Available during Prepare & Eval. Returns false if tensor is not
+ // compressed.
+ bool IsTensorCompressed(const TfLiteNode* node, int tensor_idx) override;
+
+ // Only available during Prepare. The kernel is responsible for storing the
+ // scratch buffer handle.
+ int AllocateDecompressionScratchBuffer(const TfLiteNode* node,
+ int tensor_idx) override;
+
+ // Available during Prepare & Eval. Returns nullptr if tensor is not
+ // compressed.
+ const CompressionTensorData* GetTensorCompressionData(
+ const TfLiteNode* node, int tensor_idx) override;
+
+#endif // USE_TFLM_COMPRESSION
+
private:
static constexpr int kNumScratchBuffers_ = 12;
@@ -62,6 +85,15 @@
SingleArenaBufferAllocator* allocator_;
+#ifdef USE_TFLM_COMPRESSION
+
+ //
+ // Compression
+ //
+ const CompressedTensorList* compressed_tensors_;
+
+#endif // USE_TFLM_COMPRESSION
+
TF_LITE_REMOVE_VIRTUAL_DELETE
};
diff --git a/tensorflow/lite/micro/kernels/kernel_runner.h b/tensorflow/lite/micro/kernels/kernel_runner.h
index 25b97c1..8dbd7f8 100644
--- a/tensorflow/lite/micro/kernels/kernel_runner.h
+++ b/tensorflow/lite/micro/kernels/kernel_runner.h
@@ -1,4 +1,4 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -36,7 +36,12 @@
KernelRunner(const TFLMRegistration& registration, TfLiteTensor* tensors,
int tensors_size, TfLiteIntArray* inputs,
TfLiteIntArray* outputs, const void* builtin_data,
- TfLiteIntArray* intermediates = nullptr);
+ TfLiteIntArray* intermediates = nullptr
+#ifdef USE_TFLM_COMPRESSION
+ ,
+ const CompressedTensorList* compressed_tensors = nullptr
+#endif // USE_TFLM_COMPRESSION
+ );
// Calls init and prepare on the kernel (i.e. TFLMRegistration) struct.
// Any exceptions will be DebugLog'd and returned as a status code.
diff --git a/tensorflow/lite/micro/kernels/kernel_util.h b/tensorflow/lite/micro/kernels/kernel_util.h
index f14c927..977ed95 100644
--- a/tensorflow/lite/micro/kernels/kernel_util.h
+++ b/tensorflow/lite/micro/kernels/kernel_util.h
@@ -1,4 +1,4 @@
-/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -91,6 +91,31 @@
: reinterpret_cast<const T*>(tensor->data.raw);
}
+#ifdef USE_TFLM_COMPRESSION
+
+// Overloads existing GetTensorData. If not compressed, this will return
+// tensor->data.
+//
+// TODO(ddavis-2015): make micro_context a const pointer
+template <typename T>
+const T* GetTensorData(MicroContext* micro_context,
+ const TfLiteEvalTensor* tensor,
+ const CompressionTensorData* compression_data,
+ int scratch_buffer_handle) {
+ if (tensor == nullptr) {
+ return nullptr;
+ }
+ if (compression_data == nullptr) {
+ return reinterpret_cast<const T*>(tensor->data.data);
+ }
+
+ void* uncompressed_data = micro_context->DecompressTensorToScratchBuffer(
+ *tensor, *compression_data, scratch_buffer_handle);
+ return reinterpret_cast<const T*>(uncompressed_data);
+}
+
+#endif // USE_TFLM_COMPRESSION
+
// Returns the shape of a TfLiteEvalTensor struct.
const RuntimeShape GetTensorShape(const TfLiteEvalTensor* tensor);
diff --git a/tensorflow/lite/micro/micro_context.h b/tensorflow/lite/micro/micro_context.h
index 2dd3233..33cad89 100644
--- a/tensorflow/lite/micro/micro_context.h
+++ b/tensorflow/lite/micro/micro_context.h
@@ -1,4 +1,4 @@
-/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -19,6 +19,12 @@
#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/micro/micro_graph.h"
+#ifdef USE_TFLM_COMPRESSION
+
+#include "tensorflow/lite/micro/compression.h"
+
+#endif // USE_TFLM_COMPRESSION
+
namespace tflite {
// TODO(b/149795762): kTfLiteAbort cannot be part of the tflite TfLiteStatus.
const TfLiteStatus kTfLiteAbort = static_cast<TfLiteStatus>(15);
@@ -95,6 +101,30 @@
virtual MicroGraph& graph() = 0;
+#ifdef USE_TFLM_COMPRESSION
+
+ // Available during Prepare & Eval. Returns false if tensor is not
+ // compressed.
+ virtual bool IsTensorCompressed(const TfLiteNode* node, int tensor_idx) = 0;
+
+ // Only available during Prepare. The kernel is responsible for storing the
+ // scratch buffer handle.
+ virtual int AllocateDecompressionScratchBuffer(const TfLiteNode* node,
+ int tensor_idx) = 0;
+
+ // Available during Prepare & Eval. Returns nullptr if tensor is not
+ // compressed.
+ virtual const CompressionTensorData* GetTensorCompressionData(
+ const TfLiteNode* node, int tensor_idx) = 0;
+
+ // Only available during Eval. Returns nullptr on failure, otherwise returns a
+ // pointer to the scratch buffer.
+ virtual void* DecompressTensorToScratchBuffer(
+ const TfLiteEvalTensor& tensor,
+ const CompressionTensorData& compression_data, int scratch_buffer_handle);
+
+#endif // USE_TFLM_COMPRESSION
+
private:
TF_LITE_REMOVE_VIRTUAL_DELETE
};
diff --git a/tensorflow/lite/micro/micro_interpreter_context.h b/tensorflow/lite/micro/micro_interpreter_context.h
index 5986dc3..7b336aa 100644
--- a/tensorflow/lite/micro/micro_interpreter_context.h
+++ b/tensorflow/lite/micro/micro_interpreter_context.h
@@ -1,4 +1,4 @@
-/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
@@ -106,6 +106,31 @@
// housekeeping in MicroInterpreterContext.
void SetScratchBufferHandles(ScratchBufferHandle* scratch_buffer_handles);
+#ifdef USE_TFLM_COMPRESSION
+
+ // Available during Prepare & Eval. Returns false if tensor is not
+ // compressed.
+ bool IsTensorCompressed(const TfLiteNode* node, int tensor_idx) override;
+
+ // Only available during Prepare. The kernel is responsible for storing the
+ // scratch buffer handle.
+ int AllocateDecompressionScratchBuffer(const TfLiteNode* node,
+ int tensor_idx) override;
+
+ // Available during Prepare & Eval. Returns nullptr if tensor is not
+ // compressed.
+ const CompressionTensorData* GetTensorCompressionData(
+ const TfLiteNode* node, int tensor_idx) override;
+
+ // Only available during Eval. Returns nullptr on failure, otherwise returns a
+ // pointer to the scratch buffer.
+ void* DecompressTensorToScratchBuffer(
+ const TfLiteEvalTensor& tensor,
+ const CompressionTensorData& compression_data,
+ int scratch_buffer_handle) override;
+
+#endif // USE_TFLM_COMPRESSION
+
private:
MicroAllocator& allocator_;
MicroInterpreterGraph& graph_;
