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_