Add managed_pointer that is compatible with STL.

testing/cuda/managed_memory_pointer.cu

@@ -0,0 +1,141 @@
+#include <thrust/detail/config.h>
+
+#if THRUST_CPP_DIALECT >= 2011
+
+#  include <unittest/unittest.h>
+
+#  include <thrust/allocate_unique.h>
+#  include <thrust/memory/detail/device_system_resource.h>
+#  include <thrust/mr/allocator.h>
+#  include <thrust/type_traits/is_contiguous_iterator.h>
+
+#  include <numeric>
+#  include <vector>
+
+namespace
+{
+
+template <typename T>
+using allocator =
+  thrust::mr::stateless_resource_allocator<T, thrust::universal_memory_resource>;
+
+// The managed_memory_pointer class should be identified as a
+// contiguous_iterator
+THRUST_STATIC_ASSERT(
+  thrust::is_contiguous_iterator<allocator<int>::pointer>::value);
+
+template <typename T>
+struct some_object {
+  some_object(T data)
+      : m_data(data)
+  {}
+
+  void setter(T data) { m_data = data; }
+  T getter() const { return m_data; }
+
+private:
+  T m_data;
+};
+
+} // namespace
+
+template <typename T>
+void TestAllocateUnique()
+{
+  // Simple test to ensure that pointers created with universal_memory_resource
+  // can be dereferenced and used with STL code. This is necessary as some
+  // STL implementations break when using fancy references that overload
+  // `operator&`, so universal_memory_resource uses a special pointer type that
+  // returns regular C++ references that can be safely used host-side.
+
+  // These operations fail to compile with fancy references:
+  auto pRaw = thrust::allocate_unique<T>(allocator<T>{}, 42);
+  auto pObj =
+    thrust::allocate_unique<some_object<T> >(allocator<some_object<T> >{}, 42);
+
+  static_assert(
+    std::is_same<decltype(pRaw.get()),
+                 thrust::system::cuda::detail::managed_memory_pointer<T> >::value,
+    "Unexpected pointer returned from unique_ptr::get.");
+  static_assert(
+    std::is_same<decltype(pObj.get()),
+                 thrust::system::cuda::detail::managed_memory_pointer<
+                   some_object<T> > >::value,
+    "Unexpected pointer returned from unique_ptr::get.");
+
+  ASSERT_EQUAL(*pRaw, T(42));
+  ASSERT_EQUAL(*pRaw.get(), T(42));
+  ASSERT_EQUAL(pObj->getter(), T(42));
+  ASSERT_EQUAL((*pObj).getter(), T(42));
+  ASSERT_EQUAL(pObj.get()->getter(), T(42));
+  ASSERT_EQUAL((*pObj.get()).getter(), T(42));
+}
+DECLARE_GENERIC_UNITTEST(TestAllocateUnique);
+
+template <typename T>
+void TestIterationRaw()
+{
+  auto array = thrust::allocate_unique_n<T>(allocator<T>{}, 6, 42);
+
+  static_assert(
+    std::is_same<decltype(array.get()),
+                 thrust::system::cuda::detail::managed_memory_pointer<T> >::value,
+    "Unexpected pointer returned from unique_ptr::get.");
+
+  for (auto iter = array.get(), end = array.get() + 6; iter < end; ++iter)
+  {
+    ASSERT_EQUAL(*iter, T(42));
+    ASSERT_EQUAL(*iter.get(), T(42));
+  }
+}
+DECLARE_GENERIC_UNITTEST(TestIterationRaw);
+
+template <typename T>
+void TestIterationObj()
+{
+  auto array =
+    thrust::allocate_unique_n<some_object<T> >(allocator<some_object<T> >{},
+                                               6,
+                                               42);
+
+  static_assert(
+    std::is_same<decltype(array.get()),
+                 thrust::system::cuda::detail::managed_memory_pointer<
+                   some_object<T> > >::value,
+    "Unexpected pointer returned from unique_ptr::get.");
+
+  for (auto iter = array.get(), end = array.get() + 6; iter < end; ++iter)
+  {
+    ASSERT_EQUAL(iter->getter(), T(42));
+    ASSERT_EQUAL((*iter).getter(), T(42));
+    ASSERT_EQUAL(iter.get()->getter(), T(42));
+    ASSERT_EQUAL((*iter.get()).getter(), T(42));
+  }
+}
+DECLARE_GENERIC_UNITTEST(TestIterationObj);
+
+template <typename T>
+void TestStdVector()
+{
+  // Verify that a std::vector using the universal allocator will work with
+  // STL algorithms.
+  std::vector<T, allocator<T> > v0;
+
+  static_assert(
+    std::is_same<typename std::decay<decltype(v0)>::type::pointer,
+                 thrust::system::cuda::detail::managed_memory_pointer<
+                   T > >::value,
+    "Unexpected pointer returned from unique_ptr::get.");
+
+  v0.resize(6);
+  std::iota(v0.begin(), v0.end(), 0);
+  ASSERT_EQUAL(v0[0], T(0));
+  ASSERT_EQUAL(v0[1], T(1));
+  ASSERT_EQUAL(v0[2], T(2));
+  ASSERT_EQUAL(v0[3], T(3));
+  ASSERT_EQUAL(v0[4], T(4));
+  ASSERT_EQUAL(v0[5], T(5));
+}
+DECLARE_GENERIC_UNITTEST(TestStdVector);
+
+#endif // C++11

