Index: parallel-libs/trunk/streamexecutor/include/streamexecutor/Executor.h =================================================================== --- parallel-libs/trunk/streamexecutor/include/streamexecutor/Executor.h +++ parallel-libs/trunk/streamexecutor/include/streamexecutor/Executor.h @@ -41,7 +41,11 @@ /// 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 @@ /// 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. Index: parallel-libs/trunk/streamexecutor/lib/unittests/ExecutorTest.cpp =================================================================== --- parallel-libs/trunk/streamexecutor/lib/unittests/ExecutorTest.cpp +++ parallel-libs/trunk/streamexecutor/lib/unittests/ExecutorTest.cpp @@ -54,6 +54,14 @@ 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 @@ 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) {