[Parallel_libs-commits] [parallel-libs] r279658 - [StreamExecutor] Fix allocateDeviceMemory

[Jason Henline via Parallel_libs-commits]

Author: jhen
Date: Wed Aug 24 14:42:03 2016
New Revision: 279658

URL: http://llvm.org/viewvc/llvm-project?rev=279658&view=rev
Log:
[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

Modified:
    parallel-libs/trunk/streamexecutor/include/streamexecutor/Executor.h
    parallel-libs/trunk/streamexecutor/lib/unittests/ExecutorTest.cpp

Modified: parallel-libs/trunk/streamexecutor/include/streamexecutor/Executor.h
URL: http://llvm.org/viewvc/llvm-project/parallel-libs/trunk/streamexecutor/include/streamexecutor/Executor.h?rev=279658&r1=279657&r2=279658&view=diff
==============================================================================
--- parallel-libs/trunk/streamexecutor/include/streamexecutor/Executor.h (original)
+++ parallel-libs/trunk/streamexecutor/include/streamexecutor/Executor.h Wed Aug 24 14:42:03 2016
@@ -41,7 +41,11 @@ public:
   /// Allocates an array of ElementCount entries of type T in device memory.
   template <typename T>
   Expected<GlobalDeviceMemory<T>> allocateDeviceMemory(size_t ElementCount) {
-    return PExecutor->allocateDeviceMemory(ElementCount * sizeof(T));
+    Expected<GlobalDeviceMemoryBase> MaybeBase =
+        PExecutor->allocateDeviceMemory(ElementCount * sizeof(T));
+    if (!MaybeBase)
+      return MaybeBase.takeError();
+    return GlobalDeviceMemory<T>(*MaybeBase);
   }
 
   /// Frees memory previously allocated with allocateDeviceMemory.
@@ -54,7 +58,11 @@ public:
   /// Host memory allocated by this function can be used for asynchronous memory
   /// copies on streams. See Stream::thenCopyD2H and Stream::thenCopyH2D.
   template <typename T> Expected<T *> allocateHostMemory(size_t ElementCount) {
-    return PExecutor->allocateHostMemory(ElementCount * sizeof(T));
+    Expected<void *> MaybeMemory =
+        PExecutor->allocateHostMemory(ElementCount * sizeof(T));
+    if (!MaybeMemory)
+      return MaybeMemory.takeError();
+    return static_cast<T *>(*MaybeMemory);
   }
 
   /// Frees memory previously allocated with allocateHostMemory.

Modified: parallel-libs/trunk/streamexecutor/lib/unittests/ExecutorTest.cpp
URL: http://llvm.org/viewvc/llvm-project/parallel-libs/trunk/streamexecutor/lib/unittests/ExecutorTest.cpp?rev=279658&r1=279657&r2=279658&view=diff
==============================================================================
--- parallel-libs/trunk/streamexecutor/lib/unittests/ExecutorTest.cpp (original)
+++ parallel-libs/trunk/streamexecutor/lib/unittests/ExecutorTest.cpp Wed Aug 24 14:42:03 2016
@@ -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) {