Reset Temp Allocator after each use

Differential Revision: D80191057

Pull Request resolved: #13384

Author: cmt0

runtime/executor/method.cpp

@@ -713,6 +713,9 @@ Error Method::resolve_operator(
   }
   TensorMeta* meta = allocator->allocateList<TensorMeta>(n_args);
   if (meta == nullptr) {
+    if (allocator == memory_manager_->temp_allocator()) {
+      memory_manager_->temp_allocator()->reset();
+    }
     return Error::MemoryAllocationFailed;
   }
 
@@ -726,6 +729,9 @@ Error Method::resolve_operator(
       executorch::aten::DimOrderType* dim_order_ptr =
           allocator->allocateList<executorch::aten::DimOrderType>(tensor.dim());
       if (dim_order_ptr == nullptr) {
+        if (allocator == memory_manager_->temp_allocator()) {
+          memory_manager_->temp_allocator()->reset();
+        }
         return Error::MemoryAllocationFailed;
       }
       size_t size = tensor.dim();
@@ -751,9 +757,18 @@ Error Method::resolve_operator(
         "Missing operator: [%" ET_PRIssize_t "] %s",
         static_cast<ssize_t>(op_index),
         operator_name);
+    if (allocator == memory_manager_->temp_allocator()) {
+      memory_manager_->temp_allocator()->reset();
+    }
     return op_function.error();
   }
   kernels[kernel_index] = op_function.get();
+
+  // If we used the temp allocator here, reset it.
+  if (allocator == memory_manager_->temp_allocator()) {
+    memory_manager_->temp_allocator()->reset();
+  }
+
   return Error::Ok;
 }
 
@@ -1547,6 +1562,9 @@ Error Method::execute() {
         i);
   }
   ET_LOG(Debug, "Executing method: %s.", method_meta().name());
+  if (temp_allocator_ != nullptr) {
+    temp_allocator_->reset();
+  }
 
   // Chains are executed sequentially today, but future async designs may
   // branch and run many in parallel or out of order.

runtime/executor/test/kernel_integration_test.cpp

@@ -367,8 +367,9 @@ TEST_F(KernelTempMemoryAllocatorIntegrationTest, UsingTempMemoryAllocator) {
   EXPECT_EQ(control_->total_allocated_size, 4);
   EXPECT_EQ(temp_allocator_->number_of_allocations, 1);
   EXPECT_EQ(temp_allocator_->total_allocated_size, 4);
-  // The temp allocator should have been reset after the execution.
-  EXPECT_EQ(temp_allocator_->number_of_resets, 1);
+  // The temp allocator should have been reset after the execution and before
+  // method execution.
+  EXPECT_EQ(temp_allocator_->number_of_resets, 2);
   EXPECT_EQ(temp_allocator_->currently_allocated_size, 0);
 
   control_->temp_memory_size = 8;
@@ -379,6 +380,6 @@ TEST_F(KernelTempMemoryAllocatorIntegrationTest, UsingTempMemoryAllocator) {
   EXPECT_EQ(temp_allocator_->number_of_allocations, 2);
   EXPECT_EQ(temp_allocator_->total_allocated_size, 12);
   // The temp allocator should have been reset after the execution.
-  EXPECT_EQ(temp_allocator_->number_of_resets, 2);
+  EXPECT_EQ(temp_allocator_->number_of_resets, 4);
   EXPECT_EQ(temp_allocator_->currently_allocated_size, 0);
 }