Fixes dereferencing zip_iterator over tuples of length-1 as well as copies from nested zip_iterators using OpenMP/TBB (#7882)

* Fixes dereferencing zip_iterator over tuples of length-1 as well as copies from nested zip_iterators using OpenMP/TBB

- implements is_compatible to verify that the tuple structure matches to fix conversions
- TestZipIteratorDereferenceToValue verifies that casting from the result of dereferencing to the corresponding value type works
- TestZipIteratorNestedCopy verifies that copying using nested zip_iterators works

Fixes #7855

* Simplify implementation and make it more efficient

* Bring back normalization for omp and tbb

---------

Co-authored-by: Michael Schellenberger Costa <miscco@nvidia.com>

Author: andrewcorrigan

thrust/testing/zip_iterator.cu

@@ -386,3 +386,84 @@ void TestZipIteratorCopySoAToAoS()
   ASSERT_EQUAL_QUIET(13, cuda::std::get<1>(h_soa[0]));
 };
 DECLARE_UNITTEST(TestZipIteratorCopySoAToAoS);
+
+template <typename T>
+void TestZipIteratorDereferenceToValueType(const T& t)
+{
+  thrust::device_vector<T> data(1, t);
+
+  // verify that storing the result of dereferencing a zip_iterator and then subsequently converting to its value type
+  // is handled correctly
+
+  auto a = thrust::make_zip_iterator(data.begin());
+  static_assert(cuda::std::is_same_v<cuda::std::tuple<T>, cuda::std::iter_value_t<decltype(a)>>);
+
+  auto b = a[0];
+  static_assert(
+    cuda::std::is_same_v<thrust::detail::tuple_of_iterator_references<thrust::device_reference<T>>, decltype(b)>);
+
+  // verify that the stored tuple_of_iterator_references<device_reference<T>> can be cast to tuple<T>
+  auto c = cuda::std::tuple<T>(b);
+  static_assert(cuda::std::is_same_v<cuda::std::tuple<T>, decltype(c)>);
+
+  ASSERT_EQUAL_QUIET(c, cuda::std::make_tuple(t));
+}
+
+void TestZipIteratorDereferenceToValue()
+{
+  TestZipIteratorDereferenceToValueType(1);
+  TestZipIteratorDereferenceToValueType(cuda::std::make_tuple(1));
+  TestZipIteratorDereferenceToValueType(cuda::std::make_tuple(1, cuda::std::make_tuple(1)));
+  TestZipIteratorDereferenceToValueType(cuda::std::make_tuple(1, cuda::std::make_tuple(1, 1)));
+  TestZipIteratorDereferenceToValueType(cuda::std::make_tuple(cuda::std::make_tuple(1), cuda::std::make_tuple(1, 1)));
+}
+DECLARE_UNITTEST(TestZipIteratorDereferenceToValue);
+
+void TestZipIteratorNestedCopy()
+{
+  using T = int;
+
+  {
+    thrust::device_vector<T> a(10, 1);
+    thrust::device_vector<cuda::std::tuple<cuda::std::tuple<T>>> b(a.size());
+
+    thrust::copy_n(thrust::make_zip_iterator(thrust::make_zip_iterator(a.begin())), a.size(), b.begin());
+
+    decltype(b) b_expected(b.size(), cuda::std::make_tuple(cuda::std::make_tuple(1)));
+
+    ASSERT_EQUAL_QUIET(b, b_expected);
+  }
+
+  {
+    thrust::device_vector<T> a(10, 1);
+    thrust::device_vector<cuda::std::tuple<cuda::std::tuple<T, T>, cuda::std::tuple<T, T>>> b(a.size());
+
+    thrust::copy_n(thrust::make_zip_iterator(thrust::make_zip_iterator(a.begin(), a.begin()),
+                                             thrust::make_zip_iterator(a.begin(), a.begin())),
+                   a.size(),
+                   b.begin());
+
+    decltype(b) b_expected(b.size(), cuda::std::make_tuple(cuda::std::make_tuple(1, 1), cuda::std::make_tuple(1, 1)));
+
+    ASSERT_EQUAL_QUIET(b, b_expected);
+  }
+
+  {
+    thrust::device_vector<cuda::std::tuple<T, T>> a(10, cuda::std::make_tuple(1, 1));
+    thrust::device_vector<
+      cuda::std::tuple<cuda::std::tuple<T, T>, cuda::std::tuple<T, T>, cuda::std::tuple<T, T>, cuda::std::tuple<T, T>>>
+      b(a.size());
+
+    thrust::copy_n(thrust::make_zip_iterator(a.begin(), a.begin(), a.begin(), a.begin()), a.size(), b.begin());
+
+    decltype(b) b_expected(
+      b.size(),
+      cuda::std::make_tuple(cuda::std::make_tuple(1, 1),
+                            cuda::std::make_tuple(1, 1),
+                            cuda::std::make_tuple(1, 1),
+                            cuda::std::make_tuple(1, 1)));
+
+    ASSERT_EQUAL_QUIET(b, b_expected);
+  }
+}
+DECLARE_UNITTEST(TestZipIteratorNestedCopy);

thrust/thrust/iterator/detail/tuple_of_iterator_references.h

@@ -28,7 +28,69 @@ namespace detail
 template <typename... Ts>
 class tuple_of_iterator_references;
 
