NVRTC fixes

NVIDIA · Jan 22, 2025 · b53e8ae · b53e8ae
1 parent 979ab69
commit b53e8ae
Show file tree

Hide file tree

Showing 16 changed files with 223 additions and 143 deletions.
diff --git a/thrust/thrust/detail/alignment.h b/thrust/thrust/detail/alignment.h
@@ -30,10 +30,9 @@
 #endif // no system header
 
 #include <cuda/cmath>
+#include <cuda/std/cstddef>
 #include <cuda/std/type_traits>
 
-#include <cstddef> // For `std::size_t` and `std::max_align_t`.
-
 THRUST_NAMESPACE_BEGIN
 namespace detail
 {
@@ -49,7 +48,7 @@ using alignment_of = ::cuda::std::alignment_of<T>;
 /// type whose alignment requirement is a divisor of `Align`.
 ///
 /// The behavior is undefined if `Align` is not a power of 2.
-template <std::size_t Align>
+template <::cuda::std::size_t Align>
 struct aligned_type
 {
   struct alignas(Align) type
@@ -74,7 +73,7 @@ _CCCL_HOST_DEVICE T aligned_reinterpret_cast(U u)
   return reinterpret_cast<T>(reinterpret_cast<void*>(u));
 }
 
-_CCCL_HOST_DEVICE inline std::size_t aligned_storage_size(std::size_t n, std::size_t align)
+_CCCL_HOST_DEVICE inline ::cuda::std::size_t aligned_storage_size(::cuda::std::size_t n, ::cuda::std::size_t align)
 {
   return ::cuda::ceil_div(n, align) * align;
 }

diff --git a/thrust/thrust/detail/allocator_aware_execution_policy.h b/thrust/thrust/detail/allocator_aware_execution_policy.h
@@ -28,7 +28,7 @@
 #include <thrust/detail/alignment.h>
 #include <thrust/detail/execute_with_allocator_fwd.h>
 
-#include <type_traits>
+#include <cuda/std/type_traits>
 
 THRUST_NAMESPACE_BEGIN
 
@@ -61,30 +61,33 @@ struct allocator_aware_execution_policy
   };
 
   template <typename MemoryResource>
-  typename execute_with_memory_resource_type<MemoryResource>::type operator()(MemoryResource* mem_res) const
+  _CCCL_HOST_DEVICE typename execute_with_memory_resource_type<MemoryResource>::type
+  operator()(MemoryResource* mem_res) const
   {
     return typename execute_with_memory_resource_type<MemoryResource>::type(mem_res);
   }
 
   template <typename Allocator>
-  typename execute_with_allocator_type<Allocator&>::type operator()(Allocator& alloc) const
+  _CCCL_HOST_DEVICE typename execute_with_allocator_type<Allocator&>::type operator()(Allocator& alloc) const
   {
     return typename execute_with_allocator_type<Allocator&>::type(alloc);
   }
 
   template <typename Allocator>
-  typename execute_with_allocator_type<Allocator>::type operator()(const Allocator& alloc) const
+  _CCCL_HOST_DEVICE typename execute_with_allocator_type<Allocator>::type operator()(const Allocator& alloc) const
   {
     return typename execute_with_allocator_type<Allocator>::type(alloc);
   }
 
   // just the rvalue overload
   // perfect forwarding doesn't help, because a const reference has to be turned
   // into a value by copying for the purpose of storing it in execute_with_allocator
-  template <typename Allocator, typename std::enable_if<!std::is_lvalue_reference<Allocator>::value>::type* = nullptr>
-  typename execute_with_allocator_type<Allocator>::type operator()(Allocator&& alloc) const
+  _CCCL_EXEC_CHECK_DISABLE
+  template <typename Allocator,
+            typename ::cuda::std::enable_if<!::cuda::std::is_lvalue_reference<Allocator>::value>::type* = nullptr>
+  _CCCL_HOST_DEVICE typename execute_with_allocator_type<Allocator>::type operator()(Allocator&& alloc) const
   {
-    return typename execute_with_allocator_type<Allocator>::type(std::move(alloc));
+    return typename execute_with_allocator_type<Allocator>::type(::cuda::std::move(alloc));
   }
 };
 

