Also modify indexalator and offsetalater (probably redundant, try rev…

…erting later)
rapidsai · Jul 2, 2024 · 1d3c645 · 1d3c645
1 parent b24c5eb
commit 1d3c645
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/cpp/include/cudf/detail/indexalator.cuh b/cpp/include/cudf/detail/indexalator.cuh
@@ -20,6 +20,7 @@
 #include <cudf/detail/iterator.cuh>
 #include <cudf/detail/normalizing_iterator.cuh>
 #include <cudf/scalar/scalar.hpp>
+#include <cudf/utilities/prefetch.hpp>
 #include <cudf/utilities/traits.hpp>
 
 #include <thrust/iterator/constant_iterator.h>
@@ -221,6 +222,9 @@ struct indexalator_factory {
     template <typename IndexType, CUDF_ENABLE_IF(is_index_type<IndexType>())>
     input_indexalator operator()(column_view const& indices)
     {
+      cudf::experimental::prefetch::detail::prefetch("indexalator_factor::input_indexalator_fn",
+                                                     indices.head(),
+                                                     indices.size() * size_of(indices.type()));
       return input_indexalator(indices.data<IndexType>(), indices.type());
     }
     template <typename IndexType, typename... Args, CUDF_ENABLE_IF(not is_index_type<IndexType>())>

diff --git a/cpp/include/cudf/detail/offsets_iterator_factory.cuh b/cpp/include/cudf/detail/offsets_iterator_factory.cuh
@@ -18,6 +18,7 @@
 
 #include <cudf/column/column_view.hpp>
 #include <cudf/detail/offsets_iterator.cuh>
+#include <cudf/utilities/prefetch.hpp>
 
 namespace cudf {
 namespace detail {
@@ -34,7 +35,16 @@ struct offsetalator_factory {
    */
   static input_offsetalator make_input_iterator(column_view const& offsets, size_type offset = 0)
   {
-    return input_offsetalator(offsets.head(), offsets.type(), offset);
+    if (offsets.type().id() == type_id::INT32) {
+      cudf::experimental::prefetch::detail::prefetch("offsetalator_factory::make_input_iterator",
+                                                     offsets.head(),
+                                                     offsets.size() * sizeof(int32_t));
+    } else if (offsets.type().id() == type_id::INT64) {
+      cudf::experimental::prefetch::detail::prefetch("offsetalator_factory::make_input_iterator",
+                                                     offsets.head(),
+                                                     offsets.size() * sizeof(int64_t));
+    }
+    return {offsets.head(), offsets.type(), offset};
   }
 
   /**