Port the circular buffer kernel to the new TfLiteEvalTensor API.
PiperOrigin-RevId: 322701794
Change-Id: Ib3d39f12747791760ed03d9ad39694660d05522d
diff --git a/tensorflow/lite/micro/kernels/BUILD b/tensorflow/lite/micro/kernels/BUILD
index a9b371b..c98e23a 100644
--- a/tensorflow/lite/micro/kernels/BUILD
+++ b/tensorflow/lite/micro/kernels/BUILD
@@ -585,6 +585,7 @@
"circular_buffer_test.cc",
],
deps = [
+ ":kernel_runner",
":micro_ops",
"//tensorflow/lite/c:common",
"//tensorflow/lite/micro:op_resolvers",
diff --git a/tensorflow/lite/micro/kernels/circular_buffer.cc b/tensorflow/lite/micro/kernels/circular_buffer.cc
index f588d64..876ea56 100644
--- a/tensorflow/lite/micro/kernels/circular_buffer.cc
+++ b/tensorflow/lite/micro/kernels/circular_buffer.cc
@@ -22,6 +22,7 @@
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/kernels/op_macros.h"
+#include "tensorflow/lite/micro/kernels/kernel_util.h"
/*
* The circular buffer custom operator is used to implement strided streaming
@@ -121,8 +122,10 @@
}
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
- const TfLiteTensor* input = GetInput(context, node, kInputTensor);
- TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+ const TfLiteEvalTensor* input =
+ tflite::micro::GetEvalInput(context, node, kInputTensor);
+ TfLiteEvalTensor* output =
+ tflite::micro::GetEvalOutput(context, node, kOutputTensor);
OpData* data = reinterpret_cast<OpData*>(node->user_data);
@@ -130,8 +133,8 @@
int depth = output->dims->data[3];
if (input->type == kTfLiteInt8) {
- EvalInt8(GetTensorData<int8_t>(input), num_slots, depth,
- GetTensorData<int8_t>(output));
+ EvalInt8(tflite::micro::GetTensorData<int8_t>(input), num_slots, depth,
+ tflite::micro::GetTensorData<int8_t>(output));
} else {
TF_LITE_KERNEL_LOG(context, "Type %s (%d) not supported.",
TfLiteTypeGetName(input->type), input->type);
diff --git a/tensorflow/lite/micro/kernels/circular_buffer_test.cc b/tensorflow/lite/micro/kernels/circular_buffer_test.cc
index 4c48060..770f456 100644
--- a/tensorflow/lite/micro/kernels/circular_buffer_test.cc
+++ b/tensorflow/lite/micro/kernels/circular_buffer_test.cc
@@ -16,6 +16,7 @@
#include "tensorflow/lite/c/builtin_op_data.h"
#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/micro/all_ops_resolver.h"
+#include "tensorflow/lite/micro/kernels/kernel_runner.h"
#include "tensorflow/lite/micro/kernels/micro_ops.h"
#include "tensorflow/lite/micro/testing/micro_test.h"
#include "tensorflow/lite/micro/testing/test_utils.h"
@@ -29,93 +30,6 @@
// TODO(b/149795762): Add this to TfLiteStatus enum.
constexpr int kTfLiteAbort = -9;
-TfLiteNode PrepareCircularBufferInt8(const int* input_dims_data,
- const int8_t* input_data,
- const int* output_dims_data,
- const int8_t* expected_output_data,
- int8_t* output_data) {
- const TfLiteRegistration* registration =
- ops::micro::Register_CIRCULAR_BUFFER();
-
- TfLiteNode node;
- TfLiteIntArray* input_dims = IntArrayFromInts(input_dims_data);
- TfLiteIntArray* output_dims = IntArrayFromInts(output_dims_data);
- constexpr int inputs_size = 2;
- constexpr int outputs_size = 1;
- constexpr int tensors_size = inputs_size + outputs_size;
- TfLiteTensor tensors[tensors_size] = {
- CreateQuantizedTensor(input_data, input_dims, 1, 0),
- CreateQuantizedTensor(output_data, output_dims, 1, 0),
- };
- TfLiteContext context;
- PopulateContext(tensors, tensors_size, micro_test::reporter, &context);
-
- // There is one input - tensor 0.
- const int inputs_array_data[] = {1, 0};
- TfLiteIntArray* inputs_array = IntArrayFromInts(inputs_array_data);
- // There is one output - tensor 1.
- const int outputs_array_data[] = {1, 1};
- TfLiteIntArray* outputs_array = IntArrayFromInts(outputs_array_data);
-
- node.inputs = inputs_array;
- node.outputs = outputs_array;
- node.builtin_data = nullptr;
- node.custom_initial_data = nullptr;
- node.custom_initial_data_size = 0;
-
- TF_LITE_MICRO_EXPECT_NE(nullptr, registration->prepare);
- TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->prepare(&context, &node));
- return node;
-}
-
-// Run invoke cycles_until_output times with the supplied input, expecting
-// invoke to return kTfLiteAbort until the last iteration, at which point the
-// output should match expected_output_data.
-TfLiteStatus InvokeCircularBufferInt8(const int* input_dims_data,
- const int8_t* input_data,
- const int* output_dims_data,
- const int8_t* expected_output_data,
- int8_t* output_data, TfLiteNode* node) {
- TfLiteIntArray* output_dims = IntArrayFromInts(output_dims_data);
- TfLiteIntArray* input_dims = IntArrayFromInts(input_dims_data);
-
- const int output_dims_count = ElementCount(*output_dims);
- const TfLiteRegistration* registration =
- ops::micro::Register_CIRCULAR_BUFFER();
-
- constexpr int inputs_size = 2;
- constexpr int outputs_size = 1;
- constexpr int tensors_size = inputs_size + outputs_size;
- TfLiteTensor tensors[tensors_size] = {
- CreateQuantizedTensor(input_data, input_dims, 1, 0),
- CreateQuantizedTensor(output_data, output_dims, 1, 0),
- };
- TfLiteContext context;
- PopulateContext(tensors, tensors_size, micro_test::reporter, &context);
-
- // There is one input - tensor 0.
- const int inputs_array_data[] = {1, 0};
- TfLiteIntArray* inputs_array = IntArrayFromInts(inputs_array_data);
- // There is one output - tensor 1.
- const int outputs_array_data[] = {1, 1};
- TfLiteIntArray* outputs_array = IntArrayFromInts(outputs_array_data);
-
- node->inputs = inputs_array;
- node->outputs = outputs_array;
- node->builtin_data = nullptr;
- node->custom_initial_data = nullptr;
- node->custom_initial_data_size = 0;
-
- TF_LITE_MICRO_EXPECT_NE(nullptr, registration->invoke);
-
- TfLiteStatus status = registration->invoke(&context, node);
-
- for (int i = 0; i < output_dims_count; ++i) {
- TF_LITE_MICRO_EXPECT_EQ(expected_output_data[i], output_data[i]);
- }
- return status;
-}
-
} // namespace
} // namespace testing
} // namespace tflite
@@ -125,30 +39,65 @@
TF_LITE_MICRO_TEST(OutputTensorLength4) {
constexpr int depth = 3;
constexpr int num_slots = 4;
+ int8_t input_data[depth];
int8_t output_data[depth * num_slots];
memset(output_data, 0, sizeof(output_data));
+
// There are four input dimensions - [1, 1, 1, depth].
const int input_dims[] = {4, 1, 1, 1, depth};
// There are four output dimensions - [1, num_slots, 1, depth].
const int output_dims[] = {4, 1, num_slots, 1, depth};
+ TfLiteIntArray* input_tensor_dims =
+ tflite::testing::IntArrayFromInts(input_dims);
+ TfLiteIntArray* output_tensor_dims =
+ tflite::testing::IntArrayFromInts(output_dims);
+
+ const int output_dims_count = tflite::ElementCount(*output_tensor_dims);
+
+ constexpr int inputs_size = 2;
+ constexpr int outputs_size = 1;
+ constexpr int tensors_size = inputs_size + outputs_size;
+ TfLiteTensor tensors[tensors_size] = {
+ tflite::testing::CreateQuantizedTensor(input_data, input_tensor_dims, 1,
+ 0),
+ tflite::testing::CreateQuantizedTensor(output_data, output_tensor_dims, 1,
+ 0),
+ };
+
+ // There is one input - tensor 0.
+ const int inputs_array_data[] = {1, 0};
+ TfLiteIntArray* inputs_array =
+ tflite::testing::IntArrayFromInts(inputs_array_data);
+ // There is one output - tensor 1.
+ const int outputs_array_data[] = {1, 1};
+ TfLiteIntArray* outputs_array =
+ tflite::testing::IntArrayFromInts(outputs_array_data);
+
+ const TfLiteRegistration* registration =
+ tflite::ops::micro::Register_CIRCULAR_BUFFER();
+ tflite::micro::KernelRunner runner = tflite::micro::KernelRunner(
+ *registration, tensors, tensors_size, inputs_array, outputs_array,
+ /*builtin_data=*/nullptr, micro_test::reporter);
+ TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, runner.InitAndPrepare());
+
const int8_t goldens[5][16] = {{0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3},
{0, 0, 0, 0, 0, 0, 1, 2, 3, 4, 5, 6},
{0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9},
{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12},
{4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}};
- int8_t input[depth];
- TfLiteNode node = tflite::testing::PrepareCircularBufferInt8(
- input_dims, input, output_dims, goldens[0], output_data);
// Expect the circular buffer to run every other invoke for 4xN output.
for (int i = 0; i < 5; i++) {
for (int j = 0; j < depth; j++) {
- input[j] = i * depth + j + 1;
+ input_data[j] = i * depth + j + 1;
}
- TfLiteStatus status = tflite::testing::InvokeCircularBufferInt8(
- input_dims, input, output_dims, goldens[i], output_data, &node);
+ TfLiteStatus status = runner.Invoke();
+
+ for (int j = 0; j < output_dims_count; ++j) {
+ TF_LITE_MICRO_EXPECT_EQ(goldens[i][j], output_data[j]);
+ }
// Every kRunPeriod iterations, the circular buffer should return kTfLiteOk.
if (i % tflite::testing::kRunPeriod == tflite::testing::kRunPeriod - 1) {
@@ -162,11 +111,44 @@
TF_LITE_MICRO_TEST(OutputTensorLength5) {
constexpr int depth = 4;
constexpr int num_slots = 5;
+ int8_t input_data[depth];
int8_t output_data[depth * num_slots];
memset(output_data, 0, sizeof(output_data));
const int input_dims[] = {4, 1, 1, 1, depth};
const int output_dims[] = {4, 1, num_slots, 1, depth};
+ TfLiteIntArray* input_tensor_dims =
+ tflite::testing::IntArrayFromInts(input_dims);
+ TfLiteIntArray* output_tensor_dims =
+ tflite::testing::IntArrayFromInts(output_dims);
+
+ const int output_dims_count = tflite::ElementCount(*output_tensor_dims);
+
+ constexpr int inputs_size = 2;
+ constexpr int outputs_size = 1;
+ constexpr int tensors_size = inputs_size + outputs_size;
+ TfLiteTensor tensors[tensors_size] = {
+ tflite::testing::CreateQuantizedTensor(input_data, input_tensor_dims, 1,
+ 0),
+ tflite::testing::CreateQuantizedTensor(output_data, output_tensor_dims, 1,
+ 0),
+ };
+
+ // There is one input - tensor 0.
+ const int inputs_array_data[] = {1, 0};
+ TfLiteIntArray* inputs_array =
+ tflite::testing::IntArrayFromInts(inputs_array_data);
+ // There is one output - tensor 1.
+ const int outputs_array_data[] = {1, 1};
+ TfLiteIntArray* outputs_array =
+ tflite::testing::IntArrayFromInts(outputs_array_data);
+
+ const TfLiteRegistration* registration =
+ tflite::ops::micro::Register_CIRCULAR_BUFFER();
+ tflite::micro::KernelRunner runner = tflite::micro::KernelRunner(
+ *registration, tensors, tensors_size, inputs_array, outputs_array,
+ /*builtin_data=*/nullptr, micro_test::reporter);
+ TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, runner.InitAndPrepare());
const int8_t goldens[6][20] = {
{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4},
@@ -177,18 +159,16 @@
{5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
15, 16, 17, 18, 19, 20, 21, 22, 23, 24}};
- int8_t input[depth];
- TfLiteNode node = tflite::testing::PrepareCircularBufferInt8(
- input_dims, input, output_dims, goldens[0], output_data);
// Expect circular buffer to run every cycle for 5xN output.
for (int i = 0; i < 6; i++) {
for (int j = 0; j < depth; j++) {
- input[j] = i * depth + j + 1;
+ input_data[j] = i * depth + j + 1;
}
- TF_LITE_MICRO_EXPECT_EQ(
- kTfLiteOk,
- tflite::testing::InvokeCircularBufferInt8(
- input_dims, input, output_dims, goldens[i], output_data, &node));
+ TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, runner.Invoke());
+
+ for (int j = 0; j < output_dims_count; ++j) {
+ TF_LITE_MICRO_EXPECT_EQ(goldens[i][j], output_data[j]);
+ }
}
}