diff --git a/thrust/thrust/detail/execute_with_allocator_fwd.h b/thrust/thrust/detail/execute_with_allocator_fwd.h
@@ -26,7 +26,9 @@
 #  pragma system_header
 #endif // no system header
 
-#include <thrust/detail/execute_with_dependencies.h>
+#if !_CCCL_COMPILER(NVRTC)
+#  include <thrust/detail/execute_with_dependencies.h>
+#endif // !_CCCL_COMPILER(NVRTC)
 #include <thrust/detail/type_traits.h>
 
 THRUST_NAMESPACE_BEGIN
@@ -53,11 +55,12 @@ struct execute_with_allocator : BaseSystem<execute_with_allocator<Allocator, Bas
       : alloc(alloc_)
   {}
 
-  ::cuda::std::remove_reference_t<Allocator>& get_allocator()
+  _CCCL_HOST_DEVICE ::cuda::std::remove_reference_t<Allocator>& get_allocator()
   {
     return alloc;
   }
 
+#if !_CCCL_COMPILER(NVRTC)
   template <typename... Dependencies>
   _CCCL_HOST execute_with_allocator_and_dependencies<Allocator, BaseSystem, Dependencies...>
   after(Dependencies&&... dependencies) const
@@ -97,6 +100,7 @@ struct execute_with_allocator : BaseSystem<execute_with_allocator<Allocator, Bas
   {
     return {alloc, capture_as_dependency(std::move(dependencies))};
   }
+#endif // !_CCCL_COMPILER(NVRTC)
 };
 
 } // namespace detail

diff --git a/thrust/thrust/detail/type_deduction.h b/thrust/thrust/detail/type_deduction.h
@@ -18,29 +18,29 @@
 #endif // no system header
 #include <thrust/detail/preprocessor.h>
 
-#include <type_traits>
-#include <utility>
+#include <cuda/std/type_traits>
+#include <cuda/std/utility>
 
 ///////////////////////////////////////////////////////////////////////////////
 
 /// \def THRUST_FWD(x)
 /// \brief Performs universal forwarding of a universal reference.
 ///
-#define THRUST_FWD(x) ::std::forward<decltype(x)>(x)
+#define THRUST_FWD(x) ::cuda::std::forward<decltype(x)>(x)
 
 /// \def THRUST_MVCAP(x)
 /// \brief Capture `x` into a lambda by moving.
 ///
-#define THRUST_MVCAP(x) x = ::std::move(x)
+#define THRUST_MVCAP(x) x = ::cuda::std::move(x)
 
 /// \def THRUST_RETOF(invocable, ...)
 /// \brief Expands to the type returned by invoking an instance of the invocable
 ///        type \a invocable with parameters of type \c __VA_ARGS__. Must
 ///        be called with 1 or fewer parameters to the invocable.
 ///
 #define THRUST_RETOF(...)   THRUST_PP_DISPATCH(THRUST_RETOF, __VA_ARGS__)
-#define THRUST_RETOF1(C)    decltype(::std::declval<C>()())
-#define THRUST_RETOF2(C, V) decltype(::std::declval<C>()(::std::declval<V>()))
+#define THRUST_RETOF1(C)    decltype(::cuda::std::declval<C>()())
+#define THRUST_RETOF2(C, V) decltype(::cuda::std::declval<C>()(::cuda::std::declval<V>()))
 
 /// \def THRUST_RETURNS(...)
 /// \brief Expands to a function definition that returns the expression
@@ -88,11 +88,11 @@
     }                                  \
     /**/
 #else
-#  define THRUST_DECLTYPE_RETURNS_WITH_SFINAE_CONDITION(condition, ...)                              \
-    noexcept(noexcept(__VA_ARGS__))->typename std::enable_if<condition, decltype(__VA_ARGS__)>::type \
-    {                                                                                                \
-      return (__VA_ARGS__);                                                                          \
-    }                                                                                                \
+#  define THRUST_DECLTYPE_RETURNS_WITH_SFINAE_CONDITION(condition, ...)                         \
+    noexcept(noexcept(__VA_ARGS__))->::cuda::std::enable_if_t<condition, decltype(__VA_ARGS__)> \
+    {                                                                                           \
+      return (__VA_ARGS__);                                                                     \
+    }                                                                                           \
     /**/
 #endif
 

