[SE] Don't pack raw device mem args

[jhen]

Summary

Step 4 of getting GlobalDeviceMemory to own its handle.

Take out code to pack untyped device memory types as kernel arguments.
When GlobalDeviceMemory owns its handle, users will never touch untyped
device memory types, so they will never pass them as kernel args.

[jlebar]

streamexecutor/lib/unittests/PackedKernelArgumentArrayTest.cpp

104 ↗

(On Diff #70120)

As a separate patch, I wonder if these classes should have a getByteCount() so you don't have to do this error-prone calculation.

[jhen]

streamexecutor/lib/unittests/PackedKernelArgumentArrayTest.cpp

104 ↗

(On Diff #70120)

Good idea. I'll do that.

parallel-libs/trunk/streamexecutor/include/streamexecutor/PackedKernelArgumentArray.h

private:
   // Pack a const GlobalDeviceMemory<T> pointer argument.
   template <typename T>
   void PackOneArgument(size_t Index, const GlobalDeviceMemory<T> *Argument) {
     Addresses[Index] = Argument->getHandle();
     Sizes[Index] = sizeof(void *);
     Types[Index] = KernelArgumentType::GLOBAL_DEVICE_MEMORY;
   }
 
-  // Pack a SharedDeviceMemoryBase argument.
-  void PackOneArgument(size_t Index, const SharedDeviceMemoryBase &Argument) {
-    ++SharedCount;
-    Addresses[Index] = nullptr;
-    Sizes[Index] = Argument.getByteCount();
-    Types[Index] = KernelArgumentType::SHARED_DEVICE_MEMORY;
-  }
-
-  // Pack a SharedDeviceMemoryBase pointer argument.
-  void PackOneArgument(size_t Index, SharedDeviceMemoryBase *Argument) {
-    ++SharedCount;
-    Addresses[Index] = nullptr;
-    Sizes[Index] = Argument->getByteCount();
-    Types[Index] = KernelArgumentType::SHARED_DEVICE_MEMORY;
-  }
-
-  // Pack a const SharedDeviceMemoryBase pointer argument.
-  void PackOneArgument(size_t Index, const SharedDeviceMemoryBase *Argument) {
-    ++SharedCount;
-    Addresses[Index] = nullptr;
-    Sizes[Index] = Argument->getByteCount();
-    Types[Index] = KernelArgumentType::SHARED_DEVICE_MEMORY;
-  }
-
   // Pack a SharedDeviceMemory argument.
   template <typename T>
   void PackOneArgument(size_t Index, const SharedDeviceMemory<T> &Argument) {
     ++SharedCount;
     Addresses[Index] = nullptr;
-    Sizes[Index] = Argument.getByteCount();
+    Sizes[Index] = Argument.getElementCount() * sizeof(T);
     Types[Index] = KernelArgumentType::SHARED_DEVICE_MEMORY;
   }
 
   // Pack a SharedDeviceMemory pointer argument.
   template <typename T>
   void PackOneArgument(size_t Index, SharedDeviceMemory<T> *Argument) {
     ++SharedCount;
     Addresses[Index] = nullptr;
-    Sizes[Index] = Argument->getByteCount();
+    Sizes[Index] = Argument->getElementCount() * sizeof(T);
     Types[Index] = KernelArgumentType::SHARED_DEVICE_MEMORY;
   }
 
   // Pack a const SharedDeviceMemory pointer argument.
   template <typename T>
   void PackOneArgument(size_t Index, const SharedDeviceMemory<T> *Argument) {
     ++SharedCount;
     Addresses[Index] = nullptr;
-    Sizes[Index] = Argument->getByteCount();
+    Sizes[Index] = Argument->getElementCount() * sizeof(T);
     Types[Index] = KernelArgumentType::SHARED_DEVICE_MEMORY;
   }
 
   std::array<const void *, sizeof...(ParameterTs)> Addresses;
   std::array<size_t, sizeof...(ParameterTs)> Sizes;
   std::array<KernelArgumentType, sizeof...(ParameterTs)> Types;
 };
 

parallel-libs/trunk/streamexecutor/lib/unittests/PackedKernelArgumentArrayTest.cpp

 //===-- PackedKernelArgumentArrayTest.cpp - tests for kernel arg packing --===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 ///
 /// \file
 /// Unit tests for kernel argument packing.
 ///
 //===----------------------------------------------------------------------===//
 
+#include "SimpleHostPlatformDevice.h"
+#include "streamexecutor/Device.h"
 #include "streamexecutor/DeviceMemory.h"
 #include "streamexecutor/PackedKernelArgumentArray.h"
+#include "streamexecutor/PlatformInterfaces.h"
 
 #include "llvm/ADT/Twine.h"
 
 #include "gtest/gtest.h"
 
 namespace {
 
 namespace se = ::streamexecutor;
 
 using Type = se::KernelArgumentType;
 
 // Test fixture class for testing argument packing.
 //
 // Basically defines a bunch of types to be packed so they don't have to be
 // defined separately in each test.
 class DeviceMemoryPackingTest : public ::testing::Test {
 public:
   DeviceMemoryPackingTest()
-      : Value(42), Handle(&Value), ByteCount(15), ElementCount(5),
-        UntypedGlobal(Handle, ByteCount),
-        TypedGlobal(se::GlobalDeviceMemory<int>::makeFromElementCount(
-            Handle, ElementCount)),
-        UntypedShared(ByteCount),
+      : Device(&PDevice), Value(42), Handle(&Value), ByteCount(15),
+        ElementCount(5),
+        TypedGlobal(getOrDie(Device.allocateDeviceMemory<int>(ElementCount))),
         TypedShared(
             se::SharedDeviceMemory<int>::makeFromElementCount(ElementCount)) {}
 
+  se::test::SimpleHostPlatformDevice PDevice;
+  se::Device Device;
   int Value;
   void *Handle;
   size_t ByteCount;
   size_t ElementCount;
-  se::GlobalDeviceMemoryBase UntypedGlobal;
   se::GlobalDeviceMemory<int> TypedGlobal;
-  se::SharedDeviceMemoryBase UntypedShared;
   se::SharedDeviceMemory<int> TypedShared;
 };
 
 // Utility method to check the expected address, size, and type for a packed
 // argument at the given index of a PackedKernelArgumentArray.
 template <typename... ParameterTs>
 static void
 ExpectEqual(const void *ExpectedAddress, size_t ExpectedSize, Type ExpectedType,

 
 TEST_F(DeviceMemoryPackingTest, SingleValue) {
   auto Array = se::make_kernel_argument_pack(Value);
   ExpectEqual(&Value, sizeof(Value), Type::VALUE, Array, 0);
   EXPECT_EQ(1u, Array.getArgumentCount());
   EXPECT_EQ(0u, Array.getSharedCount());
 }
 
-TEST_F(DeviceMemoryPackingTest, SingleUntypedGlobal) {
-  auto Array = se::make_kernel_argument_pack(UntypedGlobal);
-  ExpectEqual(Handle, sizeof(void *), Type::GLOBAL_DEVICE_MEMORY, Array, 0);
-  EXPECT_EQ(1u, Array.getArgumentCount());
-  EXPECT_EQ(0u, Array.getSharedCount());
-}
-
-TEST_F(DeviceMemoryPackingTest, SingleUntypedGlobalPointer) {
-  auto Array = se::make_kernel_argument_pack(&UntypedGlobal);
-  ExpectEqual(Handle, sizeof(void *), Type::GLOBAL_DEVICE_MEMORY, Array, 0);
-  EXPECT_EQ(1u, Array.getArgumentCount());
-  EXPECT_EQ(0u, Array.getSharedCount());
-}
-
-TEST_F(DeviceMemoryPackingTest, SingleConstUntypedGlobalPointer) {
-  const se::GlobalDeviceMemoryBase *ConstPointer = &UntypedGlobal;
-  auto Array = se::make_kernel_argument_pack(ConstPointer);
-  ExpectEqual(Handle, sizeof(void *), Type::GLOBAL_DEVICE_MEMORY, Array, 0);
-  EXPECT_EQ(1u, Array.getArgumentCount());
-  EXPECT_EQ(0u, Array.getSharedCount());
-}
-
 TEST_F(DeviceMemoryPackingTest, SingleTypedGlobal) {
   auto Array = se::make_kernel_argument_pack(TypedGlobal);
-  ExpectEqual(Handle, sizeof(void *), Type::GLOBAL_DEVICE_MEMORY, Array, 0);
+  ExpectEqual(TypedGlobal.getHandle(), sizeof(void *),
+              Type::GLOBAL_DEVICE_MEMORY, Array, 0);
   EXPECT_EQ(1u, Array.getArgumentCount());
   EXPECT_EQ(0u, Array.getSharedCount());
 }
 
 TEST_F(DeviceMemoryPackingTest, SingleTypedGlobalPointer) {
   auto Array = se::make_kernel_argument_pack(&TypedGlobal);
-  ExpectEqual(Handle, sizeof(void *), Type::GLOBAL_DEVICE_MEMORY, Array, 0);
+  ExpectEqual(TypedGlobal.getHandle(), sizeof(void *),
+              Type::GLOBAL_DEVICE_MEMORY, Array, 0);
   EXPECT_EQ(1u, Array.getArgumentCount());
   EXPECT_EQ(0u, Array.getSharedCount());
 }
 
 TEST_F(DeviceMemoryPackingTest, SingleConstTypedGlobalPointer) {
   const se::GlobalDeviceMemory<int> *ArgumentPointer = &TypedGlobal;
   auto Array = se::make_kernel_argument_pack(ArgumentPointer);
-  ExpectEqual(Handle, sizeof(void *), Type::GLOBAL_DEVICE_MEMORY, Array, 0);
+  ExpectEqual(TypedGlobal.getHandle(), sizeof(void *),
+              Type::GLOBAL_DEVICE_MEMORY, Array, 0);
   EXPECT_EQ(1u, Array.getArgumentCount());
   EXPECT_EQ(0u, Array.getSharedCount());
 }
 
-TEST_F(DeviceMemoryPackingTest, SingleUntypedShared) {
-  auto Array = se::make_kernel_argument_pack(UntypedShared);
-  ExpectEqual(nullptr, UntypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
-              Array, 0);
-  EXPECT_EQ(1u, Array.getArgumentCount());
-  EXPECT_EQ(1u, Array.getSharedCount());
-}
-
-TEST_F(DeviceMemoryPackingTest, SingleUntypedSharedPointer) {
-  auto Array = se::make_kernel_argument_pack(&UntypedShared);
-  ExpectEqual(nullptr, UntypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
-              Array, 0);
-  EXPECT_EQ(1u, Array.getArgumentCount());
-  EXPECT_EQ(1u, Array.getSharedCount());
-}
-
-TEST_F(DeviceMemoryPackingTest, SingleConstUntypedSharedPointer) {
-  const se::SharedDeviceMemoryBase *ArgumentPointer = &UntypedShared;
-  auto Array = se::make_kernel_argument_pack(ArgumentPointer);
-  ExpectEqual(nullptr, UntypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
-              Array, 0);
-  EXPECT_EQ(1u, Array.getArgumentCount());
-  EXPECT_EQ(1u, Array.getSharedCount());
-}
-
 TEST_F(DeviceMemoryPackingTest, SingleTypedShared) {
   auto Array = se::make_kernel_argument_pack(TypedShared);
-  ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
-              Array, 0);
+  ExpectEqual(nullptr, TypedShared.getElementCount() * sizeof(int),
+              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.getByteCount(), Type::SHARED_DEVICE_MEMORY,
-              Array, 0);
+  ExpectEqual(nullptr, TypedShared.getElementCount() * sizeof(int),
+              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.getByteCount(), Type::SHARED_DEVICE_MEMORY,
-              Array, 0);
+  ExpectEqual(nullptr, TypedShared.getElementCount() * sizeof(int),
+              Type::SHARED_DEVICE_MEMORY, Array, 0);
   EXPECT_EQ(1u, Array.getArgumentCount());
   EXPECT_EQ(1u, Array.getSharedCount());
 }
 
 TEST_F(DeviceMemoryPackingTest, PackSeveralArguments) {
-  const se::GlobalDeviceMemoryBase *UntypedGlobalPointer = &UntypedGlobal;
   const se::GlobalDeviceMemory<int> *TypedGlobalPointer = &TypedGlobal;
-  const se::SharedDeviceMemoryBase *UntypedSharedPointer = &UntypedShared;
   const se::SharedDeviceMemory<int> *TypedSharedPointer = &TypedShared;
-  auto Array = se::make_kernel_argument_pack(
-      Value, UntypedGlobal, &UntypedGlobal, UntypedGlobalPointer, TypedGlobal,
-      &TypedGlobal, TypedGlobalPointer, UntypedShared, &UntypedShared,
-      UntypedSharedPointer, TypedShared, &TypedShared, TypedSharedPointer);
+  auto Array = se::make_kernel_argument_pack(Value, TypedGlobal, &TypedGlobal,
+                                             TypedGlobalPointer, TypedShared,
+                                             &TypedShared, TypedSharedPointer);
   ExpectEqual(&Value, sizeof(Value), Type::VALUE, Array, 0);
-  ExpectEqual(Handle, sizeof(void *), Type::GLOBAL_DEVICE_MEMORY, Array, 1);
-  ExpectEqual(Handle, sizeof(void *), Type::GLOBAL_DEVICE_MEMORY, Array, 2);
-  ExpectEqual(Handle, sizeof(void *), Type::GLOBAL_DEVICE_MEMORY, Array, 3);
-  ExpectEqual(Handle, sizeof(void *), Type::GLOBAL_DEVICE_MEMORY, Array, 4);
-  ExpectEqual(Handle, sizeof(void *), Type::GLOBAL_DEVICE_MEMORY, Array, 5);
-  ExpectEqual(Handle, sizeof(void *), Type::GLOBAL_DEVICE_MEMORY, Array, 6);
-  ExpectEqual(nullptr, UntypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
-              Array, 7);
-  ExpectEqual(nullptr, UntypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
-              Array, 8);
-  ExpectEqual(nullptr, UntypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
-              Array, 9);
-  ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
-              Array, 10);
-  ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
-              Array, 11);
-  ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
-              Array, 12);
-  EXPECT_EQ(13u, Array.getArgumentCount());
-  EXPECT_EQ(6u, Array.getSharedCount());
+  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);
+  EXPECT_EQ(7u, Array.getArgumentCount());
+  EXPECT_EQ(3u, Array.getSharedCount());
 }
 
 } // namespace