[SE] RegisteredHostMemory for async device copies

[jhen]

Summary

Improve the error-prone interface that allows users to pass host
pointers that haven't been registered to asynchronous copy methods. In
CUDA, this is an extremely easy error to make, and instead of failing at
runtime, it succeeds and gives the right answers by turning the async
copy into a sync copy. So, you silently get a huge performance
degradation if you misuse the old interface. This new interface should
prevent that.

[jhen]

• Remove allocateHostMemory

[jhen]

After an offline chat with jlebar, we decided to get rid of the allocateHostMemory functions in the device interface. We thought it was confusing that the first patch in this review made a RegisteredHostMemory object that only sometimes owned its underlying pointer.

By removing allocateHostMemory the object never has to own its underlying pointer. Possible concerns with removing this function:

• The interface may be slightly more annoying for users because they have to allocate their own memory.
• Performance may be worse if users have to make two calls instead of just one to get registered host memory.
• Memory that users allocate themselves may not have optimal alignment for the registration process.

If any of these problems are later shown to be a bigger issue, we can re-introduce the host allocation functions with a different interface than the proposal from the first patch.

There are a couple of options that might work for the new interface:

• Have allocateHostMemory return a pair<RegisteredHostMemory, std::unique_ptr> so that the unique_ptr is responsible for freeing the memory and the RegisteredHostMemory is responsible for un-registering it. (Maybe a problem if the unique_ptr goes out of scope first.)
• Create a new type OwnedRegisteredHostMemory that always owns the memory and the registration of the memory. Have methods to create a slice out of this new type, and support this type as an argument to all the thenCopy Stream methods (std::enable_if should make that pretty easy).

[jhen]

For the record, in addition to using std::enable_if to get the Stream thenCopy functions to accept either mutable or immutable slice arguments, I also tried providing an implicit ctor to make immutable slices from mutable slices and designing the functions to take only immutable slices. Result: as usual, template type inference does not consider implicit type conversions, so the user would have to specify the ElemT type parameter in order to get that to work. That's the reason I stuck with std::enable_if.

[jprice]

streamexecutor/include/streamexecutor/Device.h

86 ↗