testing/cuda/managed_memory_pointer.mk

@@ -0,0 +1 @@
+CUDACC_FLAGS += -rdc=true

testing/vector.cu

@@ -52,24 +52,27 @@ DECLARE_VECTOR_UNITTEST(TestVectorFrontBack);
 template <class Vector>
 void TestVectorData(void)
 {
+    typedef typename Vector::pointer PointerT;
+    typedef typename Vector::const_pointer PointerConstT;
+
     Vector v(3);
     v[0] = 0; v[1] = 1; v[2] = 2;
 
     ASSERT_EQUAL(0,          *v.data());
     ASSERT_EQUAL(1,          *(v.data() + 1));
     ASSERT_EQUAL(2,          *(v.data() + 2));
-    ASSERT_EQUAL(&v.front(),  v.data());
-    ASSERT_EQUAL(&*v.begin(), v.data());
-    ASSERT_EQUAL(&v[0],       v.data());
+    ASSERT_EQUAL(PointerT(&v.front()),  v.data());
+    ASSERT_EQUAL(PointerT(&*v.begin()), v.data());
+    ASSERT_EQUAL(PointerT(&v[0]),       v.data());
 
     const Vector &c_v = v;
 
     ASSERT_EQUAL(0,            *c_v.data());
     ASSERT_EQUAL(1,            *(c_v.data() + 1));
     ASSERT_EQUAL(2,            *(c_v.data() + 2));
-    ASSERT_EQUAL(&c_v.front(),  c_v.data());
-    ASSERT_EQUAL(&*c_v.begin(), c_v.data());
-    ASSERT_EQUAL(&c_v[0],       c_v.data());
+    ASSERT_EQUAL(PointerConstT(&c_v.front()),  c_v.data());
+    ASSERT_EQUAL(PointerConstT(&*c_v.begin()), c_v.data());
+    ASSERT_EQUAL(PointerConstT(&c_v[0]),       c_v.data());
 }
 DECLARE_VECTOR_UNITTEST(TestVectorData);
 

thrust/detail/pointer.inl

@@ -16,6 +16,7 @@
 
 #include <thrust/detail/config.h>
 #include <thrust/detail/pointer.h>
+#include <thrust/detail/type_traits.h>
 
 
 namespace thrust
@@ -109,14 +110,43 @@ template<typename Element, typename Tag, typename Reference, typename Derived>
   return static_cast<derived_type&>(*this);
 } // end pointer::operator=
 
+namespace detail
+{
+
+// Implementation for dereference() when Reference is Element&,
+// e.g. cuda's managed_memory_pointer
+template <typename Reference, typename Derived>
+__host__ __device__
+Reference pointer_dereference_impl(const Derived& ptr,
+                                   thrust::detail::true_type /* is_cpp_ref */)
+{
+  return *ptr.get();
+}
+
+// Implementation for pointers with proxy references:
+template <typename Reference, typename Derived>
+__host__ __device__
+Reference pointer_dereference_impl(const Derived& ptr,
+                                   thrust::detail::false_type /* is_cpp_ref */)
+{
+  return Reference(ptr);
+}
+
+} // namespace detail
 
 template<typename Element, typename Tag, typename Reference, typename Derived>
   __host__ __device__
   typename pointer<Element,Tag,Reference,Derived>::super_t::reference
-    pointer<Element,Tag,Reference,Derived>
-      ::dereference() const
+  pointer<Element,Tag,Reference,Derived>
+    ::dereference() const
 {
-  return typename super_t::reference(static_cast<const derived_type&>(*this));
+  // Need to handle cpp refs and fancy refs differently:
+  typedef typename super_t::reference RefT;
+  typedef typename thrust::detail::is_reference<RefT>::type IsCppRef;
+
+  const derived_type& derivedPtr = static_cast<const derived_type&>(*this);
+
+  return detail::pointer_dereference_impl<RefT>(derivedPtr, IsCppRef());
 } // end pointer::dereference
 
 

thrust/detail/vector_base.inl

@@ -540,15 +540,15 @@ template<typename T, typename Alloc>
     vector_base<T,Alloc>
       ::data(void)
 {
-  return &front();
+  return pointer(&front());
 } // end vector_base::data()
 
 template<typename T, typename Alloc>
   typename vector_base<T,Alloc>::const_pointer
     vector_base<T,Alloc>
       ::data(void) const
 {
-  return &front();
+  return const_pointer(&front());
 } // end vector_base::data()
 
 template<typename T, typename Alloc>

thrust/mr/allocator.h

@@ -22,6 +22,7 @@
 
 #include <limits>
 
+#include <thrust/detail/config/exec_check_disable.h>
 #include <thrust/detail/type_traits/pointer_traits.h>
 
 #include <thrust/mr/detail/config.h>
@@ -93,6 +94,7 @@ class allocator : private validator<MR>
      *
      *  \returns the maximum value of \p std::size_t, divided by the size of \p T.
      */
+    __thrust_exec_check_disable__
     __host__ __device__
     size_type max_size() const
     {