[Hexagon] Add fix for vtcm allocation searches

src/runtime/hexagon/hexagon_vtcm_pool.cc

@@ -85,7 +85,8 @@ void* HexagonVtcmPool::Allocate(size_t nbytes) {
 
   auto entry_to_allocate = free_.begin();
   for (auto it = free_.begin(); it != free_.end(); it++) {
-    if ((it->second < entry_to_allocate->second) && (it->second >= nbytes)) {
+    if ((entry_to_allocate->second < nbytes || it->second < entry_to_allocate->second) &&
+        it->second >= nbytes) {
       entry_to_allocate = it;
       if (entry_to_allocate->second == nbytes) {
         break;

tests/cpp-runtime/hexagon/hexagon_vtcm_pool_tests.cc

@@ -34,6 +34,7 @@ class HexagonVtcmPoolTest : public ::testing::Test {
  public:
   HexagonVtcmPool* vtcm_pool;
   size_t max_bytes;
+  size_t four_k_block = 4096;
   size_t two_k_block = 2048;
   size_t one_k_block = 1024;
   size_t min_bytes = 128;
@@ -164,6 +165,85 @@ TEST_F(HexagonVtcmPoolTest, free_alloc_combinations) {
   vtcm_pool->Free(ptr4, max_bytes);
 }
 
+TEST_F(HexagonVtcmPoolTest, find_allocation) {
+  void* ptr1;
+  void* ptr2;
+  void* ptr3;
+
+  ptr1 = vtcm_pool->Allocate(two_k_block);
+  ptr2 = vtcm_pool->Allocate(two_k_block);
+
+  // Free the first allocation
+  vtcm_pool->Free(ptr1, two_k_block);
+
+  // Allocate a new larger block to initiate search and ensure
+  // it succeeds despite there not being a match in the first free block.
+  ptr3 = vtcm_pool->Allocate(four_k_block);
+
+  // Clean up the ptrs
+  vtcm_pool->Free(ptr2, two_k_block);
+  vtcm_pool->Free(ptr3, four_k_block);
+
+  // Make sure at the end we have the full amount available again
+  ptr1 = vtcm_pool->Allocate(max_bytes);
+  vtcm_pool->Free(ptr1, max_bytes);
+}
+
+TEST_F(HexagonVtcmPoolTest, find_smallest_allocation_combinations) {
+  void* ptr1;
+  void* ptr2;
+  void* ptr3;
+  void* ptr4;
+  void* new_ptr;
+
+  ptr1 = vtcm_pool->Allocate(two_k_block);
+  ptr2 = vtcm_pool->Allocate(two_k_block);
+  ptr3 = vtcm_pool->Allocate(four_k_block);
+  ptr4 = vtcm_pool->Allocate(four_k_block);
+
+  // Fragment memory allocations.
+  vtcm_pool->Free(ptr2, two_k_block);
+  vtcm_pool->Free(ptr3, four_k_block);
+
+  // Reallocate memory allocations and ensure that the smallest free allocations are used.
+  new_ptr = vtcm_pool->Allocate(two_k_block);
+  CHECK(new_ptr == ptr2);
+
+  new_ptr = vtcm_pool->Allocate(two_k_block);
+  CHECK(new_ptr == ptr3);
+
+  vtcm_pool->Free(ptr1, two_k_block);
+  vtcm_pool->Free(ptr2, two_k_block);
+  vtcm_pool->Free(ptr3, two_k_block);
+  vtcm_pool->Free(ptr4, four_k_block);
+
+  // Rerun the same test for non 2k aligned allocations.
+  ptr1 = vtcm_pool->Allocate(min_bytes);
+  ptr2 = vtcm_pool->Allocate(min_bytes);
+  ptr3 = vtcm_pool->Allocate(one_k_block);
+  ptr4 = vtcm_pool->Allocate(one_k_block);
+
+  // Fragment memory allocations.
+  vtcm_pool->Free(ptr2, min_bytes);
+  vtcm_pool->Free(ptr3, one_k_block);
+
+  // Reallocate memory allocations and ensure that the smallest free allocations are used.
+  new_ptr = vtcm_pool->Allocate(min_bytes);
+  CHECK(new_ptr == ptr2);
+
+  new_ptr = vtcm_pool->Allocate(one_k_block);
+  CHECK(new_ptr == ptr3);
+
+  vtcm_pool->Free(ptr1, min_bytes);
+  vtcm_pool->Free(ptr2, min_bytes);
+  vtcm_pool->Free(ptr3, one_k_block);
+  vtcm_pool->Free(ptr4, one_k_block);
+
+  // Make sure at the end we have the full amount available again
+  ptr4 = vtcm_pool->Allocate(max_bytes);
+  vtcm_pool->Free(ptr4, max_bytes);
+}
+
 // Test alignment edge cases allocating through HexagonBuffer
 TEST_F(HexagonVtcmPoolTest, vtcm_alignment) {
   std::unique_ptr<HexagonBufferManager> test_hexbuffs = std::make_unique<HexagonBufferManager>();