[Hexagon] Support both 1-d and 2-d VTCM allocations

src/runtime/hexagon/hexagon/hexagon_buffer.cc

@@ -129,7 +129,7 @@ std::unique_ptr<Allocation> Allocator<HexagonBuffer::StorageScope::kVTCM>(size_t
 }
 
 HexagonBuffer::HexagonBuffer(size_t nbytes, size_t alignment, Optional<String> scope)
-    : nallocs_(1), nbytes_(nbytes) {
+    : ndim_(1), nbytes_per_allocation_(nbytes) {
   SetStorageScope(scope);
 
   std::unique_ptr<Allocation> alloca = nullptr;
@@ -145,7 +145,7 @@ HexagonBuffer::HexagonBuffer(size_t nbytes, size_t alignment, Optional<String> s
 
 HexagonBuffer::HexagonBuffer(size_t nallocs, size_t nbytes, size_t alignment,
                              Optional<String> scope)
-    : nallocs_(nallocs), nbytes_(nallocs * nbytes) {
+    : ndim_(2), nbytes_per_allocation_(nbytes) {
   SetStorageScope(scope);
   for (size_t i = 0; i < nallocs; ++i) {
     std::unique_ptr<Allocation> alloca = nullptr;
@@ -161,7 +161,7 @@ HexagonBuffer::HexagonBuffer(size_t nallocs, size_t nbytes, size_t alignment,
 }
 
 HexagonBuffer::HexagonBuffer(void* data, size_t nbytes, Optional<String> scope)
-    : nallocs_(1), nbytes_(nbytes) {
+    : ndim_(1), nbytes_per_allocation_(nbytes) {
   SetStorageScope(scope);
   // disallow external VTCM allocations
   CHECK(GetStorageScope() != HexagonBuffer::StorageScope::kVTCM);
@@ -170,11 +170,19 @@ HexagonBuffer::HexagonBuffer(void* data, size_t nbytes, Optional<String> scope)
 
 HexagonBuffer::~HexagonBuffer() { managed_allocations_.clear(); }
 
-void** HexagonBuffer::GetPointer() {
-  if (!allocations_.size()) {
+void* HexagonBuffer::GetPointer() {
+  ICHECK(allocations_.size())
+      << "Internal failure, allocations_ should be set in HexagonBuffer constructor";
+
+  if (ndim_ == 1) {
+    ICHECK_EQ(allocations_.size(), 1);
+    return allocations_[0];
+  } else if (ndim_ == 2) {
+    return allocations_.data();
+  } else {
+    LOG(FATAL) << "HexagonBuffer should be either 1-d or 2-d, not " << ndim_ << "-d";
     return nullptr;
   }
-  return allocations_.data();
 }
 
 HexagonBuffer::StorageScope HexagonBuffer::GetStorageScope() const { return storage_scope_; }
@@ -195,68 +203,67 @@ void HexagonBuffer::SetStorageScope(Optional<String> scope) {
 }
 
 void HexagonBuffer::CopyTo(void* data, size_t nbytes) const {
-  CHECK_LE(nbytes, nbytes_);
+  CHECK_LE(nbytes, TotalBytes());
   CHECK(managed_allocations_.size() && "CopyTo not supported on unmanaged `external` allocations");
 
   size_t copied = 0;
-  for (size_t i = 0; i < nallocs_; ++i) {
-    size_t bytes_to_copy = std::min(nbytes - copied, managed_allocations_[i]->allocation_nbytes_);
+  for (const auto& managed_alloc : managed_allocations_) {
+    size_t bytes_to_copy = std::min(nbytes - copied, managed_alloc->allocation_nbytes_);
     if (bytes_to_copy == 0) break;
 
     void* data_plus_copied = static_cast<void*>((static_cast<char*>(data) + copied));
-    int status =
-        hexagon_user_dma_1d_sync(data_plus_copied, managed_allocations_[i]->data_, bytes_to_copy);
+    int status = hexagon_user_dma_1d_sync(data_plus_copied, managed_alloc->data_, bytes_to_copy);
     CHECK_EQ(status, 0);
 
     copied += bytes_to_copy;
   }
 }
 
 void HexagonBuffer::CopyFrom(void* data, size_t nbytes) {
-  CHECK_LE(nbytes, nbytes_);
+  CHECK_LE(nbytes, TotalBytes());
   CHECK(managed_allocations_.size() &&
         "CopyFrom not supported on unmanaged `external` allocations");
 
   size_t copied = 0;
-  for (size_t i = 0; i < nallocs_; ++i) {
-    size_t bytes_to_copy = std::min(nbytes - copied, managed_allocations_[i]->allocation_nbytes_);
+  for (const auto& managed_alloc : managed_allocations_) {
+    size_t bytes_to_copy = std::min(nbytes - copied, managed_alloc->allocation_nbytes_);
     if (bytes_to_copy == 0) break;
 
     void* data_plus_copied = static_cast<void*>((static_cast<char*>(data) + copied));
-    int status =
-        hexagon_user_dma_1d_sync(managed_allocations_[i]->data_, data_plus_copied, bytes_to_copy);
+    int status = hexagon_user_dma_1d_sync(managed_alloc->data_, data_plus_copied, bytes_to_copy);
     CHECK_EQ(status, 0);
 
     copied += bytes_to_copy;
   }
 }
 
 void HexagonBuffer::CopyFrom(const HexagonBuffer& other, size_t nbytes) {
-  CHECK_LE(nbytes, nbytes_);
-  CHECK_LE(nbytes, other.nbytes_);
+  CHECK_LE(nbytes, TotalBytes());
+  CHECK_LE(nbytes, other.TotalBytes());
   CHECK(managed_allocations_.size() &&
         "CopyFrom not supported on unmanaged `external` allocations");
   CHECK(other.managed_allocations_.size() &&
         "CopyFrom not supported on unmanaged `external` allocations");
 
-  if (nallocs_ == other.nallocs_) {
+  if (managed_allocations_.size() == other.managed_allocations_.size()) {
     size_t copied = 0;
-    for (size_t i = 0; i < nallocs_; ++i) {
-      size_t bytes_to_copy = std::min(nbytes - copied, managed_allocations_[i]->allocation_nbytes_);
+    for (size_t i = 0; i < managed_allocations_.size(); ++i) {
+      const auto& this_alloc = managed_allocations_[i];
+      const auto& other_alloc = other.managed_allocations_[i];
+
+      size_t bytes_to_copy = std::min(nbytes - copied, this_alloc->allocation_nbytes_);
       if (bytes_to_copy == 0) break;
 
-      CHECK_LE(other.managed_allocations_[i]->allocation_nbytes_,
-               managed_allocations_[i]->allocation_nbytes_);
+      CHECK_LE(other_alloc->allocation_nbytes_, this_alloc->allocation_nbytes_);
 
-      int status = hexagon_user_dma_1d_sync(managed_allocations_[i]->data_,
-                                            other.managed_allocations_[i]->data_, bytes_to_copy);
+      int status = hexagon_user_dma_1d_sync(this_alloc->data_, other_alloc->data_, bytes_to_copy);
       CHECK_EQ(status, 0);
 
       copied += bytes_to_copy;
     }
-  } else if (nallocs_ == 1) {
+  } else if (managed_allocations_.size() == 1) {
     return other.CopyTo(managed_allocations_[0]->data_, nbytes);
-  } else if (other.nallocs_ == 1) {
+  } else if (other.managed_allocations_.size() == 1) {
     return CopyFrom(other.managed_allocations_[0]->data_, nbytes);
   } else {
     CHECK(false) << "To copy between Hexagon Buffers they must either have the same number of "

src/runtime/hexagon/hexagon/hexagon_buffer.h

@@ -94,8 +94,24 @@ class HexagonBuffer {
   //! \brief Prevent move assignment.
   HexagonBuffer& operator=(HexagonBuffer&&) = delete;
 
-  //! \brief Return pointer to allocations.
-  void** GetPointer();
+  /*! \brief Return data pointer
+   *
+   * The return type depends on the buffer being
+   */
+  void* GetPointer();
+
+  // //! \brief Return number of allocations.
+  // size_t GetNumAllocs() { return nallocs_; }
+
+  /*! \brief Return dimensionality of buffer
+   *
+   * Will always be either 1 or 2.  If the buffer is 1-d, allocation
+   * returns a pointer to data, which is accessed by
+   * ((value_type*)ptr)[i].  If the buffer is 2-d, allocation returns
+   * a pointer to an array of pointers, which is accessed by
+   * ((value_type**)ptr)[i][j].
+   */
+  size_t GetBufferDimension() { return ndim_; }
 
   //! \brief Memory scopes managed by a Hexagon Buffer.
   enum class StorageScope {
@@ -135,6 +151,9 @@ class HexagonBuffer {
   void CopyFrom(const HexagonBuffer& other, size_t nbytes);
 
  private:
+  //! \brief Return the total number of bytes in this buffer
+  size_t TotalBytes() const { return nbytes_per_allocation_ * allocations_.size(); }
+
   //! \brief Assign a storage scope to the buffer.
   void SetStorageScope(Optional<String> scope);
   /*! \brief Array of raw pointer allocations required by the buffer.
@@ -150,8 +169,8 @@ class HexagonBuffer {
   /*! \brief The underlying storage type in which the allocation
    *  resides.
    */
-  size_t nallocs_;
-  size_t nbytes_;
+  size_t ndim_;
+  size_t nbytes_per_allocation_;
   StorageScope storage_scope_;
 };
 

src/runtime/hexagon/hexagon/hexagon_common.cc

@@ -95,9 +95,8 @@ PackedFunc WrapPackedFunc(TVMBackendPackedCFunc faddr, const ObjectPtr<Object>&
       if (args.type_codes[i] == kTVMDLTensorHandle) {
         DLTensor* tensor = static_cast<DLTensor*>(arg_values[i].v_handle);
         buffer_args.emplace_back(i, static_cast<HexagonBuffer*>(tensor->data));
-        // Assumes a single contiguous allocation
-        // TODO(Straw): Enable discontiguous allocation
-        tensor->data = buffer_args.back().second->GetPointer()[0];
+        HexagonBuffer* hexbuf = buffer_args.back().second;
+        tensor->data = hexbuf->GetPointer();
       }
     }
     int ret = (*faddr)(const_cast<TVMValue*>(args.values), const_cast<int*>(args.type_codes),

src/runtime/hexagon/hexagon/hexagon_device_api_v2.cc

@@ -59,23 +59,31 @@ void HexagonDeviceAPIv2::GetAttr(Device dev, DeviceAttrKind kind, TVMRetValue* r
 // DataSpace: static allocations for Hexagon
 void* HexagonDeviceAPIv2::AllocDataSpace(Device dev, int ndim, const int64_t* shape,
                                          DLDataType dtype, Optional<String> mem_scope) {
+  if (!mem_scope.defined() || mem_scope.value() == "global") {
+    return DeviceAPI::AllocDataSpace(dev, ndim, shape, dtype, mem_scope);
+  }
+
   CHECK(TVMDeviceExtType(dev.device_type) == kDLHexagon) << "dev.device_type: " << dev.device_type;
 
-  // Forcing contiguous allocation, for now
-  // TODO(Straw): Enable discontiguous allocation
-  size_t nallocs = 1;
-  size_t nbytes = 1;
-  for (int i = 0; i < ndim; ++i) {
-    nbytes *= shape[i];
-  }
   size_t typesize = (dtype.bits / 8) * dtype.lanes;
-  nbytes *= typesize;
 
-  size_t alignment = typesize;
+  size_t alignment = shape[ndim - 1] * typesize;
   if (alignment < kHexagonAllocAlignment) {
     alignment = kHexagonAllocAlignment;
   }
-  return new HexagonBuffer(nallocs, nbytes, alignment, mem_scope);
+
+  if (ndim == 1) {
+    size_t nbytes = shape[0] * typesize;
+    return new HexagonBuffer(nbytes, alignment, mem_scope);
+  } else if (ndim == 2) {
+    size_t nallocs = shape[0];
+    size_t nbytes = shape[1] * typesize;
+    return new HexagonBuffer(nallocs, nbytes, alignment, mem_scope);
+  } else {
+    LOG(FATAL) << "Hexagon Device API supports only 1d and 2d allocations, but received ndim = "
+               << ndim;
+    return nullptr;
+  }
 }
 
 void* HexagonDeviceAPIv2::AllocDataSpace(Device dev, size_t nbytes, size_t alignment,
@@ -109,9 +117,7 @@ void* HexagonDeviceAPIv2::AllocWorkspace(Device dev, size_t size, DLDataType typ
   auto* hexbuf = static_cast<HexagonBuffer*>(
       dmlc::ThreadLocalStore<HexagonWorkspacePool>::Get()->AllocWorkspace(dev, size));
 
-  // Assumes a single contiguous allocation
-  // TODO(Straw): Enable discontiguous allocation
-  void* ptr = hexbuf->GetPointer()[0];
+  void* ptr = hexbuf->GetPointer();
   workspace_allocations_.insert({ptr, hexbuf});
   return ptr;
 }
@@ -128,9 +134,7 @@ void HexagonDeviceAPIv2::FreeWorkspace(Device dev, void* data) {
 void* HexagonDeviceAPIv2::AllocVtcmWorkspace(Device dev, int ndim, const int64_t* shape,
                                              DLDataType dtype, Optional<String> mem_scope) {
   CHECK(TVMDeviceExtType(dev.device_type) == kDLHexagon) << "dev.device_type: " << dev.device_type;
-  // Forcing contiguous allocation, for now
-  // TODO(Straw): Enable discontiguous allocation
-  CHECK_EQ(ndim, 1);
+  CHECK((ndim == 1 || ndim == 2) && "Hexagon Device API supports only 1d and 2d allocations");
   return AllocDataSpace(dev, ndim, shape, dtype, mem_scope);
 }
 
@@ -193,9 +197,7 @@ TVM_REGISTER_GLOBAL("device_api.hexagon.alloc_nd").set_body([](TVMArgs args, TVM
   std::string scope = args[4];
   CHECK(scope.find("global.vtcm") != std::string::npos);
   int64_t ndim = args[5];
-  // Forcing contiguous allocation, for now
-  // TODO(Straw): Enable discontiguous allocation
-  CHECK_EQ(ndim, 1);
+  CHECK((ndim == 1 || ndim == 2) && "Hexagon Device API supports only 1d and 2d allocations");
   int64_t* shape = static_cast<int64_t*>(static_cast<void*>(args[6]));
 
   Device dev;
@@ -211,9 +213,7 @@ TVM_REGISTER_GLOBAL("device_api.hexagon.alloc_nd").set_body([](TVMArgs args, TVM
   HexagonBuffer* hexbuf = reinterpret_cast<HexagonBuffer*>(
       hexapi->AllocVtcmWorkspace(dev, ndim, shape, type_hint, String(scope)));
 
-  // Assumes a single contiguous allocation
-  // TODO(Straw): Enable discontiguous allocation
-  void* ptr = hexbuf->GetPointer()[0];
+  void* ptr = hexbuf->GetPointer();
   vtcmallocs[ptr] = hexbuf;
   *rv = ptr;
 });

src/runtime/hexagon/hexagon/hexagon_device_api_v2.h

@@ -75,6 +75,15 @@ class HexagonDeviceAPIv2 final : public DeviceAPI {
 
   /*!
    * \brief Allocate an Nd data space on device with memory scope support.
+   *
+   * If mem_scope is undefined or is "global", treat shape as the
+   * tensor shape, to be flattened into an allocation of 1-d physical
+   * memory.  This is done to maintain the semantics expected by callers of
+   * DeviceAPI::AllocDataSpace, in cases where it has a valid return value.
+   *
+   * For other values of mem_scope, the shape is the N-d physical
+   * shape of the allocation.
+   *
    * \param dev The device to perform the operation.
    * \param ndim The number of dimensions of allocated tensor.
    * \param shape The shape of allocated tensor.

tests/cpp/runtime/hexagon_buffer.cc

@@ -54,7 +54,7 @@ TEST(HexagonBuffer, copy_from) {
   std::vector<uint8_t> data{0, 1, 2, 3, 4, 5, 6, 7};
   hb.CopyFrom(data.data(), data.size());
 
-  uint8_t* ptr = static_cast<uint8_t*>(hb.GetPointer()[0]);
+  uint8_t* ptr = static_cast<uint8_t*>(hb.GetPointer());
   for (size_t i = 0; i < data.size(); ++i) {
     EXPECT_EQ(ptr[i], data[i]);
   }
@@ -103,17 +103,15 @@ TEST(HexagonBuffer, nd_copy_from) {
   std::vector<uint8_t> data{0, 1, 2, 3, 4, 5, 6, 7};
   hb.CopyFrom(data.data(), data.size());
 
-  uint8_t* ptr = static_cast<uint8_t*>(hb.GetPointer()[0]);
-  EXPECT_EQ(ptr[0], data[0]);
-  EXPECT_EQ(ptr[1], data[1]);
-  EXPECT_EQ(ptr[2], data[2]);
-  EXPECT_EQ(ptr[3], data[3]);
-
-  ptr = static_cast<uint8_t*>(hb.GetPointer()[1]);
-  EXPECT_EQ(ptr[0], data[4]);
-  EXPECT_EQ(ptr[1], data[5]);
-  EXPECT_EQ(ptr[2], data[6]);
-  EXPECT_EQ(ptr[3], data[7]);
+  uint8_t** ptr = static_cast<uint8_t**>(hb.GetPointer());
+  EXPECT_EQ(ptr[0][0], data[0]);
+  EXPECT_EQ(ptr[0][1], data[1]);
+  EXPECT_EQ(ptr[0][2], data[2]);
+  EXPECT_EQ(ptr[0][3], data[3]);
+  EXPECT_EQ(ptr[1][0], data[4]);
+  EXPECT_EQ(ptr[1][1], data[5]);
+  EXPECT_EQ(ptr[1][2], data[6]);
+  EXPECT_EQ(ptr[1][3], data[7]);
 }
 
 TEST(HexagonBuffer, 1d_copy_from_1d) {
@@ -127,7 +125,7 @@ TEST(HexagonBuffer, 1d_copy_from_1d) {
   from.CopyFrom(data.data(), data.size());
   to.CopyFrom(from, 8);
 
-  uint8_t* ptr = static_cast<uint8_t*>(to.GetPointer()[0]);
+  uint8_t* ptr = static_cast<uint8_t*>(to.GetPointer());
   for (size_t i = 0; i < data.size(); ++i) {
     EXPECT_EQ(ptr[i], data[i]);
   }
@@ -144,17 +142,15 @@ TEST(HexagonBuffer, 2d_copy_from_1d) {
   hb1d.CopyFrom(data.data(), data.size());
   hb2d.CopyFrom(hb1d, 8);
 
-  uint8_t* ptr = static_cast<uint8_t*>(hb2d.GetPointer()[0]);
-  EXPECT_EQ(ptr[0], data[0]);
-  EXPECT_EQ(ptr[1], data[1]);
-  EXPECT_EQ(ptr[2], data[2]);
-  EXPECT_EQ(ptr[3], data[3]);
-
-  ptr = static_cast<uint8_t*>(hb2d.GetPointer()[1]);
-  EXPECT_EQ(ptr[0], data[4]);
-  EXPECT_EQ(ptr[1], data[5]);
-  EXPECT_EQ(ptr[2], data[6]);
-  EXPECT_EQ(ptr[3], data[7]);
+  uint8_t** ptr = static_cast<uint8_t**>(hb2d.GetPointer());
+  EXPECT_EQ(ptr[0][0], data[0]);
+  EXPECT_EQ(ptr[0][1], data[1]);
+  EXPECT_EQ(ptr[0][2], data[2]);
+  EXPECT_EQ(ptr[0][3], data[3]);
+  EXPECT_EQ(ptr[1][0], data[4]);
+  EXPECT_EQ(ptr[1][1], data[5]);
+  EXPECT_EQ(ptr[1][2], data[6]);
+  EXPECT_EQ(ptr[1][3], data[7]);
 }
 
 TEST(HexagonBuffer, 1d_copy_from_2d) {
@@ -168,7 +164,7 @@ TEST(HexagonBuffer, 1d_copy_from_2d) {
   hb2d.CopyFrom(data.data(), data.size());
   hb1d.CopyFrom(hb2d, 8);
 
-  uint8_t* ptr = static_cast<uint8_t*>(hb1d.GetPointer()[0]);
+  uint8_t* ptr = static_cast<uint8_t*>(hb1d.GetPointer());
   for (size_t i = 0; i < data.size(); ++i) {
     EXPECT_EQ(ptr[i], data[i]);
   }
@@ -245,12 +241,12 @@ TEST(HexagonBuffer, external) {
 
   Optional<String> def;
   HexagonBuffer hb_default(data.data(), data.size(), def);
-  EXPECT_EQ(hb_default.GetPointer()[0], data.data());
+  EXPECT_EQ(hb_default.GetPointer(), data.data());
   EXPECT_EQ(hb_default.GetStorageScope(), HexagonBuffer::StorageScope::kDDR);
 
   Optional<String> global("global");
   HexagonBuffer hb_global(data.data(), data.size(), global);
-  EXPECT_EQ(hb_global.GetPointer()[0], data.data());
+  EXPECT_EQ(hb_global.GetPointer(), data.data());
   EXPECT_EQ(hb_global.GetStorageScope(), HexagonBuffer::StorageScope::kDDR);
 
   Optional<String> vtcm("global.vtcm");