[Parallel_libs-commits] [parallel-libs] r280563 - [SE] Add getByteCount methods for device memory

[Jason Henline via Parallel_libs-commits]

Author: jhen
Date: Fri Sep  2 19:32:07 2016
New Revision: 280563

URL: http://llvm.org/viewvc/llvm-project?rev=280563&view=rev
Log:
[SE] Add getByteCount methods for device memory

Summary:
Simple utility methods will prevent users from making mistakes when
converting element counts to byte counts.

Reviewers: jlebar

Subscribers: jprice, parallel_libs-commits

Differential Revision: https://reviews.llvm.org/D24197

Modified:
    parallel-libs/trunk/streamexecutor/include/streamexecutor/DeviceMemory.h
    parallel-libs/trunk/streamexecutor/lib/unittests/PackedKernelArgumentArrayTest.cpp

Modified: parallel-libs/trunk/streamexecutor/include/streamexecutor/DeviceMemory.h
URL: http://llvm.org/viewvc/llvm-project/parallel-libs/trunk/streamexecutor/include/streamexecutor/DeviceMemory.h?rev=280563&r1=280562&r2=280563&view=diff
==============================================================================
--- parallel-libs/trunk/streamexecutor/include/streamexecutor/DeviceMemory.h (original)
+++ parallel-libs/trunk/streamexecutor/include/streamexecutor/DeviceMemory.h Fri Sep  2 19:32:07 2016
@@ -71,6 +71,9 @@ public:
   /// Gets the number of elements in this slice.
   size_t getElementCount() const { return ElementCount; }
 
+  /// Returns the number of bytes that can fit in this slice.
+  size_t getByteCount() const { return ElementCount * sizeof(ElemT); }
+
   /// Creates a slice of the memory with the first DropCount elements removed.
   GlobalDeviceMemorySlice<ElemT> drop_front(size_t DropCount) const {
     assert(DropCount <= ElementCount &&
@@ -175,6 +178,9 @@ public:
   /// allocation.
   size_t getElementCount() const { return ByteCount / sizeof(ElemT); }
 
+  /// Returns the number of bytes that can fit in this memory buffer.
+  size_t getByteCount() const { return ByteCount; }
+
   /// Converts this memory object into a slice.
   GlobalDeviceMemorySlice<ElemT> asSlice() const {
     return GlobalDeviceMemorySlice<ElemT>(*this);
@@ -224,10 +230,13 @@ public:
   /// Copy-assignable because it is just an array size.
   SharedDeviceMemory &operator=(const SharedDeviceMemory &) = default;
 
-  /// Returns the number of elements of type ElemT that can fit this memory
+  /// Returns the number of elements of type ElemT that can fit in this memory
   /// buffer.
   size_t getElementCount() const { return ElementCount; }
 
+  /// Returns the number of bytes that can fit in this memory buffer.
+  size_t getByteCount() const { return ElementCount * sizeof(ElemT); }
+
   /// Returns whether this is a single-element memory buffer.
   bool isScalar() const { return getElementCount() == 1; }
 

Modified: parallel-libs/trunk/streamexecutor/lib/unittests/PackedKernelArgumentArrayTest.cpp
URL: http://llvm.org/viewvc/llvm-project/parallel-libs/trunk/streamexecutor/lib/unittests/PackedKernelArgumentArrayTest.cpp?rev=280563&r1=280562&r2=280563&view=diff
==============================================================================
--- parallel-libs/trunk/streamexecutor/lib/unittests/PackedKernelArgumentArrayTest.cpp (original)
+++ parallel-libs/trunk/streamexecutor/lib/unittests/PackedKernelArgumentArrayTest.cpp Fri Sep  2 19:32:07 2016
@@ -101,16 +101,16 @@ TEST_F(DeviceMemoryPackingTest, SingleCo
 
 TEST_F(DeviceMemoryPackingTest, SingleTypedShared) {
   auto Array = se::make_kernel_argument_pack(TypedShared);
-  ExpectEqual(nullptr, TypedShared.getElementCount() * sizeof(int),
-              Type::SHARED_DEVICE_MEMORY, Array, 0);
+  ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
+              Array, 0);
   EXPECT_EQ(1u, Array.getArgumentCount());
   EXPECT_EQ(1u, Array.getSharedCount());
 }
 
 TEST_F(DeviceMemoryPackingTest, SingleTypedSharedPointer) {
   auto Array = se::make_kernel_argument_pack(&TypedShared);
-  ExpectEqual(nullptr, TypedShared.getElementCount() * sizeof(int),
-              Type::SHARED_DEVICE_MEMORY, Array, 0);
+  ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
+              Array, 0);
   EXPECT_EQ(1u, Array.getArgumentCount());
   EXPECT_EQ(1u, Array.getSharedCount());
 }
@@ -118,8 +118,8 @@ TEST_F(DeviceMemoryPackingTest, SingleTy
 TEST_F(DeviceMemoryPackingTest, SingleConstTypedSharedPointer) {
   const se::SharedDeviceMemory<int> *ArgumentPointer = &TypedShared;
   auto Array = se::make_kernel_argument_pack(ArgumentPointer);
-  ExpectEqual(nullptr, TypedShared.getElementCount() * sizeof(int),
-              Type::SHARED_DEVICE_MEMORY, Array, 0);
+  ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
+              Array, 0);
   EXPECT_EQ(1u, Array.getArgumentCount());
   EXPECT_EQ(1u, Array.getSharedCount());
 }
@@ -137,12 +137,12 @@ TEST_F(DeviceMemoryPackingTest, PackSeve
               Type::GLOBAL_DEVICE_MEMORY, Array, 2);
   ExpectEqual(TypedGlobal.getHandle(), sizeof(void *),
               Type::GLOBAL_DEVICE_MEMORY, Array, 3);
-  ExpectEqual(nullptr, TypedShared.getElementCount() * sizeof(int),
-              Type::SHARED_DEVICE_MEMORY, Array, 4);
-  ExpectEqual(nullptr, TypedShared.getElementCount() * sizeof(int),
-              Type::SHARED_DEVICE_MEMORY, Array, 5);
-  ExpectEqual(nullptr, TypedShared.getElementCount() * sizeof(int),
-              Type::SHARED_DEVICE_MEMORY, Array, 6);
+  ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
+              Array, 4);
+  ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
+              Array, 5);
+  ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
+              Array, 6);
   EXPECT_EQ(7u, Array.getArgumentCount());
   EXPECT_EQ(3u, Array.getSharedCount());
 }