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

   1 //===-- PackedKernelArgumentArrayTest.cpp - tests for kernel arg packing --===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file is distributed under the University of Illinois Open Source
   6 // License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9 ///
  10 /// \file
  11 /// Unit tests for kernel argument packing.
  12 ///
  13 //===----------------------------------------------------------------------===//
  14
  15 #include "SimpleHostPlatformDevice.h"
  16 #include "streamexecutor/Device.h"
  17 #include "streamexecutor/DeviceMemory.h"
  18 #include "streamexecutor/PackedKernelArgumentArray.h"
  19 #include "streamexecutor/PlatformDevice.h"
  20
  21 #include "llvm/ADT/Twine.h"
  22
  23 #include "gtest/gtest.h"
  24
  25 namespace {
  26
  27 namespace se = ::streamexecutor;
  28
  29 using Type = se::KernelArgumentType;
  30
  31 // Test fixture class for testing argument packing.
  32 //
  33 // Basically defines a bunch of types to be packed so they don't have to be
  34 // defined separately in each test.
  35 class DeviceMemoryPackingTest : public ::testing::Test {
  36 public:
  37   DeviceMemoryPackingTest()
  38       : Device(&PDevice), Value(42), Handle(&Value), ByteCount(15),
  39         ElementCount(5),
  40         TypedGlobal(getOrDie(Device.allocateDeviceMemory<int>(ElementCount))),
  41         TypedShared(
  42             se::SharedDeviceMemory<int>::makeFromElementCount(ElementCount)) {}
  43
  44   se::test::SimpleHostPlatformDevice PDevice;
  45   se::Device Device;
  46   int Value;
  47   void *Handle;
  48   size_t ByteCount;
  49   size_t ElementCount;
  50   se::GlobalDeviceMemory<int> TypedGlobal;
  51   se::SharedDeviceMemory<int> TypedShared;
  52 };
  53
  54 // Utility method to check the expected address, size, and type for a packed
  55 // argument at the given index of a PackedKernelArgumentArray.
  56 template <typename... ParameterTs>
  57 static void
  58 ExpectEqual(const void *ExpectedAddress, size_t ExpectedSize, Type ExpectedType,
  59             const se::PackedKernelArgumentArray<ParameterTs...> &Observed,
  60             size_t Index) {
  61   SCOPED_TRACE(("Index = " + llvm::Twine(Index)).str());
  62   EXPECT_EQ(ExpectedAddress, Observed.getAddress(Index));
  63   EXPECT_EQ(ExpectedAddress, Observed.getAddresses()[Index]);
  64   EXPECT_EQ(ExpectedSize, Observed.getSize(Index));
  65   EXPECT_EQ(ExpectedSize, Observed.getSizes()[Index]);
  66   EXPECT_EQ(ExpectedType, Observed.getType(Index));
  67   EXPECT_EQ(ExpectedType, Observed.getTypes()[Index]);
  68 }
  69
  70 TEST_F(DeviceMemoryPackingTest, SingleValue) {
  71   auto Array = se::make_kernel_argument_pack(Value);
  72   ExpectEqual(&Value, sizeof(Value), Type::VALUE, Array, 0);
  73   EXPECT_EQ(1u, Array.getArgumentCount());
  74   EXPECT_EQ(0u, Array.getSharedCount());
  75 }
  76
  77 TEST_F(DeviceMemoryPackingTest, SingleTypedGlobal) {
  78   auto Array = se::make_kernel_argument_pack(TypedGlobal);
  79   ExpectEqual(TypedGlobal.getHandle(), sizeof(void *),
  80               Type::GLOBAL_DEVICE_MEMORY, Array, 0);
  81   EXPECT_EQ(1u, Array.getArgumentCount());
  82   EXPECT_EQ(0u, Array.getSharedCount());
  83 }
  84
  85 TEST_F(DeviceMemoryPackingTest, SingleTypedGlobalPointer) {
  86   auto Array = se::make_kernel_argument_pack(&TypedGlobal);
  87   ExpectEqual(TypedGlobal.getHandle(), sizeof(void *),
  88               Type::GLOBAL_DEVICE_MEMORY, Array, 0);
  89   EXPECT_EQ(1u, Array.getArgumentCount());
  90   EXPECT_EQ(0u, Array.getSharedCount());
  91 }
  92
  93 TEST_F(DeviceMemoryPackingTest, SingleConstTypedGlobalPointer) {
  94   const se::GlobalDeviceMemory<int> *ArgumentPointer = &TypedGlobal;
  95   auto Array = se::make_kernel_argument_pack(ArgumentPointer);
  96   ExpectEqual(TypedGlobal.getHandle(), sizeof(void *),
  97               Type::GLOBAL_DEVICE_MEMORY, Array, 0);
  98   EXPECT_EQ(1u, Array.getArgumentCount());
  99   EXPECT_EQ(0u, Array.getSharedCount());
 100 }
 101
 102 TEST_F(DeviceMemoryPackingTest, SingleTypedShared) {
 103   auto Array = se::make_kernel_argument_pack(TypedShared);
 104   ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
 105               Array, 0);
 106   EXPECT_EQ(1u, Array.getArgumentCount());
 107   EXPECT_EQ(1u, Array.getSharedCount());
 108 }
 109
 110 TEST_F(DeviceMemoryPackingTest, SingleTypedSharedPointer) {
 111   auto Array = se::make_kernel_argument_pack(&TypedShared);
 112   ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
 113               Array, 0);
 114   EXPECT_EQ(1u, Array.getArgumentCount());
 115   EXPECT_EQ(1u, Array.getSharedCount());
 116 }
 117
 118 TEST_F(DeviceMemoryPackingTest, SingleConstTypedSharedPointer) {
 119   const se::SharedDeviceMemory<int> *ArgumentPointer = &TypedShared;
 120   auto Array = se::make_kernel_argument_pack(ArgumentPointer);
 121   ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
 122               Array, 0);
 123   EXPECT_EQ(1u, Array.getArgumentCount());
 124   EXPECT_EQ(1u, Array.getSharedCount());
 125 }
 126
 127 TEST_F(DeviceMemoryPackingTest, PackSeveralArguments) {
 128   const se::GlobalDeviceMemory<int> *TypedGlobalPointer = &TypedGlobal;
 129   const se::SharedDeviceMemory<int> *TypedSharedPointer = &TypedShared;
 130   auto Array = se::make_kernel_argument_pack(Value, TypedGlobal, &TypedGlobal,
 131                                              TypedGlobalPointer, TypedShared,
 132                                              &TypedShared, TypedSharedPointer);
 133   ExpectEqual(&Value, sizeof(Value), Type::VALUE, Array, 0);
 134   ExpectEqual(TypedGlobal.getHandle(), sizeof(void *),
 135               Type::GLOBAL_DEVICE_MEMORY, Array, 1);
 136   ExpectEqual(TypedGlobal.getHandle(), sizeof(void *),
 137               Type::GLOBAL_DEVICE_MEMORY, Array, 2);
 138   ExpectEqual(TypedGlobal.getHandle(), sizeof(void *),
 139               Type::GLOBAL_DEVICE_MEMORY, Array, 3);
 140   ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
 141               Array, 4);
 142   ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
 143               Array, 5);
 144   ExpectEqual(nullptr, TypedShared.getByteCount(), Type::SHARED_DEVICE_MEMORY,
 145               Array, 6);
 146   EXPECT_EQ(7u, Array.getArgumentCount());
 147   EXPECT_EQ(3u, Array.getSharedCount());
 148 }
 149
 150 } // namespace