[StreamExecutor] Fix allocateDeviceMemory

Summary:
The return value from PlatformExecutor::allocateDeviceMemory needs to be
converted from Expected<GlobalDeviceMemoryBase> to
Expected<GlobalDeviceMemory<T>> in Executor::allocateDeviceMemory.

A similar bug is also fixed for Executor::allocateHostMemory.

Thanks to jprice for identifying this bug.

Reviewers: jprice, jlebar

Subscribers: parallel_libs-commits

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

llvm-svn: 279658

1 files changed, 27 insertions, 0 deletions

diff --git a/parallel-libs/streamexecutor/lib/unittests/ExecutorTest.cpp b/parallel-libs/streamexecutor/lib/unittests/ExecutorTest.cpp
index d2d03fb6c88..b6719d303ec 100644
--- a/parallel-libs/streamexecutor/lib/unittests/ExecutorTest.cpp
+++ b/parallel-libs/streamexecutor/lib/unittests/ExecutorTest.cpp
@@ -54,6 +54,14 @@ public:
     return se::Error::success();
   }
 
+  se::Error registerHostMemory(void *, size_t) override {
+    return se::Error::success();
+  }
+
+  se::Error unregisterHostMemory(void *) override {
+    return se::Error::success();
+  }
+
   se::Error synchronousCopyD2H(const se::GlobalDeviceMemoryBase &DeviceSrc,
                                size_t SrcByteOffset, void *HostDst,
                                size_t DstByteOffset,
@@ -131,6 +139,25 @@ public:
 using llvm::ArrayRef;
 using llvm::MutableArrayRef;
 
+TEST_F(ExecutorTest, AllocateAndFreeDeviceMemory) {
+  se::Expected<se::GlobalDeviceMemory<int>> MaybeMemory =
+      Executor.allocateDeviceMemory<int>(10);
+  EXPECT_TRUE(static_cast<bool>(MaybeMemory));
+  EXPECT_NO_ERROR(Executor.freeDeviceMemory(*MaybeMemory));
+}
+
+TEST_F(ExecutorTest, AllocateAndFreeHostMemory) {
+  se::Expected<int *> MaybeMemory = Executor.allocateHostMemory<int>(10);
+  EXPECT_TRUE(static_cast<bool>(MaybeMemory));
+  EXPECT_NO_ERROR(Executor.freeHostMemory(*MaybeMemory));
+}
+
+TEST_F(ExecutorTest, RegisterAndUnregisterHostMemory) {
+  std::vector<int> Data(10);
+  EXPECT_NO_ERROR(Executor.registerHostMemory(Data.data(), 10));
+  EXPECT_NO_ERROR(Executor.unregisterHostMemory(Data.data()));
+}
+
 // D2H tests
 
 TEST_F(ExecutorTest, SyncCopyD2HToMutableArrayRefByCount) {