feat: add fixed-capacity, statically-allocated type tflite::StaticVector (#2642)
Add a type, tflite::StaticVector, which behaves like std::vector, but
which avoids heap memory allocation.
BUG=#2636
diff --git a/tensorflow/lite/micro/BUILD b/tensorflow/lite/micro/BUILD
index 527b85a..97d6897 100644
--- a/tensorflow/lite/micro/BUILD
+++ b/tensorflow/lite/micro/BUILD
@@ -408,6 +408,15 @@
)
cc_library(
+ name = "static_vector",
+ hdrs = ["static_vector.h"],
+ copts = micro_copts(),
+ deps = [
+ "//tensorflow/lite/kernels:op_macros",
+ ],
+)
+
+cc_library(
name = "system_setup",
srcs = [
"system_setup.cc",
@@ -616,6 +625,18 @@
],
)
+cc_test(
+ name = "static_vector_test",
+ size = "small",
+ srcs = [
+ "static_vector_test.cc",
+ ],
+ deps = [
+ ":static_vector",
+ "//tensorflow/lite/micro/testing:micro_test",
+ ],
+)
+
bzl_library(
name = "build_def_bzl",
srcs = ["build_def.bzl"],
diff --git a/tensorflow/lite/micro/static_vector.h b/tensorflow/lite/micro/static_vector.h
new file mode 100644
index 0000000..8b9e063
--- /dev/null
+++ b/tensorflow/lite/micro/static_vector.h
@@ -0,0 +1,83 @@
+// 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_STATIC_VECTOR_H_
+#define TENSORFLOW_LITE_MICRO_STATIC_VECTOR_H_
+
+#include <array>
+#include <cassert>
+#include <cstddef>
+
+#include "tensorflow/lite/kernels/op_macros.h" // for TF_LITE_ASSERT
+
+namespace tflite {
+
+template <typename T, std::size_t MaxSize>
+class StaticVector {
+ // A staticlly-allocated vector. Add to the interface as needed.
+
+ private:
+ std::array<T, MaxSize> array_;
+ std::size_t size_{0};
+
+ public:
+ using iterator = typename decltype(array_)::iterator;
+ using const_iterator = typename decltype(array_)::const_iterator;
+ using pointer = typename decltype(array_)::pointer;
+ using reference = typename decltype(array_)::reference;
+ using const_reference = typename decltype(array_)::const_reference;
+
+ StaticVector() {}
+
+ StaticVector(std::initializer_list<T> values) {
+ for (const T& v : values) {
+ push_back(v);
+ }
+ }
+
+ static constexpr std::size_t max_size() { return MaxSize; }
+ std::size_t size() const { return size_; }
+ bool full() const { return size() == max_size(); }
+ iterator begin() { return array_.begin(); }
+ const_iterator begin() const { return array_.begin(); }
+ iterator end() { return begin() + size(); }
+ const_iterator end() const { return begin() + size(); }
+ pointer data() { return array_.data(); }
+ reference operator[](int i) { return array_[i]; }
+ const_reference operator[](int i) const { return array_[i]; }
+ void clear() { size_ = 0; }
+
+ template <std::size_t N>
+ bool operator==(const StaticVector<T, N>& other) const {
+ return std::equal(begin(), end(), other.begin(), other.end());
+ }
+
+ template <std::size_t N>
+ bool operator!=(const StaticVector<T, N>& other) const {
+ return !(*this == other);
+ }
+
+ void push_back(const T& t) {
+ TF_LITE_ASSERT(!full());
+ *end() = t;
+ ++size_;
+ }
+};
+
+template <typename T, typename... U>
+StaticVector(T, U...) -> StaticVector<T, 1 + sizeof...(U)>;
+
+} // end namespace tflite
+
+#endif // TENSORFLOW_LITE_MICRO_STATIC_VECTOR_H_
diff --git a/tensorflow/lite/micro/static_vector_test.cc b/tensorflow/lite/micro/static_vector_test.cc
new file mode 100644
index 0000000..6d601bc
--- /dev/null
+++ b/tensorflow/lite/micro/static_vector_test.cc
@@ -0,0 +1,82 @@
+// 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.
+
+#include "tensorflow/lite/micro/static_vector.h"
+
+#include "tensorflow/lite/micro/testing/micro_test.h"
+
+using tflite::StaticVector;
+
+TF_LITE_MICRO_TESTS_BEGIN
+
+TF_LITE_MICRO_TEST(StaticVectorPushBack) {
+ StaticVector<int, 4> a;
+ TF_LITE_MICRO_EXPECT(a.max_size() == 4);
+ TF_LITE_MICRO_EXPECT(a.size() == 0);
+
+ a.push_back(1);
+ TF_LITE_MICRO_EXPECT(a.size() == 1);
+ TF_LITE_MICRO_EXPECT(a[0] == 1);
+
+ a.push_back(2);
+ TF_LITE_MICRO_EXPECT(a.size() == 2);
+ TF_LITE_MICRO_EXPECT(a[1] == 2);
+
+ a.push_back(3);
+ TF_LITE_MICRO_EXPECT(a.size() == 3);
+ TF_LITE_MICRO_EXPECT(a[2] == 3);
+}
+
+TF_LITE_MICRO_TEST(StaticVectorInitializationPartial) {
+ const StaticVector<int, 4> a{1, 2, 3};
+ TF_LITE_MICRO_EXPECT(a.max_size() == 4);
+ TF_LITE_MICRO_EXPECT(a.size() == 3);
+ TF_LITE_MICRO_EXPECT(a[0] == 1);
+ TF_LITE_MICRO_EXPECT(a[1] == 2);
+ TF_LITE_MICRO_EXPECT(a[2] == 3);
+}
+
+TF_LITE_MICRO_TEST(StaticVectorInitializationFull) {
+ const StaticVector b{1, 2, 3};
+ TF_LITE_MICRO_EXPECT(b.max_size() == 3);
+ TF_LITE_MICRO_EXPECT(b.size() == 3);
+}
+
+TF_LITE_MICRO_TEST(StaticVectorEquality) {
+ const StaticVector a{1, 2, 3};
+ const StaticVector b{1, 2, 3};
+ TF_LITE_MICRO_EXPECT(a == b);
+ TF_LITE_MICRO_EXPECT(!(a != b));
+}
+
+TF_LITE_MICRO_TEST(StaticVectorInequality) {
+ const StaticVector a{1, 2, 3};
+ const StaticVector b{3, 2, 1};
+ TF_LITE_MICRO_EXPECT(a != b);
+ TF_LITE_MICRO_EXPECT(!(a == b));
+}
+
+TF_LITE_MICRO_TEST(StaticVectorSizeInequality) {
+ const StaticVector a{1, 2};
+ const StaticVector b{1, 2, 3};
+ TF_LITE_MICRO_EXPECT(a != b);
+}
+
+TF_LITE_MICRO_TEST(StaticVectorPartialSizeInequality) {
+ const StaticVector<int, 3> a{1, 2};
+ const StaticVector<int, 3> b{1, 2, 3};
+ TF_LITE_MICRO_EXPECT(a != b);
+}
+
+TF_LITE_MICRO_TESTS_END
