[StreamExecutor] Add DeviceMemory and kernel arg packing

Summary:
Add types for device memory and add the code that knows how to pack these
device memory types if they are passed as arguments to kernel launches.

Reviewers: jlebar, tra

Subscribers: parallel_libs-commits

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

llvm-svn: 278021

2 files changed, 212 insertions, 0 deletions

diff --git a/parallel-libs/streamexecutor/lib/unittests/CMakeLists.txt b/parallel-libs/streamexecutor/lib/unittests/CMakeLists.txt
index d88332a2ee7..efe5f76eb26 100644
--- a/parallel-libs/streamexecutor/lib/unittests/CMakeLists.txt
+++ b/parallel-libs/streamexecutor/lib/unittests/CMakeLists.txt
@@ -17,3 +17,13 @@ target_link_libraries(
     ${GTEST_BOTH_LIBRARIES}
     ${CMAKE_THREAD_LIBS_INIT})
 add_test(KernelSpecTest kernel_spec_test)
+
+add_executable(
+    packed_kernel_argument_array_test
+    PackedKernelArgumentArrayTest.cpp)
+target_link_libraries(
+    packed_kernel_argument_array_test
+    ${llvm_libs}
+    ${GTEST_BOTH_LIBRARIES}
+    ${CMAKE_THREAD_LIBS_INIT})
+add_test(PackedKernelArgumentArrayTest packed_kernel_argument_array_test)
diff --git a/parallel-libs/streamexecutor/lib/unittests/PackedKernelArgumentArrayTest.cpp b/parallel-libs/streamexecutor/lib/unittests/PackedKernelArgumentArrayTest.cpp
new file mode 100644
index 00000000000..7e87b59f656
--- /dev/null
+++ b/parallel-libs/streamexecutor/lib/unittests/PackedKernelArgumentArrayTest.cpp
@@ -0,0 +1,202 @@
+//===-- 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 "streamexecutor/DeviceMemory.h"
+#include "streamexecutor/PackedKernelArgumentArray.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),
+        TypedShared(
+            se::SharedDeviceMemory<int>::makeFromElementCount(ElementCount)) {}
+
+  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,
+            const se::PackedKernelArgumentArray<ParameterTs...> &Observed,
+            size_t Index) {
+  SCOPED_TRACE(("Index = " + llvm::Twine(Index)).str());
+  EXPECT_EQ(ExpectedAddress, Observed.getAddress(Index));
+  EXPECT_EQ(ExpectedAddress, Observed.getAddresses()[Index]);
+  EXPECT_EQ(ExpectedSize, Observed.getSize(Index));
+  EXPECT_EQ(ExpectedSize, Observed.getSizes()[Index]);
+  EXPECT_EQ(ExpectedType, Observed.getType(Index));
+  EXPECT_EQ(ExpectedType, Observed.getTypes()[Index]);
+}
+
+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);
+  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);
+  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);
+  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);
+  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);
+  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);
+  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);
+  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());
+}
+
+} // namespace