[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

host

streamexecutor/include/streamexecutor/HostMemory.h

106

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

108–110

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

33

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

161

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.

[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

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

161

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

[jlebar]

streamexecutor/include/streamexecutor/HostMemory.h

46

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.

46

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.

49

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.

100

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

streamexecutor/include/streamexecutor/Stream.h

169

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

175

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

46

Switching back to mutable array.

49

Switching back to mutable array.

100

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

169

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

175

Using mutable array now.

[jlebar]

Ship it.

streamexecutor/include/streamexecutor/HostMemory.h

95

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

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<float>(RegisteredX, X)
+      .thenCopyH2D<float>(RegisteredY, Y)
       .thenLaunch(ArraySize, 1, Kernel, A, X, Y)
-      .thenCopyD2H<float>(X, HostX);
+      .thenCopyD2H<float>(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]
 }

streamexecutor/include/streamexecutor/Device.h

 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef STREAMEXECUTOR_DEVICE_H
 #define STREAMEXECUTOR_DEVICE_H
 
 #include <type_traits>
 
+#include "streamexecutor/HostMemory.h"
 #include "streamexecutor/KernelSpec.h"
 #include "streamexecutor/PlatformDevice.h"
 #include "streamexecutor/Utils/Error.h"
 
 namespace streamexecutor {
 
 class Stream;
 
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 a RegisteredHostMemoryBase may unregister host memory.
+  friend RegisteredHostMemoryBase;
+
+  Error unregisterHostMemory(const RegisteredHostMemoryBase &Memory) {
+    return PDevice->unregisterHostMemory(Memory.getUntypedPointer());
+  }
+
   PlatformDevice *PDevice;
 };
 
 } // namespace streamexecutor
 
 #endif // STREAMEXECUTOR_DEVICE_H

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;
+
+/// An immutable slice of registered host memory.
+///
+/// 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:
+  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:
+  friend RegisteredHostMemory<ElemT>;
+
+  RegisteredHostMemorySlice(llvm::ArrayRef<ElemT> ArrayRef)
+      : ArrayRef(ArrayRef) {}
+
+  llvm::ArrayRef<ElemT> ArrayRef;
+};
+
+/// A mutable slice of registered host memory.
+///
+/// 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:
+  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:
+  friend RegisteredHostMemory<ElemT>;
+
+  MutableRegisteredHostMemorySlice(llvm::MutableArrayRef<ElemT> MutableArrayRef)
+      : MutableArrayRef(MutableArrayRef) {}
+
+  llvm::MutableArrayRef<ElemT> MutableArrayRef;
+};
+
+/// Base class for registered host memory that knows how to unregister itself
+/// upon destruction.
+///
+/// This class does not keep track of the data types of the underlying memory.
+/// See RegisteredHostMemory<T> for the subclass that does keep track of types.
+///
+/// Can be created by registering previously allocated host memory and will
+/// unregister that memory upon destruction.
+class RegisteredHostMemoryBase {
+public:
+  RegisteredHostMemoryBase(const RegisteredHostMemoryBase &) = delete;
+  RegisteredHostMemoryBase &
+  operator=(const RegisteredHostMemoryBase &) = delete;
+
+  RegisteredHostMemoryBase(RegisteredHostMemoryBase &&Other)
+      : TheDevice(Other.TheDevice), Pointer(Other.Pointer),
+        ByteCount(Other.ByteCount) {
+    Other.TheDevice = nullptr;
+    Other.Pointer = nullptr;
+  };
+
+  RegisteredHostMemoryBase &operator=(RegisteredHostMemoryBase &&Other) {
+    TheDevice = Other.TheDevice;
+    Pointer = Other.Pointer;
+    ByteCount = Other.ByteCount;
+    Other.TheDevice = nullptr;
+    Other.Pointer = nullptr;
+    return *this;
+  }
+
+  const void *getUntypedPointer() const { return Pointer; }
+  size_t getByteCount() const { return ByteCount; }
+
+protected:
+  RegisteredHostMemoryBase(Device *TheDevice, void *Pointer, size_t ByteCount)
+      : TheDevice(TheDevice), Pointer(Pointer), ByteCount(ByteCount) {
+    assert(TheDevice != nullptr && "cannot construct a "
+                                   "RegisteredHostMemoryBase with a null "
+                                   "platform device");
+  }
+
+  ~RegisteredHostMemoryBase();
+
+  Device *TheDevice;
+  void *Pointer;
+  size_t ByteCount;
+};
+
+/// A subclass of RegisteredHostMemoryBase that keeps track of the types of
+/// elements stored in the host buffer.
+///
+/// ElemT is the type of element stored in the host buffer.
+template <typename ElemT>
+class RegisteredHostMemory : RegisteredHostMemoryBase {
+public:
+  RegisteredHostMemory(Device *TheDevice, ElemT *Pointer, size_t ElementCount)
+      : RegisteredHostMemoryBase(TheDevice, Pointer,
+                                 ElementCount * sizeof(ElemT)) {}
+
+  RegisteredHostMemory(RegisteredHostMemory &&) = default;
+  RegisteredHostMemory &operator=(RegisteredHostMemory &&) = default;
+
+  ElemT *getPointer() { return static_cast<ElemT *>(Pointer); }
+  size_t getElementCount() const { return ByteCount / sizeof(ElemT); }
+
+  /// Creates an immutable slice for the entire contents of this memory.
+  RegisteredHostMemorySlice<ElemT> asSlice() const {
+    return RegisteredHostMemorySlice<ElemT>(llvm::ArrayRef<ElemT>(
+        static_cast<ElemT *>(Pointer), getElementCount()));
+  }
+
+  /// Creates a mutable slice for the entire contents of this memory.
+  MutableRegisteredHostMemorySlice<ElemT> asSlice() {
+    return MutableRegisteredHostMemorySlice<ElemT>(llvm::MutableArrayRef<ElemT>(
+        static_cast<ElemT *>(Pointer), getElementCount()));
+  }
+};
+
+/// \name Type trait unifying  RegisteredHostMemorySlice<T> and
+/// MutableRegisteredHostMemorySlice<T>.
+///
+/// In many cases, it doesn't matter if a slice is mutable or not, so this type
+/// trait makes it easy to write templated functions that take either kind of
+/// slice.
+///
+/// is_registered_host_memory_slice<T, ElemT> is std::true_type if T is either
+/// RegisteredHostMemorySlice<ElemT> or MutableRegisteredHostMemorySlice<T>.
+/// Otherwise, it is std::false_type.
+///
+/// @{
+
+template <typename T, typename ElemT>
+struct is_registered_host_memory_slice : public std::false_type {};
+
+template <typename ElemT>
+struct is_registered_host_memory_slice<RegisteredHostMemorySlice<ElemT>, ElemT>
+    : public std::true_type {};
+
+template <typename ElemT>
+struct is_registered_host_memory_slice<MutableRegisteredHostMemorySlice<ElemT>,
+                                       ElemT> : public std::true_type {};
+
+/// @}
+
+} // namespace streamexecutor
+
+#endif // STREAMEXECUTOR_HOSTMEMORY_H

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

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/HostMemory.h"
 #include "streamexecutor/Kernel.h"
 #include "streamexecutor/LaunchDimensions.h"
 #include "streamexecutor/PackedKernelArgumentArray.h"
 #include "streamexecutor/PlatformDevice.h"
 #include "streamexecutor/Utils/Error.h"
 
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/Twine.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.
   ///@{
 
   template <typename T>
   Stream &thenCopyD2H(GlobalDeviceMemorySlice<T> Src,
-                      llvm::MutableArrayRef<T> Dst, size_t ElementCount) {
+                      MutableRegisteredHostMemorySlice<T> Dst,
+                      size_t ElementCount) {
     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())
