diff --git a/include/af/array.h b/include/af/array.h index e1048dd114..ccf99ed0a4 100644 --- a/include/af/array.h +++ b/include/af/array.h @@ -108,6 +108,7 @@ namespace af dim_t dims(unsigned dim) const; unsigned numdims() const; size_t bytes() const; + size_t allocated() const; array copy() const; bool isempty() const; bool isscalar() const; @@ -586,6 +587,12 @@ namespace af */ size_t bytes() const; + /** + Get the size of the array in memory. This will return the parent's + bytes() if the array is indexed. + */ + size_t allocated() const; + /** Perform deep copy of the array */ diff --git a/include/af/internal.h b/include/af/internal.h index 53002929c3..c441919107 100644 --- a/include/af/internal.h +++ b/include/af/internal.h @@ -176,6 +176,17 @@ extern "C" AFAPI af_err af_is_owner(bool *result, const af_array arr); #endif +#if AF_API_VERSION >= 35 + /** + \param[out] bytes the size of the physical allocated bytes. This will return the size + of the parent/owner if the \p arr is an indexed array. + \param[in] arr the input array. + + \ingroup internal_func_allocatedbytes + */ + AFAPI af_err af_get_allocated_bytes(size_t *bytes, const af_array arr); +#endif + #ifdef __cplusplus } #endif diff --git a/src/api/c/internal.cpp b/src/api/c/internal.cpp index 47c62c6478..5c9d6e3797 100644 --- a/src/api/c/internal.cpp +++ b/src/api/c/internal.cpp @@ -168,3 +168,32 @@ af_err af_is_owner(bool *result, const af_array arr) CATCHALL; return AF_SUCCESS; } + +af_err af_get_allocated_bytes(size_t *bytes, const af_array arr) +{ + try { + af_dtype ty = getInfo(arr).getType(); + + size_t res = 0; + + switch (ty) { + case f32: res = getArray(arr).getAllocatedBytes(); break; + case f64: res = getArray(arr).getAllocatedBytes(); break; + case c32: res = getArray(arr).getAllocatedBytes(); break; + case c64: res = getArray(arr).getAllocatedBytes(); break; + case u32: res = getArray(arr).getAllocatedBytes(); break; + case s32: res = getArray(arr).getAllocatedBytes(); break; + case u64: res = getArray(arr).getAllocatedBytes(); break; + case s64: res = getArray(arr).getAllocatedBytes(); break; + case u16: res = getArray(arr).getAllocatedBytes(); break; + case s16: res = getArray(arr).getAllocatedBytes(); break; + case b8 : res = getArray(arr).getAllocatedBytes(); break; + case u8 : res = getArray(arr).getAllocatedBytes(); break; + default: TYPE_ERROR(6, ty); + } + + std::swap(*bytes, res); + } + CATCHALL; + return AF_SUCCESS; +} diff --git a/src/api/cpp/array.cpp b/src/api/cpp/array.cpp index a4a7937523..8420582f33 100644 --- a/src/api/cpp/array.cpp +++ b/src/api/cpp/array.cpp @@ -17,6 +17,7 @@ #include #include #include +#include #include "error.hpp" namespace af @@ -266,6 +267,13 @@ namespace af return nElements * getSizeOf(type()); } + size_t array::allocated() const + { + size_t result = 0; + AF_THROW(af_get_allocated_bytes(&result, get())); + return result; + } + array array::copy() const { af_array other = 0; @@ -597,6 +605,7 @@ namespace af MEM_FUNC(dim4 , dims) MEM_FUNC(unsigned , numdims) MEM_FUNC(size_t , bytes) + MEM_FUNC(size_t , allocated) MEM_FUNC(array , copy) MEM_FUNC(bool , isempty) MEM_FUNC(bool , isscalar) diff --git a/src/api/unified/internal.cpp b/src/api/unified/internal.cpp index b9ac0ac277..3a2d51cca5 100644 --- a/src/api/unified/internal.cpp +++ b/src/api/unified/internal.cpp @@ -52,3 +52,9 @@ af_err af_is_owner(bool *result, const af_array arr) CHECK_ARRAYS(arr); return CALL(result, arr); } + +af_err af_get_allocated_bytes(size_t *bytes, const af_array arr) +{ + CHECK_ARRAYS(arr); + return CALL(bytes, arr); +} diff --git a/src/backend/cpu/Array.hpp b/src/backend/cpu/Array.hpp index 8762e7e54e..65ef5f1434 100644 --- a/src/backend/cpu/Array.hpp +++ b/src/backend/cpu/Array.hpp @@ -188,6 +188,11 @@ namespace cpu data_dims = new_dims; } + size_t getAllocatedBytes() const + { + return data_dims.elements() * sizeof(T); + } + T* device(); T* device() const diff --git a/src/backend/cuda/Array.hpp b/src/backend/cuda/Array.hpp index 6f8acc1ca6..127e0344e1 100644 --- a/src/backend/cuda/Array.hpp +++ b/src/backend/cuda/Array.hpp @@ -184,6 +184,11 @@ namespace cuda data_dims = new_dims; } + size_t getAllocatedBytes() const + { + return data_dims.elements() * sizeof(T); + } + T* device(); T* device() const diff --git a/src/backend/opencl/Array.hpp b/src/backend/opencl/Array.hpp index 9dc9b427f7..b3a168efa2 100644 --- a/src/backend/opencl/Array.hpp +++ b/src/backend/opencl/Array.hpp @@ -211,6 +211,11 @@ namespace opencl data_dims = new_dims; } + size_t getAllocatedBytes() const + { + return data_dims.elements() * sizeof(T); + } + operator Param() const { KParam info = {{dims()[0], dims()[1], dims()[2], dims()[3]}, diff --git a/test/internal.cpp b/test/internal.cpp index 75fa54fdb9..f60c570d99 100644 --- a/test/internal.cpp +++ b/test/internal.cpp @@ -122,3 +122,36 @@ TEST(Internal, Linear) ASSERT_EQ(isOwner(c), false); } } + +TEST(Internal, Allocated) +{ + af::array a = af::randu(10, 8); + const size_t aBytes = a.bytes(); + + // b is just pointing to same underlying data + // b is an owner; + af::array b = a; + ASSERT_EQ(b.allocated(), aBytes); + ASSERT_EQ(b.bytes(), aBytes); + + // C is considered sub array + // C will not be an owner + af::array c = a(af::span); + ASSERT_EQ(c.allocated(), aBytes); + ASSERT_EQ(c.bytes(), aBytes); + + af::array d = a.col(1); + ASSERT_EQ(d.allocated(), aBytes); + ASSERT_EQ(d.bytes(), (size_t)10 * 4); + + a = af::randu(20); + b = af::randu(20); + + // Even though a, b are reallocated and c, d are not owners + // the allocated and bytes should remain the same + ASSERT_EQ(c.allocated(), aBytes); + ASSERT_EQ(c.bytes(), aBytes); + + ASSERT_EQ(d.allocated(), aBytes); + ASSERT_EQ(d.bytes(), (size_t)10 * 4); +}