[StreamExecutor] Add basic Stream operations

[jhen]

Summary

Add the Stream class and a few of the operations it supports.

[jlebar]

streamexecutor/include/streamexecutor/StreamExecutor.h

57 ↗

(On Diff #67413)

It looks like Stream is concrete and movable, so can we get rid of the unique_ptr?

streamexecutor/lib/Stream.cpp

36 ↗

(On Diff #67413)

So...it's illegal to move a stream if there are any outstanding operations? (Among other reasons, we're not taking a mutex here.) Maybe we should have a comment.

38 ↗

(On Diff #67413)

Is it important for us to have an API that would let consumers capture this error, if they wanted to?

[jlebar]

streamexecutor/include/streamexecutor/Stream.h

17 ↗

(On Diff #67413)

Extra space

21 ↗

(On Diff #67413)

Missing a comma before "and".

25 ↗

(On Diff #67413)

, and

26 ↗

(On Diff #67413)

to be

27 ↗

(On Diff #67413)

Perhaps: Multiple Stream instances can be created for the same StreamExecutor. This allows ...

72 ↗

(On Diff #67413)

This is intriguing...

Is the idea that you don't want to make SE a friend of Stream, because that's too much of a layering violation?

I wonder if you could make SEKey a (private?) type inside Stream.

109 ↗

(On Diff #67413)

LLVM style is to leave off curly braces where possible.

113 ↗

(On Diff #67413)

Maybe s/ for the invocation//?

122 ↗

(On Diff #67413)

Do you want Params&&... and std::forward(Arguments...) below?

124 ↗

(On Diff #67413)

Do we want std::move(ArgumentArray)? (At the moment it doesn't make sense because thenRawLaunch takes a const ref to the pack, but see below...)

132 ↗

(On Diff #67413)

StreamExecutor::registerHostMemory might be a better name, but we can cross that bridge.

136 ↗

(On Diff #67413)

These default args seem weird...would this be better as two functions, one that does a full copy and one that takes an ElementCount?

203 ↗

(On Diff #67413)

"up to this point"?

216 ↗

(On Diff #67413)

Could we store an Error internally, instead of a string?

[jhen]

• Respond to jlebar's comments

streamexecutor/include/streamexecutor/Stream.h

72 ↗

(On Diff #67413)

Is the idea that you don't want to make SE a friend of Stream, because that's too much of a layering violation?

Yes, I thought it would be nice to clarify the interface between StreamExecutor and Stream. Maybe this is overkill, but I thought it seemed lightweight enough to be OK.

I wonder if you could make SEKey a (private?) type inside Stream.

I put SEKey into Stream, but I don't know if there is any way to make it a private type. I think StreamExecutor has to be able to see the definition so it can create an instance of SEKey, but maybe I'm missing something.

122 ↗

(On Diff #67413)

Ugh, thanks for catching that ugly copy. I changed the arguments to const Params &... Arguments instead of Params &&... for the following horrifying reason:

When I tried to use Params &&..., I got a compiler error when I tried to call thenLaunch with one int argument and one GlobalDeviceMemory<int> argument. The error said the compiler could not infer the types for Params because the TypedKernel<Params...> argument made it infer int, GlobalDeviceMemory<int> and the Params &&... argument made it infer int, GlobalDeviceMemory<int> &.

So, the && seems to make the compiler infer a reference type for some arguments and that conflicts with the types from TypedKernel<Params...>.

Actually, I think const ref is the right signature for this function anyway, so I don't think we need to figure out a solution for the conflicting type inference. What do you think?

124 ↗

(On Diff #67413)

Maybe I missed it, but I didn't see the comment below about thenRawLaunch. I didn't change anything for now until we clear up the confusion.

136 ↗

(On Diff #67413)

Thanks for the idea. I was struggling with that design. Multiple functions is much nicer.

216 ↗

(On Diff #67413)

The problem I had with trying to use an Error is that the method to get out the error string destroys the error object. This is a side-effect of the way LLVM Error enforces the law that users cannot ignore their errors. So it looked like I was going to have to keep destroying and rebuilding the saved Error every time somebody queried it.

streamexecutor/include/streamexecutor/StreamExecutor.h

57 ↗

(On Diff #67413)

Oops, I was experimenting with moving streams and forgot to delete the stuff. Streams shouldn't be movable, so we'll stick with a pointer here.

streamexecutor/lib/Stream.cpp

36 ↗

(On Diff #67413)

Sorry I forgot to delete this garbage before. It's gone now and Streams are no longer movable.

38 ↗

(On Diff #67413)

I think we can do that in a sane way. I change StreamExecutor::deallocateStream to return void and added a comment that says it will record an error if it fails. I left it as a TODO for now to design the interface for fetching stream destruction errors.

[jprice]

I think I understand the motivation for restricting the use of StreamExecutor::launch/memcpy etc to only the Stream class via the StreamKey struct, but I'm not sure I see the need for the StreamExecutorKey struct. It seems like this is currently just protecting the Stream constructor, but is it really a problem if someone can write new Stream(Executor, nullptr) as well as Executor->createStream()? Don't have strong opinions about this, just curious.

The StreamExecutor class seems to be missing a virtual destructor (warned by Clang).

[jhen]

• Respond to jprice's comments

[jhen]

In D23333#511039, @jprice wrote:

I think I understand the motivation for restricting the use of StreamExecutor::launch/memcpy etc to only the Stream class via the StreamKey struct, but I'm not sure I see the need for the StreamExecutorKey struct. It seems like this is currently just protecting the Stream constructor, but is it really a problem if someone can write new Stream(Executor, nullptr) as well as Executor->createStream()? Don't have strong opinions about this, just curious.

My goal was to make sure that Streams always have non-null implementation pointers. Basically, I didn't want any half-formed Stream objects floating around. This goal was really undercut when I originally left in the code to allow Streams to be moved. After being moved from, a Stream would be half-formed. But now I've removed the move stuff for Streams, and I think that leaves us with a useful invariant that Streams will always be fully-formed.

I added a comment to the Stream constructor to try to clarify this because I think you're right that it's a bit confusing, and I think others will have the same question.

The StreamExecutor class seems to be missing a virtual destructor (warned by Clang).

Thanks for catching that. I have added in a virtual destructor now.

[jlebar]

streamexecutor/include/streamexecutor/Stream.h

90 ↗

(On Diff #67540)

If that's the only point, could we just assert that Implementation is not null?

One advantage of getting rid of SEKey would be, that (maybe?) we could move StreamKey into Stream itself. Right now we can't do that because you can't forward-declare a nested type, and so there's no way to have StreamKey inside Stream and SEKey inside SE and have them both depend on each other via #includes.

123 ↗

(On Diff #67540)

Discussed IRL, it sounds like the point is to pass references, not ownership. So this is good.

125 ↗

(On Diff #67540)

Sorry, I meant to add another comment. But it's moot if we're taking const refs.

217 ↗

(On Diff #67540)

Hm. Okay, I don't see a better way to do this at the moment. If we had a more complex error class it would be worth worrying about, but as-is, it's fine.

272 ↗

(On Diff #67540)

Nit, I think you want llvm::optional<std::string>.

[jhen]

• Respond to jlebar's comments (take 2)

streamexecutor/include/streamexecutor/Stream.h

90 ↗

(On Diff #67540)

Thinking more about this, I realized that the more important point is that I don't want users dealing with the StreamInterface class at all. It should be private to StreamExecutor. So I think it is good to mark this interface as only callable by StreamExecutor so user's don't get the wrong idea about StreamInterface. I updated the constructor comment to say this instead of talking about null pointers.

I agree that I don't like the way the StreamKey is part of the StreamExecutor class and vice versa. To deal with this, I factored the whole key object pattern out into a header called Passkey<T> that makes it possible for any class to make its own key instance. I think this cleans up the design significantly.

[jlebar]

streamexecutor/include/streamexecutor/Passkey.h

1 ↗

(On Diff #67572)

Move into Util?

41 ↗

(On Diff #67572)

Should we take Passkey<B> by value instead of by reference?

55 ↗

(On Diff #67572)

We call it a "reference" in the comment and here, but apparently it's a pointer; that should be harmonized. I kind of like the idea of keeping a proper C++ reference, so it's clear that it's never null.

62 ↗

(On Diff #67572)

Not sure about the "&". In fact I would expect that if you have

void foo(Passkey<Bar>& b);

and then have

void Bar::DoSomething() {
  foo(this);
}

that won't compile because we're trying to bind a temporary Passkey to b.

65 ↗

(On Diff #67572)

I think you'll need a copy-constructor if you want to pass these by value. I guess we could take them by rvalue reference if you really wanted to disallow copying. But even then we'd still need to allow moves, which is almost as bad from an access-control perspective.

So I'm not sure disallowing copying is worthwhile.

streamexecutor/include/streamexecutor/Stream.h

91 ↗

(On Diff #67572)

I realized that the more important point is that I don't want users dealing with the StreamInterface class at all. It should be private to StreamExecutor.

Can't SE simply avoid passing out references to StreamInterfaces, then?

For the "private-to-Stream" functions in SE, I'm wondering something similar: It seems like maybe SE should be split into two types. One type is visible to end-users, and it contains only functions they'd call. The other type is visible only to e.g. Stream, and it actually does work. Stream would take a pointer to that thing in its constructor.

Since SE isn't going to give out pointers to this "internal SE" except when creating Streams, we get access control for free.

[jhen]

• Respond to jlebar's comments (take 3)

streamexecutor/include/streamexecutor/Passkey.h

1 ↗

(On Diff #67572)

I got rid of this class (see comment below), so I'm marking all the comments on this class as "done". Thanks for suggesting a better way to handle all this.

streamexecutor/include/streamexecutor/Stream.h

91 ↗

(On Diff #67572)

Yes, I think the design you outline here is much better. I created a class called StreamParent to hold the interface that a Stream uses to call its parent StreamExecutor. I got rid of the Passkey class and all the key arguments.

This seems much cleaner. Thanks for the advice!

[jhen]

• Fixes to breaks from last patch

[jhen]

Previous patch accidentally left things in a broken state. Should be correct now.

[jlebar]

streamexecutor/include/streamexecutor/Stream.h

64 ↗

(On Diff #67638)

explicit?

110 ↗

(On Diff #67638)

LLVM style is to leave off curly braces where possible.

149 ↗

(On Diff #67638)

Leave off curly braces? (I know, it's awful, but consistency is good too...) Same below.

streamexecutor/lib/Stream.cpp

51 ↗

(On Diff #67638)

I wonder if we need any of these *raw* functions anymore, now that they're just forwarding to the parent.

[jhen]

• Remove StreamParent class

[jhen]

This version gets rid of the StreamParent because it was obscuring the logic about which functions are overridden by each platform. I renamed the file Interfaces.h to PlatformInterfaces.h to try to clarify its purpose as the interfaces that each platform must implement. I also added the PlatformStreamExecutor interface class in PlatformInterfaces.h, and that class takes the place of what StreamParent was trying to do.

Thanks to jlebar for talking this over with me offline to try to get the design right.

streamexecutor/include/streamexecutor/Stream.h

64 ↗

(On Diff #67638)

It wasn't meant to be and it's not anymore.

streamexecutor/lib/Stream.cpp

51 ↗

(On Diff #67638)

I got rid of them for now, but I might bring them back later if users need to include Stream.h and I don't want them to see the PlatformInterface.h file.

[jprice]

FWIW, this design feels a lot cleaner to me :-)

streamexecutor/include/streamexecutor/PlatformInterfaces.h

48 ↗

(On Diff #67791)

Missing virtual destructor in this class.

60 ↗

(On Diff #67791)

Should there be a corresponding platform-specific allocateStream()?

[jhen]

• Add PlatformStream

streamexecutor/include/streamexecutor/PlatformInterfaces.h

60 ↗

(On Diff #67791)

Really good point. In response, I changed this part of the design. Now there is a PlatformStream base class that is responsible for cleaning up its resources during destruction, so there is no more deallocateStream method.

I also changed the methods in PlatformStreamExecutor to accept PlatformStream pointers rather than Stream pointers. This seems to fit better with the design of having the Platform class methods take raw handles rather than public wrappers.

[jprice]

One last nit from me.

streamexecutor/include/streamexecutor/Stream.h

51 ↗

(On Diff #67859)

This declaration isn't needed anymore.

[jhen]

• Remove StreamParent forward decl

[jhen]

streamexecutor/include/streamexecutor/Stream.h

52 ↗

(On Diff #67871)

Thanks for catching that.

[jlebar]

Sorry to keep nit'ing on the design here. I think each revision has been an improvement, though; hopefully we don't start going in circles.

streamexecutor/include/streamexecutor/PlatformInterfaces.h

42 ↗

(On Diff #67871)

If Stream is still calling into PlatformSE to do its work (so this class is just an RAII thing), maybe PlatformStreamHandle would be a better name?

streamexecutor/include/streamexecutor/Stream.h

64 ↗

(On Diff #67871)

Now we're kind of taking redundant information here. PStream's PE must match PlatformExecutor, right?

It doesn't seem so bad to me if PStream exposed a pointer back to its PExecutor...

streamexecutor/include/streamexecutor/StreamExecutor.h

30 ↗

(On Diff #67871)

What does it mean to have an SE without a backing PlatformSE?

31 ↗

(On Diff #67871)

Presumably you want explicit?

[jhen]

• Tweaks to PlatformStream

streamexecutor/include/streamexecutor/PlatformInterfaces.h

42 ↗

(On Diff #67871)

Yes, that is a better name. I made the change.

streamexecutor/include/streamexecutor/StreamExecutor.h

30 ↗

(On Diff #67871)

I just forgot to delete this ctor during development. It is gone now.

[jhen]

In D23333#514158, @jlebar wrote:

Sorry to keep nit'ing on the design here. I think each revision has been an improvement, though; hopefully we don't start going in circles.

No worries from my end. I definitely feel like things are getting better each time. Thanks for your help!

[jlebar]

streamexecutor/include/streamexecutor/PlatformInterfaces.h

45 ↗

(On Diff #67890)

Why do we allow a null executor?

57 ↗

(On Diff #67890)

Not sold on the "parent" part of the name. At best it seems redundant. Would it be confusing if we just called it "executor"?

streamexecutor/lib/StreamExecutor.cpp

35 ↗

(On Diff #67890)

Hm, I think this should be the PE's responsibility, since otherwise anyone who calls PE::createStream will have to remember to do this. We can assert it here if you want.

[jhen]

• Make platform exec set stream exec during creation

streamexecutor/include/streamexecutor/PlatformInterfaces.h

45 ↗

(On Diff #67890)

This was to allow a default construction followed by a call to setParentExecutor. It's removed now because it is not needed since I moved the responsibility of setting the executor pointer to the platform-specific implementation (as per jlebar's comment below).

parallel-libs/trunk/streamexecutor/CMakeLists.txt

if(STREAM_EXECUTOR_STANDALONE)
 
     find_package(LLVM REQUIRED CONFIG)
     message(STATUS "Found LLVM ${LLVM_PACKAGE_VERSION}")
     message(STATUS "Using LLVMConfig.cmake in: ${LLVM_DIR}")
 
     include_directories(${LLVM_INCLUDE_DIRS})
     add_definitions(${LLVM_DEFINITIONS})
 
+    # Get the LLVM cxxflags by using llvm-config.
+    #
+    # This is necessary to get -fno-rtti if LLVM is compiled that way.
+    execute_process(
+        COMMAND
+        "${LLVM_BINARY_DIR}/bin/llvm-config"
+        --cxxflags
+        OUTPUT_VARIABLE
+        LLVM_CXXFLAGS
+        OUTPUT_STRIP_TRAILING_WHITESPACE)
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${LLVM_CXXFLAGS}")
+
     # Find the libraries that correspond to the LLVM components
     # that we wish to use
     llvm_map_components_to_libnames(llvm_libs support symbolize)
 
     if(STREAM_EXECUTOR_UNIT_TESTS)
         enable_testing()
         find_package(GTest REQUIRED)
         include_directories(${GTEST_INCLUDE_DIRS})

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

-//===-- Interfaces.h - Interfaces to platform-specific impls ----*- C++ -*-===//
-//
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
-//
-//===----------------------------------------------------------------------===//
-///
-/// \file
-/// Interfaces to platform-specific StreamExecutor type implementations.
-///
-//===----------------------------------------------------------------------===//
-
-#ifndef STREAMEXECUTOR_INTERFACES_H
-#define STREAMEXECUTOR_INTERFACES_H
-
-namespace streamexecutor {
-
-/// Methods supported by device kernel function objects on all platforms.
-class KernelInterface {
-  // TODO(jhen): Add methods.
-};
-
-// TODO(jhen): Add other interfaces such as Stream.
-
-} // namespace streamexecutor
-
-#endif // STREAMEXECUTOR_INTERFACES_H

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

+//===-- LaunchDimensions.h - Kernel block and grid sizes --------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// Structures to hold sizes for blocks and grids which are used as parameters
+/// for kernel launches.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef STREAMEXECUTOR_LAUNCHDIMENSIONS_H
+#define STREAMEXECUTOR_LAUNCHDIMENSIONS_H
+
+namespace streamexecutor {
+
+/// The dimensions of a device block of execution.
+///
+/// A block is made up of an array of X by Y by Z threads.
+struct BlockDimensions {
+  BlockDimensions(unsigned X = 1, unsigned Y = 1, unsigned Z = 1)
+      : X(X), Y(Y), Z(Z) {}
+
+  unsigned X;
+  unsigned Y;
+  unsigned Z;
+};
+
+/// The dimensions of a device grid of execution.
+///
+/// A grid is made up of an array of X by Y by Z blocks.
+struct GridDimensions {
+  GridDimensions(unsigned X = 1, unsigned Y = 1, unsigned Z = 1)
+      : X(X), Y(Y), Z(Z) {}
+
+  unsigned X;
+  unsigned Y;
+  unsigned Z;
+};
+
+} // namespace streamexecutor
+
+#endif // STREAMEXECUTOR_LAUNCHDIMENSIONS_H

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

 ///   * array of argument address,
 ///   * array of argument sizes,
 ///   * array of argument types, and
 ///   * shared pointer count.
 /// This information should be enough to allow any platform to launch the kernel
 /// efficiently, although it is probably more information than is needed for any
 /// specific platform.
 ///
+/// The PackedKernelArgumentArrayBase class has no template parameters, so it
+/// does not benefit from compile-time type checking. However, since it has no
+/// template parameters, it can be passed as an argument to virtual functions,
+/// and this allows it to be passed to functions that use virtual function
+/// overloading to handle platform-specific kernel launching.
+///
 //===----------------------------------------------------------------------===//
 
 #ifndef STREAMEXECUTOR_PACKEDKERNELARGUMENTARRAY_H
 #define STREAMEXECUTOR_PACKEDKERNELARGUMENTARRAY_H
 
 #include <array>
 
 #include "streamexecutor/DeviceMemory.h"
 
 namespace streamexecutor {
 
 enum class KernelArgumentType {
   VALUE,                /// Non-device-memory argument.
   GLOBAL_DEVICE_MEMORY, /// Non-shared device memory argument.
   SHARED_DEVICE_MEMORY  /// Shared device memory argument.
 };
 
-/// An array of packed kernel arguments.
-template <typename... ParameterTs> class PackedKernelArgumentArray {
+/// An array of packed kernel arguments without compile-time type information.
+///
+/// This un-templated base class is useful because packed kernel arguments must
+/// at some point be passed to a virtual function that performs
+/// platform-specific kernel launches. Such a virtual function cannot be
+/// templated to handle all specializations of the
+/// PackedKernelArgumentArray<...> class template, so, instead, references to
+/// PackedKernelArgumentArray<...> are passed as references to this base class.
+class PackedKernelArgumentArrayBase {
 public:
-  /// Constructs an instance by packing the specified arguments.
-  PackedKernelArgumentArray(const ParameterTs &... Arguments)
-      : SharedCount(0u) {
-    PackArguments(0, Arguments...);
-  }
+  virtual ~PackedKernelArgumentArrayBase();
 
   /// Gets the number of packed arguments.
-  size_t getArgumentCount() const { return sizeof...(ParameterTs); }
+  size_t getArgumentCount() const { return ArgumentCount; }
 
   /// Gets the address of the argument at the given index.
-  const void *getAddress(size_t Index) const { return Addresses[Index]; }
+  const void *getAddress(size_t Index) const { return AddressesData[Index]; }
 
   /// Gets the size of the argument at the given index.
-  size_t getSize(size_t Index) const { return Sizes[Index]; }
+  size_t getSize(size_t Index) const { return SizesData[Index]; }
 
   /// Gets the type of the argument at the given index.
-  KernelArgumentType getType(size_t Index) const { return Types[Index]; }
+  KernelArgumentType getType(size_t Index) const { return TypesData[Index]; }
 
   /// Gets a pointer to the address array.
-  const void *const *getAddresses() const { return Addresses.data(); }
+  const void *const *getAddresses() const { return AddressesData; }
 
   /// Gets a pointer to the sizes array.
-  const size_t *getSizes() const { return Sizes.data(); }
+  const size_t *getSizes() const { return SizesData; }
 
   /// Gets a pointer to the types array.
-  const KernelArgumentType *getTypes() const { return Types.data(); }
+  const KernelArgumentType *getTypes() const { return TypesData; }
 
   /// Gets the number of shared device memory arguments.
   size_t getSharedCount() const { return SharedCount; }
 
+protected:
+  PackedKernelArgumentArrayBase(size_t ArgumentCount)
+      : ArgumentCount(ArgumentCount), SharedCount(0u) {}
+
+  size_t ArgumentCount;
+  size_t SharedCount;
+  const void *const *AddressesData;
+  size_t *SizesData;
+  KernelArgumentType *TypesData;
+};
+
+/// An array of packed kernel arguments with compile-time type information.
+///
+/// This is used by the platform-independent StreamExecutor code to pack
+/// arguments in a compile-time type-safe way. In order to actually launch a
+/// kernel on a specific platform, however, a reference to this class will have
+/// to be passed to a virtual, platform-specific kernel launch function. Such a
+/// reference will be passed as a reference to the base class rather than a
+/// reference to this subclass itself because a virtual function cannot be
+/// templated in such a way to maintain the template parameter types of the
+/// subclass.
+template <typename... ParameterTs>
+class PackedKernelArgumentArray : public PackedKernelArgumentArrayBase {
+public:
+  /// Constructs an instance by packing the specified arguments.
+  ///
+  /// Rather than using this constructor directly, consider using the
+  /// make_kernel_argument_pack function instead, to get the compiler to infer
+  /// the parameter types for you.
+  PackedKernelArgumentArray(const ParameterTs &... Arguments)
+      : PackedKernelArgumentArrayBase(sizeof...(ParameterTs)) {
+    AddressesData = Addresses.data();
+    SizesData = Sizes.data();
+    TypesData = Types.data();
+    PackArguments(0, Arguments...);
+  }
+
+  ~PackedKernelArgumentArray() override = default;
+
 private:
   // Base case for PackArguments when there are no arguments to pack.
   void PackArguments(size_t) {}
 
   // Induction step for PackArguments.
   template <typename T, typename... RemainingParameterTs>
   void PackArguments(size_t Index, const T &Argument,
                      const RemainingParameterTs &... RemainingArguments) {
void PackOneArgument(size_t Index, const SharedDeviceMemory<T> *Argument) {
     Addresses[Index] = nullptr;
     Sizes[Index] = Argument->getByteCount();
     Types[Index] = KernelArgumentType::SHARED_DEVICE_MEMORY;
   }
 
   std::array<const void *, sizeof...(ParameterTs)> Addresses;
   std::array<size_t, sizeof...(ParameterTs)> Sizes;
   std::array<KernelArgumentType, sizeof...(ParameterTs)> Types;
-  size_t SharedCount;
 };
 
 // Utility template function to call the PackedKernelArgumentArray constructor
 // with the template arguments matching the types of the arguments passed to
 // this function.
 template <typename... ParameterTs>
 PackedKernelArgumentArray<ParameterTs...>
 make_kernel_argument_pack(const ParameterTs &... Arguments) {
   return PackedKernelArgumentArray<ParameterTs...>(Arguments...);
 }
 
 } // namespace streamexecutor
 
 #endif // STREAMEXECUTOR_PACKEDKERNELARGUMENTARRAY_H

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

+//===-- PlatformInterfaces.h - Interfaces to platform impls -----*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// Interfaces to platform-specific implementations.
+///
+/// The general pattern is that the functions in these interfaces take raw
+/// handle types as parameters. This means that these types and functions are
+/// not intended for public use. Instead, corresponding methods in public types
+/// like Stream, StreamExecutor, and Kernel use C++ templates to create
+/// type-safe public interfaces. Those public functions do the type-unsafe work
+/// of extracting raw handles from their arguments and forwarding those handles
+/// to the methods defined in this file in the proper format.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef STREAMEXECUTOR_PLATFORMINTERFACES_H
+#define STREAMEXECUTOR_PLATFORMINTERFACES_H
+
+#include "streamexecutor/DeviceMemory.h"
+#include "streamexecutor/Kernel.h"
+#include "streamexecutor/LaunchDimensions.h"
+#include "streamexecutor/PackedKernelArgumentArray.h"
+#include "streamexecutor/Utils/Error.h"
+
+namespace streamexecutor {
+
+class PlatformStreamExecutor;
+
+/// Methods supported by device kernel function objects on all platforms.
+class KernelInterface {
+  // TODO(jhen): Add methods.
+};
+
+/// Platform-specific stream handle.
+class PlatformStreamHandle {
+public:
+  explicit PlatformStreamHandle(PlatformStreamExecutor *Executor)
+      : Executor(Executor) {}
+
+  virtual ~PlatformStreamHandle();
+
+  PlatformStreamExecutor *getExecutor() { return Executor; }
+
+private:
+  PlatformStreamExecutor *Executor;
+};
+
+/// Raw executor methods that must be implemented by each platform.
+///
+/// This class defines the platform interface that supports executing work on a
+/// device.
+///
+/// The public StreamExecutor and Stream classes have the type-safe versions of
+/// the functions in this interface.
+class PlatformStreamExecutor {
+public:
+  virtual ~PlatformStreamExecutor();
+
+  virtual std::string getName() const = 0;
+
+  /// Creates a platform-specific stream.
+  virtual Expected<std::unique_ptr<PlatformStreamHandle>> createStream() = 0;
+
+  /// Launches a kernel on the given stream.
+  virtual Error launch(PlatformStreamHandle *S, BlockDimensions BlockSize,
+                       GridDimensions GridSize, const KernelBase &Kernel,
+                       const PackedKernelArgumentArrayBase &ArgumentArray) {
+    return make_error("launch not implemented for platform " + getName());
+  }
+
+  /// Copies data from the device to the host.
+  virtual Error memcpyD2H(PlatformStreamHandle *S,
+                          const GlobalDeviceMemoryBase &DeviceSrc,
+                          void *HostDst, size_t ByteCount) {
+    return make_error("memcpyD2H not implemented for platform " + getName());
+  }
+
+  /// Copies data from the host to the device.
+  virtual Error memcpyH2D(PlatformStreamHandle *S, const void *HostSrc,
+                          GlobalDeviceMemoryBase *DeviceDst, size_t ByteCount) {
+    return make_error("memcpyH2D not implemented for platform " + getName());
+  }
+
+  /// Copies data from one device location to another.
+  virtual Error memcpyD2D(PlatformStreamHandle *S,
+                          const GlobalDeviceMemoryBase &DeviceSrc,
+                          GlobalDeviceMemoryBase *DeviceDst, size_t ByteCount) {
+    return make_error("memcpyD2D 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());
+  }
+};
+
+} // namespace streamexecutor
+
+#endif // STREAMEXECUTOR_PLATFORMINTERFACES_H

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

+//===-- Stream.h - A stream of execution ------------------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+///
+/// A Stream instance represents a queue of sequential, host-asynchronous work
+/// to be performed on a device.
+///
+/// To enqueue work on a device, first create a StreamExecutor instance for a
+/// given device and then use that StreamExecutor to create a Stream instance.
+/// The Stream instance will perform its work on the device managed by the
+/// StreamExecutor that created it.
+///
+/// The various "then" methods of the Stream object, such as thenMemcpyH2D and
+/// thenLaunch, may be used to enqueue work on the Stream, and the
+/// blockHostUntilDone() method may be used to block the host code until the
+/// Stream has completed all its work.
+///
+/// Multiple Stream instances can be created for the same StreamExecutor. This
+/// allows several independent streams of computation to be performed
+/// simultaneously on a single device.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef STREAMEXECUTOR_STREAM_H
+#define STREAMEXECUTOR_STREAM_H
+
+#include <cassert>
+#include <memory>
+#include <string>
+
+#include "streamexecutor/DeviceMemory.h"
+#include "streamexecutor/Kernel.h"
+#include "streamexecutor/LaunchDimensions.h"
+#include "streamexecutor/PackedKernelArgumentArray.h"
+#include "streamexecutor/PlatformInterfaces.h"
+#include "streamexecutor/Utils/Error.h"
+
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/Twine.h"
+#include "llvm/Support/RWMutex.h"
+
+namespace streamexecutor {
+
+/// Represents a stream of dependent computations on a device.
+///
+/// The operations within a stream execute sequentially and asynchronously until
+/// blockHostUntilDone() is invoked, which synchronously joins host code with
+/// the execution of the stream.
+///
+/// If any given operation fails when entraining work for the stream, isOK()
+/// will indicate that an error has occurred and getStatus() will get the first
+/// error that occurred on the stream. There is no way to clear the error state
+/// of a stream once it is in an error state.
+class Stream {
+public:
+  explicit Stream(std::unique_ptr<PlatformStreamHandle> PStream);
+
+  ~Stream();
+
+  /// Returns whether any error has occurred while entraining work on this
+  /// stream.
+  bool isOK() const {
+    llvm::sys::ScopedReader ReaderLock(ErrorMessageMutex);
+    return !ErrorMessage;
+  }
+
+  /// Returns the status created by the first error that occurred while
+  /// entraining work on this stream.
+  Error getStatus() const {
+    llvm::sys::ScopedReader ReaderLock(ErrorMessageMutex);
+    if (ErrorMessage)
+      return make_error(*ErrorMessage);
+    else
+      return Error::success();
+  };
+
+  /// Entrains onto the stream of operations a kernel launch with the given
+  /// arguments.
+  ///
+  /// These arguments can be device memory types like GlobalDeviceMemory<T> and
+  /// SharedDeviceMemory<T>, or they can be primitive types such as int. The
+  /// allowable argument types are determined by the template parameters to the
+  /// TypedKernel argument.
+  template <typename... ParameterTs>
+  Stream &thenLaunch(BlockDimensions BlockSize, GridDimensions GridSize,
+                     const TypedKernel<ParameterTs...> &Kernel,
+                     const ParameterTs &... Arguments) {
+    auto ArgumentArray =
+        make_kernel_argument_pack<ParameterTs...>(Arguments...);
+    setError(PlatformExecutor->launch(ThePlatformStream.get(), BlockSize,
+                                      GridSize, Kernel, ArgumentArray));
+    return *this;
+  }
+
+  /// Entrain onto the stream a memcpy of a given number of elements from a
+  /// device source to a host destination.
+  ///
+  /// HostDst must be a pointer to host memory allocated by
+  /// StreamExecutor::allocateHostMemory or otherwise allocated and then
+  /// registered with StreamExecutor::registerHostMemory.
+  template <typename T>
+  Stream &thenMemcpyD2H(const GlobalDeviceMemory<T> &DeviceSrc,
+                        llvm::MutableArrayRef<T> HostDst, size_t ElementCount) {
+    if (ElementCount > DeviceSrc.getElementCount())
+      setError("copying too many elements, " + llvm::Twine(ElementCount) +
+               ", from device memory array of size " +
+               llvm::Twine(DeviceSrc.getElementCount()));
+    else if (ElementCount > HostDst.size())
+      setError("copying too many elements, " + llvm::Twine(ElementCount) +
+               ", to host array of size " + llvm::Twine(HostDst.size()));
+    else
+      setError(PlatformExecutor->memcpyD2H(ThePlatformStream.get(), DeviceSrc,
+                                           HostDst.data(),
+                                           ElementCount * sizeof(T)));
+    return *this;
+  }
+
+  /// Same as thenMemcpyD2H above, but copies the entire source to the
+  /// destination.
+  template <typename T>
+  Stream &thenMemcpyD2H(const GlobalDeviceMemory<T> &DeviceSrc,
+                        llvm::MutableArrayRef<T> HostDst) {
+    return thenMemcpyD2H(DeviceSrc, HostDst, DeviceSrc.getElementCount());
+  }
+
+  /// Entrain onto the stream a memcpy of a given number of elements from a host
+  /// source to a device destination.
+  ///
+  /// HostSrc must be a pointer to host memory allocated by
+  /// StreamExecutor::allocateHostMemory or otherwise allocated and then
+  /// registered with StreamExecutor::registerHostMemory.
+  template <typename T>
+  Stream &thenMemcpyH2D(llvm::ArrayRef<T> HostSrc,
+                        GlobalDeviceMemory<T> *DeviceDst, size_t ElementCount) {
+    if (ElementCount > HostSrc.size())
+      setError("copying too many elements, " + llvm::Twine(ElementCount) +
+               ", from host array of size " + llvm::Twine(HostSrc.size()));
+    else if (ElementCount > DeviceDst->getElementCount())
+      setError("copying too many elements, " + llvm::Twine(ElementCount) +
+               ", to device memory array of size " +
+               llvm::Twine(DeviceDst->getElementCount()));
+    else
+      setError(PlatformExecutor->memcpyH2D(ThePlatformStream.get(),
+                                           HostSrc.data(), DeviceDst,
+                                           ElementCount * sizeof(T)));
+    return *this;
+  }
+
+  /// Same as thenMemcpyH2D above, but copies the entire source to the
+  /// destination.
+  template <typename T>
+  Stream &thenMemcpyH2D(llvm::ArrayRef<T> HostSrc,
+                        GlobalDeviceMemory<T> *DeviceDst) {
+    return thenMemcpyH2D(HostSrc, DeviceDst, HostSrc.size());
+  }
+
+  /// Entrain onto the stream a memcpy of a given number of elements from a
+  /// device source to a device destination.
+  template <typename T>
+  Stream &thenMemcpyD2D(const GlobalDeviceMemory<T> &DeviceSrc,
+                        GlobalDeviceMemory<T> *DeviceDst, size_t ElementCount) {
+    if (ElementCount > DeviceSrc.getElementCount())
+      setError("copying too many elements, " + llvm::Twine(ElementCount) +
+               ", from device memory array of size " +
+               llvm::Twine(DeviceSrc.getElementCount()));
+    else if (ElementCount > DeviceDst->getElementCount())
+      setError("copying too many elements, " + llvm::Twine(ElementCount) +
+               ", to device memory array of size " +
+               llvm::Twine(DeviceDst->getElementCount()));
+    else
+      setError(PlatformExecutor->memcpyD2D(ThePlatformStream.get(), DeviceSrc,
+                                           DeviceDst,
+                                           ElementCount * sizeof(T)));
+    return *this;
+  }
+
+  /// Same as thenMemcpyD2D above, but copies the entire source to the
+  /// destination.
+  template <typename T>
+  Stream &thenMemcpyD2D(const GlobalDeviceMemory<T> &DeviceSrc,
+                        GlobalDeviceMemory<T> *DeviceDst) {
+    return thenMemcpyD2D(DeviceSrc, DeviceDst, DeviceSrc.getElementCount());
+  }
+
+  /// Blocks the host code, waiting for the operations entrained on the stream
+  /// (enqueued up to this point in program execution) to complete.
+  ///
+  /// Returns true if there are no errors on the stream.
+  bool blockHostUntilDone() {
+    Error E = PlatformExecutor->blockHostUntilDone(ThePlatformStream.get());
+    bool returnValue = static_cast<bool>(E);
+    setError(std::move(E));
+    return returnValue;
+  }
+
+private:
+  /// Sets the error state from an Error object.
+  ///
+  /// Does not overwrite the error if it is already set.
+  void setError(Error &&E) {
+    if (E) {
+      llvm::sys::ScopedWriter WriterLock(ErrorMessageMutex);
+      if (!ErrorMessage)
+        ErrorMessage = consumeAndGetMessage(std::move(E));
+    }
+  }
+
+  /// Sets the error state from an error message.
+  ///
+  /// Does not overwrite the error if it is already set.
+  void setError(llvm::Twine Message) {
+    llvm::sys::ScopedWriter WriterLock(ErrorMessageMutex);
+    if (!ErrorMessage)
+      ErrorMessage = Message.str();
+  }
+
+  /// The PlatformStreamExecutor that supports the operations of this stream.
+  PlatformStreamExecutor *PlatformExecutor;
+
+  /// The platform-specific stream handle for this instance.
+  std::unique_ptr<PlatformStreamHandle> ThePlatformStream;
+
+  /// Mutex that guards the error state flags.
+  ///
+  /// Mutable so that it can be obtained via const reader lock.
+  mutable llvm::sys::RWMutex ErrorMessageMutex;
+
+  /// First error message for an operation in this stream or empty if there have
+  /// been no errors.
+  llvm::Optional<std::string> ErrorMessage;
+
+  Stream(const Stream &) = delete;
+  void operator=(const Stream &) = delete;
+};
+
+} // namespace streamexecutor
+
+#endif // STREAMEXECUTOR_STREAM_H

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

 /// The StreamExecutor class which represents a single device of a specific
 /// platform.
 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef STREAMEXECUTOR_STREAMEXECUTOR_H
 #define STREAMEXECUTOR_STREAMEXECUTOR_H
 
+#include "streamexecutor/KernelSpec.h"
 #include "streamexecutor/Utils/Error.h"
 
 namespace streamexecutor {
 
 class KernelInterface;
+class PlatformStreamExecutor;
+class Stream;
 
 class StreamExecutor {
 public:
+  explicit StreamExecutor(PlatformStreamExecutor *PlatformExecutor);
+  virtual ~StreamExecutor();
+
   /// Gets the kernel implementation for the underlying platform.
   virtual Expected<std::unique_ptr<KernelInterface>>
   getKernelImplementation(const MultiKernelLoaderSpec &Spec) {
     // TODO(jhen): Implement this.
     return nullptr;
   }
 
-  // TODO(jhen): Add other methods.
+  Expected<std::unique_ptr<Stream>> createStream();
+
+private:
+  PlatformStreamExecutor *PlatformExecutor;
 };
 
 } // namespace streamexecutor
 
 #endif // STREAMEXECUTOR_STREAMEXECUTOR_H

parallel-libs/trunk/streamexecutor/lib/CMakeLists.txt

 add_library(
     utils
     OBJECT
     Utils/Error.cpp)
 
 add_library(
     streamexecutor
     $<TARGET_OBJECTS:utils>
     Kernel.cpp
-    KernelSpec.cpp)
+    KernelSpec.cpp
+    PackedKernelArgumentArray.cpp
+    PlatformInterfaces.cpp
+    Stream.cpp
+    StreamExecutor.cpp)
 target_link_libraries(streamexecutor ${llvm_libs})
 
 if(STREAM_EXECUTOR_UNIT_TESTS)
     add_subdirectory(unittests)
 endif()

parallel-libs/trunk/streamexecutor/lib/Kernel.cpp

 //===-- Kernel.cpp - General kernel implementation ------------------------===//
 //
 //                     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 implementation details for kernel types.
 ///
 //===----------------------------------------------------------------------===//
 
 #include "streamexecutor/Kernel.h"
-#include "streamexecutor/Interfaces.h"
+#include "streamexecutor/PlatformInterfaces.h"
 #include "streamexecutor/StreamExecutor.h"
 
 #include "llvm/DebugInfo/Symbolize/Symbolize.h"
 
 namespace streamexecutor {
 
 KernelBase::KernelBase(StreamExecutor *ParentExecutor, const std::string &Name,
                        const std::string &DemangledName,

parallel-libs/trunk/streamexecutor/lib/PackedKernelArgumentArray.cpp

+//===-- PackedKernelArgumentArray.cpp - Packed argument array impl --------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// Implementation details for classes from PackedKernelArgumentArray.h.
+///
+//===----------------------------------------------------------------------===//
+
+#include "streamexecutor/PackedKernelArgumentArray.h"
+
+namespace streamexecutor {
+
+PackedKernelArgumentArrayBase::~PackedKernelArgumentArrayBase() = default;
+
+} // namespace streamexecutor

parallel-libs/trunk/streamexecutor/lib/PlatformInterfaces.cpp

+//===-- PlatformInterfaces.cpp - Platform interface implementations -------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// Implementation file for PlatformInterfaces.h.
+///
+//===----------------------------------------------------------------------===//
+
+#include "streamexecutor/PlatformInterfaces.h"
+
+namespace streamexecutor {
+
+PlatformStreamHandle::~PlatformStreamHandle() = default;
+
+PlatformStreamExecutor::~PlatformStreamExecutor() = default;
+
+} // namespace streamexecutor

parallel-libs/trunk/streamexecutor/lib/Stream.cpp

+//===-- Stream.cpp - General stream implementation ------------------------===//
+//
+//                     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 implementation details for a general stream object.
+///
+//===----------------------------------------------------------------------===//
+
+#include "streamexecutor/Stream.h"
+
+namespace streamexecutor {
+
+Stream::Stream(std::unique_ptr<PlatformStreamHandle> PStream)
+    : PlatformExecutor(PStream->getExecutor()),
+      ThePlatformStream(std::move(PStream)) {}
+
+Stream::~Stream() = default;
+
+} // namespace streamexecutor

parallel-libs/trunk/streamexecutor/lib/StreamExecutor.cpp

+//===-- StreamExecutor.cpp - StreamExecutor implementation ----------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// Implementation of StreamExecutor class internals.
+///
+//===----------------------------------------------------------------------===//
+
+#include "streamexecutor/StreamExecutor.h"
+
+#include <cassert>
+
+#include "streamexecutor/PlatformInterfaces.h"
+#include "streamexecutor/Stream.h"
+
+#include "llvm/ADT/STLExtras.h"
+
+namespace streamexecutor {
+
+StreamExecutor::StreamExecutor(PlatformStreamExecutor *PlatformExecutor)
+    : PlatformExecutor(PlatformExecutor) {}
+
+StreamExecutor::~StreamExecutor() = default;
+
+Expected<std::unique_ptr<Stream>> StreamExecutor::createStream() {
+  Expected<std::unique_ptr<PlatformStreamHandle>> MaybePlatformStream =
+      PlatformExecutor->createStream();
+  if (!MaybePlatformStream) {
+    return MaybePlatformStream.takeError();
+  }
+  assert((*MaybePlatformStream)->getExecutor() == PlatformExecutor &&
+         "an executor created a stream with a different stored executor");
+  return llvm::make_unique<Stream>(std::move(*MaybePlatformStream));
+}
+
+} // namespace streamexecutor

parallel-libs/trunk/streamexecutor/lib/unittests/CMakeLists.txt

target_link_libraries(
     ${CMAKE_THREAD_LIBS_INIT})
 add_test(KernelSpecTest kernel_spec_test)
 
 add_executable(
     packed_kernel_argument_array_test
     PackedKernelArgumentArrayTest.cpp)
 target_link_libraries(
     packed_kernel_argument_array_test
+    streamexecutor
     ${llvm_libs}
     ${GTEST_BOTH_LIBRARIES}
     ${CMAKE_THREAD_LIBS_INIT})
 add_test(PackedKernelArgumentArrayTest packed_kernel_argument_array_test)
+
+add_executable(
+    stream_test
+    StreamTest.cpp)
+target_link_libraries(
+    stream_test
+    streamexecutor
+    ${llvm_libs}
+    ${GTEST_BOTH_LIBRARIES}
+    ${CMAKE_THREAD_LIBS_INIT})
+add_test(StreamTest stream_test)

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

 //===-- KernelTest.cpp - Tests for Kernel objects -------------------------===//
 //
 //                     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 the code in Kernel.
 ///
 //===----------------------------------------------------------------------===//
 
 #include <cassert>
 
-#include "streamexecutor/Interfaces.h"
 #include "streamexecutor/Kernel.h"
 #include "streamexecutor/KernelSpec.h"
+#include "streamexecutor/PlatformInterfaces.h"
 #include "streamexecutor/StreamExecutor.h"
 
 #include "llvm/ADT/STLExtras.h"
 
 #include "gtest/gtest.h"
 
 namespace {
 

 // getKernelImplementation method and will thereby take ownership of the
 // unique_ptr, but the copy of the raw pointer will stay behind in this mock
 // object. The raw pointer copy can then be used to identify the unique_ptr in
 // its new location (by comparing the raw pointer with unique_ptr::get), to
 // verify that the unique_ptr ended up where it was supposed to be.
 class MockStreamExecutor : public se::StreamExecutor {
 public:
   MockStreamExecutor()
-      : Unique(llvm::make_unique<se::KernelInterface>()), Raw(Unique.get()) {}
+      : se::StreamExecutor(nullptr),
+        Unique(llvm::make_unique<se::KernelInterface>()), Raw(Unique.get()) {}
 
   // Moves the unique pointer into the returned se::Expected instance.
   //
   // Asserts that it is not called again after the unique pointer has been moved
   // out.
   se::Expected<std::unique_ptr<se::KernelInterface>>
   getKernelImplementation(const se::MultiKernelLoaderSpec &) override {
     assert(Unique && "MockStreamExecutor getKernelImplementation should not be "

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 "streamexecutor/Kernel.h"
+#include "streamexecutor/KernelSpec.h"
+#include "streamexecutor/PlatformInterfaces.h"
+#include "streamexecutor/Stream.h"
+#include "streamexecutor/StreamExecutor.h"
+
+#include "gtest/gtest.h"
+
+namespace {
+
+namespace se = ::streamexecutor;
+
+/// Mock PlatformStreamExecutor that performs asynchronous memcpy operations by
+/// ignoring the stream argument and calling std::memcpy on device memory
+/// handles.
+class MockPlatformStreamExecutor : public se::PlatformStreamExecutor {
+public:
+  ~MockPlatformStreamExecutor() override {}
+
+  std::string getName() const override { return "MockPlatformStreamExecutor"; }
+
+  se::Expected<std::unique_ptr<se::PlatformStreamHandle>>
+  createStream() override {
+    return nullptr;
+  }
+
+  se::Error memcpyD2H(se::PlatformStreamHandle *,
+                      const se::GlobalDeviceMemoryBase &DeviceSrc,
+                      void *HostDst, size_t ByteCount) override {
+    std::memcpy(HostDst, DeviceSrc.getHandle(), ByteCount);
+    return se::Error::success();
+  }
+
+  se::Error memcpyH2D(se::PlatformStreamHandle *, const void *HostSrc,
+                      se::GlobalDeviceMemoryBase *DeviceDst,
+                      size_t ByteCount) override {
+    std::memcpy(const_cast<void *>(DeviceDst->getHandle()), HostSrc, ByteCount);
+    return se::Error::success();
+  }
+
+  se::Error memcpyD2D(se::PlatformStreamHandle *,
+                      const se::GlobalDeviceMemoryBase &DeviceSrc,
+                      se::GlobalDeviceMemoryBase *DeviceDst,
+                      size_t ByteCount) override {
+    std::memcpy(const_cast<void *>(DeviceDst->getHandle()),
+                DeviceSrc.getHandle(), ByteCount);
+    return se::Error::success();
+  }
+};
+
+/// Test fixture to hold objects used by tests.
+class StreamTest : public ::testing::Test {
+public:
+  StreamTest()
+      : DeviceA(se::GlobalDeviceMemory<int>::makeFromElementCount(HostA, 10)),
+        DeviceB(se::GlobalDeviceMemory<int>::makeFromElementCount(HostB, 10)),
+        Stream(llvm::make_unique<se::PlatformStreamHandle>(&PlatformExecutor)) {
+  }
+
+protected:
+  // Device memory is backed by host arrays.
+  int HostA[10];
+  se::GlobalDeviceMemory<int> DeviceA;
+  int HostB[10];
+  se::GlobalDeviceMemory<int> DeviceB;
+
+  // Host memory to be used as actual host memory.
+  int Host[10];
+
+  MockPlatformStreamExecutor PlatformExecutor;
+  se::Stream Stream;
+};
+
+TEST_F(StreamTest, MemcpyCorrectSize) {
+  Stream.thenMemcpyH2D(llvm::ArrayRef<int>(Host), &DeviceA);
+  EXPECT_TRUE(Stream.isOK());
+
+  Stream.thenMemcpyD2H(DeviceA, llvm::MutableArrayRef<int>(Host));
+  EXPECT_TRUE(Stream.isOK());
+
+  Stream.thenMemcpyD2D(DeviceA, &DeviceB);
+  EXPECT_TRUE(Stream.isOK());
+}
+
+TEST_F(StreamTest, MemcpyH2DTooManyElements) {
+  Stream.thenMemcpyH2D(llvm::ArrayRef<int>(Host), &DeviceA, 20);
+  EXPECT_FALSE(Stream.isOK());
+}
+
+TEST_F(StreamTest, MemcpyD2HTooManyElements) {
+  Stream.thenMemcpyD2H(DeviceA, llvm::MutableArrayRef<int>(Host), 20);
+  EXPECT_FALSE(Stream.isOK());
+}
+
+TEST_F(StreamTest, MemcpyD2DTooManyElements) {
+  Stream.thenMemcpyD2D(DeviceA, &DeviceB, 20);
+  EXPECT_FALSE(Stream.isOK());
+}
+
+} // namespace