+    else if (ElementCount > Dst.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(Dst.getElementCount()));
     else
       setError(PDevice->copyD2H(PlatformStreamHandle,
                                 Src.getBaseMemory().getHandle(),
-                                Src.getElementOffset() * sizeof(T), Dst.data(),
-                                0, ElementCount * sizeof(T)));
+                                Src.getElementOffset() * sizeof(T),
+                                Dst.getPointer(), 0, ElementCount * sizeof(T)));
     return *this;
   }
 
   template <typename T>
   Stream &thenCopyD2H(GlobalDeviceMemorySlice<T> Src,
-                      llvm::MutableArrayRef<T> Dst) {
-    if (Src.getElementCount() != Dst.size())
+                      RegisteredHostMemory<T> &Dst, size_t ElementCount) {
+    return thenCopyD2H(Src, Dst.asSlice(), ElementCount);
+  }
+
+  template <typename T>
+  Stream &thenCopyD2H(GlobalDeviceMemorySlice<T> Src,
+                      MutableRegisteredHostMemorySlice<T> Dst) {
+    if (Src.getElementCount() != Dst.getElementCount())
       setError("array size mismatch for D2H, device source has element count " +
                llvm::Twine(Src.getElementCount()) +
                " but host destination has element count " +
-               llvm::Twine(Dst.size()));
+               llvm::Twine(Dst.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;
+  Stream &thenCopyD2H(GlobalDeviceMemorySlice<T> Src,
+                      RegisteredHostMemory<T> &Dst) {
+    return thenCopyD2H(Src, Dst.asSlice());
   }
 
   template <typename T>
   Stream &thenCopyD2H(const GlobalDeviceMemory<T> &Src,
-                      llvm::MutableArrayRef<T> Dst, size_t ElementCount) {
-    thenCopyD2H(Src.asSlice(), Dst, ElementCount);
-    return *this;
+                      MutableRegisteredHostMemorySlice<T> Dst,
+                      size_t ElementCount) {
+    return thenCopyD2H(Src.asSlice(), Dst, ElementCount);
   }
 
   template <typename T>
   Stream &thenCopyD2H(const GlobalDeviceMemory<T> &Src,
-                      llvm::MutableArrayRef<T> Dst) {
-    thenCopyD2H(Src.asSlice(), Dst);
-    return *this;
+                      RegisteredHostMemory<T> &Dst, size_t ElementCount) {
+    return thenCopyD2H(Src.asSlice(), Dst.asSlice(), ElementCount);
   }
 
   template <typename T>
-  Stream &thenCopyD2H(const GlobalDeviceMemory<T> &Src, T *Dst,
-                      size_t ElementCount) {
-    thenCopyD2H(Src.asSlice(), Dst, ElementCount);
-    return *this;
+  Stream &thenCopyD2H(const GlobalDeviceMemory<T> &Src,
+                      MutableRegisteredHostMemorySlice<T> Dst) {
+    return thenCopyD2H(Src.asSlice(), Dst);
   }
 
   template <typename T>
-  Stream &thenCopyH2D(llvm::ArrayRef<T> Src, GlobalDeviceMemorySlice<T> Dst,
-                      size_t ElementCount) {
-    if (ElementCount > Src.size())
+  Stream &thenCopyD2H(const GlobalDeviceMemory<T> &Src,
+                      RegisteredHostMemory<T> &Dst) {
+    return thenCopyD2H(Src.asSlice(), Dst.asSlice());
+  }
+
+  template <typename T, typename SrcSliceT>
+  typename std::enable_if<is_registered_host_memory_slice<SrcSliceT, T>::value,
+                          Stream>::type &
+  thenCopyH2D(SrcSliceT Src, GlobalDeviceMemorySlice<T> Dst,
+    if (ElementCount > Src.getElementCount())
       setError("copying too many elements, " + llvm::Twine(ElementCount) +
                ", from a host array of element count " +
-               llvm::Twine(Src.size()));
+               llvm::Twine(Src.getElementCount()));
     else if (ElementCount > Dst.getElementCount())
       setError("copying too many elements, " + llvm::Twine(ElementCount) +
                ", to a device array of element count " +
                llvm::Twine(Dst.getElementCount()));
     else
-      setError(PDevice->copyH2D(
-          PlatformStreamHandle, Src.data(), 0, Dst.getBaseMemory().getHandle(),
-          Dst.getElementOffset() * sizeof(T), ElementCount * sizeof(T)));
+      setError(PDevice->copyH2D(PlatformStreamHandle, Src.getPointer(), 0,
+                                Dst.getBaseMemory().getHandle(),
+                                Dst.getElementOffset() * sizeof(T),
+                                ElementCount * sizeof(T)));
     return *this;
   }
 
   template <typename T>
-  Stream &thenCopyH2D(llvm::ArrayRef<T> Src, GlobalDeviceMemorySlice<T> Dst) {
-    if (Src.size() != Dst.getElementCount())
+  Stream &thenCopyH2D(const RegisteredHostMemory<T> &Src,
+                      GlobalDeviceMemorySlice<T> Dst, size_t ElementCount) {
+    return thenCopyH2D(Src.asSlice(), Dst, ElementCount);
+  }
+
+  template <typename T, typename SrcSliceT>
+  typename std::enable_if<is_registered_host_memory_slice<SrcSliceT, T>::value,
+                          Stream>::type &
+  thenCopyH2D(SrcSliceT Src, GlobalDeviceMemorySlice<T> Dst) {
+    if (Src.getElementCount() != Dst.getElementCount())
       setError("array size mismatch for H2D, host source has element count " +
-               llvm::Twine(Src.size()) +
+               llvm::Twine(Src.getElementCount()) +
                " but device destination has element count " +
                llvm::Twine(Dst.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;
+  Stream &thenCopyH2D(const RegisteredHostMemory<T> &Src,
+                      GlobalDeviceMemorySlice<T> Dst) {
+    return thenCopyH2D(Src.asSlice(), Dst);
   }
 
-  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, typename SrcSliceT>
+  typename std::enable_if<is_registered_host_memory_slice<SrcSliceT, T>::value,
+                          Stream>::type &
+  thenCopyH2D(SrcSliceT Src, GlobalDeviceMemory<T> &Dst, size_t ElementCount) {
+    return thenCopyH2D(Src, Dst.asSlice(), ElementCount);
   }
 
   template <typename T>
-  Stream &thenCopyH2D(llvm::ArrayRef<T> Src, GlobalDeviceMemory<T> &Dst) {
-    thenCopyH2D(Src, Dst.asSlice());
-    return *this;
+  Stream &thenCopyH2D(const RegisteredHostMemory<T> &Src,
+                      GlobalDeviceMemory<T> &Dst, size_t ElementCount) {
+    return thenCopyH2D(Src.asSlice(), Dst, ElementCount);
+  }
+
+  template <typename T, typename SrcSliceT>
+  typename std::enable_if<is_registered_host_memory_slice<SrcSliceT, T>::value,
+                          Stream>::type &
+  thenCopyH2D(SrcSliceT Src, GlobalDeviceMemory<T> &Dst) {
+    return thenCopyH2D(Src, Dst.asSlice());
   }
 
   template <typename T>
-  Stream &thenCopyH2D(T *Src, GlobalDeviceMemory<T> &Dst, size_t ElementCount) {
-    thenCopyH2D(Src, Dst.asSlice(), ElementCount);
-    return *this;
+  Stream &thenCopyH2D(const RegisteredHostMemory<T> &Src,
+                      GlobalDeviceMemory<T> &Dst) {
+    return thenCopyH2D(Src.asSlice(), Dst);
   }
 
   template <typename T>
   Stream &thenCopyD2D(GlobalDeviceMemorySlice<T> Src,
                       GlobalDeviceMemorySlice<T> Dst, size_t ElementCount) {
     if (ElementCount > Src.getElementCount())
       setError("copying too many elements, " + llvm::Twine(ElementCount) +
                ", from a device array of element count " +

streamexecutor/lib/CMakeLists.txt

 add_library(
     utils
     OBJECT
     Utils/Error.cpp)
 
 add_library(
     streamexecutor
     $<TARGET_OBJECTS:utils>
     Device.cpp
     DeviceMemory.cpp
+    HostMemory.cpp
     Kernel.cpp
     KernelSpec.cpp
     PackedKernelArgumentArray.cpp
     Platform.cpp
     PlatformDevice.cpp
     PlatformManager.cpp
     Stream.cpp)
 target_link_libraries(streamexecutor ${llvm_libs})
 
 install(TARGETS streamexecutor DESTINATION lib)
 
 if(STREAM_EXECUTOR_UNIT_TESTS)
     add_subdirectory(unittests)
 endif()

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

streamexecutor/lib/unittests/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) {

streamexecutor/lib/unittests/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 {

streamexecutor/lib/unittests/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