(On Diff #70734)

host

streamexecutor/include/streamexecutor/HostMemory.h

105 ↗

(On Diff #70734)

"newly allocating registered host memory" isn't an option anymore

107–109 ↗

(On Diff #70734)

Likewise - this last sentence isn't valid anymore

[jhen]

• Fix broken comments found by jprice

[jhen]

Thanks for catching those comment errors, jprice!

[jlebar]

I talked to Chandler about this slicing thing...

His suggestion was to give up on deducing T in e.g. thenCopyH2D. Instead, we'd write function signatures like the following:

template <typename HostTy, typename DeviceTy>
thenCopyH2D(HostTy &&Src, DeviceTy &&Dst);

We have constraints on HostTy and DeviceTy that we still need to enforce. Maybe our constraint is that they both expose a typedef named ElementTy and they are convertible to RegisteredHostMemorySlice<ElementTy> and GlobalMemorySlice<ElementTy> respectively. In this case we'd write something like

static_assert(std::is_same<HostTy::ElementTy, DeviceTy::ElementTy>::value, "...");
RegisteredHostMemorySlice<HostTy::ElemTy> HostSlice(Src);
GlobalMemorySlice<DeviceTy::ElemTy> DeviceSlice(Dst);

If this compiles, we're good to go. If it doesn't compile, that means you passed bad types to this function.

One nice thing about this is that it would let us unify the mutable / immutable slices into a single type (I think):

RegisteredHostMemorySlice<const HostTy::ElemTy> ImmutableSlice;
RegisteredHostMemorySlice<HostTy::ElemTy> MutableSlice;

This would solve our combinatorial explosion problem in Stream.h, and wouldn't require complex enable_if shenanigans. (Although it involves its own shenanigans, I suppose.)

What do you think?

streamexecutor/include/streamexecutor/HostMemory.h

161 ↗

(On Diff #70766)

It doesn't matter much, but this base class is kind of annoying, partially because we have the added overhead of this division (shift) every time we get the size, but more because it's another class kicking around.

It seems like the main reason we need the base class is to make the out-of-line destructor work without including device.h here. But we can do something similar without all the fuss of a base class. We could define an out-of-line function in a "detail" namespace, or we could have a base class with just one function, defined out-of-line, that we call from our destructor.

32 ↗

(On Diff #70720)

Can we define in what way the memory is "registered" here?

[jhen]

• implement chanderc's suggestion
• Get rid of mutable host array slice

[jhen]

I really like the idea of using completely templated Src and Dst arguments. I've implemented that in this patch, and I think everyone can agree that it really cleaned up the code.

I'm not 100% sure that combining the mutable and immutable host slices in one class is the right choice, but I've also made that change so we can take a look and see what folks think.

streamexecutor/include/streamexecutor/HostMemory.h

33 ↗

(On Diff #70766)

I added comments to each of these classes to say the memory is registered in the sense of Device::registerHostMemory.

161 ↗

(On Diff #70766)

I made a function called internal::destroyRegisteredHostMemoryInternals and got rid of the base class.

[jlebar]

streamexecutor/include/streamexecutor/HostMemory.h

45 ↗

(On Diff #70889)

I'm not sure that a const_cast is the right thing to do. Instead, could we have two constructors, one which takes a const RHM<ElemT> and one which takes a non-const RHM<ElemT>?

Also, how does this work when ElemT = "const int" and we get an RHM<ElemT>? I would expect you need some const-decay here.

45 ↗

(On Diff #70889)

If we go with this approach, we should make sure we can convert from RHMSlice<const T> to RHMSlice<T>, which I don't think works atm.

48 ↗

(On Diff #70889)

Hmmm...

In a vacuum, I think we'd want just one overload,

ElemT *getPointer() const;

This matches what C++17's std::slice (nee array_view) does. And it sort of makes sense: We already have constness in the template type to determine whether or not we can modify the underlying ElemTs.

(By the way, if we stick with this approach, we should explain in a comment about putting const T in the template param.)

But I see that we have the two overloads on RHM, and there they make more sense. So...I am not sure what we should do here.

One option is we could also make RHM unconditionally give out non-const pointers (and mutable slices). That would make it behave like unique_ptr -- if I have a const unique_ptr<int[]>, the get() function still gives me a non-const pointer back.

I guess one argument against merging the slice classes is that it seems like we're having difficulty keeping them straight within Stream.h. So I think I'm OK splitting them back out for now. It shouldn't complicate the Stream interface any.

99 ↗

(On Diff #70889)

Nit, clang-format is being conservative and not re-merging strings, but this should probably be reflowed.

streamexecutor/include/streamexecutor/Stream.h

172 ↗

(On Diff #70889)

Is it worth merging this with the overload that explicitly takes slices, or do you think that would be too confusing?

178 ↗

(On Diff #70889)

There should be const in some of these?

[jhen]

• Remove allocateHostMemory
• Fix broken comments found by jprice
• implement chanderc's suggestion
• Combine more thenCopy functions as templates

streamexecutor/include/streamexecutor/HostMemory.h

45 ↗

(On Diff #70889)

Switching back to mutable array.

48 ↗

(On Diff #70889)

Switching back to mutable array.

99 ↗

(On Diff #70889)

Actually, clang-format really seems to like this. It won't seem to break after the && even if I manually combine the string literals into one long one. So I'll just let clang-format have its way here.

streamexecutor/include/streamexecutor/Stream.h

172 ↗

(On Diff #70889)

Yes, I did it now. I don't know why I didn't before.

178 ↗

(On Diff #70889)

Using mutable array now.

[jlebar]

Ship it.

streamexecutor/include/streamexecutor/HostMemory.h

94 ↗

(On Diff #70934)

While I'm thinking about it, we should add LLVM_ATTRIBUTE_UNUSED_RESULT to these (and other similar functions) in a separate patch.

parallel-libs/trunk/streamexecutor/examples/CUDASaxpy.cpp

if (Platform->getDeviceCount() == 0) {
     return EXIT_FAILURE;
   }
   se::Device *Device = getOrDie(Platform->getDevice(0));
 
   // Load the kernel onto the device.
   cg::SaxpyKernel Kernel =
       getOrDie(Device->createKernel<cg::SaxpyKernel>(cg::SaxpyLoaderSpec));
 
+  se::RegisteredHostMemory<float> RegisteredX =
+      getOrDie(Device->registerHostMemory<float>(HostX));
+  se::RegisteredHostMemory<float> RegisteredY =
+      getOrDie(Device->registerHostMemory<float>(HostY));
+
   // Allocate memory on the device.
   se::GlobalDeviceMemory<float> X =
       getOrDie(Device->allocateDeviceMemory<float>(ArraySize));
   se::GlobalDeviceMemory<float> Y =
       getOrDie(Device->allocateDeviceMemory<float>(ArraySize));
 
   // Run operations on a stream.
   se::Stream Stream = getOrDie(Device->createStream());
-  Stream.thenCopyH2D<float>(HostX, X)
-      .thenCopyH2D<float>(HostY, Y)
+  Stream.thenCopyH2D(RegisteredX, X)
+      .thenCopyH2D(RegisteredY, Y)
       .thenLaunch(ArraySize, 1, Kernel, A, X, Y)
-      .thenCopyD2H<float>(X, HostX);
+      .thenCopyD2H(X, RegisteredX);
   // Wait for the stream to complete.
   se::dieIfError(Stream.blockHostUntilDone());
 
   // Process output data in HostX.
   std::vector<float> ExpectedX = {4, 47, 90, 133};
   assert(std::equal(ExpectedX.begin(), ExpectedX.end(), HostX.begin()));
   /// [Example saxpy host main]
 }

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

 //===----------------------------------------------------------------------===//
 
 #ifndef STREAMEXECUTOR_DEVICE_H
 #define STREAMEXECUTOR_DEVICE_H
 
 #include <type_traits>
 
 #include "streamexecutor/Error.h"
+#include "streamexecutor/HostMemory.h"
 #include "streamexecutor/KernelSpec.h"
 #include "streamexecutor/PlatformDevice.h"
 
 namespace streamexecutor {
 
 class Stream;
 
 class Device {
public:
   Expected<GlobalDeviceMemory<T>> allocateDeviceMemory(size_t ElementCount) {
     Expected<void *> MaybeMemory =
         PDevice->allocateDeviceMemory(ElementCount * sizeof(T));
     if (!MaybeMemory)
       return MaybeMemory.takeError();
     return GlobalDeviceMemory<T>(this, *MaybeMemory, ElementCount);
   }
 
-  /// 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 =
-        PDevice->allocateHostMemory(ElementCount * sizeof(T));
-    if (!MaybeMemory)
-      return MaybeMemory.takeError();
-    return static_cast<T *>(*MaybeMemory);
-  }
-
-  /// Frees memory previously allocated with allocateHostMemory.
-  template <typename T> Error freeHostMemory(T *Memory) {
-    return PDevice->freeHostMemory(Memory);
-  }
-
   /// Registers a previously allocated host array of type T for asynchronous
   /// memory operations.
   ///
   /// Host memory registered by this function can be used for asynchronous
   /// memory copies on streams. See Stream::thenCopyD2H and Stream::thenCopyH2D.
   template <typename T>
-  Error registerHostMemory(T *Memory, size_t ElementCount) {
-    return PDevice->registerHostMemory(Memory, ElementCount * sizeof(T));
+  Expected<RegisteredHostMemory<T>>
+  registerHostMemory(llvm::MutableArrayRef<T> Memory) {
+    if (Error E = PDevice->registerHostMemory(Memory.data(),
+                                              Memory.size() * sizeof(T))) {
+      return std::move(E);
     }
-
-  /// Unregisters host memory previously registered by registerHostMemory.
-  template <typename T> Error unregisterHostMemory(T *Memory) {
-    return PDevice->unregisterHostMemory(Memory);
+    return RegisteredHostMemory<T>(this, Memory.data(), Memory.size());
   }
 
   /// \anchor DeviceHostSyncCopyGroup
   /// \name Host-synchronous device memory copying functions
   ///
   /// These methods block the calling host thread while copying data to or from
   /// device memory. On the device side, these methods do not block any ongoing
   /// device calls.
   ///
   /// There are no restrictions on the host memory that is used as a source or
-  /// destination in these copy methods, so there is no need to allocate that
-  /// host memory using allocateHostMemory or register it with
-  /// registerHostMemory.
+  /// destination in these copy methods, so there is no need to register that
+  /// host memory with registerHostMemory.
   ///
   /// Each of these methods has a single template parameter, T, that specifies
   /// the type of data being copied. The ElementCount arguments specify the
   /// number of objects of type T to be copied.
   ///
   /// For ease of use, each of the methods is overloaded to take either a
   /// GlobalDeviceMemorySlice or a GlobalDeviceMemory argument in the device
   /// memory argument slots, and the GlobalDeviceMemory arguments are just

 
 private:
   // Only a GlobalDeviceMemoryBase may free device memory.
   friend GlobalDeviceMemoryBase;
   Error freeDeviceMemory(const GlobalDeviceMemoryBase &Memory) {
     return PDevice->freeDeviceMemory(Memory.getHandle());
   }
 
+  // Only destroyRegisteredHostMemoryInternals may unregister host memory.
+  friend void internal::destroyRegisteredHostMemoryInternals(Device *, void *);
+  Error unregisterHostMemory(const void *Pointer) {
+    return PDevice->unregisterHostMemory(Pointer);
+  }
+
   PlatformDevice *PDevice;
 };
 
 } // namespace streamexecutor
 
 #endif // STREAMEXECUTOR_DEVICE_H

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

 template <typename ElemT> class GlobalDeviceMemory;
 
 /// Reference to a slice of device memory.
 ///
 /// Contains a base memory handle, an element count offset into that base
 /// memory, and an element count for the size of the slice.
 template <typename ElemT> class GlobalDeviceMemorySlice {
 public:
+  using ElementTy = ElemT;
+
   /// Intentionally implicit so GlobalDeviceMemory<T> can be passed to functions
   /// expecting GlobalDeviceMemorySlice<T> arguments.
   GlobalDeviceMemorySlice(const GlobalDeviceMemory<ElemT> &Memory)
       : BaseMemory(Memory), ElementOffset(0),
         ElementCount(Memory.getElementCount()) {}
 
   GlobalDeviceMemorySlice(const GlobalDeviceMemory<ElemT> &BaseMemory,
                           size_t ElementOffset, size_t ElementCount)

 /// Typed wrapper around the "void *"-like GlobalDeviceMemoryBase class.
 ///
 /// For example, GlobalDeviceMemory<int> is a simple wrapper around
 /// GlobalDeviceMemoryBase that represents a buffer of integers stored in global
 /// device memory.
 template <typename ElemT>
 class GlobalDeviceMemory : public GlobalDeviceMemoryBase {
 public:
+  using ElementTy = ElemT;
+
   GlobalDeviceMemory(GlobalDeviceMemory &&Other) = default;
   GlobalDeviceMemory &operator=(GlobalDeviceMemory &&Other) = default;
 
   /// Returns the number of elements of type ElemT that constitute this
   /// allocation.
   size_t getElementCount() const { return ByteCount / sizeof(ElemT); }
 
   /// Returns the number of bytes that can fit in this memory buffer.

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

+//===-- HostMemory.h - Types for registered host memory ---------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+///
+/// This file defines types that represent registered host memory buffers. Host
+/// memory must be registered to participate in asynchronous copies to or from
+/// device memory.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef STREAMEXECUTOR_HOSTMEMORY_H
+#define STREAMEXECUTOR_HOSTMEMORY_H
+
+#include <cassert>
+#include <cstddef>
+#include <type_traits>
+
+#include "llvm/ADT/ArrayRef.h"
+
+namespace streamexecutor {
+
+class Device;
+template <typename ElemT> class RegisteredHostMemory;
+
+/// A mutable slice of registered host memory.
+///
+/// The memory is registered in the sense of
+/// streamexecutor::Device::registerHostMemory.
+///
+/// Holds a reference to an underlying registered host memory buffer. Must not
+/// be used after the underlying buffer is freed or unregistered.
+template <typename ElemT> class MutableRegisteredHostMemorySlice {
+public:
+  using ElementTy = ElemT;
+
+  MutableRegisteredHostMemorySlice(RegisteredHostMemory<ElemT> &Registered)
+      : MutableArrayRef(Registered.getPointer(), Registered.getElementCount()) {
+  }
+
+  ElemT *getPointer() const { return MutableArrayRef.data(); }
+  size_t getElementCount() const { return MutableArrayRef.size(); }
+
+  /// Chops off the first N elements of the slice.
+  MutableRegisteredHostMemorySlice slice(size_t N) const {
+    return MutableRegisteredHostMemorySlice(MutableArrayRef.slice(N));
+  }
+
+  /// Chops off the first N elements of the slice and keeps the next M elements.
+  MutableRegisteredHostMemorySlice slice(size_t N, size_t M) const {
+    return MutableRegisteredHostMemorySlice(MutableArrayRef.slice(N, M));
+  }
+
+  /// Chops off the last N elements of the slice.
+  MutableRegisteredHostMemorySlice drop_back(size_t N) const {
+    return MutableRegisteredHostMemorySlice(MutableArrayRef.drop_back(N));
+  }
+
+private:
+  MutableRegisteredHostMemorySlice(llvm::MutableArrayRef<ElemT> MutableArrayRef)
+      : MutableArrayRef(MutableArrayRef) {}
+
+  llvm::MutableArrayRef<ElemT> MutableArrayRef;
+};
+
+/// An immutable slice of registered host memory.
+///
+/// The memory is registered in the sense of
+/// streamexecutor::Device::registerHostMemory.
+///
+/// Holds a reference to an underlying registered host memory buffer. Must not
+/// be used after the underlying buffer is freed or unregistered.
+template <typename ElemT> class RegisteredHostMemorySlice {
+public:
+  using ElementTy = ElemT;
+
+  RegisteredHostMemorySlice(const RegisteredHostMemory<ElemT> &Registered)
+      : ArrayRef(Registered.getPointer(), Registered.getElementCount()) {}
+
+  RegisteredHostMemorySlice(
+      MutableRegisteredHostMemorySlice<ElemT> MutableSlice)
+      : ArrayRef(MutableSlice.getPointer(), MutableSlice.getElementCount()) {}
+
+  const ElemT *getPointer() const { return ArrayRef.data(); }
+  size_t getElementCount() const { return ArrayRef.size(); }
+
+  /// Chops off the first N elements of the slice.
+  RegisteredHostMemorySlice slice(size_t N) const {
+    return RegisteredHostMemorySlice(ArrayRef.slice(N));
+  }
+
+  /// Chops off the first N elements of the slice and keeps the next M elements.
+  RegisteredHostMemorySlice slice(size_t N, size_t M) const {
+    return RegisteredHostMemorySlice(ArrayRef.slice(N, M));
+  }
+
+  /// Chops off the last N elements of the slice.
+  RegisteredHostMemorySlice drop_back(size_t N) const {
+    return RegisteredHostMemorySlice(ArrayRef.drop_back(N));
+  }
+
+private:
+  llvm::ArrayRef<ElemT> ArrayRef;
+};
+
+namespace internal {
+
+/// Helper function to unregister host memory.
+///
+/// This is a thin wrapper around streamexecutor::Device::unregisterHostMemory.
+/// It is defined so this operation can be performed from the destructor of the
+/// template class RegisteredHostMemory without including Device.h in this
+/// header and creating a header inclusion cycle.
+void destroyRegisteredHostMemoryInternals(Device *TheDevice, void *Pointer);
+
+} // namespace internal
+
+/// Registered host memory that knows how to unregister itself upon destruction.
+///
+/// The memory is registered in the sense of
+/// streamexecutor::Device::registerHostMemory.
+///
+/// ElemT is the type of element stored in the host buffer.
+template <typename ElemT> class RegisteredHostMemory {
+public:
+  using ElementTy = ElemT;
+
+  RegisteredHostMemory(Device *TheDevice, ElemT *Pointer, size_t ElementCount)
+      : TheDevice(TheDevice), Pointer(Pointer), ElementCount(ElementCount) {
+    assert(TheDevice != nullptr && "cannot construct a "
+                                   "RegisteredHostMemoryBase with a null "
+                                   "platform device");
+  }
+
+  RegisteredHostMemory(const RegisteredHostMemory &) = delete;
+  RegisteredHostMemory &operator=(const RegisteredHostMemory &) = delete;
+
+  RegisteredHostMemory(RegisteredHostMemory &&Other)
+      : TheDevice(Other.TheDevice), Pointer(Other.Pointer),
+        ElementCount(Other.ElementCount) {
+    Other.TheDevice = nullptr;
+    Other.Pointer = nullptr;
+  }
+
+  RegisteredHostMemory &operator=(RegisteredHostMemory &&Other) {
+    TheDevice = Other.TheDevice;
+    Pointer = Other.Pointer;
+    ElementCount = Other.ElementCount;
+    Other.TheDevice = nullptr;
+    Other.Pointer = nullptr;
+  }
+
+  ~RegisteredHostMemory() {
+    internal::destroyRegisteredHostMemoryInternals(TheDevice, Pointer);
+  }
+
+  ElemT *getPointer() { return static_cast<ElemT *>(Pointer); }
+  const ElemT *getPointer() const { return static_cast<ElemT *>(Pointer); }
+  size_t getElementCount() const { return ElementCount; }
+
+  /// Creates an immutable slice for the entire contents of this memory.
+  RegisteredHostMemorySlice<ElemT> asSlice() const {
+    return RegisteredHostMemorySlice<ElemT>(*this);
+  }
+
+  /// Creates a mutable slice for the entire contents of this memory.
+  MutableRegisteredHostMemorySlice<ElemT> asSlice() {
+    return MutableRegisteredHostMemorySlice<ElemT>(*this);
+  }
+
+private:
+  Device *TheDevice;
+  void *Pointer;
+  size_t ElementCount;
+};
+
+} // namespace streamexecutor
+
+#endif // STREAMEXECUTOR_HOSTMEMORY_H

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

virtual Error launch(const void *PlatformStreamHandle,
                        BlockDimensions BlockSize, GridDimensions GridSize,
                        const void *PKernelHandle,
                        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.
+  /// HostDst should have been registered with registerHostMemory.
   virtual Error copyD2H(const void *PlatformStreamHandle,
                         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.
+  /// HostSrc should have been registered with registerHostMemory.
   virtual Error copyH2D(const void *PlatformStreamHandle, const void *HostSrc,
                         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(const void *PlatformStreamHandle,
public:
   }
 
   /// Frees device memory previously allocated by allocateDeviceMemory.
   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.
-  virtual Expected<void *> allocateHostMemory(size_t ByteCount) {
-    return make_error("allocateHostMemory not implemented for platform " +
-                      getName());
-  }
-
-  /// Frees host memory allocated by allocateHostMemory.
-  virtual Error freeHostMemory(void *Memory) {
-    return make_error("freeHostMemory not implemented for platform " +
-                      getName());
-  }
-
   /// Registers previously allocated host memory so it can be used with copyH2D
   /// and copyD2H.
   virtual Error registerHostMemory(void *Memory, size_t ByteCount) {
     return make_error("registerHostMemory not implemented for platform " +
                       getName());
   }
 
   /// Unregisters host memory previously registered with registerHostMemory.
-  virtual Error unregisterHostMemory(void *Memory) {
+  virtual Error unregisterHostMemory(const void *Memory) {
     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.
+  /// any host memory, not just registered host memory. Does not block any
+  /// ongoing device calls.
   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 void *, size_t, void

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

 //===----------------------------------------------------------------------===//
 
 #ifndef STREAMEXECUTOR_STREAM_H
 #define STREAMEXECUTOR_STREAM_H
 
 #include <cassert>
 #include <memory>
 #include <string>
+#include <type_traits>
 
 #include "streamexecutor/DeviceMemory.h"
 #include "streamexecutor/Error.h"
+#include "streamexecutor/HostMemory.h"
 #include "streamexecutor/Kernel.h"
 #include "streamexecutor/LaunchDimensions.h"
 #include "streamexecutor/PackedKernelArgumentArray.h"
 #include "streamexecutor/PlatformDevice.h"
 
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/RWMutex.h"
Stream &thenLaunch(BlockDimensions BlockSize, GridDimensions GridSize,
     return *this;
   }
 
   /// \name Device memory copying functions
   ///
   /// These methods enqueue a device memory copy operation on the stream and
   /// return without waiting for the operation to complete.
   ///
-  /// Any host memory used as a source or destination for one of these
-  /// operations must be allocated with Device::allocateHostMemory or registered
-  /// with Device::registerHostMemory. Otherwise, the enqueuing operation may
-  /// block until the copy operation is fully complete.
-  ///
   /// The arguments and bounds checking for these methods match the API of the
   /// \ref DeviceHostSyncCopyGroup
   /// "host-synchronous device memory copying functions" of Device.
+  ///
+  /// The template types SrcTy and DstTy must match the following constraints:
+  ///   * Must define typename ElementTy (the type of element stored in the
+  ///   memory);
+  ///   * ElementTy for the source argument must be the same as ElementTy for
+  ///     the destination argument;
+  ///   * Must be convertible to the correct slice type:
+  ///     * GlobalDeviceMemorySlice<ElementTy> for device memory arguments,
+  ///     * RegisteredHostMemorySlice<ElementTy> for host memory source
+  ///       arguments,
+  ///     * MutableRegisteredHostMemorySlice<ElementT> for host memory
+  ///       destination arguments.
   ///@{
 
-  template <typename T>
-  Stream &thenCopyD2H(GlobalDeviceMemorySlice<T> Src,
-                      llvm::MutableArrayRef<T> Dst, size_t ElementCount) {
+  // D2H
+
+  template <typename SrcTy, typename DstTy>
+  Stream &thenCopyD2H(SrcTy &&Src, DstTy &&Dst, size_t ElementCount) {
+    using SrcElemTy = typename std::remove_reference<SrcTy>::type::ElementTy;
+    using DstElemTy = typename std::remove_reference<DstTy>::type::ElementTy;
+    static_assert(std::is_same<SrcElemTy, DstElemTy>::value,
+                  "src/dst element type mismatch for thenCopyD2H");
+    GlobalDeviceMemorySlice<SrcElemTy> SrcSlice(Src);
+    MutableRegisteredHostMemorySlice<DstElemTy> DstSlice(Dst);
     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())
+               llvm::Twine(SrcSlice.getElementCount()));
+    else if (ElementCount > DstSlice.getElementCount())
       setError("copying too many elements, " + llvm::Twine(ElementCount) +
-               ", to a host array of element count " + llvm::Twine(Dst.size()));
+               ", to a host array of element count " +
+               llvm::Twine(DstSlice.getElementCount()));
     else
-      setError(PDevice->copyD2H(PlatformStreamHandle,
-                                Src.getBaseMemory().getHandle(),
-                                Src.getElementOffset() * sizeof(T), Dst.data(),
-                                0, ElementCount * sizeof(T)));
+      setError(PDevice->copyD2H(
+          PlatformStreamHandle, SrcSlice.getBaseMemory().getHandle(),
+          SrcSlice.getElementOffset() * sizeof(SrcElemTy),
+          DstSlice.getPointer(), 0, ElementCount * sizeof(DstElemTy)));
     return *this;
   }
 
-  template <typename T>
-  Stream &thenCopyD2H(GlobalDeviceMemorySlice<T> Src,
-                      llvm::MutableArrayRef<T> Dst) {
-    if (Src.getElementCount() != Dst.size())
+  template <typename SrcTy, typename DstTy>
+  Stream &thenCopyD2H(SrcTy &&Src, DstTy &&Dst) {
+    using SrcElemTy = typename std::remove_reference<SrcTy>::type::ElementTy;
+    using DstElemTy = typename std::remove_reference<DstTy>::type::ElementTy;
+    static_assert(std::is_same<SrcElemTy, DstElemTy>::value,
+                  "src/dst element type mismatch for thenCopyD2H");
+    GlobalDeviceMemorySlice<SrcElemTy> SrcSlice(Src);
+    MutableRegisteredHostMemorySlice<DstElemTy> DstSlice(Dst);
+    if (SrcSlice.getElementCount() != DstSlice.getElementCount())
       setError("array size mismatch for D2H, device source has element count " +
-               llvm::Twine(Src.getElementCount()) +
+               llvm::Twine(SrcSlice.getElementCount()) +
                " but host destination has element count " +
-               llvm::Twine(Dst.size()));
+               llvm::Twine(DstSlice.getElementCount()));
     else
-      thenCopyD2H(Src, Dst, Src.getElementCount());
-    return *this;
-  }
-
-  template <typename T>
-  Stream &thenCopyD2H(GlobalDeviceMemorySlice<T> Src, T *Dst,
-                      size_t ElementCount) {
-    thenCopyD2H(Src, llvm::MutableArrayRef<T>(Dst, ElementCount), ElementCount);
-    return *this;
-  }
-
-  template <typename T>
-  Stream &thenCopyD2H(const GlobalDeviceMemory<T> &Src,
-                      llvm::MutableArrayRef<T> Dst, size_t ElementCount) {
-    thenCopyD2H(Src.asSlice(), Dst, ElementCount);
-    return *this;
-  }
-
-  template <typename T>
-  Stream &thenCopyD2H(const GlobalDeviceMemory<T> &Src,
-                      llvm::MutableArrayRef<T> Dst) {
-    thenCopyD2H(Src.asSlice(), Dst);
+      thenCopyD2H(SrcSlice, DstSlice, SrcSlice.getElementCount());
     return *this;
   }
 
-  template <typename T>
-  Stream &thenCopyD2H(const GlobalDeviceMemory<T> &Src, T *Dst,
-                      size_t ElementCount) {
-    thenCopyD2H(Src.asSlice(), Dst, ElementCount);
-    return *this;
-  }
+  // H2D
 
-  template <typename T>
-  Stream &thenCopyH2D(llvm::ArrayRef<T> Src, GlobalDeviceMemorySlice<T> Dst,
-                      size_t ElementCount) {
-    if (ElementCount > Src.size())
+  template <typename SrcTy, typename DstTy>
+  Stream &thenCopyH2D(SrcTy &&Src, DstTy &&Dst, size_t ElementCount) {
+    using SrcElemTy = typename std::remove_reference<SrcTy>::type::ElementTy;
+    using DstElemTy = typename std::remove_reference<DstTy>::type::ElementTy;
+    static_assert(std::is_same<SrcElemTy, DstElemTy>::value,
+                  "src/dst element type mismatch for thenCopyH2D");
+    RegisteredHostMemorySlice<SrcElemTy> SrcSlice(Src);
+    GlobalDeviceMemorySlice<DstElemTy> DstSlice(Dst);
+    if (ElementCount > SrcSlice.getElementCount())
       setError("copying too many elements, " + llvm::Twine(ElementCount) +
                ", from a host array of element count " +
-               llvm::Twine(Src.size()));
-    else if (ElementCount > Dst.getElementCount())
+               llvm::Twine(SrcSlice.getElementCount()));
+    else if (ElementCount > DstSlice.getElementCount())
       setError("copying too many elements, " + llvm::Twine(ElementCount) +
                ", to a device array of element count " +
-               llvm::Twine(Dst.getElementCount()));
+               llvm::Twine(DstSlice.getElementCount()));
     else
-      setError(PDevice->copyH2D(
-          PlatformStreamHandle, Src.data(), 0, Dst.getBaseMemory().getHandle(),
-          Dst.getElementOffset() * sizeof(T), ElementCount * sizeof(T)));
+      setError(PDevice->copyH2D(PlatformStreamHandle, SrcSlice.getPointer(), 0,
+                                DstSlice.getBaseMemory().getHandle(),
+                                DstSlice.getElementOffset() * sizeof(DstElemTy),
+                                ElementCount * sizeof(SrcElemTy)));
     return *this;
   }
 
-  template <typename T>
-  Stream &thenCopyH2D(llvm::ArrayRef<T> Src, GlobalDeviceMemorySlice<T> Dst) {
-    if (Src.size() != Dst.getElementCount())
+  template <typename SrcTy, typename DstTy>
+  Stream &thenCopyH2D(SrcTy &&Src, DstTy &&Dst) {
+    using SrcElemTy = typename std::remove_reference<SrcTy>::type::ElementTy;
+    using DstElemTy = typename std::remove_reference<DstTy>::type::ElementTy;
+    static_assert(std::is_same<SrcElemTy, DstElemTy>::value,
+                  "src/dst element type mismatch for thenCopyH2D");
+    RegisteredHostMemorySlice<SrcElemTy> SrcSlice(Src);
+    GlobalDeviceMemorySlice<DstElemTy> DstSlice(Dst);
+    if (SrcSlice.getElementCount() != DstSlice.getElementCount())
       setError("array size mismatch for H2D, host source has element count " +
-               llvm::Twine(Src.size()) +
+               llvm::Twine(SrcSlice.getElementCount()) +
                " but device destination has element count " +
-               llvm::Twine(Dst.getElementCount()));
+               llvm::Twine(DstSlice.getElementCount()));
     else
-      thenCopyH2D(Src, Dst, Dst.getElementCount());
-    return *this;
-  }
-
-  template <typename T>
-  Stream &thenCopyH2D(T *Src, GlobalDeviceMemorySlice<T> Dst,
-                      size_t ElementCount) {
-    thenCopyH2D(llvm::ArrayRef<T>(Src, ElementCount), Dst, ElementCount);
-    return *this;
-  }
-
-  template <typename T>
-  Stream &thenCopyH2D(llvm::ArrayRef<T> Src, GlobalDeviceMemory<T> &Dst,
-                      size_t ElementCount) {
-    thenCopyH2D(Src, Dst.asSlice(), ElementCount);
-    return *this;
-  }
-
-  template <typename T>
-  Stream &thenCopyH2D(llvm::ArrayRef<T> Src, GlobalDeviceMemory<T> &Dst) {
-    thenCopyH2D(Src, Dst.asSlice());
+      thenCopyH2D(SrcSlice, DstSlice, DstSlice.getElementCount());
     return *this;
   }
 
-  template <typename T>
-  Stream &thenCopyH2D(T *Src, GlobalDeviceMemory<T> &Dst, size_t ElementCount) {
-    thenCopyH2D(Src, Dst.asSlice(), ElementCount);
-    return *this;
-  }
+  // D2D
 
-  template <typename T>
-  Stream &thenCopyD2D(GlobalDeviceMemorySlice<T> Src,
-                      GlobalDeviceMemorySlice<T> Dst, size_t ElementCount) {
-    if (ElementCount > Src.getElementCount())
+  template <typename SrcTy, typename DstTy>
+  Stream &thenCopyD2D(SrcTy &&Src, DstTy &&Dst, size_t ElementCount) {
+    using SrcElemTy = typename std::remove_reference<SrcTy>::type::ElementTy;
+    using DstElemTy = typename std::remove_reference<DstTy>::type::ElementTy;
+    static_assert(std::is_same<SrcElemTy, DstElemTy>::value,
+                  "src/dst element type mismatch for thenCopyD2D");
+    GlobalDeviceMemorySlice<SrcElemTy> SrcSlice(Src);
+    GlobalDeviceMemorySlice<DstElemTy> DstSlice(Dst);
+    if (ElementCount > SrcSlice.getElementCount())
       setError("copying too many elements, " + llvm::Twine(ElementCount) +
                ", from a device array of element count " +
-               llvm::Twine(Src.getElementCount()));
-    else if (ElementCount > Dst.getElementCount())
+               llvm::Twine(SrcSlice.getElementCount()));
+    else if (ElementCount > DstSlice.getElementCount())
       setError("copying too many elements, " + llvm::Twine(ElementCount) +
                ", to a device array of element count " +
-               llvm::Twine(Dst.getElementCount()));
+               llvm::Twine(DstSlice.getElementCount()));
     else
-      setError(PDevice->copyD2D(
-          PlatformStreamHandle, Src.getBaseMemory().getHandle(),
-          Src.getElementOffset() * sizeof(T), Dst.getBaseMemory().getHandle(),
-          Dst.getElementOffset() * sizeof(T), ElementCount * sizeof(T)));
+      setError(PDevice->copyD2D(PlatformStreamHandle,
+                                SrcSlice.getBaseMemory().getHandle(),
+                                SrcSlice.getElementOffset() * sizeof(SrcElemTy),
+                                DstSlice.getBaseMemory().getHandle(),
+                                DstSlice.getElementOffset() * sizeof(DstElemTy),
+                                ElementCount * sizeof(SrcElemTy)));
     return *this;
   }
 
-  template <typename T>
-  Stream &thenCopyD2D(GlobalDeviceMemorySlice<T> Src,
-                      GlobalDeviceMemorySlice<T> Dst) {
-    if (Src.getElementCount() != Dst.getElementCount())
+  template <typename SrcTy, typename DstTy>
+  Stream &thenCopyD2D(SrcTy &&Src, DstTy &&Dst) {
+    using SrcElemTy = typename std::remove_reference<SrcTy>::type::ElementTy;
+    using DstElemTy = typename std::remove_reference<DstTy>::type::ElementTy;
+    static_assert(std::is_same<SrcElemTy, DstElemTy>::value,
+                  "src/dst element type mismatch for thenCopyD2D");
+    GlobalDeviceMemorySlice<SrcElemTy> SrcSlice(Src);
+    GlobalDeviceMemorySlice<DstElemTy> DstSlice(Dst);
+    if (SrcSlice.getElementCount() != DstSlice.getElementCount())
       setError("array size mismatch for D2D, device source has element count " +
-               llvm::Twine(Src.getElementCount()) +
+               llvm::Twine(SrcSlice.getElementCount()) +
                " but device destination has element count " +
-               llvm::Twine(Dst.getElementCount()));
+               llvm::Twine(DstSlice.getElementCount()));
     else
-      thenCopyD2D(Src, Dst, Src.getElementCount());
-    return *this;
-  }
-
-  template <typename T>
-  Stream &thenCopyD2D(const GlobalDeviceMemory<T> &Src,
-                      GlobalDeviceMemorySlice<T> Dst, size_t ElementCount) {
-    thenCopyD2D(Src.asSlice(), Dst, ElementCount);
-    return *this;
-  }
-
-  template <typename T>
-  Stream &thenCopyD2D(const GlobalDeviceMemory<T> &Src,
-                      GlobalDeviceMemorySlice<T> Dst) {
-    thenCopyD2D(Src.asSlice(), Dst);
-    return *this;
-  }
-
-  template <typename T>
-  Stream &thenCopyD2D(GlobalDeviceMemorySlice<T> Src,
-                      GlobalDeviceMemory<T> &Dst, size_t ElementCount) {
-    thenCopyD2D(Src, Dst.asSlice(), ElementCount);
-    return *this;
-  }
-
-  template <typename T>
-  Stream &thenCopyD2D(GlobalDeviceMemorySlice<T> Src,
-                      GlobalDeviceMemory<T> &Dst) {
-    thenCopyD2D(Src, Dst.asSlice());
-    return *this;
-  }
-
-  template <typename T>
-  Stream &thenCopyD2D(const GlobalDeviceMemory<T> &Src,
-                      GlobalDeviceMemory<T> &Dst, size_t ElementCount) {
-    thenCopyD2D(Src.asSlice(), Dst.asSlice(), ElementCount);
-    return *this;
-  }
-
-  template <typename T>
-  Stream &thenCopyD2D(const GlobalDeviceMemory<T> &Src,
-                      GlobalDeviceMemory<T> &Dst) {
-    thenCopyD2D(Src.asSlice(), Dst.asSlice());
+      thenCopyD2D(SrcSlice, DstSlice, SrcSlice.getElementCount());
     return *this;
   }
 
   ///@} End device memory copying functions
 
 private:
   /// Sets the error state from an Error object.
   ///

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

 macro(add_se_library name)
   add_llvm_library(${name} ${ARGN})
   set_target_properties(${name} PROPERTIES FOLDER "streamexecutor libraries")
 endmacro(add_se_library)
 
 add_se_library(
     streamexecutor
     Device.cpp
     DeviceMemory.cpp
     Error.cpp
+    HostMemory.cpp
     Kernel.cpp
     KernelSpec.cpp
     PackedKernelArgumentArray.cpp
     Platform.cpp
     PlatformDevice.cpp
     PlatformManager.cpp
     Stream.cpp
     )
 
 install(TARGETS streamexecutor DESTINATION lib)

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

+//===-- HostMemory.cpp - HostMemory 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 HostMemory internals.
+///
+//===----------------------------------------------------------------------===//
+
+#include "streamexecutor/HostMemory.h"
+#include "streamexecutor/Device.h"
+
+namespace streamexecutor {
+namespace internal {
+
+void destroyRegisteredHostMemoryInternals(Device *TheDevice, void *Pointer) {
+  // TODO(jhen): How to handle errors here?
+  if (Pointer) {
+    consumeError(TheDevice->unregisterHostMemory(Pointer));
+  }
+}
+
+} // namespace internal
+} // namespace streamexecutor

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

 }
 
 TEST_F(DeviceTest, AllocateAndFreeDeviceMemory) {
   se::Expected<se::GlobalDeviceMemory<int>> MaybeMemory =
       Device.allocateDeviceMemory<int>(10);
   EXPECT_TRUE(static_cast<bool>(MaybeMemory));
 }
 
-TEST_F(DeviceTest, AllocateAndFreeHostMemory) {
-  se::Expected<int *> MaybeMemory = Device.allocateHostMemory<int>(10);
-  EXPECT_TRUE(static_cast<bool>(MaybeMemory));
-  EXPECT_NO_ERROR(Device.freeHostMemory(*MaybeMemory));
-}
-
 TEST_F(DeviceTest, RegisterAndUnregisterHostMemory) {
   std::vector<int> Data(10);
-  EXPECT_NO_ERROR(Device.registerHostMemory(Data.data(), 10));
-  EXPECT_NO_ERROR(Device.unregisterHostMemory(Data.data()));
+  se::Expected<se::RegisteredHostMemory<int>> MaybeMemory =
+      Device.registerHostMemory<int>(Data);
+  EXPECT_TRUE(static_cast<bool>(MaybeMemory));
 }
 
 // D2H tests
 
 TEST_F(DeviceTest, SyncCopyD2HToMutableArrayRefByCount) {
   EXPECT_NO_ERROR(
       Device.synchronousCopyD2H(DeviceA5, MutableArrayRef<int>(Host5), 5));
   for (int I = 0; I < 5; ++I) {

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

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 {
+  streamexecutor::Error unregisterHostMemory(const void *Memory) override {
     return streamexecutor::Error::success();
   }
 
   streamexecutor::Error copyD2H(const void *StreamHandle,
                                 const void *DeviceHandleSrc,
                                 size_t SrcByteOffset, void *HostDst,
                                 size_t DstByteOffset,
                                 size_t ByteCount) override {

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

 class StreamTest : public ::testing::Test {
 public:
   StreamTest()
       : DummyPlatformStream(1), Device(&PDevice),
         Stream(&PDevice, &DummyPlatformStream), 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}, Host5{24, 25, 26, 27, 28},
         Host7{29, 30, 31, 32, 33, 34, 35},
+        RegisteredHost5(getOrDie(
+            Device.registerHostMemory(llvm::MutableArrayRef<int>(Host5)))),
+        RegisteredHost7(getOrDie(
+            Device.registerHostMemory(llvm::MutableArrayRef<int>(Host7)))),
         DeviceA5(getOrDie(Device.allocateDeviceMemory<int>(5))),
         DeviceB5(getOrDie(Device.allocateDeviceMemory<int>(5))),
         DeviceA7(getOrDie(Device.allocateDeviceMemory<int>(7))),
         DeviceB7(getOrDie(Device.allocateDeviceMemory<int>(7))) {
     se::dieIfError(Device.synchronousCopyH2D<int>(HostA5, DeviceA5));
     se::dieIfError(Device.synchronousCopyH2D<int>(HostB5, DeviceB5));
     se::dieIfError(Device.synchronousCopyH2D<int>(HostA7, DeviceA7));
     se::dieIfError(Device.synchronousCopyH2D<int>(HostB7, DeviceB7));
protected:
   int HostB5[5];
   int HostA7[7];
   int HostB7[7];
 
   // Host memory to be used as actual host memory.
   int Host5[5];
   int Host7[7];
 
+  se::RegisteredHostMemory<int> RegisteredHost5;
+  se::RegisteredHostMemory<int> RegisteredHost7;
+
   // Device memory.
   se::GlobalDeviceMemory<int> DeviceA5;
   se::GlobalDeviceMemory<int> DeviceB5;
   se::GlobalDeviceMemory<int> DeviceA7;
   se::GlobalDeviceMemory<int> DeviceB7;
 };
 
 using llvm::ArrayRef;
 using llvm::MutableArrayRef;
 
 // D2H tests
 
-TEST_F(StreamTest, CopyD2HToMutableArrayRefByCount) {
-  Stream.thenCopyD2H(DeviceA5, MutableArrayRef<int>(Host5), 5);
+TEST_F(StreamTest, CopyD2HToRegisteredRefByCount) {
+  Stream.thenCopyD2H(DeviceA5, RegisteredHost5, 5);
   EXPECT_TRUE(Stream.isOK());
   for (int I = 0; I < 5; ++I) {
     EXPECT_EQ(HostA5[I], Host5[I]);
   }
 
-  Stream.thenCopyD2H(DeviceB5, MutableArrayRef<int>(Host5), 2);
+  Stream.thenCopyD2H(DeviceB5, RegisteredHost5, 2);
   EXPECT_TRUE(Stream.isOK());
   for (int I = 0; I < 2; ++I) {
     EXPECT_EQ(HostB5[I], Host5[I]);
   }
 
-  Stream.thenCopyD2H(DeviceA7, MutableArrayRef<int>(Host5), 7);
-  EXPECT_FALSE(Stream.isOK());
-}
-
-TEST_F(StreamTest, CopyD2HToMutableArrayRef) {
-  Stream.thenCopyD2H(DeviceA5, MutableArrayRef<int>(Host5));
-  EXPECT_TRUE(Stream.isOK());
-  for (int I = 0; I < 5; ++I) {
-    EXPECT_EQ(HostA5[I], Host5[I]);
-  }
-
-  Stream.thenCopyD2H(DeviceA5, MutableArrayRef<int>(Host7));
+  Stream.thenCopyD2H(DeviceA7, RegisteredHost5, 7);
   EXPECT_FALSE(Stream.isOK());
 }
 
-TEST_F(StreamTest, CopyD2HToPointer) {
-  Stream.thenCopyD2H(DeviceA5, Host5, 5);
+TEST_F(StreamTest, CopyD2HToRegistered) {
+  Stream.thenCopyD2H(DeviceA5, RegisteredHost5);
   EXPECT_TRUE(Stream.isOK());
   for (int I = 0; I < 5; ++I) {
     EXPECT_EQ(HostA5[I], Host5[I]);
   }
 
-  Stream.thenCopyD2H(DeviceA5, Host7, 7);
+  Stream.thenCopyD2H(DeviceA5, RegisteredHost7);
   EXPECT_FALSE(Stream.isOK());
 }
 
-TEST_F(StreamTest, CopyD2HSliceToMutableArrayRefByCount) {
+TEST_F(StreamTest, CopyD2HSliceToRegiseredSliceByCount) {
   Stream.thenCopyD2H(DeviceA5.asSlice().drop_front(1),
-                     MutableArrayRef<int>(Host5 + 1, 4), 4);
+                     RegisteredHost5.asSlice().slice(1, 4), 4);
   EXPECT_TRUE(Stream.isOK());
   for (int I = 1; I < 5; ++I) {
     EXPECT_EQ(HostA5[I], Host5[I]);
   }
 
-  Stream.thenCopyD2H(DeviceB5.asSlice().drop_back(1),
-                     MutableArrayRef<int>(Host5), 2);
+  Stream.thenCopyD2H(DeviceB5.asSlice().drop_back(1), RegisteredHost5, 2);
   EXPECT_TRUE(Stream.isOK());
   for (int I = 0; I < 2; ++I) {
     EXPECT_EQ(HostB5[I], Host5[I]);
   }
 
-  Stream.thenCopyD2H(DeviceA5.asSlice(), MutableArrayRef<int>(Host7), 7);
+  Stream.thenCopyD2H(DeviceA5.asSlice(), RegisteredHost7, 7);
   EXPECT_FALSE(Stream.isOK());
 }
 
-TEST_F(StreamTest, CopyD2HSliceToMutableArrayRef) {
-  Stream.thenCopyD2H(DeviceA7.asSlice().slice(1, 5),
-                     MutableArrayRef<int>(Host5));
+TEST_F(StreamTest, CopyD2HSliceToRegistered) {
+  Stream.thenCopyD2H(DeviceA7.asSlice().slice(1, 5), RegisteredHost5);
   EXPECT_TRUE(Stream.isOK());
   for (int I = 0; I < 5; ++I) {
     EXPECT_EQ(HostA7[I + 1], Host5[I]);
   }
 
-  Stream.thenCopyD2H(DeviceA5.asSlice(), MutableArrayRef<int>(Host7));
-  EXPECT_FALSE(Stream.isOK());
-}
-
-TEST_F(StreamTest, CopyD2HSliceToPointer) {
-  Stream.thenCopyD2H(DeviceA5.asSlice().drop_front(1), Host5 + 1, 4);
-  EXPECT_TRUE(Stream.isOK());
-  for (int I = 1; I < 5; ++I) {
-    EXPECT_EQ(HostA5[I], Host5[I]);
-  }
-
-  Stream.thenCopyD2H(DeviceA5.asSlice(), Host7, 7);
+  Stream.thenCopyD2H(DeviceA5.asSlice(), RegisteredHost7);
   EXPECT_FALSE(Stream.isOK());
 }
 
 // H2D tests
 
-TEST_F(StreamTest, CopyH2DToArrayRefByCount) {
-  Stream.thenCopyH2D(ArrayRef<int>(Host5), DeviceA5, 5);
+TEST_F(StreamTest, CopyH2DFromRegisterdByCount) {
+  Stream.thenCopyH2D(RegisteredHost5, DeviceA5, 5);
   EXPECT_TRUE(Stream.isOK());
   for (int I = 0; I < 5; ++I) {
     EXPECT_EQ(getDeviceValue(DeviceA5, I), Host5[I]);
   }
 
-  Stream.thenCopyH2D(ArrayRef<int>(Host5), DeviceB5, 2);
+  Stream.thenCopyH2D(RegisteredHost5, DeviceB5, 2);
   EXPECT_TRUE(Stream.isOK());
   for (int I = 0; I < 2; ++I) {
     EXPECT_EQ(getDeviceValue(DeviceB5, I), Host5[I]);
   }
 
-  Stream.thenCopyH2D(ArrayRef<int>(Host7), DeviceA5, 7);
-  EXPECT_FALSE(Stream.isOK());
-}
-
-TEST_F(StreamTest, CopyH2DToArrayRef) {
-  Stream.thenCopyH2D(ArrayRef<int>(Host5), DeviceA5);
-  EXPECT_TRUE(Stream.isOK());
-  for (int I = 0; I < 5; ++I) {
-    EXPECT_EQ(getDeviceValue(DeviceA5, I), Host5[I]);
-  }
-
-  Stream.thenCopyH2D(ArrayRef<int>(Host7), DeviceA5);
+  Stream.thenCopyH2D(RegisteredHost7, DeviceA5, 7);
   EXPECT_FALSE(Stream.isOK());
 }
 
-TEST_F(StreamTest, CopyH2DToPointer) {
-  Stream.thenCopyH2D(Host5, DeviceA5, 5);
+TEST_F(StreamTest, CopyH2DFromRegistered) {
+  Stream.thenCopyH2D(RegisteredHost5, DeviceA5);
   EXPECT_TRUE(Stream.isOK());
   for (int I = 0; I < 5; ++I) {
     EXPECT_EQ(getDeviceValue(DeviceA5, I), Host5[I]);
   }
 
-  Stream.thenCopyH2D(Host7, DeviceA5, 7);
+  Stream.thenCopyH2D(RegisteredHost7, DeviceA5);
   EXPECT_FALSE(Stream.isOK());
 }
 
-TEST_F(StreamTest, CopyH2DSliceToArrayRefByCount) {
-  Stream.thenCopyH2D(ArrayRef<int>(Host5 + 1, 4),
+TEST_F(StreamTest, CopyH2DFromRegisteredSliceToSlice) {
+  Stream.thenCopyH2D(RegisteredHost5.asSlice().slice(1, 4),
                      DeviceA5.asSlice().drop_front(1), 4);
   EXPECT_TRUE(Stream.isOK());
   for (int I = 1; I < 5; ++I) {
     EXPECT_EQ(getDeviceValue(DeviceA5, I), Host5[I]);
   }
 
-  Stream.thenCopyH2D(ArrayRef<int>(Host5), DeviceB5.asSlice().drop_back(1), 2);
+  Stream.thenCopyH2D(RegisteredHost5, DeviceB5.asSlice().drop_back(1), 2);
   EXPECT_TRUE(Stream.isOK());
   for (int I = 0; I < 2; ++I) {
     EXPECT_EQ(getDeviceValue(DeviceB5, I), Host5[I]);
   }
 
-  Stream.thenCopyH2D(ArrayRef<int>(Host5), DeviceA5.asSlice(), 7);
+  Stream.thenCopyH2D(RegisteredHost5, DeviceA5.asSlice(), 7);
   EXPECT_FALSE(Stream.isOK());
 }
 
-TEST_F(StreamTest, CopyH2DSliceToArrayRef) {
-
-  Stream.thenCopyH2D(ArrayRef<int>(Host5), DeviceA5.asSlice());
+TEST_F(StreamTest, CopyH2DRegisteredToSlice) {
+  Stream.thenCopyH2D(RegisteredHost5, DeviceA5.asSlice());
   EXPECT_TRUE(Stream.isOK());
   for (int I = 0; I < 5; ++I) {
     EXPECT_EQ(getDeviceValue(DeviceA5, I), Host5[I]);
   }
 
-  Stream.thenCopyH2D(ArrayRef<int>(Host7), DeviceA5.asSlice());
-  EXPECT_FALSE(Stream.isOK());
-}
-
-TEST_F(StreamTest, CopyH2DSliceToPointer) {
-  Stream.thenCopyH2D(Host5, DeviceA5.asSlice(), 5);
-  EXPECT_TRUE(Stream.isOK());
-  for (int I = 0; I < 5; ++I) {
-    EXPECT_EQ(getDeviceValue(DeviceA5, I), Host5[I]);
-  }
-
-  Stream.thenCopyH2D(Host7, DeviceA5.asSlice(), 7);
+  Stream.thenCopyH2D(RegisteredHost7, DeviceA5.asSlice());
   EXPECT_FALSE(Stream.isOK());
 }
 
 // D2D tests
 
 TEST_F(StreamTest, CopyD2DByCount) {
   Stream.thenCopyD2D(DeviceA5, DeviceB5, 5);
   EXPECT_TRUE(Stream.isOK());
for (int I = 0; I < 2; ++I) {
     EXPECT_EQ(getDeviceValue(DeviceA7, I), getDeviceValue(DeviceB7, I));
   }
 
   Stream.thenCopyD2D(DeviceA5.asSlice(), DeviceB5, 7);
   EXPECT_FALSE(Stream.isOK());
 }
 
 TEST_F(StreamTest, CopySliceD2D) {
-
   Stream.thenCopyD2D(DeviceA7.asSlice().drop_back(2), DeviceB5);
   EXPECT_TRUE(Stream.isOK());
   for (int I = 0; I < 5; ++I) {
     EXPECT_EQ(getDeviceValue(DeviceA7, I), getDeviceValue(DeviceB5, I));
   }
 
   Stream.thenCopyD2D(DeviceA5.asSlice().drop_back(1), DeviceB7);
   EXPECT_FALSE(Stream.isOK());
for (int I = 0; I < 2; ++I) {
     EXPECT_EQ(getDeviceValue(DeviceA7, I), getDeviceValue(DeviceB7, I));
   }
 
   Stream.thenCopyD2D(DeviceA5, DeviceB7.asSlice(), 7);
   EXPECT_FALSE(Stream.isOK());
 }
 
 TEST_F(StreamTest, CopyD2DSlice) {
-
   Stream.thenCopyD2D(DeviceA5, DeviceB7.asSlice().drop_back(2));
   EXPECT_TRUE(Stream.isOK());
   for (int I = 0; I < 5; ++I) {
     EXPECT_EQ(getDeviceValue(DeviceA5, I), getDeviceValue(DeviceB7, I));
   }
 
   Stream.thenCopyD2D(DeviceA5, DeviceB7.asSlice());
   EXPECT_FALSE(Stream.isOK());
 }
 
 TEST_F(StreamTest, CopySliceD2DSliceByCount) {
-
   Stream.thenCopyD2D(DeviceA5.asSlice(), DeviceB5.asSlice(), 5);
   EXPECT_TRUE(Stream.isOK());
   for (int I = 0; I < 5; ++I) {
     EXPECT_EQ(getDeviceValue(DeviceA5, I), getDeviceValue(DeviceB5, I));
   }
 
   Stream.thenCopyD2D(DeviceA7.asSlice(), DeviceB7.asSlice(), 2);
   EXPECT_TRUE(Stream.isOK());
   for (int I = 0; I < 2; ++I) {
     EXPECT_EQ(getDeviceValue(DeviceA7, I), getDeviceValue(DeviceB7, I));
   }
 
   Stream.thenCopyD2D(DeviceA7.asSlice(), DeviceB5.asSlice(), 7);
   EXPECT_FALSE(Stream.isOK());
 }
 
 TEST_F(StreamTest, CopySliceD2DSlice) {
-
   Stream.thenCopyD2D(DeviceA5.asSlice(), DeviceB5.asSlice());
   EXPECT_TRUE(Stream.isOK());
   for (int I = 0; I < 5; ++I) {
     EXPECT_EQ(getDeviceValue(DeviceA5, I), getDeviceValue(DeviceB5, I));
   }
 
   Stream.thenCopyD2D(DeviceA5.asSlice(), DeviceB7.asSlice());
   EXPECT_FALSE(Stream.isOK());
 }
 
 } // namespace