tflite: Add tflite aligned buffer allocator helper Since sync_driver requires the buffer to be aligned when fed in to it. It would be nice to have a helper function to help create these aligned buffer, so we don't need to manually copy those buffer allocation and deallocation everywhere. BUG=b:374245241 TEST=run sync_driver_test on host with sample delegate Change-Id: Ifc7b9c9991344f4f6a8e1388fafad4e55769d25c Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/tflite/+/6367801 Reviewed-by: Shik Chen <shik@chromium.org> Commit-Queue: ChromeOS Auto Retry <chromeos-auto-retry@chromeos-bot.iam.gserviceaccount.com> Tested-by: Tommy Chiang <ototot@google.com>
diff --git a/common/BUILD.bazel b/common/BUILD.bazel index a7ad4c8..d252424 100644 --- a/common/BUILD.bazel +++ b/common/BUILD.bazel
@@ -196,6 +196,7 @@ deps = [ ":simple_model_builder", ":sync_driver", + "//common/test_util:aligned_tflite_buffer_allocator", "//common/test_util:stable_delegate_env", "@com_google_absl//absl/random", "@com_google_googletest//:gtest",
diff --git a/common/sync_driver_test.cc b/common/sync_driver_test.cc index 3cbc877..64860a0 100644 --- a/common/sync_driver_test.cc +++ b/common/sync_driver_test.cc
@@ -11,6 +11,7 @@ #include "absl/random/random.h" #include "common/simple_model_builder.h" #include "common/sync_driver.h" +#include "common/test_util/aligned_tflite_buffer_allocator.h" #include "common/test_util/stable_delegate_env.h" #include "tensorflow/lite/core/c/builtin_op_data.h" @@ -75,28 +76,12 @@ n *= dim; } - auto deallocate_tflite_aligned_tensor_buffer = [](uint8_t* ptr) { - operator delete[](ptr, std::align_val_t{tflite::kDefaultTensorAlignment}); - }; - auto allocate_tflite_aligned_tensor_buffer = [](int size) { - return new (std::align_val_t{tflite::kDefaultTensorAlignment}) - uint8_t[size]; - }; - const int num_iterations = GetParam(); for (int iteration = 0; iteration < num_iterations; ++iteration) { - auto a_buffer = std::shared_ptr<uint8_t[]>( - allocate_tflite_aligned_tensor_buffer(n * sizeof(float)), - deallocate_tflite_aligned_tensor_buffer); - auto b_buffer = std::shared_ptr<uint8_t[]>( - allocate_tflite_aligned_tensor_buffer(n * sizeof(float)), - deallocate_tflite_aligned_tensor_buffer); - auto c_buffer = std::shared_ptr<uint8_t[]>( - allocate_tflite_aligned_tensor_buffer(n * sizeof(float)), - deallocate_tflite_aligned_tensor_buffer); - auto d_buffer = std::shared_ptr<uint8_t[]>( - allocate_tflite_aligned_tensor_buffer(n * sizeof(float)), - deallocate_tflite_aligned_tensor_buffer); + auto a_buffer = AllocateTfLiteAlignedTensorBuffer(n * sizeof(float)); + auto b_buffer = AllocateTfLiteAlignedTensorBuffer(n * sizeof(float)); + auto c_buffer = AllocateTfLiteAlignedTensorBuffer(n * sizeof(float)); + auto d_buffer = AllocateTfLiteAlignedTensorBuffer(n * sizeof(float)); driver->SetInputTensorBuffer("a", a_buffer); driver->SetInputTensorBuffer("b", b_buffer); driver->SetInputTensorBuffer("c", c_buffer);
diff --git a/common/test_util/BUILD.bazel b/common/test_util/BUILD.bazel index 2a14924..b1be3ba 100644 --- a/common/test_util/BUILD.bazel +++ b/common/test_util/BUILD.bazel
@@ -30,3 +30,13 @@ "@org_tensorflow//tensorflow/lite/tools/delegates/experimental/stable_delegate:delegate_provider", ], ) + +cc_library( + name = "aligned_tflite_buffer_allocator", + testonly = True, + srcs = ["aligned_tflite_buffer_allocator.cc"], + hdrs = ["aligned_tflite_buffer_allocator.h"], + deps = [ + "@org_tensorflow//tensorflow/lite:util", + ], +)
diff --git a/common/test_util/aligned_tflite_buffer_allocator.cc b/common/test_util/aligned_tflite_buffer_allocator.cc new file mode 100644 index 0000000..a27deb8 --- /dev/null +++ b/common/test_util/aligned_tflite_buffer_allocator.cc
@@ -0,0 +1,30 @@ +/* + * Copyright 2025 The ChromiumOS Authors + * Use of this source code is governed by a BSD-style license that can be + * found in the LICENSE file. + */ + +#include "common/test_util/aligned_tflite_buffer_allocator.h" + +#include <cstdint> +#include <memory> + +#include "tensorflow/lite/util.h" + +namespace { + +void DeallocateTfLiteAlignedTensorBufferHelper(uint8_t* ptr) { + operator delete[](ptr, std::align_val_t{tflite::kDefaultTensorAlignment}); +} + +uint8_t* AllocateTfLiteAlignedTensorBufferHelper(int bytes) { + return new (std::align_val_t{tflite::kDefaultTensorAlignment}) uint8_t[bytes]; +} + +} // namespace + +std::shared_ptr<uint8_t[]> AllocateTfLiteAlignedTensorBuffer(size_t bytes) { + return std::shared_ptr<uint8_t[]>( + AllocateTfLiteAlignedTensorBufferHelper(bytes), + DeallocateTfLiteAlignedTensorBufferHelper); +}
diff --git a/common/test_util/aligned_tflite_buffer_allocator.h b/common/test_util/aligned_tflite_buffer_allocator.h new file mode 100644 index 0000000..d3408c6 --- /dev/null +++ b/common/test_util/aligned_tflite_buffer_allocator.h
@@ -0,0 +1,17 @@ +/* + * Copyright 2025 The ChromiumOS Authors + * Use of this source code is governed by a BSD-style license that can be + * found in the LICENSE file. + */ + +#ifndef COMMON_TEST_UTIL_ALIGNED_TFLITE_BUFFER_ALLOCATOR_H_ +#define COMMON_TEST_UTIL_ALIGNED_TFLITE_BUFFER_ALLOCATOR_H_ + +#include <cstdint> +#include <memory> + +// Allocate a buffer with the size of |bytes| and aligned to the alignment +// requirement of tflite. +std::shared_ptr<uint8_t[]> AllocateTfLiteAlignedTensorBuffer(size_t bytes); + +#endif // COMMON_TEST_UTIL_ALIGNED_TFLITE_BUFFER_ALLOCATOR_H_