[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

llvm-svn: 280563
This commit is contained in:
Jason Henline 2016-09-03 00:32:07 +00:00
parent 6098fd1989
commit 3956b2840b
2 changed files with 22 additions and 13 deletions

View File

@ -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; }

View File

@ -101,16 +101,16 @@ TEST_F(DeviceMemoryPackingTest, SingleConstTypedGlobalPointer) {
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, SingleTypedSharedPointer) {
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, PackSeveralArguments) {
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());
}