diff --git a/thrust/thrust/iterator/counting_iterator.h b/thrust/thrust/iterator/counting_iterator.h
@@ -142,7 +142,7 @@ class _CCCL_DECLSPEC_EMPTY_BASES counting_iterator
    */
   using super_t = typename detail::counting_iterator_base<Incrementable, System, Traversal, Difference>::type;
 
-  friend class thrust::iterator_core_access;
+  friend class iterator_core_access;
 
 public:
   using reference       = typename super_t::reference;

diff --git a/thrust/thrust/iterator/detail/any_system_tag.h b/thrust/thrust/iterator/detail/any_system_tag.h
@@ -34,7 +34,7 @@ struct any_system_tag : thrust::execution_policy<any_system_tag>
   // allow any_system_tag to convert to any type at all
   // XXX make this safer using enable_if<is_tag<T>> upon c++11
   template <typename T>
-  operator T() const
+  _CCCL_HOST_DEVICE operator T() const
   {
     return T();
   }

diff --git a/thrust/thrust/iterator/detail/discard_iterator_base.h b/thrust/thrust/iterator/detail/discard_iterator_base.h
@@ -29,7 +29,7 @@
 #include <thrust/iterator/detail/any_assign.h>
 #include <thrust/iterator/iterator_adaptor.h>
 
-#include <cstddef> // for std::ptrdiff_t
+#include <cuda/std/cstddef>
 
 THRUST_NAMESPACE_BEGIN
 
@@ -47,7 +47,7 @@ struct discard_iterator_base
   //     but this interferes with zip_iterator<discard_iterator>
   using value_type    = any_assign;
   using reference     = any_assign&;
-  using incrementable = std::ptrdiff_t;
+  using incrementable = ::cuda::std::ptrdiff_t;
 
   using base_iterator = typename thrust::counting_iterator<incrementable, System, thrust::random_access_traversal_tag>;
 

diff --git a/thrust/thrust/iterator/detail/iterator_facade_category.h b/thrust/thrust/iterator/detail/iterator_facade_category.h
@@ -94,6 +94,7 @@ struct iterator_facade_default_category;
 // Instead, it simply assumes that if is_convertible<Traversal, single_pass_traversal_tag>,
 // then the category is input_iterator_tag
 
+#if !_CCCL_COMPILER(NVRTC)
 // this is the function for standard system iterators
 template <typename Traversal, typename ValueParam, typename Reference>
 struct iterator_facade_default_category_std
@@ -110,6 +111,7 @@ struct iterator_facade_default_category_std
           thrust::detail::identity_<std::input_iterator_tag>,
           thrust::detail::identity_<Traversal>>>
 {}; // end iterator_facade_default_category_std
+#endif // !_CCCL_COMPILER(NVRTC)
 
 // this is the function for host system iterators
 template <typename Traversal, typename ValueParam, typename Reference>
@@ -127,7 +129,6 @@ struct iterator_facade_default_category_host
           thrust::detail::identity_<thrust::input_host_iterator_tag>,
           thrust::detail::identity_<Traversal>>>
 {}; // end iterator_facade_default_category_host
-
 // this is the function for device system iterators
 template <typename Traversal, typename ValueParam, typename Reference>
 struct iterator_facade_default_category_device
@@ -146,6 +147,7 @@ struct iterator_facade_default_category_device
           thrust::detail::identity_<Traversal>>>
 {}; // end iterator_facade_default_category_device
 
+#if !_CCCL_COMPILER(NVRTC)
 // this is the function for any system iterators
 template <typename Traversal, typename ValueParam, typename Reference>
 struct iterator_facade_default_category_any
@@ -155,29 +157,41 @@ struct iterator_facade_default_category_any
     thrust::any_system_tag,
     Traversal>;
 }; // end iterator_facade_default_category_any
