Return the number of dtype entries from memory::size().

libocca · Oct 20, 2023 · 0f2f339 · 0f2f339
1 parent ca23135
commit 0f2f339
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 28 deletions.
diff --git a/include/occa/core/memory.hpp b/include/occa/core/memory.hpp
@@ -196,38 +196,35 @@ namespace occa {
      * @startDoc{size}
      *
      * Description:
-     *   Get the byte size of the allocated memory
+     *   Returns the number of elements of size [[dtype_t]] held in this memory.
+     *   If no type was given during [[allocation|device.malloc]], returns the
+     *   size of the storage in bytes.
      *
      * @endDoc
      */
     udim_t size() const;
 
     /**
-     * @startDoc{length[0]}
+     * @startDoc{length}
      *
      * Description:
-     *   Get the length of the memory object, using its underlying [[dtype_t]].
-     *   This [[dtype_t]] can be fetched through the [[memory.dtype]] method
-     *
-     *   If no type was given during [[allocation|device.malloc]] or was ever set
-     *   through [[casting it|memory.cast]], it will return the bytes just like [[memory.size]].
+     *   Returns the number of elements of size [[dtype_t]] held in this memory.
+     *   If no type was given during [[allocation|device.malloc]], returns the
+     *   size of the storage in bytes.
      *
      * @endDoc
      */
     udim_t length() const;
 
     /**
-     * @startDoc{length[1]}
+     * @startDoc{byte_size}
      *
-     * Overloaded Description:
-     *   Same as above but explicitly chose the type (`T`)
+     * Description:
+     *   Get the size of the allocated memory in bytes.
      *
      * @endDoc
      */
-    template <class T>
-    udim_t length() const {
-      return size() / sizeof(T);
-    }
+    udim_t byte_size() const;
 
     /**
      * @startDoc{operator_equals[0]}

diff --git a/include/occa/functional/array.hpp b/include/occa/functional/array.hpp
@@ -319,8 +319,8 @@ namespace occa {
     }
 
     array concat(const array &other) const {
-      const udim_t bytes1 = memory_.size();
-      const udim_t bytes2 = other.memory_.size();
+      const udim_t bytes1 = memory_.byte_size();
+      const udim_t bytes2 = other.memory_.byte_size();
 
       occa::memory ret = getDevice().template malloc<T>(length() + other.length());
       ret.copyFrom(memory_, bytes1, 0);

diff --git a/src/core/memory.cpp b/src/core/memory.cpp
@@ -132,10 +132,7 @@ namespace occa {
   }
 
   udim_t memory::size() const {
-    if (modeMemory == NULL) {
-      return 0;
-    }
-    return modeMemory->size;
+    return length();
   }
 
   udim_t memory::length() const {
@@ -145,6 +142,11 @@ namespace occa {
     return modeMemory->size / modeMemory->dtype_->bytes();
   }
 
+  udim_t memory::byte_size() const {
+    if (modeMemory == NULL) return 0;
+    else return modeMemory->size;
+  }
+
   bool memory::operator == (const occa::memory &other) const {
     return (modeMemory == other.modeMemory);
   }
@@ -172,12 +174,12 @@ namespace occa {
                                       ? (length() - offset)
                                       : count);
 
-    OCCA_ERROR("Trying to allocate negative bytes (" << bytes << ")",
+    OCCA_ERROR("Trying to allocate negative elements (" << count << ")",
                bytes >= 0);
 
-    OCCA_ERROR("Cannot have offset and bytes greater than the memory size ("
-               << offset_ << " + " << bytes << " > " << size() << ")",
-               (offset_ + (dim_t) bytes) <= (dim_t) size());
+    OCCA_ERROR("Memory size is less than offset + count ("
+                << size() << " <" << offset << " + " << count << ")",
+               (offset + (dim_t) count) <= (dim_t) size());
 
     occa::memory m(modeMemory->slice(offset_, bytes));
     m.setDtype(dtype());
@@ -320,7 +322,7 @@ namespace occa {
 
     occa::memory mem = (
       occa::device(modeMemory->getModeDevice())
-      .malloc(size(), *this, properties())
+      .malloc(byte_size(), *this, properties())
     );
     mem.setDtype(dtype());
 

diff --git a/tests/src/core/device.cpp b/tests/src/core/device.cpp
@@ -110,19 +110,22 @@ void testWrapMemory() {
   int *hostPtr = &value;
 
   occa::memory mem = device.wrapMemory((void*) hostPtr, bytes);
+  mem.setDtype(occa::dtype::int_);
   ASSERT_EQ(mem.ptr<int>()[0], value);
   ASSERT_EQ(mem.ptr<int>(), hostPtr);
-  ASSERT_EQ((int) mem.length<int>(), 1);
+  ASSERT_EQ((int) mem.length(), 1);
 
   mem = device.wrapMemory(hostPtr, 1);
+  mem.setDtype(occa::dtype::int_);
   ASSERT_EQ(mem.ptr<int>()[0], value);
   ASSERT_EQ(mem.ptr<int>(), hostPtr);
-  ASSERT_EQ((int) mem.length<int>(), 1);
+  ASSERT_EQ((int) mem.length(), 1);
 
   mem = device.wrapMemory(hostPtr, 1, {{"use_host_pointer", false}});
+  mem.setDtype(occa::dtype::int_);
   ASSERT_EQ(mem.ptr<int>()[0], value);
   ASSERT_EQ(mem.ptr<int>(), hostPtr);
-  ASSERT_EQ((int) mem.length<int>(), 1);
+  ASSERT_EQ((int) mem.length(), 1);
 }
 
 void testUnwrap() {