Clear the cache as shape propagation is unpredictable, and make sure the size is updated through tensor buffer requirements

And enhance the testing to run full inference

LiteRT-PiperOrigin-RevId: 825765979

Author: ai-edge-bot

litert/cc/litert_compiled_model.cc

@@ -94,13 +94,11 @@ Expected<TensorBuffer> CompiledModel::CreateInputOutputBuffer(
       is_input ? model_.GetInputTensorType(signature_index, tensor_name)
                : model_.GetOutputTensorType(signature_index, tensor_name);
   LITERT_ASSIGN_OR_RETURN(RankedTensorType tensor_type, tensor_type_expected);
-  Expected<TensorBufferRequirements> buffer_requirements_expected =
-      is_input ? GetInputBufferRequirements(signature_index, tensor_name)
-               : GetOutputBufferRequirements(signature_index, tensor_name);
-
-  LITERT_ASSIGN_OR_RETURN(const TensorBufferRequirements& buffer_requirements,
-                          buffer_requirements_expected);
+  LITERT_ASSIGN_OR_RETURN(auto env, GetEnvironment());
   if (is_input) {
+    LITERT_ASSIGN_OR_RETURN(
+        TensorBufferRequirements buffer_requirements,
+        GetInputBufferRequirements(signature_index, tensor_name));
     LITERT_ASSIGN_OR_RETURN(size_t tensor_index,
                             FindInputIndex(signature_index, tensor_name));
     LiteRtLayout input_layout;
@@ -111,20 +109,28 @@ Expected<TensorBuffer> CompiledModel::CreateInputOutputBuffer(
       tensor_type = RankedTensorType(tensor_type.ElementType(),
                                      std::move(runtime_layout));
     }
+    return CreateBufferImpl(env, buffer_requirements, tensor_type);
   } else {
     const auto& dims = tensor_type.Layout().Dimensions();
     if (absl::c_find(dims, -1) != dims.end()) {
       LITERT_ASSIGN_OR_RETURN(size_t tensor_index,
                               FindOutputIndex(signature_index, tensor_name));
       LITERT_ASSIGN_OR_RETURN(
-          std::vector<Layout> output_layouts,
+          std::vector<Layout> runtime_layouts,
           GetOutputTensorLayouts(signature_index, /*update_allocation=*/true));
       tensor_type = RankedTensorType(tensor_type.ElementType(),
-                                     std::move(output_layouts[tensor_index]));
+                                     std::move(runtime_layouts[tensor_index]));
+      LITERT_ASSIGN_OR_RETURN(
+          const TensorBufferRequirements& refreshed_requirements,
+          GetOutputBufferRequirements(signature_index, tensor_name));
+      return CreateBufferImpl(env, refreshed_requirements, tensor_type);
+    } else {
+      LITERT_ASSIGN_OR_RETURN(
+          const TensorBufferRequirements& requirements,
+          GetOutputBufferRequirements(signature_index, tensor_name));
+      return CreateBufferImpl(env, requirements, tensor_type);
     }
   }
-  LITERT_ASSIGN_OR_RETURN(auto env, GetEnvironment());
-  return CreateBufferImpl(env, buffer_requirements, tensor_type);
 }
 
 Expected<std::vector<TensorBuffer>> CompiledModel::CreateInputOutputBuffers(

litert/cc/litert_compiled_model_test.cc

@@ -141,16 +141,33 @@ TEST(CompiledModelTest,
   absl::string_view signature_key = model.DefaultSignatureKey();
 
   const std::vector<int> resized_dims = {4, 2, 3};
-  LITERT_ASSERT_OK(compiled_model.ResizeInputTensor(
-      signature_key, "arg0", absl::MakeConstSpan(resized_dims)));
+  LITERT_ASSERT_OK_AND_ASSIGN(const auto& input_names,
+                              model.GetSignatureInputNames(signature_key));
+  absl::flat_hash_map<absl::string_view, TensorBuffer> input_map;
+  for (const auto& input_name : input_names) {
+    LITERT_ASSERT_OK(compiled_model.ResizeInputTensor(
+        signature_key, input_name, absl::MakeConstSpan(resized_dims)));
+    LITERT_ASSERT_OK_AND_ASSIGN(
+        TensorBuffer input_buffer,
+        compiled_model.CreateInputBuffer(signature_key, input_name));
+    LITERT_ASSERT_OK_AND_ASSIGN(RankedTensorType buffer_type,
+                                input_buffer.TensorType());
+    EXPECT_THAT(buffer_type.Layout().Dimensions(),
+                ElementsAre(resized_dims[0], resized_dims[1], resized_dims[2]));
+    input_map[input_name] = std::move(input_buffer);
+  }
 
   LITERT_ASSERT_OK_AND_ASSIGN(
-      TensorBuffer input_buffer,
-      compiled_model.CreateInputBuffer(signature_key, "arg0"));
+      TensorBuffer output_buffer,
+      compiled_model.CreateOutputBuffer(signature_key, "tfl.add"));
   LITERT_ASSERT_OK_AND_ASSIGN(RankedTensorType buffer_type,
-                              input_buffer.TensorType());
+                              output_buffer.TensorType());
   EXPECT_THAT(buffer_type.Layout().Dimensions(),
               ElementsAre(resized_dims[0], resized_dims[1], resized_dims[2]));
+  absl::flat_hash_map<absl::string_view, TensorBuffer> output_map;
+  output_map["tfl.add"] = std::move(output_buffer);
+
+  LITERT_ASSERT_OK(compiled_model.Run(input_map, output_map));
 }
 
 TEST(CompiledModelTest, BasicSignatureIndex) {

litert/runtime/compiled_model.cc

@@ -1451,11 +1451,9 @@ litert::Expected<void> LiteRtCompiledModelT::ResizeInputTensor(
                               "Failed to resize input tensor");
   }
 
-  // Clear cached buffer requirements for this tensor
-  LITERT_ASSIGN_OR_RETURN(const auto tensor_id,
-                          GetTensorIdentifier(*interp_, input_tensor));
-  cpu_buffer_requirements_.erase(tensor_id);
-
+  // Clear cached buffer requirements for all tensors since output and
+  // intermediate tensors may change shape after an explicit resize.
+  cpu_buffer_requirements_.clear();
   return {};
 }