* GetDataSize is moved to DeviceAPI and memory_manager uses this interface.

Author: srkreddy1238

include/tvm/runtime/device_api.h

@@ -95,6 +95,14 @@ class TVM_DLL DeviceAPI {
    */
   virtual void GetAttr(Device dev, DeviceAttrKind kind, TVMRetValue* rv) = 0;
 
+  /*!
+   * \brief Get the physical memory size required.
+   * \param arr the tensor object.
+   * \param mem_scope the memory scope if any
+   * \return the memory size.
+   */
+  virtual size_t GetDataSize(const DLTensor& arr, Optional<String> mem_scope = NullOpt);
+
   /*!
    * \brief Query the device for specified properties.
    *

src/runtime/c_runtime_api.cc

@@ -152,6 +152,20 @@ static size_t GetDataAlignment(const DLDataType dtype) {
   return align;
 }
 
+size_t DeviceAPI::GetDataSize(const DLTensor& arr, Optional<String> mem_scope) {
+  if (!mem_scope.defined() || mem_scope.value().empty() || mem_scope.value() == "global") {
+    size_t size = 1;
+    for (tvm_index_t i = 0; i < arr.ndim; ++i) {
+      size *= static_cast<size_t>(arr.shape[i]);
+    }
+    size *= (arr.dtype.bits * arr.dtype.lanes + 7) / 8;
+    return size;
+  }
+  LOG(FATAL) << "Device does not support physical mem computation with "
+             << "specified memory scope: " << mem_scope.value();
+  return 0;
+}
+
 void* DeviceAPI::AllocDataSpace(Device dev, int ndim, const int64_t* shape, DLDataType dtype,
                                 Optional<String> mem_scope) {
   if (!mem_scope.defined() || mem_scope.value() == "" || mem_scope.value() == "global") {

src/runtime/memory/memory_manager.cc

@@ -85,7 +85,7 @@ NDArray StorageObj::AllocNDArray(size_t offset, ShapeTuple shape, DLDataType dty
   container->dl_tensor.byte_offset = offset;
 
   container->SetDeleter(StorageObj::Deleter);
-  size_t needed_size = GetDataSize(container->dl_tensor);
+  size_t needed_size = DeviceAPI::Get(this->buffer.device)->GetDataSize(container->dl_tensor);
   this->IncRef();
   // The manager context pointer must continue to point to the storage object
   // which owns the backing memory, and keeps track of the reference count.
@@ -159,7 +159,7 @@ NDArray Allocator::Empty(ShapeTuple shape, DLDataType dtype, DLDevice dev,
   VerifyDataType(dtype);
   NDArray::Container* container = new NDArray::Container(nullptr, shape, dtype, dev);
   container->SetDeleter(BufferDeleter);
-  size_t size = GetDataSize(container->dl_tensor);
+  size_t size = DeviceAPI::Get(dev)->GetDataSize(container->dl_tensor);
   size_t alignment = GetDataAlignment(container->dl_tensor);
   Buffer* buffer = new Buffer;
   if (!mem_scope.defined() || mem_scope == "global") {
@@ -177,7 +177,7 @@ Buffer Allocator::Alloc(Device dev, ShapeTuple shape, DLDataType type_hint,
   if (mem_scope.empty() || mem_scope == "global") {
     // by default, we can always redirect to the flat memory allocations
     NDArray::Container container(nullptr, shape, type_hint, dev);
-    size_t size = GetDataSize(container.dl_tensor);
+    size_t size = DeviceAPI::Get(dev)->GetDataSize(container.dl_tensor);
     size_t alignment = GetDataAlignment(container.dl_tensor);
     return Alloc(size, alignment, type_hint);
   }

src/runtime/memory/naive_allocator.h

@@ -51,7 +51,7 @@ class NaiveAllocator final : public Allocator {
     Buffer buf;
     size_t nbytes = 1;
     buf.shape = shape;
-    for (int i = 0; i < shape.size(); ++i) {
+    for (int i = 0; i < static_cast<int>(shape.size()); ++i) {
       nbytes *= static_cast<size_t>(shape[i]);
     }
     nbytes *= (type_hint.bits * type_hint.lanes + 7) / 8;