+#endif // !_CCCL_COMPILER(NVRTC)
 
 template <typename System, typename Traversal, typename ValueParam, typename Reference>
 struct iterator_facade_default_category
+    :
+#if !_CCCL_COMPILER(NVRTC)
     // check for any system
-    : thrust::detail::eval_if<
-        ::cuda::std::is_convertible<System, thrust::any_system_tag>::value,
-        iterator_facade_default_category_any<Traversal, ValueParam, Reference>,
-
-        // check for host system
-        thrust::detail::eval_if<
-          ::cuda::std::is_convertible<System, thrust::host_system_tag>::value,
-          iterator_facade_default_category_host<Traversal, ValueParam, Reference>,
-
-          // check for device system
-          thrust::detail::eval_if<::cuda::std::is_convertible<System, thrust::device_system_tag>::value,
-                                  iterator_facade_default_category_device<Traversal, ValueParam, Reference>,
-
-                                  // if we don't recognize the system, get a standard iterator category
-                                  // and combine it with System & Traversal
-                                  thrust::detail::identity_<thrust::detail::iterator_category_with_system_and_traversal<
-                                    typename iterator_facade_default_category_std<Traversal, ValueParam, Reference>::type,
-                                    System,
-                                    Traversal>>>>>
+    thrust::detail::eval_if<
+      ::cuda::std::is_convertible<System, thrust::any_system_tag>::value,
+      iterator_facade_default_category_any<Traversal, ValueParam, Reference>,
+#endif // !_CCCL_COMPILER(NVRTC)
+       // check for host system
+      thrust::detail::eval_if<
+        ::cuda::std::is_convertible<System, thrust::host_system_tag>::value,
+        iterator_facade_default_category_host<Traversal, ValueParam, Reference>,
+
+        // check for device system
+        thrust::detail::eval_if<::cuda::std::is_convertible<System, thrust::device_system_tag>::value,
+                                iterator_facade_default_category_device<Traversal, ValueParam, Reference>,
+
+                                // if we don't recognize the system, get a standard iterator category
+                                // and combine it with System & Traversal (for NVRTC fall back to host categories)
+                                thrust::detail::identity_<thrust::detail::iterator_category_with_system_and_traversal<
+                                  typename
+#if _CCCL_COMPILER(NVRTC)
+                                  iterator_facade_default_category_host
+#else // _CCCL_COMPILER(NVRTC)
+                                iterator_facade_default_category_std
+#endif // _CCCL_COMPILER(NVRTC)
+                                  <Traversal, ValueParam, Reference>::type,
+                                  System,
+                                  Traversal>>>>
+#if !_CCCL_COMPILER(NVRTC)
+      >
+#endif // !_CCCL_COMPILER(NVRTC)
 {};
 
 template <typename System, typename Traversal, typename ValueParam, typename Reference>

diff --git a/thrust/thrust/iterator/discard_iterator.h b/thrust/thrust/iterator/discard_iterator.h
@@ -102,7 +102,7 @@ class discard_iterator : public detail::discard_iterator_base<System>::type
 {
   /*! \cond
    */
-  friend class thrust::iterator_core_access;
+  friend class iterator_core_access;
   using super_t       = typename detail::discard_iterator_base<System>::type;
   using incrementable = typename detail::discard_iterator_base<System>::incrementable;
   using base_iterator = typename detail::discard_iterator_base<System>::base_iterator;

diff --git a/thrust/thrust/iterator/iterator_adaptor.h b/thrust/thrust/iterator/iterator_adaptor.h
@@ -55,6 +55,8 @@ THRUST_NAMESPACE_BEGIN
  *  \{
  */
 
+class iterator_core_access;
+
 /*! \p iterator_adaptor is an iterator which adapts an existing type of iterator to create a new type of
  *  iterator. Most of Thrust's fancy iterators are defined via inheritance from \p iterator_adaptor.
  *  While composition of these existing Thrust iterators is often sufficient for expressing the desired
@@ -132,7 +134,7 @@ class _CCCL_DECLSPEC_EMPTY_BASES iterator_adaptor
   /*! \cond
    */
 
-  friend class thrust::iterator_core_access;
+  friend class iterator_core_access;
 
 protected:
   using super_t =