From 8a0aba2e951f44fb48ce40f4f207868e90f76d9f Mon Sep 17 00:00:00 2001
From: Eric Lunderberg <Lunderberg@users.noreply.github.com>
Date: Wed, 6 Apr 2022 10:58:00 -0500
Subject: [PATCH] [Hexagon] Use single allocation to back 2-d arrays (#10903)

* [Hexagon] Use single allocation to back 2-d arrays

Currently, each allocation allocates an entire page, so even a
relatively small number of allocations can use very large amounts of
VTCM.  This commit changes calls to `AllocVtcmWorkspace` of shape
`[N,M]` from performing `N` allocations of size `M`, to 1 allocation
of size `N*M`.  Since `N` is usually much smaller than a page, this
reduces the total amount of memory required.

This is an intermediate step, where the long-term solution is to use
static planning for VTCM allocations.  This returns the same `void**`
type as the static planning eventually will, but avoids excess memory
use in the meantime.

* [Hexagon] Maintain alignment of allocations

Previously, when a single monolithic allocation is used to back a 2-d
Hexagon buffer of shape `[nallocs, nbytes_per_allocation]`, the
allocation itself is aligned, but each individual region is not.  This
commit ensures that each individual region also followed the alignment
specified.
---
 src/runtime/hexagon/hexagon/hexagon_buffer.cc | 25 ++++++++++++-------
 1 file changed, 16 insertions(+), 9 deletions(-)

diff --git a/src/runtime/hexagon/hexagon/hexagon_buffer.cc b/src/runtime/hexagon/hexagon/hexagon_buffer.cc
index b1de44df330c7..fc8cfa4efb3a1 100644
--- a/src/runtime/hexagon/hexagon/hexagon_buffer.cc
+++ b/src/runtime/hexagon/hexagon/hexagon_buffer.cc
@@ -150,17 +150,24 @@ HexagonBuffer::HexagonBuffer(size_t nallocs, size_t nbytes, size_t alignment,
                              Optional<String> scope)
     : ndim_(2), nbytes_per_allocation_(nbytes) {
   SetStorageScope(scope);
+
+  size_t nbytes_aligned = ((nbytes + (alignment - 1)) / alignment) * alignment;
+  size_t nbytes_monolithic = nallocs * nbytes_aligned;
+
+  std::unique_ptr<Allocation> alloca = nullptr;
+  if (GetStorageScope() == StorageScope::kDDR) {
+    alloca = Allocator<StorageScope::kDDR>(nbytes_monolithic, alignment);
+  } else if (GetStorageScope() == StorageScope::kVTCM) {
+    alloca = Allocator<StorageScope::kVTCM>(nbytes_monolithic, alignment);
+  }
+  CHECK(alloca) << "could not create allocation";
+
   for (size_t i = 0; i < nallocs; ++i) {
-    std::unique_ptr<Allocation> alloca = nullptr;
-    if (GetStorageScope() == StorageScope::kDDR) {
-      alloca = Allocator<StorageScope::kDDR>(nbytes, alignment);
-    } else if (GetStorageScope() == StorageScope::kVTCM) {
-      alloca = Allocator<StorageScope::kVTCM>(nbytes, alignment);
-    }
-    CHECK(alloca != nullptr);
-    allocations_.push_back(alloca->data_);
-    managed_allocations_.push_back(std::move(alloca));
+    void* alloc_offset = static_cast<unsigned char*>(alloca->data_) + i * nbytes_aligned;
+    allocations_.push_back(alloc_offset);
   }
+
+  managed_allocations_.push_back(std::move(alloca));
 }
 
 HexagonBuffer::HexagonBuffer(void* data, size_t nbytes, Optional<String> scope)