Support dynamic tensor shape in LiteRt Inference Runner.

LiteRT-PiperOrigin-RevId: 815896283

Author: ai-edge-bot

litert/c/BUILD

@@ -442,6 +442,7 @@ cc_library(
     deps = [
         ":litert_common",
         ":litert_environment",
+        ":litert_layout",
         ":litert_logging",
         ":litert_metrics",
         ":litert_model",
@@ -468,11 +469,13 @@ cc_test(
         ":litert_compiled_model",
         ":litert_environment",
         ":litert_environment_options",
+        ":litert_layout",
         ":litert_logging",
         ":litert_model",
         ":litert_options",
         ":litert_tensor_buffer",
         "//litert/test:common",
+        "//litert/test:matchers",
         "//litert/test:simple_model",
         "@com_google_absl//absl/log:absl_log",
         "@com_google_absl//absl/strings:string_view",
@@ -493,6 +496,7 @@ cc_test(
         ":litert_environment_options",
         ":litert_runtime_c_api_shared_lib",
         "//litert/test:common",
+        "//litert/test:matchers",
         "//litert/test:simple_model",
         "@com_google_absl//absl/log:absl_log",
         "@com_google_absl//absl/strings:string_view",

litert/c/litert_compiled_model.cc

@@ -27,6 +27,7 @@
 #include "absl/types/span.h"  // from @com_google_absl
 #include "litert/c/litert_common.h"
 #include "litert/c/litert_environment.h"
+#include "litert/c/litert_layout.h"
 #include "litert/c/litert_logging.h"
 #include "litert/c/litert_metrics.h"
 #include "litert/c/litert_model.h"
@@ -96,6 +97,23 @@ LiteRtStatus LiteRtGetCompiledModelEnvironment(
   return kLiteRtStatusOk;
 }
 
+LiteRtStatus LiteRtGetCompiledModelOutputTensorLayouts(
+    LiteRtCompiledModel compiled_model, LiteRtParamIndex signature_index,
+    size_t num_layouts, LiteRtLayout* layouts, bool update_allocation) {
+  if (!compiled_model || !layouts) {
+    return kLiteRtStatusErrorInvalidArgument;
+  }
+  absl::Span<LiteRtLayout> output_layouts(layouts, num_layouts);
+  LITERT_RETURN_IF_ERROR(compiled_model->GetOutputTensorShapes(
+      signature_index, output_layouts, update_allocation));
+  size_t tensors_size = output_layouts.size();
+  if (tensors_size == 0) {
+    LITERT_LOG(LITERT_WARNING, "No output tensors found for signature index.");
+    return kLiteRtStatusErrorInvalidArgument;
+  }
+  return kLiteRtStatusOk;
+}
+
 LiteRtStatus LiteRtRunCompiledModel(LiteRtCompiledModel compiled_model,
                                     LiteRtParamIndex signature_index,
                                     size_t num_input_buffers,

litert/c/litert_compiled_model.h

@@ -18,6 +18,7 @@
 #include <stddef.h>
 
 #include "litert/c/litert_common.h"
+#include "litert/c/litert_layout.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -85,6 +86,23 @@ LiteRtStatus LiteRtGetCompiledModelOutputBufferRequirements(
     LiteRtParamIndex output_index,
     LiteRtTensorBufferRequirements* buffer_requirements);
 
+// Returns the tensor layouts for all output tensors.
+//
+// Parameters:
+// - compiled_model: the target `LiteRtCompiledModel` object.
+// - signature_index: the index of the signature in `LiteRtModel`.
+// - num_layouts: the number of output tensor layouts.
+// - layouts: user allocated memory to store `LiteRtLayout` for tensor outputs.
+// - update_allocation: whether to update the tensor allocation. Set to true
+//   for dynamic models after resize input tensors.
+//
+// Note: This function usually should be called after resizing input tensors
+// to get the new output tensor layouts. User should be responsible for
+// allocation and deallocating of the layouts memory.
+LiteRtStatus LiteRtGetCompiledModelOutputTensorLayouts(
+    LiteRtCompiledModel compiled_model, LiteRtParamIndex signature_index,
+    size_t num_layouts, LiteRtLayout* layouts, bool update_allocation);
+
 // Returns the associated environment of the given compiled model.
 LiteRtStatus LiteRtGetCompiledModelEnvironment(
     LiteRtCompiledModel compiled_model, LiteRtEnvironment* environment);
@@ -183,7 +201,9 @@ LiteRtStatus LiteRtCompiledModelGetProfiler(LiteRtCompiledModel compiled_model,
 // - dims: A span containing the new dimensions for the input tensor.
 //
 // Note: After resizing, the previously obtained buffer requirements may be
-// invalidated. Callers should re-query buffer requirements if needed.
+// invalidated. Callers should re-query buffer requirements if needed. After
+// resizing, LiteRtGetCompiledModelAllOutputTensorLayouts can be used to get
+// the new output tensor layouts.
 //
 // Returns:
 // - kLiteRtStatusOk: Success.