[StreamExecutor] Dev handles in platform interface

[jhen]

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.

[jlebar]

I like factoring out this test platform.

streamexecutor/include/streamexecutor/PlatformInterfaces.h

108 ↗

(On Diff #69933)

Do we want the dst handles to be non-const void*s? Or is the idea that the handle is always a const thing because you never write to it from host code?

134 ↗

(On Diff #69933)

Same for the handle we're freeing here. (Although, as was pointed out to me recently, it's perfectly OK to call delete on a C++ const pointer.)

[jhen]

streamexecutor/include/streamexecutor/PlatformInterfaces.h

108 ↗

(On Diff #69933)

Yeah, I wanted it to be const because I think that's more appropriate for an opaque handle. I'll keep it that way unless others want it changed.

134 ↗

(On Diff #69933)

My idea is that the platform-specific implementation Will probably have to cast this handle to something else before using it anyway, so throwing on a const_cast too, if necessary, should seem natural.

parallel-libs/trunk/streamexecutor/include/streamexecutor/Device.h

public:
   }
 
   /// Creates a stream object for this device.
   Expected<Stream> createStream();
 
   /// Allocates an array of ElementCount entries of type T in device memory.
   template <typename T>
   Expected<GlobalDeviceMemory<T>> allocateDeviceMemory(size_t ElementCount) {
-    Expected<GlobalDeviceMemoryBase> MaybeBase =
+    Expected<void *> MaybeMemory =
         PDevice->allocateDeviceMemory(ElementCount * sizeof(T));
-    if (!MaybeBase)
-      return MaybeBase.takeError();
-    return GlobalDeviceMemory<T>(*MaybeBase);
+    if (!MaybeMemory)
+      return MaybeMemory.takeError();
+    return GlobalDeviceMemory<T>::makeFromElementCount(*MaybeMemory,
+                                                       ElementCount);
   }
 
   /// Frees memory previously allocated with allocateDeviceMemory.
   template <typename T> Error freeDeviceMemory(GlobalDeviceMemory<T> Memory) {
-    return PDevice->freeDeviceMemory(Memory);
+    return PDevice->freeDeviceMemory(Memory.getHandle());
   }
 
   /// Allocates an array of ElementCount entries of type T in host memory.
   ///
   /// 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) {
     Expected<void *> MaybeMemory =
if (ElementCount > Src.getElementCount())
       return make_error("copying too many elements, " +
                         llvm::Twine(ElementCount) +
                         ", from a device array of element count " +
                         llvm::Twine(Src.getElementCount()));
     if (ElementCount > Dst.size())
       return make_error(
           "copying too many elements, " + llvm::Twine(ElementCount) +
           ", to a host array of element count " + llvm::Twine(Dst.size()));
-    return PDevice->synchronousCopyD2H(Src.getBaseMemory(),
+    return PDevice->synchronousCopyD2H(Src.getBaseMemory().getHandle(),
                                        Src.getElementOffset() * sizeof(T),
                                        Dst.data(), 0, ElementCount * sizeof(T));
   }
 
   template <typename T>
   Error synchronousCopyD2H(GlobalDeviceMemorySlice<T> Src,
                            llvm::MutableArrayRef<T> Dst) {
     if (Src.getElementCount() != Dst.size())
if (ElementCount > Src.size())
       return make_error(
           "copying too many elements, " + llvm::Twine(ElementCount) +
           ", from a host array of element count " + llvm::Twine(Src.size()));
     if (ElementCount > Dst.getElementCount())
       return make_error("copying too many elements, " +
                         llvm::Twine(ElementCount) +
                         ", to a device array of element count " +
                         llvm::Twine(Dst.getElementCount()));
-    return PDevice->synchronousCopyH2D(Src.data(), 0, Dst.getBaseMemory(),
-                                       Dst.getElementOffset() * sizeof(T),
-                                       ElementCount * sizeof(T));
+    return PDevice->synchronousCopyH2D(
+        Src.data(), 0, Dst.getBaseMemory().getHandle(),
+        Dst.getElementOffset() * sizeof(T), ElementCount * sizeof(T));
   }
 
   template <typename T>
   Error synchronousCopyH2D(llvm::ArrayRef<T> Src,
                            GlobalDeviceMemorySlice<T> Dst) {
     if (Src.size() != Dst.getElementCount())
       return make_error(
           "array size mismatch for H2D, host source has element count " +
if (ElementCount > Src.getElementCount())
                         ", from a device array of element count " +
                         llvm::Twine(Src.getElementCount()));
     if (ElementCount > Dst.getElementCount())
       return make_error("copying too many elements, " +
                         llvm::Twine(ElementCount) +
                         ", to a device array of element count " +
                         llvm::Twine(Dst.getElementCount()));
     return PDevice->synchronousCopyD2D(
-        Src.getBaseMemory(), Src.getElementOffset() * sizeof(T),
-        Dst.getBaseMemory(), Dst.getElementOffset() * sizeof(T),
+        Src.getBaseMemory().getHandle(), Src.getElementOffset() * sizeof(T),
+        Dst.getBaseMemory().getHandle(), Dst.getElementOffset() * sizeof(T),
         ElementCount * sizeof(T));
   }
 
   template <typename T>
   Error synchronousCopyD2D(GlobalDeviceMemorySlice<T> Src,
                            GlobalDeviceMemorySlice<T> Dst) {
     if (Src.getElementCount() != Dst.getElementCount())
       return make_error(

parallel-libs/trunk/streamexecutor/include/streamexecutor/PlatformInterfaces.h

virtual Error launch(PlatformStreamHandle *S, BlockDimensions BlockSize,
                        const PackedKernelArgumentArrayBase &ArgumentArray) {
     return make_error("launch not implemented for platform " + getName());
   }
 
   /// Copies data from the device to the host.
   ///
   /// HostDst should have been allocated by allocateHostMemory or registered
   /// with registerHostMemory.
-  virtual Error copyD2H(PlatformStreamHandle *S,
-                        const GlobalDeviceMemoryBase &DeviceSrc,
+  virtual Error copyD2H(PlatformStreamHandle *S, const void *DeviceSrcHandle,
                         size_t SrcByteOffset, void *HostDst,
                         size_t DstByteOffset, size_t ByteCount) {
     return make_error("copyD2H not implemented for platform " + getName());
   }
 
   /// Copies data from the host to the device.
   ///
   /// HostSrc should have been allocated by allocateHostMemory or registered
   /// with registerHostMemory.
   virtual Error copyH2D(PlatformStreamHandle *S, const void *HostSrc,
-                        size_t SrcByteOffset, GlobalDeviceMemoryBase DeviceDst,
+                        size_t SrcByteOffset, const void *DeviceDstHandle,
                         size_t DstByteOffset, size_t ByteCount) {
     return make_error("copyH2D not implemented for platform " + getName());
   }
 
   /// Copies data from one device location to another.
-  virtual Error copyD2D(PlatformStreamHandle *S,
-                        const GlobalDeviceMemoryBase &DeviceSrc,
-                        size_t SrcByteOffset, GlobalDeviceMemoryBase DeviceDst,
+  virtual Error copyD2D(PlatformStreamHandle *S, const void *DeviceSrcHandle,
+                        size_t SrcByteOffset, const void *DeviceDstHandle,
                         size_t DstByteOffset, size_t ByteCount) {
     return make_error("copyD2D not implemented for platform " + getName());
   }
 
   /// Blocks the host until the given stream completes all the work enqueued up
   /// to the point this function is called.
   virtual Error blockHostUntilDone(PlatformStreamHandle *S) {
     return make_error("blockHostUntilDone not implemented for platform " +
                       getName());
   }
 
   /// Allocates untyped device memory of a given size in bytes.
-  virtual Expected<GlobalDeviceMemoryBase>
-  allocateDeviceMemory(size_t ByteCount) {
+  virtual Expected<void *> allocateDeviceMemory(size_t ByteCount) {
     return make_error("allocateDeviceMemory not implemented for platform " +
                       getName());
   }
 
   /// Frees device memory previously allocated by allocateDeviceMemory.
-  virtual Error freeDeviceMemory(GlobalDeviceMemoryBase Memory) {
+  virtual Error freeDeviceMemory(const void *Handle) {
     return make_error("freeDeviceMemory not implemented for platform " +
                       getName());
   }
 
   /// Allocates untyped host memory of a given size in bytes.
   ///
   /// Host memory allocated via this method is suitable for use with copyH2D and
   /// copyD2H.
return make_error("unregisterHostMemory not implemented for platform " +
                       getName());
   }
 
   /// Copies the given number of bytes from device memory to host memory.
   ///
   /// Blocks the calling host thread until the copy is completed. Can operate on
   /// any host memory, not just registered host memory or host memory allocated
   /// by allocateHostMemory. Does not block any ongoing device calls.
-  virtual Error synchronousCopyD2H(const GlobalDeviceMemoryBase &DeviceSrc,
+  virtual Error synchronousCopyD2H(const void *DeviceSrcHandle,
                                    size_t SrcByteOffset, void *HostDst,
                                    size_t DstByteOffset, size_t ByteCount) {
     return make_error("synchronousCopyD2H not implemented for platform " +
                       getName());
   }
 
-  /// Similar to synchronousCopyD2H(const GlobalDeviceMemoryBase &, size_t, void
+  /// Similar to synchronousCopyD2H(const void *, size_t, void
   /// *, size_t, size_t), but copies memory from host to device rather than
   /// device to host.
   virtual Error synchronousCopyH2D(const void *HostSrc, size_t SrcByteOffset,
-                                   GlobalDeviceMemoryBase DeviceDst,
+                                   const void *DeviceDstHandle,
                                    size_t DstByteOffset, size_t ByteCount) {
     return make_error("synchronousCopyH2D not implemented for platform " +
                       getName());
   }
 
-  /// Similar to synchronousCopyD2H(const GlobalDeviceMemoryBase &, size_t, void
+  /// Similar to synchronousCopyD2H(const void *, size_t, void
   /// *, size_t, size_t), but copies memory from one location in device memory
   /// to another rather than from device to host.
-  virtual Error synchronousCopyD2D(GlobalDeviceMemoryBase DeviceDst,
+  virtual Error synchronousCopyD2D(const void *DeviceDstHandle,
                                    size_t DstByteOffset,
-                                   const GlobalDeviceMemoryBase &DeviceSrc,
+                                   const void *DeviceSrcHandle,
                                    size_t SrcByteOffset, size_t ByteCount) {
     return make_error("synchronousCopyD2D not implemented for platform " +
                       getName());
   }
 };
 
 } // namespace streamexecutor
 
 #endif // STREAMEXECUTOR_PLATFORMINTERFACES_H

parallel-libs/trunk/streamexecutor/include/streamexecutor/Stream.h

Stream &thenCopyD2H(GlobalDeviceMemorySlice<T> Src,
     if (ElementCount > Src.getElementCount())
       setError("copying too many elements, " + llvm::Twine(ElementCount) +
                ", from a device array of element count " +
                llvm::Twine(Src.getElementCount()));
     else if (ElementCount > Dst.size())
       setError("copying too many elements, " + llvm::Twine(ElementCount) +
                ", to a host array of element count " + llvm::Twine(Dst.size()));
     else
-      setError(PDevice->copyD2H(ThePlatformStream.get(), Src.getBaseMemory(),
+      setError(PDevice->copyD2H(ThePlatformStream.get(),
+                                Src.getBaseMemory().getHandle(),
                                 Src.getElementOffset() * sizeof(T), Dst.data(),
                                 0, ElementCount * sizeof(T)));
     return *this;
   }
 
   template <typename T>
   Stream &thenCopyD2H(GlobalDeviceMemorySlice<T> Src,
                       llvm::MutableArrayRef<T> Dst) {
if (ElementCount > Src.size())
       setError("copying too many elements, " + llvm::Twine(ElementCount) +
                ", from a host array of element count " +
                llvm::Twine(Src.size()));
     else if (ElementCount > Dst.getElementCount())
       setError("copying too many elements, " + llvm::Twine(ElementCount) +
                ", to a device array of element count " +
                llvm::Twine(Dst.getElementCount()));
     else
-      setError(PDevice->copyH2D(
-          ThePlatformStream.get(), Src.data(), 0, Dst.getBaseMemory(),
-          Dst.getElementOffset() * sizeof(T), ElementCount * sizeof(T)));
+      setError(PDevice->copyH2D(ThePlatformStream.get(), Src.data(), 0,
+                                Dst.getBaseMemory().getHandle(),
+                                Dst.getElementOffset() * sizeof(T),
+                                ElementCount * sizeof(T)));
     return *this;
   }
 
   template <typename T>
   Stream &thenCopyH2D(llvm::ArrayRef<T> Src, GlobalDeviceMemorySlice<T> Dst) {
     if (Src.size() != Dst.getElementCount())
       setError("array size mismatch for H2D, host source has element count " +
                llvm::Twine(Src.size()) +
if (ElementCount > Src.getElementCount())
                ", from a device array of element count " +
                llvm::Twine(Src.getElementCount()));
     else if (ElementCount > Dst.getElementCount())
       setError("copying too many elements, " + llvm::Twine(ElementCount) +
                ", to a device array of element count " +
                llvm::Twine(Dst.getElementCount()));
     else
       setError(PDevice->copyD2D(
-          ThePlatformStream.get(), Src.getBaseMemory(),
-          Src.getElementOffset() * sizeof(T), Dst.getBaseMemory(),
+          ThePlatformStream.get(), Src.getBaseMemory().getHandle(),
+          Src.getElementOffset() * sizeof(T), Dst.getBaseMemory().getHandle(),
           Dst.getElementOffset() * sizeof(T), ElementCount * sizeof(T)));
     return *this;
   }
 
   template <typename T>
   Stream &thenCopyD2D(GlobalDeviceMemorySlice<T> Src,
                       GlobalDeviceMemorySlice<T> Dst) {
     if (Src.getElementCount() != Dst.getElementCount())

parallel-libs/trunk/streamexecutor/lib/unittests/DeviceTest.cpp

 /// \file
 /// This file contains the unit tests for Device code.
 ///
 //===----------------------------------------------------------------------===//
 
 #include <cstdlib>
 #include <cstring>
 
+#include "SimpleHostPlatformDevice.h"
 #include "streamexecutor/Device.h"
 #include "streamexecutor/PlatformInterfaces.h"
 
 #include "gtest/gtest.h"
 
 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:
   DeviceTest()
       : HostA5{0, 1, 2, 3, 4}, HostB5{5, 6, 7, 8, 9},
         HostA7{10, 11, 12, 13, 14, 15, 16}, HostB7{17, 18, 19, 20, 21, 22, 23},
         DeviceA5(se::GlobalDeviceMemory<int>::makeFromElementCount(HostA5, 5)),
         DeviceB5(se::GlobalDeviceMemory<int>::makeFromElementCount(HostB5, 5)),
public:
   se::GlobalDeviceMemory<int> DeviceB5;
   se::GlobalDeviceMemory<int> DeviceA7;
   se::GlobalDeviceMemory<int> DeviceB7;
 
   // Host memory to be used as actual host memory.
   int Host5[5];
   int Host7[7];
 
-  MockPlatformDevice PDevice;
+  SimpleHostPlatformDevice PDevice;
   se::Device Device;
 };
 
 #define EXPECT_NO_ERROR(E) EXPECT_FALSE(static_cast<bool>(E))
 #define EXPECT_ERROR(E)                                                        \
   do {                                                                         \
     se::Error E__ = E;                                                         \
     EXPECT_TRUE(static_cast<bool>(E__));                                       \

parallel-libs/trunk/streamexecutor/lib/unittests/SimpleHostPlatformDevice.h

+//===-- 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

parallel-libs/trunk/streamexecutor/lib/unittests/StreamTest.cpp

 //===-- StreamTest.cpp - Tests for Stream ---------------------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 ///
 /// \file
 /// This file contains the unit tests for Stream code.
 ///
 //===----------------------------------------------------------------------===//
 
 #include <cstring>
 
+#include "SimpleHostPlatformDevice.h"
 #include "streamexecutor/Device.h"
 #include "streamexecutor/Kernel.h"
 #include "streamexecutor/KernelSpec.h"
 #include "streamexecutor/PlatformInterfaces.h"
 #include "streamexecutor/Stream.h"
 
 #include "gtest/gtest.h"
 
 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:
   StreamTest()
       : HostA5{0, 1, 2, 3, 4}, HostB5{5, 6, 7, 8, 9},
         HostA7{10, 11, 12, 13, 14, 15, 16}, HostB7{17, 18, 19, 20, 21, 22, 23},
         DeviceA5(se::GlobalDeviceMemory<int>::makeFromElementCount(HostA5, 5)),
         DeviceB5(se::GlobalDeviceMemory<int>::makeFromElementCount(HostB5, 5)),
protected:
   se::GlobalDeviceMemory<int> DeviceB5;
   se::GlobalDeviceMemory<int> DeviceA7;
   se::GlobalDeviceMemory<int> DeviceB7;
 
   // Host memory to be used as actual host memory.
   int Host5[5];
   int Host7[7];
 
-  MockPlatformDevice PDevice;
+  SimpleHostPlatformDevice PDevice;
   se::Stream Stream;
 };
 
 using llvm::ArrayRef;
 using llvm::MutableArrayRef;
 
 // D2H tests