Skip to content

Commit 6c1c1a8

Browse files
ddavis-2015ddavis-2015
authored
Support for DECODE operator (#3213)
* Support for DECODE operator @tensorflow/micro Add support for alternate decompression memory to DECODE operator. Additional unit tests. Update generic benchmark application and Makefile. bug=fixes #3212 * fixes. * Changes as per review.
1 parent b8c7ecc commit 6c1c1a8

15 files changed

+288
-95
lines changed

‎tensorflow/lite/micro/kernels/decode.cc‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,37 @@ limitations under the License.
1818
#include "tensorflow/lite/kernels/kernel_util.h"
1919
#include "tensorflow/lite/micro/kernels/decode_state.h"
2020
#include "tensorflow/lite/micro/kernels/kernel_util.h"
21+
#include "tensorflow/lite/micro/micro_arena_constants.h"
2122
#include "tensorflow/lite/micro/micro_context.h"
2223
#include "tensorflow/lite/micro/micro_log.h"
2324

2425
namespace tflite {
2526
namespace {
2627

28+
TfLiteStatus SetOutputTensorData(TfLiteContext* context, const TfLiteNode* node,
29+
size_t tensor_output_index,
30+
TfLiteTensor* output) {
31+
if (output->data.data != nullptr) {
32+
// If memory has already been assigned to the tensor, leave it be
33+
return kTfLiteOk;
34+
}
35+
36+
// If alternate decompression memory is available, set the tensor data
37+
// pointer now to preclude allocation by the memory planner.
38+
void* alternate_decompress_mem =
39+
GetMicroContext(context)->AllocateDecompressionMemory(
40+
output->bytes, MicroArenaBufferAlignment());
41+
if (alternate_decompress_mem != nullptr) {
42+
TfLiteEvalTensor* output_eval =
43+
tflite::micro::GetEvalOutput(context, node, tensor_output_index);
44+
TF_LITE_ENSURE(context, output_eval != nullptr);
45+
output_eval->data.data = alternate_decompress_mem;
46+
output->data.data = alternate_decompress_mem;
47+
}
48+
49+
return kTfLiteOk;
50+
}
51+
2752
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
2853
const size_t num_inputs = NumInputs(node);
2954
const size_t num_outputs = NumOutputs(node);
@@ -43,6 +68,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
4368
TfLiteTensor* output = nullptr;
4469
TfLiteStatus status = kTfLiteOk;
4570

71+
micro_context->ResetDecompressionMemoryAllocations();
72+
4673
for (size_t i = 0; i < num_inputs; i += 2) {
4774
input = micro_context->AllocateTempInputTensor(node, i);
4875
if (input == nullptr) {
@@ -95,6 +122,11 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
95122
break;
96123
}
97124

125+
status = SetOutputTensorData(context, node, i / 2, output);
126+
if (status != kTfLiteOk) {
127+
break;
128+
}
129+
98130
if (dsp != nullptr) {
99131
status = dsp->Setup(*input, *ancillary, *output);
100132
if (status != kTfLiteOk) {

‎tensorflow/lite/micro/kernels/decode_state_huffman_test.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -271,7 +271,7 @@ TF_LITE_MICRO_TEST(DecodeHuffmanTable16BitsInt16Fail) {
271271
tflite::testing::TestDecode<encodes.size() + ancillaries.size(),
272272
outputs.size()>(
273273
encodes, ancillaries, outputs, expected, tflite::Register_DECODE(),
274-
kTfLiteError);
274+
nullptr, kTfLiteError);
275275
}
276276

277277
TF_LITE_MICRO_TEST(DecodeHuffmanTable32BitsInt8) {

‎tensorflow/lite/micro/kernels/decode_state_prune_test.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -575,7 +575,7 @@ TF_LITE_MICRO_TEST(DecodePruneQuantizedInvalidZeroPointInt16) {
575575
tflite::testing::TestDecode<kEncodes.size() + kAncillaries.size(),
576576
kOutputs.size()>(
577577
kEncodes, kAncillaries, kOutputs, kExpected, tflite::Register_DECODE(),
578-
kTfLiteError);
578+
nullptr, kTfLiteError);
579579
}
580580

581581
TF_LITE_MICRO_TESTS_END

‎tensorflow/lite/micro/kernels/decode_test.cc‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,4 +196,54 @@ TF_LITE_MICRO_TEST(DecodeTwoTensorsLUT) {
196196
encodes, ancillaries, outputs, expected, tflite::Register_DECODE());
197197
}
198198

199+
TF_LITE_MICRO_TEST(DecodeWithAltDecompressionMemory) {
200+
// Align the tensor data the same as a Buffer in the TfLite schema
201+
alignas(16) int8_t output_data[std::size(kExpectLUT0)] = {};
202+
alignas(16) const AncillaryData<int8_t, std::size(kAncillaryDataLUT0)>
203+
kAncillaryData = {{kDcmLUT0}, {kAncillaryDataLUT0}};
204+
205+
constexpr int kAncillaryShapeLUT[] = {1, sizeof(kAncillaryData)};
206+
207+
const TfLiteIntArray* const encoded_dims =
208+
tflite::testing::IntArrayFromInts(kEncodedShapeLUT);
209+
static const TensorInDatum tid_encode = {
210+
kEncodedLUT,
211+
*encoded_dims,
212+
};
213+
static constexpr std::initializer_list<const TensorInDatum*> encodes = {
214+
&tid_encode,
215+
};
216+
217+
const TfLiteIntArray* const ancillary_dims =
218+
tflite::testing::IntArrayFromInts(kAncillaryShapeLUT);
219+
static const TensorInDatum tid_ancillary = {
220+
&kAncillaryData,
221+
*ancillary_dims,
222+
};
223+
static constexpr std::initializer_list<const TensorInDatum*> ancillaries = {
224+
&tid_ancillary};
225+
226+
const TfLiteIntArray* const output_dims =
227+
tflite::testing::IntArrayFromInts(kOutputShapeLUT);
228+
constexpr int kOutputZeroPointsData[] = {0};
229+
const TfLiteIntArray* const kOutputZeroPoints =
230+
tflite::testing::IntArrayFromInts(kOutputZeroPointsData);
231+
const TfLiteFloatArray kOutputScales = {kOutputZeroPoints->size};
232+
static const TensorOutDatum tod = {
233+
nullptr, // using alternate decompression memory
234+
*output_dims, kTfLiteInt8, kOutputScales, *kOutputZeroPoints, 0, {},
235+
};
236+
static constexpr std::initializer_list<const TensorOutDatum*> outputs = {
237+
&tod};
238+
239+
const std::initializer_list<const void*> expected = {kExpectLUT0};
240+
241+
std::initializer_list<tflite::MicroContext::AlternateMemoryRegion> amr = {
242+
{output_data, sizeof(output_data)}};
243+
244+
tflite::testing::TestDecode<encodes.size() + ancillaries.size(),
245+
outputs.size()>(
246+
encodes, ancillaries, outputs, expected, tflite::Register_DECODE(), &amr);
247+
}
248+
199249
TF_LITE_MICRO_TESTS_END

‎tensorflow/lite/micro/kernels/decode_test_helpers.h‎

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,9 @@ TfLiteStatus CheckOutput(const TfLiteTensor& output,
8383
template <size_t kNumInputs, size_t kNumOutputs>
8484
TfLiteStatus ExecuteDecodeTest(
8585
TfLiteTensor* tensors, const TFLMRegistration& registration,
86-
const std::initializer_list<const void*>& expected) {
86+
const std::initializer_list<const void*>& expected,
87+
const std::initializer_list<MicroContext::AlternateMemoryRegion>* amr =
88+
nullptr) {
8789
int kInputArrayData[kNumInputs + 1] = {kNumInputs};
8890
for (size_t i = 0; i < kNumInputs; i++) {
8991
kInputArrayData[i + 1] = i;
@@ -99,6 +101,10 @@ TfLiteStatus ExecuteDecodeTest(
99101
micro::KernelRunner runner(registration, tensors, kNumInputs + kNumOutputs,
100102
inputs_array, outputs_array, nullptr);
101103

104+
if (amr != nullptr) {
105+
runner.GetFakeMicroContext()->SetDecompressionMemory(*amr);
106+
}
107+
102108
if (runner.InitAndPrepare() != kTfLiteOk || runner.Invoke() != kTfLiteOk) {
103109
return kTfLiteError;
104110
}
@@ -135,12 +141,15 @@ TfLiteStatus ExecuteDecodeTest(
135141
}
136142

137143
template <size_t kNumInputs, size_t kNumOutputs>
138-
void TestDecode(const std::initializer_list<const TensorInDatum*>& encodes,
139-
const std::initializer_list<const TensorInDatum*>& ancillaries,
140-
const std::initializer_list<const TensorOutDatum*>& outputs,
141-
const std::initializer_list<const void*>& expected,
142-
const TFLMRegistration& registration,
143-
const TfLiteStatus expected_status = kTfLiteOk) {
144+
void TestDecode(
145+
const std::initializer_list<const TensorInDatum*>& encodes,
146+
const std::initializer_list<const TensorInDatum*>& ancillaries,
147+
const std::initializer_list<const TensorOutDatum*>& outputs,
148+
const std::initializer_list<const void*>& expected,
149+
const TFLMRegistration& registration,
150+
const std::initializer_list<MicroContext::AlternateMemoryRegion>* amr =
151+
nullptr,
152+
const TfLiteStatus expected_status = kTfLiteOk) {
144153
TfLiteTensor tensors[kNumInputs + kNumOutputs] = {};
145154

146155
for (size_t i = 0; i < kNumInputs; i += 2) {
@@ -173,7 +182,7 @@ void TestDecode(const std::initializer_list<const TensorInDatum*>& encodes,
173182
}
174183

175184
TfLiteStatus s = ExecuteDecodeTest<kNumInputs, kNumOutputs>(
176-
tensors, registration, expected);
185+
tensors, registration, expected, amr);
177186
TF_LITE_MICRO_EXPECT_EQ(s, expected_status);
178187
}
179188

‎tensorflow/lite/micro/kernels/kernel_runner.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,9 @@ class KernelRunner {
6767
// to stub out MicroGraph methods and track invocations on each subgraph.
6868
MockMicroGraph* GetMockGraph() { return &mock_micro_graph_; }
6969

70+
// Returns a pointer to the internal FakeMicroContext.
71+
FakeMicroContext* GetFakeMicroContext() { return &fake_micro_context_; }
72+
7073
// Returns true if all temp buffer in tests are deallocated.
7174
// TODO(b/209453859): move this function to private after deallocation checks
7275
// are enabled for all kernel tests.

‎tensorflow/lite/micro/micro_context.cc‎

Lines changed: 38 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,13 @@ limitations under the License.
1515

1616
#include "tensorflow/lite/micro/micro_context.h"
1717

18+
#include <algorithm>
1819
#include <cstdarg>
1920
#include <cstddef>
2021

2122
#include "tensorflow/lite/kernels/internal/compatibility.h"
2223
#include "tensorflow/lite/micro/kernels/decompress.h"
24+
#include "tensorflow/lite/micro/memory_helpers.h"
2325
#include "tensorflow/lite/micro/micro_common.h"
2426
#include "tensorflow/lite/micro/micro_log.h"
2527
#include "tensorflow/lite/micro/micro_utils.h"
@@ -125,18 +127,50 @@ void* MicroContext::DecompressTensorToBuffer(
125127
return nullptr;
126128
}
127129

130+
#endif // USE_TFLM_COMPRESSION
131+
128132
TfLiteStatus MicroContext::SetDecompressionMemory(
129133
const std::initializer_list<AlternateMemoryRegion>& regions) {
130-
return kTfLiteError;
134+
if (decompress_regions_ != nullptr) {
135+
return kTfLiteError;
136+
}
137+
138+
decompress_regions_ = &regions;
139+
decompress_regions_allocations_ = static_cast<size_t*>(
140+
AllocatePersistentBuffer(sizeof(size_t) * regions.size()));
141+
if (decompress_regions_allocations_ == nullptr) {
142+
return kTfLiteError;
143+
}
144+
ResetDecompressionMemoryAllocations();
145+
146+
return kTfLiteOk;
131147
}
132148

133149
void* MicroContext::AllocateDecompressionMemory(size_t bytes,
134150
size_t alignment) {
151+
if (decompress_regions_ != nullptr) {
152+
for (size_t i = 0; i < decompress_regions_->size(); i++) {
153+
const AlternateMemoryRegion* region = &decompress_regions_->begin()[i];
154+
uint8_t* start = static_cast<uint8_t*>(region->address) +
155+
decompress_regions_allocations_[i];
156+
uint8_t* aligned_start = AlignPointerUp(start, alignment);
157+
size_t total = bytes + (aligned_start - start);
158+
if (total + decompress_regions_allocations_[i] <= region->bytes) {
159+
decompress_regions_allocations_[i] += total;
160+
return aligned_start;
161+
}
162+
}
163+
}
164+
135165
return nullptr;
136166
}
137167

138-
void MicroContext::ResetDecompressionMemoryAllocations() {}
139-
140-
#endif // USE_TFLM_COMPRESSION
168+
void MicroContext::ResetDecompressionMemoryAllocations() {
169+
if (decompress_regions_ == nullptr) {
170+
return;
171+
}
172+
TFLITE_DCHECK(decompress_regions_allocations_ != nullptr);
173+
std::fill_n(decompress_regions_allocations_, decompress_regions_->size(), 0);
174+
}
141175

142176
} // namespace tflite

‎tensorflow/lite/micro/micro_context.h‎

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,15 @@ limitations under the License.
1616
#ifndef TENSORFLOW_LITE_MICRO_MICRO_CONTEXT_H_
1717
#define TENSORFLOW_LITE_MICRO_MICRO_CONTEXT_H_
1818

19+
#include <cstddef>
20+
#include <initializer_list>
21+
1922
#include "tensorflow/lite/c/common.h"
2023
#include "tensorflow/lite/micro/micro_graph.h"
2124
#include "tensorflow/lite/micro/micro_profiler_interface.h"
2225

2326
#ifdef USE_TFLM_COMPRESSION
2427

25-
#include <initializer_list>
26-
2728
#include "tensorflow/lite/micro/compression.h"
2829

2930
#endif // USE_TFLM_COMPRESSION
@@ -126,6 +127,8 @@ class MicroContext {
126127
const TfLiteEvalTensor& tensor,
127128
const CompressionTensorData& compression_data, void* buffer);
128129

130+
#endif // USE_TFLM_COMPRESSION
131+
129132
// Used for configuring alternate decompression memory
130133
struct AlternateMemoryRegion {
131134
void* address;
@@ -140,14 +143,13 @@ class MicroContext {
140143
// Return a pointer to memory that can be used for decompression.
141144
// The pointer will be aligned to the <alignment> value.
142145
// Return nullptr if the requested size is not available.
143-
// Can be called during kPrepare and kInvoke states.
146+
// Can be called during kPrepare state.
144147
virtual void* AllocateDecompressionMemory(size_t bytes, size_t alignment);
145148

146-
// reset all allocation tracking
149+
// Reset all allocation tracking.
150+
// Can be called during kPrepare state.
147151
virtual void ResetDecompressionMemoryAllocations();
148152

149-
#endif // USE_TFLM_COMPRESSION
150-
151153
// Set the alternate MicroProfilerInterface.
152154
// This can be used to profile subsystems simultaneously with the profiling
153155
// of kernels during the Eval phase. See (b/379584353).
@@ -168,6 +170,11 @@ class MicroContext {
168170
}
169171

170172
private:
173+
const std::initializer_list<AlternateMemoryRegion>* decompress_regions_ =
174+
nullptr;
175+
// array of size_t elements with length equal to decompress_regions_.size()
176+
size_t* decompress_regions_allocations_ = nullptr;
177+
171178
TF_LITE_REMOVE_VIRTUAL_DELETE
172179
};
173180

‎tensorflow/lite/micro/micro_interpreter.cc‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -339,14 +339,10 @@ TfLiteStatus MicroInterpreter::SetAlternateProfiler(
339339
return micro_context_.SetAlternateProfiler(alt_profiler);
340340
}
341341

342-
#ifdef USE_TFLM_COMPRESSION
343-
344342
TfLiteStatus MicroInterpreter::SetDecompressionMemory(
345343
const std::initializer_list<MicroInterpreterContext::AlternateMemoryRegion>&
346344
regions) {
347345
return micro_context_.SetDecompressionMemory(regions);
348346
}
349347

350-
#endif // USE_TFLM_COMPRESSION
351-
352348
} // namespace tflite

‎tensorflow/lite/micro/micro_interpreter.h‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -160,17 +160,13 @@ class MicroInterpreter {
160160
// decompression subsystem.
161161
TfLiteStatus SetAlternateProfiler(MicroProfilerInterface* alt_profiler);
162162

163-
#ifdef USE_TFLM_COMPRESSION
164-
165163
// Set the alternate decompression memory regions.
166164
// Can only be called during the MicroInterpreter kInit state (i.e. must
167165
// be called before MicroInterpreter::AllocateTensors).
168166
TfLiteStatus SetDecompressionMemory(
169167
const std::initializer_list<MicroContext::AlternateMemoryRegion>&
170168
regions);
171169

172-
#endif // USE_TFLM_COMPRESSION
173-
174170
protected:
175171
const MicroAllocator& allocator() const { return allocator_; }
176172
const TfLiteContext& context() const { return context_; }

0 commit comments

Comments
 (0)