-template <class U, class T>
+// is_compatible_tuple_normalize:
+//   device_reference<T> --> T
+//   tuple_of_iterator_references<Ts...> --> tuple<Ts...>
+//   T& --> T
+
+template <typename T>
+struct is_compatible_tuple_normalize
+{
+  using type = T;
+};
+
+template <typename... Ts>
+struct is_compatible_tuple_normalize<tuple_of_iterator_references<Ts...>>
+{
+  using type = ::cuda::std::tuple<Ts...>;
+};
+
+template <typename T>
+struct is_compatible_tuple_normalize<thrust::device_reference<T>>
+{
+  using type = T;
+};
+
+template <typename T>
+struct is_compatible_tuple_normalize<T&>
+{
+  using type = T;
+};
+
+template <typename T>
+using is_compatible_tuple_normalize_t = typename is_compatible_tuple_normalize<T>::type;
+
+// is_compatible_tuple_v:
+//  - checks if the tuple structure matches
+//  - rather than just testing the top-level size, this handles nesting with length-1 tuples,
+
+// is_compatible_tuple_v:
+//  - case of two non-tuple types are compatible
+//  - case of mixing tuples is not compatible
+template <typename U, typename T>
+inline constexpr bool is_compatible_tuple_v = ::cuda::std::__tuple_like<U> == ::cuda::std::__tuple_like<T>;
+
+// is_compatible_tuple_helper_v: verifies that the outer-most tuple_size matches prior to recursing further
+//  - case1: non-viable, sizes don't even match, do not recurse
+template <typename U, typename T, bool TupleSizeMatches>
+inline constexpr bool is_compatible_tuple_helper_v = false;
+
+// is_compatible_tuple_helper_v: viable, sizes match, recurse further but unwrap references
+template <template <class...> class Tuple1, template <class...> class Tuple2, typename... Ts, typename... Us>
+inline constexpr bool is_compatible_tuple_helper_v<Tuple1<Us...>, Tuple2<Ts...>, true> =
+  (is_compatible_tuple_v<is_compatible_tuple_normalize_t<Us>, is_compatible_tuple_normalize_t<Ts>> && ...);
+
+// is_compatible_tuple_v: recurse via is_compatible_tuple_helper_v to see if the two tuples are compatible
+template <template <class...> class Tuple1, template <class...> class Tuple2, typename... Ts, typename... Us>
+inline constexpr bool is_compatible_tuple_v<Tuple1<Us...>, Tuple2<Ts...>> =
+  is_compatible_tuple_helper_v<Tuple1<Us...>, Tuple2<Ts...>, sizeof...(Us) == sizeof...(Ts)>;
+
+// is_compatible_tuple_v: recurse via is_compatible_tuple_helper_v to see if the two tuples are compatible
+template <typename... Us, typename... Ts>
+inline constexpr bool is_compatible_tuple_v<::cuda::std::tuple<Us...>, ::cuda::std::tuple<Ts...>> =
+  is_compatible_tuple_helper_v<::cuda::std::tuple<Us...>, ::cuda::std::tuple<Ts...>, sizeof...(Us) == sizeof...(Ts)>;
+
+template <class U, class T, class Enable = void>
 struct maybe_unwrap_nested
 {
   _CCCL_HOST_DEVICE U operator()(const T& t) const
@@ -38,7 +100,10 @@ struct maybe_unwrap_nested
 };
 
 template <class... Us, class... Ts>
-struct maybe_unwrap_nested<::cuda::std::tuple<Us...>, tuple_of_iterator_references<Ts...>>
+struct maybe_unwrap_nested<
+  ::cuda::std::tuple<Us...>,
+  tuple_of_iterator_references<Ts...>,
+  ::cuda::std::enable_if_t<is_compatible_tuple_v<::cuda::std::tuple<Us...>, ::cuda::std::tuple<Ts...>>, int>>
 {
   _CCCL_HOST_DEVICE ::cuda::std::tuple<Us...> operator()(const tuple_of_iterator_references<Ts...>& t) const
   {
@@ -98,7 +163,9 @@ class tuple_of_iterator_references : public ::cuda::std::tuple<Ts...>
     return *this;
   }
 
-  template <class... Us, ::cuda::std::enable_if_t<sizeof...(Us) == sizeof...(Ts), int> = 0>
+  template <
+    class... Us,
+    ::cuda::std::enable_if_t<is_compatible_tuple_v<::cuda::std::tuple<Us...>, ::cuda::std::tuple<Ts...>>, int> = 0>
   _CCCL_HOST_DEVICE constexpr operator ::cuda::std::tuple<Us...>() const
   {
     return __to_tuple<Us...>(typename ::cuda::std::__make_tuple_indices<sizeof...(Ts)>::type{});
@@ -112,7 +179,10 @@ class tuple_of_iterator_references : public ::cuda::std::tuple<Ts...>
     x.swap(y);
   }
 
-  template <class... Us, size_t... Id>
+  template <
+    class... Us,
+    size_t... Id,
+    ::cuda::std::enable_if_t<is_compatible_tuple_v<::cuda::std::tuple<Us...>, ::cuda::std::tuple<Ts...>>, int> = 0>
   _CCCL_HOST_DEVICE constexpr ::cuda::std::tuple<Us...> __to_tuple(::cuda::std::__tuple_indices<Id...>) const
   {
     return {maybe_unwrap_nested<Us, Ts>{}(::cuda::std::get<Id>(*this))...};