Update CompiledModel::CreateInputBuffers() and CompiledModel::CreateOutputBuffers()

Fix them to respect alignment in the TensorBufferRequirements by using
LiteRtCreateManagedTensorBufferFromRequirements() C API.

This PR fixes #5373.

LiteRT-PiperOrigin-RevId: 893695151

Author: terryheo

litert/cc/BUILD

@@ -143,8 +143,8 @@ cc_library(
     hdrs = ["litert_expected.h"],
     compatible_with = get_compatible_with_portable(),
     deps = [
+        ":litert_common",
         "//litert/c:litert_common",
-        "//litert/cc:litert_common",
         "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
@@ -210,11 +210,11 @@ cc_library(
     srcs = ["litert_opaque_options.cc"],
     hdrs = ["litert_opaque_options.h"],
     deps = [
+        ":litert_common",
         ":litert_expected",
         ":litert_macros",
         "//litert/c:litert_common",
         "//litert/c:litert_opaque_options",
-        "//litert/cc:litert_common",
         "//litert/cc/internal:litert_handle",
         "@com_google_absl//absl/strings:string_view",
     ],
@@ -482,6 +482,15 @@ litert_device_test(
 #         "requires-gpu-nvidia",
 #     ],
 #     deps = [
+#         ":litert_common",
+#         ":litert_compiled_model",
+#         ":litert_environment",
+#         ":litert_event",
+#         ":litert_expected",
+#         ":litert_macros",
+#         ":litert_model",
+#         ":litert_options",
+#         ":litert_tensor_buffer",
 #         "@com_google_googletest//:gtest_main",
 #         "@com_google_absl//absl/debugging:leak_check",
 #         "@com_google_absl//absl/log:absl_log",
@@ -493,15 +502,6 @@ litert_device_test(
 #         "//litert/c:litert_event_type",
 #         "//litert/c:litert_profiler_event",
 #         "//litert/c:litert_tensor_buffer_types",
-#         "//litert/cc:litert_common",
-#         "//litert/cc:litert_compiled_model",
-#         "//litert/cc:litert_environment",
-#         "//litert/cc:litert_event",
-#         "//litert/cc:litert_expected",
-#         "//litert/cc:litert_macros",
-#         "//litert/cc:litert_model",
-#         "//litert/cc:litert_options",
-#         "//litert/cc:litert_tensor_buffer",
 #         "//litert/cc/internal:litert_platform_support",
 #         "//litert/cc/options:litert_gpu_options",
 #         "//litert/cc/options:litert_runtime_options",
@@ -629,14 +629,14 @@ cc_library(
         # copybara:comment_end
     ],
     deps = [
+        ":litert_common",
         ":litert_expected",
         ":litert_layout",
         "//litert/c:litert_common",
         "//litert/c:litert_custom_op_kernel",
         "//litert/c:litert_layout",
         "//litert/c:litert_tensor_buffer",
         "//litert/c/internal:litert_logging",
-        "//litert/cc:litert_common",
         "//litert/cc/internal:litert_handle",
         "//litert/cc/internal:litert_tensor_buffer_without_registry",
         "@com_google_absl//absl/types:span",
@@ -780,14 +780,14 @@ cc_library(
         # copybara:comment_end
     ],
     deps = [
+        ":litert_environment",
         ":litert_expected",
         ":litert_macros",
         "//litert/c:litert_common",
         "//litert/c:litert_event_type",
         "//litert/c:litert_gl_types",
         "//litert/c:litert_opencl_types",
         "//litert/c:litert_profiler_event",
-        "//litert/cc:litert_environment",
         "//litert/cc/internal:litert_handle",
     ],
 )
@@ -797,9 +797,9 @@ cc_test(
     srcs = ["litert_event_test.cc"],
     linkopts = litert_android_linkopts() + gles_linkopts(),
     deps = [
+        ":litert_environment",
         ":litert_event",
         "//litert/c:litert_event_type",
-        "//litert/cc:litert_environment",
         "//litert/cc/internal:litert_platform_support",
         "//litert/test:matchers",
         "@com_google_googletest//:gtest_main",
@@ -951,7 +951,6 @@ cc_library(
         ":litert_common",
         ":litert_expected",
         ":litert_tensor_buffer_types",
-        "//litert/c:litert_common",
         "//litert/c/internal:litert_logging",
         "@com_google_absl//absl/types:span",
     ],
@@ -984,11 +983,13 @@ cc_library(
         # copybara:comment_end
     ],
     deps = [
+        ":litert_common",
         ":litert_environment",
         ":litert_event",
         ":litert_expected",
         ":litert_macros",
         ":litert_ranked_tensor_type",
+        ":litert_tensor_buffer_requirements",
         ":litert_tensor_buffer_types",
         "//litert/c:litert_common",
         "//litert/c:litert_custom_tensor_buffer",
@@ -1000,7 +1001,6 @@ cc_library(
         "//litert/c:litert_tensor_buffer_types",
         "//litert/c:litert_webgpu_types",
         "//litert/c/internal:litert_tensor_buffer_registry",  # buildcleaner: keep
-        "//litert/cc:litert_common",
         "//litert/cc/internal:litert_handle",
         "//litert/cc/internal:litert_tensor_buffer_without_registry",
         "@com_google_absl//absl/cleanup",
@@ -1042,6 +1042,7 @@ cc_test(
 #         ":litert_model",
 #         ":litert_ranked_tensor_type",
 #         ":litert_tensor_buffer",
+#         ":litert_tensor_buffer_requirements",
 #         ":litert_tensor_buffer_types",
 #         "@com_google_googletest//:gtest_main",
 #         "@com_google_absl//absl/debugging:leak_check",

litert/cc/dynamic_runtime/BUILD

@@ -469,6 +469,7 @@ cc_library(
         "//litert/cc:litert_expected",
         "//litert/cc:litert_macros",
         "//litert/cc:litert_ranked_tensor_type",
+        "//litert/cc:litert_tensor_buffer_requirements",
         "//litert/cc:litert_tensor_buffer_types",
         "//litert/cc/internal:litert_handle",
         "@com_google_absl//absl/cleanup",
@@ -507,6 +508,7 @@ cc_test(
         "//litert/cc:litert_layout",
         "//litert/cc:litert_macros",
         "//litert/cc:litert_ranked_tensor_type",
+        "//litert/cc:litert_tensor_buffer_requirements",
         "//litert/cc:litert_tensor_buffer_types",
         "//litert/cc/internal:litert_handle",
         "//litert/cc/internal:litert_platform_support",

litert/cc/internal/BUILD

@@ -95,6 +95,7 @@ cc_library(
         "//litert/cc:litert_expected",
         "//litert/cc:litert_macros",
         "//litert/cc:litert_ranked_tensor_type",
+        "//litert/cc:litert_tensor_buffer_requirements",
         "//litert/cc:litert_tensor_buffer_types",
         "@com_google_absl//absl/cleanup",
         "@com_google_absl//absl/log:absl_check",

litert/cc/litert_compiled_model.cc

@@ -73,6 +73,10 @@ Expected<TensorBufferRequirements> ToTensorBufferRequirements(
   LITERT_RETURN_IF_ERROR(env.runtime->GetTensorBufferRequirementsBufferSize(
       litert_requirements, &buffer_size));
 
+  size_t alignment;
+  LITERT_RETURN_IF_ERROR(env.runtime->GetTensorBufferRequirementsAlignment(
+      litert_requirements, &alignment));
+
   int num_strides;
   const uint32_t* strides_ptr;
   LITERT_RETURN_IF_ERROR(env.runtime->GetTensorBufferRequirementsStrides(
@@ -82,10 +86,11 @@ Expected<TensorBufferRequirements> ToTensorBufferRequirements(
     strides.assign(strides_ptr, strides_ptr + num_strides);
   }
 
-  size_t alignment;
-  LITERT_RETURN_IF_ERROR(env.runtime->GetTensorBufferRequirementsAlignment(
-      litert_requirements, &alignment));
-
+  if (num_strides == 0 || strides[0] == 0) {
+    // Strides are not specified.
+    return TensorBufferRequirements::CreateWithAlignment(
+        absl::MakeConstSpan(supported_types), buffer_size, alignment);
+  }
   return TensorBufferRequirements::CreateWithAlignment(
       absl::MakeConstSpan(supported_types), buffer_size, alignment,
       absl::MakeConstSpan(strides));
@@ -347,20 +352,8 @@ Expected<size_t> CompiledModel::FindOutputIndex(
 Expected<TensorBuffer> CompiledModel::CreateBufferImpl(
     const Environment& env, const TensorBufferRequirements& buffer_requirements,
     const RankedTensorType& tensor_type) {
-  LITERT_ASSIGN_OR_RETURN(const std::vector<TensorBufferType>& supported_types,
-                          buffer_requirements.SupportedTypes());
-  if (supported_types.empty()) {
-    return Unexpected(Status::kErrorRuntimeFailure,
-                      "Input doesn't support any tensor buffer types");
-  }
-  // For simplicity we just pick the first supported tensor buffer type.
-  TensorBufferType tensor_buffer_type = supported_types[0];
-  LITERT_ASSIGN_OR_RETURN(size_t buffer_size, buffer_requirements.BufferSize());
-
-  LITERT_ASSIGN_OR_RETURN(TensorBuffer buffer, TensorBuffer::CreateManaged(
-                                                   env, tensor_buffer_type,
-                                                   tensor_type, buffer_size));
-  return buffer;
+  return TensorBuffer::CreateManagedFromRequirements(env, tensor_type,
+                                                     buffer_requirements);
 }
 
 Expected<TensorBuffer> CompiledModel::CreateInputOutputBuffer(

litert/cc/litert_compiled_model_test.cc

@@ -1153,8 +1153,11 @@ TEST(CompiledModelTest, GetBufferRequirementsDetailed) {
   EXPECT_GE(input_alignment, 0);
 
   LITERT_ASSERT_OK_AND_ASSIGN(auto input_strides, input_requirements.Strides());
-  EXPECT_EQ(input_strides.size(), 1);
-  EXPECT_EQ(input_strides[0], 0);
+
+  EXPECT_LE(input_strides.size(), 1);
+  if (input_strides.size() == 1) {
+    EXPECT_EQ(input_strides[0], 0);
+  }
 
   // Check output buffer requirements.
   LITERT_ASSERT_OK_AND_ASSIGN(
@@ -1175,8 +1178,10 @@ TEST(CompiledModelTest, GetBufferRequirementsDetailed) {
 
   LITERT_ASSERT_OK_AND_ASSIGN(auto output_strides,
                               output_requirements.Strides());
-  EXPECT_EQ(output_strides.size(), 1);
-  EXPECT_EQ(output_strides[0], 0);
+  EXPECT_LE(output_strides.size(), 1);
+  if (output_strides.size() == 1) {
+    EXPECT_EQ(output_strides[0], 0);
+  }
 }
 
 }  // namespace

litert/cc/litert_tensor_buffer.cc

@@ -15,7 +15,9 @@
 #include "litert/cc/litert_tensor_buffer.h"
 
 #include <cstddef>
+#include <vector>
 
+#include "absl/cleanup/cleanup.h"  // from @com_google_absl
 #include "litert/c/litert_common.h"
 #include "litert/c/litert_gl_types.h"
 #include "litert/c/litert_model_types.h"
@@ -28,10 +30,44 @@
 #include "litert/cc/litert_expected.h"
 #include "litert/cc/litert_macros.h"
 #include "litert/cc/litert_ranked_tensor_type.h"
+#include "litert/cc/litert_tensor_buffer_requirements.h"
 #include "litert/cc/litert_tensor_buffer_types.h"
 
 namespace litert {
 
+namespace {
+
+// Converts a `TensorBufferRequirements` C++ object to a
+// `LiteRtTensorBufferRequirements` C object. In compiled_model.cc, there is a
+// function named `ToTensorBufferRequirements` which converts a
+// `LiteRtTensorBufferRequirements` C object to a `TensorBufferRequirements`
+// C++ object.
+Expected<LiteRtTensorBufferRequirements> ToLiteRtTensorBufferRequirements(
+    const internal::EnvironmentHolder& env,
+    const TensorBufferRequirements& requirements) {
+  LITERT_ASSIGN_OR_RETURN(const auto supported_types,
+                          requirements.SupportedTypes());
+  LITERT_ASSIGN_OR_RETURN(const auto buffer_size, requirements.BufferSize());
+  LITERT_ASSIGN_OR_RETURN(const auto alignment, requirements.Alignment());
+  LITERT_ASSIGN_OR_RETURN(const auto strides, requirements.Strides());
+
+  std::vector<LiteRtTensorBufferType> litert_buffer_types;
+  litert_buffer_types.reserve(supported_types.size());
+  for (auto type : supported_types) {
+    litert_buffer_types.push_back(static_cast<LiteRtTensorBufferType>(type));
+  }
+
+  LiteRtTensorBufferRequirements litert_requirements;
+  LITERT_RETURN_IF_ERROR(
+      env.runtime->CreateTensorBufferRequirementsWithAlignment(
+          litert_buffer_types.size(), litert_buffer_types.data(), buffer_size,
+          strides.size(), strides.data(), alignment, &litert_requirements));
+
+  return litert_requirements;
+}
+
+}  // namespace
+
 Expected<TensorBuffer> TensorBuffer::Duplicate() const {
   if (!IsOwned()) {
     return Unexpected(Status::kErrorInvalidArgument,
@@ -53,6 +89,28 @@ Expected<TensorBuffer> TensorBuffer::CreateManaged(
   return TensorBuffer(env_holder, tensor_buffer, OwnHandle::kYes);
 }
 
+Expected<TensorBuffer> TensorBuffer::CreateManagedFromRequirements(
+    const Environment& env, const RankedTensorType& tensor_type,
+    const TensorBufferRequirements& requirements) {
+  LiteRtTensorBuffer tensor_buffer;
+  auto litert_tensor_type = static_cast<LiteRtRankedTensorType>(tensor_type);
+  auto env_holder = env.GetHolder();
+
+  LITERT_ASSIGN_OR_RETURN(
+      LiteRtTensorBufferRequirements litert_requirements,
+      ToLiteRtTensorBufferRequirements(env_holder, requirements));
+
+  auto cleanup = absl::MakeCleanup([&env_holder, litert_requirements] {
+    env_holder.runtime->DestroyTensorBufferRequirements(litert_requirements);
+  });
+
+  LITERT_RETURN_IF_ERROR(
+      env_holder.runtime->CreateManagedTensorBufferFromRequirements(
+          env_holder.handle, &litert_tensor_type, litert_requirements,
+          &tensor_buffer));
+  return TensorBuffer(env_holder, tensor_buffer, OwnHandle::kYes);
+}
+
 // TODO(terryheo): make this function not depend on Environment.
 Expected<TensorBuffer> TensorBuffer::CreateManagedHostMemory(
     const RankedTensorType& tensor_type, size_t buffer_size) {

litert/cc/litert_tensor_buffer.h

@@ -44,6 +44,8 @@
 
 namespace litert {
 
+class TensorBufferRequirements;
+
 /// @brief A C++ wrapper for `LiteRtTensorBuffer`, representing a tensor and
 /// its associated backing buffer.
 class TensorBuffer : public internal::BaseHandle<LiteRtTensorBuffer> {
@@ -64,6 +66,14 @@ class TensorBuffer : public internal::BaseHandle<LiteRtTensorBuffer> {
       const Environment& env, TensorBufferType buffer_type,
       const RankedTensorType& tensor_type, size_t buffer_size);
 
+  /// @brief Creates a managed `TensorBuffer` from requirements.
+  ///
+  /// The returned object is owned by the caller. It automatically selects
+  /// the best buffer type and applies required alignment and padding.
+  static Expected<TensorBuffer> CreateManagedFromRequirements(
+      const Environment& env, const RankedTensorType& tensor_type,
+      const TensorBufferRequirements& requirements);
+
   /// @brief Creates a managed host memory `TensorBuffer` using the default
   /// environment (if applicable).
   ///

litert/cc/litert_tensor_buffer_requirements.h

@@ -24,7 +24,6 @@
 
 #include "absl/types/span.h"  // from @com_google_absl
 #include "litert/c/internal/litert_logging.h"
-#include "litert/c/litert_common.h"
 #include "litert/cc/litert_common.h"
 #include "litert/cc/litert_expected.h"
 #include "litert/cc/litert_tensor_buffer_types.h"
@@ -84,6 +83,10 @@ class TensorBufferRequirements {
 
   Expected<size_t> BufferSize() const { return buffer_size_; }
 
+  /// @brief Returns the strides of the tensor buffer requirements.
+  ///
+  /// If the strides are not specified, either an empty span or a span with a
+  /// single element of 0 is returned, which is equivalent to no strides.
   Expected<absl::Span<const uint32_t>> Strides() const {
     return absl::MakeConstSpan(strides_);
   }