[StreamExecutor] Dev handles in platform interface

Summary:
This is the first in a series of patches that will convert
GlobalDeviceMemory to own its device memory handle. The first step is to
remove GlobalDeviceMemoryBase from the PlatformInterface interfaces and
use raw handles there instead. This is useful because
GlobalDeviceMemoryBase is going to lose its importance in this process.

Reviewers: jlebar

Subscribers: jprice, parallel_libs-commits

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

llvm-svn: 280401

3 files changed, 139 insertions, 121 deletions

diff --git a/parallel-libs/streamexecutor/lib/unittests/DeviceTest.cpp b/parallel-libs/streamexecutor/lib/unittests/DeviceTest.cpp
index cb34b8b92d5..93d378fff5a 100644
--- a/parallel-libs/streamexecutor/lib/unittests/DeviceTest.cpp
+++ b/parallel-libs/streamexecutor/lib/unittests/DeviceTest.cpp
@@ -15,6 +15,7 @@
 #include <cstdlib>
 #include <cstring>
 
+#include "SimpleHostPlatformDevice.h"
 #include "streamexecutor/Device.h"
 #include "streamexecutor/PlatformInterfaces.h"
 
@@ -24,79 +25,6 @@ namespace {
 
 namespace se = ::streamexecutor;
 
-class MockPlatformDevice : public se::PlatformDevice {
-public:
-  ~MockPlatformDevice() override {}
-
-  std::string getName() const override { return "MockPlatformDevice"; }
-
-  se::Expected<std::unique_ptr<se::PlatformStreamHandle>>
-  createStream() override {
-    return se::make_error("not implemented");
-  }
-
-  se::Expected<se::GlobalDeviceMemoryBase>
-  allocateDeviceMemory(size_t ByteCount) override {
-    return se::GlobalDeviceMemoryBase(std::malloc(ByteCount));
-  }
-
-  se::Error freeDeviceMemory(se::GlobalDeviceMemoryBase Memory) override {
-    std::free(const_cast<void *>(Memory.getHandle()));
-    return se::Error::success();
-  }
-
-  se::Expected<void *> allocateHostMemory(size_t ByteCount) override {
-    return std::malloc(ByteCount);
-  }
-
-  se::Error freeHostMemory(void *Memory) override {
-    std::free(Memory);
-    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,
-                               size_t ByteCount) override {
-    std::memcpy(static_cast<char *>(HostDst) + DstByteOffset,
-                static_cast<const char *>(DeviceSrc.getHandle()) +
-                    SrcByteOffset,
-                ByteCount);
-    return se::Error::success();
-  }
-
-  se::Error synchronousCopyH2D(const void *HostSrc, size_t SrcByteOffset,
-                               se::GlobalDeviceMemoryBase DeviceDst,
-                               size_t DstByteOffset,
-                               size_t ByteCount) override {
-    std::memcpy(static_cast<char *>(const_cast<void *>(DeviceDst.getHandle())) +
-                    DstByteOffset,
-                static_cast<const char *>(HostSrc) + SrcByteOffset, ByteCount);
-    return se::Error::success();
-  }
-
-  se::Error synchronousCopyD2D(se::GlobalDeviceMemoryBase DeviceDst,
-                               size_t DstByteOffset,
-                               const se::GlobalDeviceMemoryBase &DeviceSrc,
-                               size_t SrcByteOffset,
-                               size_t ByteCount) override {
-    std::memcpy(static_cast<char *>(const_cast<void *>(DeviceDst.getHandle())) +
-                    DstByteOffset,
-                static_cast<const char *>(DeviceSrc.getHandle()) +
-                    SrcByteOffset,
-                ByteCount);
-    return se::Error::success();
-  }
-};
-
 /// Test fixture to hold objects used by tests.
 class DeviceTest : public ::testing::Test {
 public:
@@ -124,7 +52,7 @@ public:
   int Host5[5];
   int Host7[7];
 
-  MockPlatformDevice PDevice;
+  SimpleHostPlatformDevice PDevice;
   se::Device Device;
 };
 
diff --git a/parallel-libs/streamexecutor/lib/unittests/SimpleHostPlatformDevice.h b/parallel-libs/streamexecutor/lib/unittests/SimpleHostPlatformDevice.h
new file mode 100644
index 00000000000..a2dd3c8738f
--- /dev/null
+++ b/parallel-libs/streamexecutor/lib/unittests/SimpleHostPlatformDevice.h
@@ -0,0 +1,135 @@
+//===-- SimpleHostPlatformDevice.h - Host device for testing ----*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// The SimpleHostPlatformDevice class is a streamexecutor::PlatformDevice that
+/// is really just the host processor and memory. It is useful for testing
+/// because no extra device platform is required.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef STREAMEXECUTOR_LIB_UNITTESTS_SIMPLEHOSTPLATFORMDEVICE_H
+#define STREAMEXECUTOR_LIB_UNITTESTS_SIMPLEHOSTPLATFORMDEVICE_H
+
+#include <cstdlib>
+#include <cstring>
+
+#include "streamexecutor/PlatformInterfaces.h"
+
+/// A streamexecutor::PlatformDevice that simply forwards all operations to the
+/// host platform.
+///
+/// The allocate and copy methods are simple wrappers for std::malloc and
+/// std::memcpy.
+class SimpleHostPlatformDevice : public streamexecutor::PlatformDevice {
+  std::string getName() const override { return "SimpleHostPlatformDevice"; }
+
+  streamexecutor::Expected<
+      std::unique_ptr<streamexecutor::PlatformStreamHandle>>
+  createStream() override {
+    return nullptr;
+  }
+
+  streamexecutor::Expected<void *>
+  allocateDeviceMemory(size_t ByteCount) override {
+    return std::malloc(ByteCount);
+  }
+
+  streamexecutor::Error freeDeviceMemory(const void *Handle) override {
+    std::free(const_cast<void *>(Handle));
+    return streamexecutor::Error::success();
+  }
+
+  streamexecutor::Expected<void *>
+  allocateHostMemory(size_t ByteCount) override {
+    return std::malloc(ByteCount);
+  }
+
+  streamexecutor::Error freeHostMemory(void *Memory) override {
+    std::free(const_cast<void *>(Memory));
+    return streamexecutor::Error::success();
+  }
+
+  streamexecutor::Error registerHostMemory(void *Memory,
+                                           size_t ByteCount) override {
+    return streamexecutor::Error::success();
+  }
+
+  streamexecutor::Error unregisterHostMemory(void *Memory) override {
+    return streamexecutor::Error::success();
+  }
+
+  streamexecutor::Error copyD2H(streamexecutor::PlatformStreamHandle *S,
+                                const void *DeviceHandleSrc,
+                                size_t SrcByteOffset, void *HostDst,
+                                size_t DstByteOffset,
+                                size_t ByteCount) override {
+    std::memcpy(static_cast<char *>(HostDst) + DstByteOffset,
+                static_cast<const char *>(DeviceHandleSrc) + SrcByteOffset,
+                ByteCount);
+    return streamexecutor::Error::success();
+  }
+
+  streamexecutor::Error copyH2D(streamexecutor::PlatformStreamHandle *S,
+                                const void *HostSrc, size_t SrcByteOffset,
+                                const void *DeviceHandleDst,
+                                size_t DstByteOffset,
+                                size_t ByteCount) override {
+    std::memcpy(static_cast<char *>(const_cast<void *>(DeviceHandleDst)) +
+                    DstByteOffset,
+                static_cast<const char *>(HostSrc) + SrcByteOffset, ByteCount);
+    return streamexecutor::Error::success();
+  }
+
+  streamexecutor::Error
+  copyD2D(streamexecutor::PlatformStreamHandle *S, const void *DeviceHandleSrc,
+          size_t SrcByteOffset, const void *DeviceHandleDst,
+          size_t DstByteOffset, size_t ByteCount) override {
+    std::memcpy(static_cast<char *>(const_cast<void *>(DeviceHandleDst)) +
+                    DstByteOffset,
+                static_cast<const char *>(DeviceHandleSrc) + SrcByteOffset,
+                ByteCount);
+    return streamexecutor::Error::success();
+  }
+
+  streamexecutor::Error synchronousCopyD2H(const void *DeviceHandleSrc,
+                                           size_t SrcByteOffset, void *HostDst,
+                                           size_t DstByteOffset,
+                                           size_t ByteCount) override {
+    std::memcpy(static_cast<char *>(HostDst) + DstByteOffset,
+                static_cast<const char *>(DeviceHandleSrc) + SrcByteOffset,
+                ByteCount);
+    return streamexecutor::Error::success();
+  }
+
+  streamexecutor::Error synchronousCopyH2D(const void *HostSrc,
+                                           size_t SrcByteOffset,
+                                           const void *DeviceHandleDst,
+                                           size_t DstByteOffset,
+                                           size_t ByteCount) override {
+    std::memcpy(static_cast<char *>(const_cast<void *>(DeviceHandleDst)) +
+                    DstByteOffset,
+                static_cast<const char *>(HostSrc) + SrcByteOffset, ByteCount);
+    return streamexecutor::Error::success();
+  }
+
+  streamexecutor::Error synchronousCopyD2D(const void *DeviceHandleSrc,
+                                           size_t SrcByteOffset,
+                                           const void *DeviceHandleDst,
+                                           size_t DstByteOffset,
+                                           size_t ByteCount) override {
+    std::memcpy(static_cast<char *>(const_cast<void *>(DeviceHandleDst)) +
+                    DstByteOffset,
+                static_cast<const char *>(DeviceHandleSrc) + SrcByteOffset,
+                ByteCount);
+    return streamexecutor::Error::success();
+  }
+};
+
+#endif // STREAMEXECUTOR_LIB_UNITTESTS_SIMPLEHOSTPLATFORMDEVICE_H
diff --git a/parallel-libs/streamexecutor/lib/unittests/StreamTest.cpp b/parallel-libs/streamexecutor/lib/unittests/StreamTest.cpp
index d05c928dcb3..b194bf02082 100644
--- a/parallel-libs/streamexecutor/lib/unittests/StreamTest.cpp
+++ b/parallel-libs/streamexecutor/lib/unittests/StreamTest.cpp
@@ -14,6 +14,7 @@
 
 #include <cstring>
 
+#include "SimpleHostPlatformDevice.h"
 #include "streamexecutor/Device.h"
 #include "streamexecutor/Kernel.h"
 #include "streamexecutor/KernelSpec.h"
@@ -26,52 +27,6 @@ namespace {
 
 namespace se = ::streamexecutor;
 
-/// Mock PlatformDevice that performs asynchronous memcpy operations by
-/// ignoring the stream argument and calling std::memcpy on device memory
-/// handles.
-class MockPlatformDevice : public se::PlatformDevice {
-public:
-  ~MockPlatformDevice() override {}
-
-  std::string getName() const override { return "MockPlatformDevice"; }
-
-  se::Expected<std::unique_ptr<se::PlatformStreamHandle>>
-  createStream() override {
-    return nullptr;
-  }
-
-  se::Error copyD2H(se::PlatformStreamHandle *S,
-                    const se::GlobalDeviceMemoryBase &DeviceSrc,
-                    size_t SrcByteOffset, void *HostDst, size_t DstByteOffset,
-                    size_t ByteCount) override {
-    std::memcpy(HostDst, static_cast<const char *>(DeviceSrc.getHandle()) +
-                             SrcByteOffset,
-                ByteCount);
-    return se::Error::success();
-  }
-
-  se::Error copyH2D(se::PlatformStreamHandle *S, const void *HostSrc,
-                    size_t SrcByteOffset, se::GlobalDeviceMemoryBase DeviceDst,
-                    size_t DstByteOffset, size_t ByteCount) override {
-    std::memcpy(static_cast<char *>(const_cast<void *>(DeviceDst.getHandle())) +
-                    DstByteOffset,
-                HostSrc, ByteCount);
-    return se::Error::success();
-  }
-
-  se::Error copyD2D(se::PlatformStreamHandle *S,
-                    const se::GlobalDeviceMemoryBase &DeviceSrc,
-                    size_t SrcByteOffset, se::GlobalDeviceMemoryBase DeviceDst,
-                    size_t DstByteOffset, size_t ByteCount) override {
-    std::memcpy(static_cast<char *>(const_cast<void *>(DeviceDst.getHandle())) +
-                    DstByteOffset,
-                static_cast<const char *>(DeviceSrc.getHandle()) +
-                    SrcByteOffset,
-                ByteCount);
-    return se::Error::success();
-  }
-};
-
 /// Test fixture to hold objects used by tests.
 class StreamTest : public ::testing::Test {
 public:
@@ -100,7 +55,7 @@ protected:
   int Host5[5];
   int Host7[7];
 
-  MockPlatformDevice PDevice;
+  SimpleHostPlatformDevice PDevice;
   se::Stream Stream;
 };