TFLM compression changes (3rd) (#2658) @tensorflow/micro Updates to support TFLM compression: MicroContext MicroInterpreterContext FakeMicroContext KernelRunner bug=#2657
diff --git a/tensorflow/lite/micro/fake_micro_context.cc b/tensorflow/lite/micro/fake_micro_context.cc index 5787ffd..8874798 100644 --- a/tensorflow/lite/micro/fake_micro_context.cc +++ b/tensorflow/lite/micro/fake_micro_context.cc
@@ -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. @@ -23,10 +23,23 @@ namespace tflite { -FakeMicroContext::FakeMicroContext(TfLiteTensor* tensors, - SingleArenaBufferAllocator* allocator, - MicroGraph* micro_graph) - : graph_(*micro_graph), tensors_(tensors), allocator_(allocator) {} +FakeMicroContext::FakeMicroContext( + TfLiteTensor* tensors, SingleArenaBufferAllocator* allocator, + MicroGraph* micro_graph +#ifdef USE_TFLM_COMPRESSION + , + const CompressedTensorList* compressed_tensors +#endif // USE_TFLM_COMPRESSION + ) + : graph_(*micro_graph), + tensors_(tensors), + allocator_(allocator) +#ifdef USE_TFLM_COMPRESSION + , + compressed_tensors_(compressed_tensors) +#endif // USE_TFLM_COMPRESSION +{ +} TfLiteTensor* FakeMicroContext::AllocateTempTfLiteTensor(int tensor_index) { allocated_temp_count_++; @@ -112,4 +125,60 @@ MicroGraph& FakeMicroContext::graph() { return graph_; } +#ifdef USE_TFLM_COMPRESSION + +// Available during Prepare & Eval. Returns false if tensor is not +// compressed. +bool FakeMicroContext::IsTensorCompressed(const TfLiteNode* node, + int tensor_idx) { + if (compressed_tensors_ != nullptr && tensor_idx < node->inputs->size) { + int index = node->inputs->data[tensor_idx]; + if (index >= 0 && compressed_tensors_->tensors[index] != nullptr) { + return true; + } + } + + return false; +} + +// Only available during Prepare. The kernel is responsible for storing the +// scratch buffer handle. +int FakeMicroContext::AllocateDecompressionScratchBuffer(const TfLiteNode* node, + int tensor_idx) { + if (compressed_tensors_ == nullptr || tensor_idx >= node->inputs->size) { + return -1; + } + int index = node->inputs->data[tensor_idx]; + if (index < 0 || compressed_tensors_->tensors[index] == nullptr) { + return -1; + } + TfLiteTensor* tensor = &tensors_[index]; + int scratch_index = -1; + TfLiteStatus result = + RequestScratchBufferInArena(tensor->bytes, &scratch_index); + if (result != kTfLiteOk) { + return -1; + } + + return scratch_index; +} + +// Available during Prepare & Eval. Returns nullptr if tensor is not +// compressed. +const CompressionTensorData* FakeMicroContext::GetTensorCompressionData( + const TfLiteNode* node, int tensor_idx) { + if (compressed_tensors_ == nullptr || tensor_idx >= node->inputs->size) { + return nullptr; + } + + int index = node->inputs->data[tensor_idx]; + if (index < 0) { + return nullptr; + } + + return compressed_tensors_->tensors[index]; +} + +#endif // USE_TFLM_COMPRESSION + } // namespace tflite
diff --git a/tensorflow/lite/micro/kernels/kernel_runner.cc b/tensorflow/lite/micro/kernels/kernel_runner.cc index 602778d..79824ef 100644 --- a/tensorflow/lite/micro/kernels/kernel_runner.cc +++ b/tensorflow/lite/micro/kernels/kernel_runner.cc
@@ -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. @@ -18,7 +18,6 @@ #include "tensorflow/lite/micro/arena_allocator/single_arena_buffer_allocator.h" #include "tensorflow/lite/micro/micro_arena_constants.h" #include "tensorflow/lite/micro/micro_log.h" -#include "tensorflow/lite/micro/test_helpers.h" namespace tflite { namespace micro { @@ -38,12 +37,22 @@ TfLiteTensor* tensors, int tensors_size, TfLiteIntArray* inputs, TfLiteIntArray* outputs, const void* builtin_data, - TfLiteIntArray* intermediates) + TfLiteIntArray* intermediates +#ifdef USE_TFLM_COMPRESSION + , + const CompressedTensorList* compressed_tensors +#endif // USE_TFLM_COMPRESSION + ) : registration_(registration), allocator_(SingleArenaBufferAllocator::Create(kKernelRunnerBuffer_, kKernelRunnerBufferSize_)), mock_micro_graph_(allocator_), - fake_micro_context_(tensors, allocator_, &mock_micro_graph_) { + fake_micro_context_(tensors, allocator_, &mock_micro_graph_ +#ifdef USE_TFLM_COMPRESSION + , + compressed_tensors +#endif // USE_TFLM_COMPRESSION + ) { // Prepare TfLiteContext: context_.impl_ = static_cast<void*>(&fake_micro_context_); context_.ReportError = MicroContextReportOpError;
diff --git a/tensorflow/lite/micro/micro_context.cc b/tensorflow/lite/micro/micro_context.cc index 295b3c3..680dee8 100644 --- a/tensorflow/lite/micro/micro_context.cc +++ b/tensorflow/lite/micro/micro_context.cc
@@ -1,4 +1,4 @@ -/* Copyright 2023 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. @@ -18,8 +18,10 @@ #include <cstdarg> #include <cstddef> +#include "tensorflow/lite/kernels/internal/compatibility.h" #include "tensorflow/lite/micro/micro_common.h" #include "tensorflow/lite/micro/micro_log.h" +#include "tensorflow/lite/micro/micro_utils.h" namespace tflite { namespace { @@ -34,6 +36,103 @@ return -1; } +#ifdef USE_TFLM_COMPRESSION + +struct DecompressionState { + DecompressionState() = delete; + + DecompressionState(const uint8_t* compressed_indices, + const size_t count_indices, + const CompressionTensorData& comp_data, + const size_t num_channels) + : compressed_indices_(compressed_indices), + count_indices_(count_indices), + comp_data_(comp_data), + num_channels_(num_channels) {} + + template <typename T> + T* DecompressToBuffer(void* buffer); + + size_t GetNextTableIndex(); + void UpdateBufferAndChannelIndex(); + + private: + const uint8_t* compressed_indices_; + const size_t count_indices_; + const CompressionTensorData& comp_data_; + const size_t num_channels_; + const size_t compressed_bit_width_ = + comp_data_.data.lut_data->compressed_bit_width; + size_t channel_ = 0; + size_t index_in_channel_ = 0; + const size_t elements_per_channel_ = + comp_data_.data.lut_data->use_alternate_axis + ? 1 + : count_indices_ / num_channels_; + size_t buffer_index_ = 0; + size_t current_offset_ = 0; + size_t current_bits_remaining_ = 8; + uint8_t current_byte_ = compressed_indices_[0]; +}; + +template <typename T> +T* DecompressionState::DecompressToBuffer(void* buffer) { + while (buffer_index_ < count_indices_) { + const size_t table_index = GetNextTableIndex(); + static_cast<T*>(buffer)[buffer_index_] = + static_cast<const T*>(comp_data_.data.lut_data->value_table) + [table_index + + (channel_ * comp_data_.data.lut_data->value_table_channel_stride)]; + UpdateBufferAndChannelIndex(); + } + + return static_cast<T*>(buffer); +} + +size_t DecompressionState::GetNextTableIndex() { + TFLITE_DCHECK(compressed_bit_width_ <= LookupTableData::kMaxBitWidth); + TFLITE_DCHECK(compressed_bit_width_ > 0); + + size_t table_index_bits_to_fill = compressed_bit_width_; + size_t table_index = 0; + + while (table_index_bits_to_fill > 0) { + if (current_bits_remaining_ == 0) { + current_offset_++; + current_byte_ = compressed_indices_[current_offset_]; + current_bits_remaining_ = 8; + } + + const uint8_t mask_bit_count = + std::min(table_index_bits_to_fill, + std::min(compressed_bit_width_, current_bits_remaining_)); + const uint8_t current_byte_mask = (1 << mask_bit_count) - 1; + table_index <<= mask_bit_count; + table_index |= + (current_byte_ >> (current_bits_remaining_ - mask_bit_count)) & + current_byte_mask; + + table_index_bits_to_fill -= mask_bit_count; + current_bits_remaining_ -= mask_bit_count; + } + + return table_index; +} + +void DecompressionState::UpdateBufferAndChannelIndex() { + buffer_index_++; + index_in_channel_++; + if (index_in_channel_ == elements_per_channel_) { + index_in_channel_ = 0; + channel_++; + if (channel_ == num_channels_) { + channel_ = 0; + } + } +} + +#endif // USE_TFLM_COMPRESSION + } // namespace TfLiteTensor* MicroContext::AllocateTempInputTensor(const TfLiteNode* node, @@ -74,4 +173,56 @@ va_end(args); } +#ifdef USE_TFLM_COMPRESSION + +void* MicroContext::DecompressTensorToScratchBuffer( + const TfLiteEvalTensor& tensor, + const CompressionTensorData& compression_data, int scratch_buffer_handle) { + TFLITE_DCHECK(compression_data.scheme == CompressionScheme::kBinQuant); + TFLITE_DCHECK(scratch_buffer_handle != -1); + void* scratch_buffer = GetScratchBuffer(scratch_buffer_handle); + TFLITE_DCHECK(scratch_buffer != nullptr); + size_t count = ElementCount(*tensor.dims); + size_t num_channels = 1; + + if (compression_data.data.lut_data->is_per_channel_quantized) { + const size_t channel_axis = + compression_data.data.lut_data->use_alternate_axis + ? tensor.dims->size - 1 + : 0; + num_channels = tensor.dims->data[channel_axis]; + } + + DecompressionState ds(static_cast<uint8_t*>(tensor.data.data), count, + compression_data, num_channels); + + switch (tensor.type) { + case kTfLiteBool: { + return ds.DecompressToBuffer<bool>(scratch_buffer); + } break; + case kTfLiteInt8: { + return ds.DecompressToBuffer<int8_t>(scratch_buffer); + } break; + case kTfLiteInt16: { + return ds.DecompressToBuffer<int16_t>(scratch_buffer); + } break; + case kTfLiteInt32: { + return ds.DecompressToBuffer<int32_t>(scratch_buffer); + } break; + case kTfLiteInt64: { + return ds.DecompressToBuffer<int64_t>(scratch_buffer); + } break; + case kTfLiteFloat32: { + return ds.DecompressToBuffer<float>(scratch_buffer); + } break; + default: { + MicroPrintf("Unsupported decompression tensor type %d", tensor.type); + } break; + } + + return nullptr; +} + +#endif // USE_TFLM_COMPRESSION + } // namespace tflite
diff --git a/tensorflow/lite/micro/micro_interpreter_context.cc b/tensorflow/lite/micro/micro_interpreter_context.cc index 098df15..0ba461f 100644 --- a/tensorflow/lite/micro/micro_interpreter_context.cc +++ b/tensorflow/lite/micro/micro_interpreter_context.cc
@@ -1,4 +1,4 @@ -/* Copyright 2023 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. @@ -18,8 +18,28 @@ #include <cstdint> #include "tensorflow/lite/kernels/internal/compatibility.h" +#include "tensorflow/lite/micro/micro_utils.h" namespace tflite { + +namespace { + +#ifdef USE_TFLM_COMPRESSION + +int GetInputTensorIndex(const TfLiteNode* node, const int index) { + if (index >= 0 && index < node->inputs->size) { + const int tensor_index = node->inputs->data[index]; + if (tensor_index != kTfLiteOptionalTensor) { + return tensor_index; + } + } + return -1; +} + +#endif // USE_TFLM_COMPRESSION + +} // namespace + MicroInterpreterContext::MicroInterpreterContext(MicroAllocator* allocator, const Model* model, MicroInterpreterGraph* graph) @@ -106,4 +126,83 @@ return state_; } +#ifdef USE_TFLM_COMPRESSION + +// Available during Prepare & Eval. Returns false if tensor is not +// compressed. +bool MicroInterpreterContext::IsTensorCompressed(const TfLiteNode* node, + int tensor_idx) { + TFLITE_DCHECK(state_ == InterpreterState::kPrepare || + state_ == InterpreterState::kInvoke); + + const SubgraphAllocations* allocations = + &graph_.GetAllocations()[graph_.GetCurrentSubgraphIndex()]; + if (allocations->compressed.tensors == nullptr) { + return false; + } + int index = GetInputTensorIndex(node, tensor_idx); + if (index == -1) { + return false; + } + return allocations->compressed.tensors[index] != nullptr; +} + +// Only available during Prepare. The kernel is responsible for storing the +// scratch buffer handle. +int MicroInterpreterContext::AllocateDecompressionScratchBuffer( + const TfLiteNode* node, int tensor_idx) { + TFLITE_DCHECK(state_ == InterpreterState::kPrepare); + + const SubgraphAllocations* allocations = + &graph_.GetAllocations()[graph_.GetCurrentSubgraphIndex()]; + if (allocations->compressed.tensors == nullptr) { + return -1; + } + int index = GetInputTensorIndex(node, tensor_idx); + if (index == -1 || allocations->compressed.tensors[index] == nullptr) { + return -1; + } + const TfLiteEvalTensor* tensor = &allocations->tensors[index]; + const size_t byte_count = EvalTensorBytes(tensor); + int scratch_index = -1; + TfLiteStatus result = RequestScratchBufferInArena(byte_count, &scratch_index); + if (result != kTfLiteOk) { + return -1; + } + + return scratch_index; +} + +// Available during Prepare & Eval. Returns nullptr if tensor is not +// compressed. +const CompressionTensorData* MicroInterpreterContext::GetTensorCompressionData( + const TfLiteNode* node, int tensor_idx) { + TFLITE_DCHECK(state_ == InterpreterState::kPrepare || + state_ == InterpreterState::kInvoke); + + const SubgraphAllocations* allocations = + &graph_.GetAllocations()[graph_.GetCurrentSubgraphIndex()]; + if (allocations->compressed.tensors == nullptr) { + return nullptr; + } + int index = GetInputTensorIndex(node, tensor_idx); + if (index == -1) { + return nullptr; + } + return allocations->compressed.tensors[index]; +} + +// Only available during Eval. Returns nullptr on failure, otherwise returns a +// pointer to the scratch buffer. +void* MicroInterpreterContext::DecompressTensorToScratchBuffer( + const TfLiteEvalTensor& tensor, + const CompressionTensorData& compression_data, int scratch_buffer_handle) { + TFLITE_DCHECK(state_ == InterpreterState::kInvoke); + + return MicroContext::DecompressTensorToScratchBuffer(tensor, compression_data, + scratch_buffer_handle); +} + +#endif // USE_TFLM_COMPRESSION + } // namespace tflite