[SE] Add getByteCount methods for device memory

[jhen]

Summary

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

streamexecutor/include/streamexecutor/DeviceMemory.h

public:
   /// Gets the offset of this slice from the base memory.
   ///
   /// The offset is measured in elements, not bytes.
   size_t getElementOffset() const { return ElementOffset; }
 
   /// 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 &&
            "dropping more than the size of a slice");
     return GlobalDeviceMemorySlice<ElemT>(BaseMemory, ElementOffset + DropCount,
                                           ElementCount - DropCount);
   }
 

 public:
   GlobalDeviceMemory(GlobalDeviceMemory &&Other) = default;
   GlobalDeviceMemory &operator=(GlobalDeviceMemory &&Other) = default;
 
   /// Returns the number of elements of type ElemT that constitute this
   /// 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);
   }
 
 private:
   GlobalDeviceMemory(const GlobalDeviceMemory &) = delete;
   GlobalDeviceMemory &operator=(const GlobalDeviceMemory &) = delete;
public:
   }
 
   /// Copyable because it is just an array size.
   SharedDeviceMemory(const SharedDeviceMemory &) = default;
 
   /// 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; }
 
 private:
   /// Constructs a SharedDeviceMemory instance from an element count.
   ///
   /// This constructor is not public because there is a potential for confusion
   /// between the size of the buffer in bytes and the size of the buffer in

streamexecutor/lib/unittests/PackedKernelArgumentArrayTest.cpp

TEST_F(DeviceMemoryPackingTest, SingleConstTypedGlobalPointer) {
   ExpectEqual(TypedGlobal.getHandle(), sizeof(void *),
               Type::GLOBAL_DEVICE_MEMORY, Array, 0);
   EXPECT_EQ(1u, Array.getArgumentCount());
   EXPECT_EQ(0u, Array.getSharedCount());
 }
 
 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());
 }
 
 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());
 }
 
 TEST_F(DeviceMemoryPackingTest, PackSeveralArguments) {
   const se::GlobalDeviceMemory<int> *TypedGlobalPointer = &TypedGlobal;
   const se::SharedDeviceMemory<int> *TypedSharedPointer = &TypedShared;
   auto Array = se::make_kernel_argument_pack(Value, TypedGlobal, &TypedGlobal,
                                              TypedGlobalPointer, TypedShared,
                                              &TypedShared, TypedSharedPointer);
   ExpectEqual(&Value, sizeof(Value), Type::VALUE, Array, 0);
   ExpectEqual(TypedGlobal.getHandle(), sizeof(void *),
               Type::GLOBAL_DEVICE_MEMORY, Array, 1);
   ExpectEqual(TypedGlobal.getHandle(), sizeof(void *),
               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());
 }
 
 } // namespace