Use cudf::detail::make_counting_transform_iterator (#7338)

Came across some "raw" uses of `thrust::make_transform_iterator` + `thrust::make_counting_iterator` where now `cudf::detail::make_counting_transform_iterator` can be used (because of #6546). Authors: - Conor Hoekstra (@codereport) Approvers: - Devavret Makkar (@devavret) - David (@davidwendt) - Mark Harris (@harrism) URL: #7338
rapidsai · Feb 9, 2021 · c0282e6 · c0282e6
1 parent 9e8bdfc
commit c0282e6
Show file tree

Hide file tree

Showing 21 changed files with 131 additions and 134 deletions.
diff --git a/cpp/include/cudf/detail/iterator.cuh b/cpp/include/cudf/detail/iterator.cuh
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -43,6 +43,31 @@
 
 namespace cudf {
 namespace detail {
+/**
+ * @brief Convenience wrapper for creating a `thrust::transform_iterator` over a
+ * `thrust::counting_iterator`.
+ *
+ * Example:
+ * @code{.cpp}
+ * // Returns square of the value of the counting iterator
+ * auto iter = make_counting_transform_iterator(0, [](auto i){ return (i * i);});
+ * iter[0] == 0
+ * iter[1] == 1
+ * iter[2] == 4
+ * ...
+ * iter[n] == n * n
+ * @endcode
+ *
+ * @param start The starting value of the counting iterator
+ * @param f The unary function to apply to the counting iterator.
+ * @return A transform iterator that applies `f` to a counting iterator
+ */
+template <typename UnaryFunction>
+inline auto make_counting_transform_iterator(cudf::size_type start, UnaryFunction f)
+{
+  return thrust::make_transform_iterator(thrust::make_counting_iterator(start), f);
+}
+
 /**
  * @brief value accessor of column with null bitmask
  * A unary functor returns scalar value at `id`.
@@ -128,9 +153,8 @@ template <typename Element>
 auto make_null_replacement_iterator(column_device_view const& column,
                                     Element const null_replacement = Element{0})
 {
-  return thrust::make_transform_iterator(
-    thrust::counting_iterator<cudf::size_type>{0},
-    null_replaced_value_accessor<Element>{column, null_replacement});
+  return make_counting_transform_iterator(
+    0, null_replaced_value_accessor<Element>{column, null_replacement});
 }
 
 /**
@@ -174,8 +198,7 @@ auto make_pair_iterator(column_device_view const& column)
  */
 auto inline make_validity_iterator(column_device_view const& column)
 {
-  return thrust::make_transform_iterator(thrust::counting_iterator<cudf::size_type>{0},
-                                         validity_accessor{column});
+  return make_counting_transform_iterator(cudf::size_type{0}, validity_accessor{column});
 }
 
 /**
@@ -320,31 +343,5 @@ auto inline make_pair_iterator(scalar const& scalar_value)
                                          scalar_pair_accessor<Element>{scalar_value});
 }
 
-/**
- * @brief Convenience wrapper for creating a `thrust::transform_iterator` over a
- * `thrust::counting_iterator`.
- *
- * Example:
- * @code{.cpp}
- * // Returns square of the value of the counting iterator
- * auto iter = make_counting_transform_iterator(0, [](auto i){ return (i * i);});
- * iter[0] == 0
- * iter[1] == 1
- * iter[2] == 4
- * ...
- * iter[n] == n * n
- * @endcode
- *
- * @param start The starting value of the counting iterator
- * @param f The unary function to apply to the counting iterator.
- * This should be a host function and not a device function.
- * @return auto A transform iterator that applies `f` to a counting iterator
- */
-template <typename UnaryFunction>
-inline auto make_counting_transform_iterator(cudf::size_type start, UnaryFunction f)
-{
-  return thrust::make_transform_iterator(thrust::make_counting_iterator(start), f);
-}
-
 }  // namespace detail
 }  // namespace cudf
diff --git a/cpp/include/cudf/dictionary/detail/iterator.cuh b/cpp/include/cudf/dictionary/detail/iterator.cuh
@@ -15,6 +15,7 @@
  */
 
 #include <cudf/column/column_device_view.cuh>
+#include <cudf/detail/iterator.cuh>
 #include <cudf/dictionary/dictionary_column_view.hpp>
 
 namespace cudf {
@@ -58,8 +59,8 @@ auto make_dictionary_iterator(column_device_view const& dictionary_column)
 {
   CUDF_EXPECTS(is_dictionary(dictionary_column.type()),
                "Dictionary iterator is only for dictionary columns");
-  return thrust::make_transform_iterator(thrust::make_counting_iterator<size_type>(0),
-                                         dictionary_access_fn<KeyType>{dictionary_column});
+  return cudf::detail::make_counting_transform_iterator(
+    size_type{0}, dictionary_access_fn<KeyType>{dictionary_column});
 }
 
 /**
@@ -112,9 +113,8 @@ auto make_dictionary_pair_iterator(column_device_view const& dictionary_column,
 {
   CUDF_EXPECTS(is_dictionary(dictionary_column.type()),
                "Dictionary iterator is only for dictionary columns");
-  return thrust::make_transform_iterator(
-    thrust::make_counting_iterator<size_type>(0),
-    dictionary_access_pair_fn<KeyType>{dictionary_column, has_nulls});
+  return cudf::detail::make_counting_transform_iterator(
+    0, dictionary_access_pair_fn<KeyType>{dictionary_column, has_nulls});
 }
 
 }  // namespace detail

diff --git a/cpp/include/cudf/strings/detail/modify_strings.cuh b/cpp/include/cudf/strings/detail/modify_strings.cuh
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,10 +16,12 @@
 #pragma once
 
 #include <cudf/column/column_device_view.cuh>
+#include <cudf/detail/iterator.cuh>
 #include <cudf/detail/null_mask.hpp>
 #include <cudf/strings/detail/utilities.hpp>
 #include <cudf/strings/string_view.cuh>
 #include <cudf/strings/strings_column_view.hpp>
+
 #include <strings/utilities.cuh>
 
 #include <rmm/cuda_stream_view.hpp>
@@ -71,9 +73,8 @@ std::unique_ptr<column> modify_strings(strings_column_view const& strings,
   device_probe_functor d_probe_fctr{d_column, std::forward<Types>(args)...};
 
   // build offsets column -- calculate the size of each output string
-  auto offsets_transformer_itr =
-    thrust::make_transform_iterator(thrust::make_counting_iterator<size_type>(0), d_probe_fctr);
-  auto offsets_column = detail::make_offsets_child_column(
+  auto offsets_transformer_itr = cudf::detail::make_counting_transform_iterator(0, d_probe_fctr);
+  auto offsets_column          = detail::make_offsets_child_column(
     offsets_transformer_itr, offsets_transformer_itr + strings_count, stream, mr);
   auto offsets_view = offsets_column->view();
   auto d_new_offsets =

diff --git a/cpp/src/copying/contiguous_split.cu b/cpp/src/copying/contiguous_split.cu
@@ -18,6 +18,7 @@
 #include <cudf/column/column_view.hpp>
 #include <cudf/copying.hpp>
 #include <cudf/detail/copy.hpp>
+#include <cudf/detail/iterator.cuh>
 #include <cudf/detail/null_mask.hpp>
 #include <cudf/detail/nvtx/ranges.hpp>
 #include <cudf/detail/utilities/cuda.cuh>
@@ -943,10 +944,10 @@ std::vector<packed_table> contiguous_split(cudf::table_view const& input,
   // compute total size of each partition
   {
     // key is split index
-    auto keys   = thrust::make_transform_iterator(thrust::make_counting_iterator(0),
-                                                split_key_functor{static_cast<int>(num_src_bufs)});
-    auto values = thrust::make_transform_iterator(thrust::make_counting_iterator(0),
-                                                  buf_size_functor{d_dst_buf_info});
+    auto keys = cudf::detail::make_counting_transform_iterator(
+      0, split_key_functor{static_cast<int>(num_src_bufs)});
+    auto values =
+      cudf::detail::make_counting_transform_iterator(0, buf_size_functor{d_dst_buf_info});
 
     thrust::reduce_by_key(rmm::exec_policy(stream),
                           keys,
@@ -958,10 +959,11 @@ std::vector<packed_table> contiguous_split(cudf::table_view const& input,
 
   // compute start offset for each output buffer
   {
-    auto keys   = thrust::make_transform_iterator(thrust::make_counting_iterator(0),
-                                                split_key_functor{static_cast<int>(num_src_bufs)});
-    auto values = thrust::make_transform_iterator(thrust::make_counting_iterator(0),
-                                                  buf_size_functor{d_dst_buf_info});
+    auto keys = cudf::detail::make_counting_transform_iterator(
+      0, split_key_functor{static_cast<int>(num_src_bufs)});
+    auto values =
+      cudf::detail::make_counting_transform_iterator(0, buf_size_functor{d_dst_buf_info});
+
     thrust::exclusive_scan_by_key(rmm::exec_policy(stream),
                                   keys,
                                   keys + num_bufs,

diff --git a/cpp/src/copying/sample.cu b/cpp/src/copying/sample.cu
@@ -18,6 +18,7 @@
 #include <cudf/column/column_factories.hpp>
 #include <cudf/detail/gather.cuh>
 #include <cudf/detail/gather.hpp>
+#include <cudf/detail/iterator.cuh>
 #include <cudf/detail/nvtx/ranges.hpp>
 #include <cudf/table/table.hpp>
 #include <cudf/table/table_view.hpp>
@@ -56,11 +57,9 @@ std::unique_ptr<table> sample(table_view const& input,
       return dist(rng);
     };
 
-    auto begin =
-      thrust::make_transform_iterator(thrust::counting_iterator<size_type>(0), RandomGen);
-    auto end = thrust::make_transform_iterator(thrust::counting_iterator<size_type>(n), RandomGen);
+    auto begin = cudf::detail::make_counting_transform_iterator(0, RandomGen);
 
-    return detail::gather(input, begin, end, out_of_bounds_policy::DONT_CHECK, stream, mr);
+    return detail::gather(input, begin, begin + n, out_of_bounds_policy::DONT_CHECK, stream, mr);
   } else {
     auto gather_map =
       make_numeric_column(data_type{type_id::INT32}, num_rows, mask_state::UNALLOCATED, stream);

diff --git a/cpp/src/dictionary/detail/concatenate.cu b/cpp/src/dictionary/detail/concatenate.cu
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,6 +17,7 @@
 #include <cudf/column/column_factories.hpp>
 #include <cudf/detail/concatenate.cuh>
 #include <cudf/detail/indexalator.cuh>
+#include <cudf/detail/iterator.cuh>
 #include <cudf/detail/stream_compaction.hpp>
 #include <cudf/dictionary/detail/concatenate.hpp>
 #include <cudf/dictionary/dictionary_column_view.hpp>
@@ -231,10 +232,10 @@ std::unique_ptr<column> concatenate(std::vector<column_view> const& columns,
   // build a vector of values to map the old indices to the concatenated keys
   auto children_offsets = child_offsets_fn.create_children_offsets(stream);
   rmm::device_uvector<size_type> map_to_keys(indices_size, stream);
-  auto indices_itr = thrust::make_transform_iterator(thrust::make_counting_iterator<size_type>(1),
-                                                     [] __device__(size_type idx) {
-                                                       return offsets_pair{0, idx};
-                                                     });
+  auto indices_itr =
+    cudf::detail::make_counting_transform_iterator(1, [] __device__(size_type idx) {
+      return offsets_pair{0, idx};
+    });
   // the indices offsets (pair.second) are for building the map
   thrust::lower_bound(
     rmm::exec_policy(stream),

diff --git a/cpp/src/dictionary/replace.cu b/cpp/src/dictionary/replace.cu
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,6 +17,7 @@
 #include <cudf/detail/copy.hpp>
 #include <cudf/detail/copy_if_else.cuh>
 #include <cudf/detail/indexalator.cuh>
+#include <cudf/detail/iterator.cuh>
 #include <cudf/detail/unary.hpp>
 #include <cudf/dictionary/detail/encode.hpp>
 #include <cudf/dictionary/detail/replace.hpp>
@@ -74,8 +75,7 @@ struct nullable_index_accessor {
 template <bool has_nulls>
 auto make_nullable_index_iterator(column_view const& col)
 {
-  return thrust::make_transform_iterator(thrust::make_counting_iterator<size_type>(0),
-                                         nullable_index_accessor<has_nulls>{col});
+  return cudf::detail::make_counting_transform_iterator(0, nullable_index_accessor<has_nulls>{col});
 }
 
 /**

diff --git a/cpp/src/io/parquet/page_enc.cu b/cpp/src/io/parquet/page_enc.cu
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,6 +16,7 @@
 #include <io/parquet/parquet_gpu.hpp>
 #include <io/utilities/block_utils.cuh>
 
+#include <cudf/detail/iterator.cuh>
 #include <cudf/detail/utilities/cuda.cuh>
 
 #include <rmm/cuda_stream_view.hpp>
@@ -1839,10 +1840,10 @@ dremel_data get_dremel_data(column_view h_col,
     curr_rep_values_size = ends.second - output_zip_it;
 
     // Scan to get distance by which each offset value is shifted due to the insertion of empties
-    auto scan_it =
-      thrust::make_transform_iterator(thrust::make_counting_iterator(column_offsets[level]),
-                                      [off = lcv.offsets().data<size_type>()] __device__(
-                                        auto i) -> int { return off[i] == off[i + 1]; });
+    auto scan_it = cudf::detail::make_counting_transform_iterator(
+      column_offsets[level], [off = lcv.offsets().data<size_type>()] __device__(auto i) -> int {
+        return off[i] == off[i + 1];
+      });
     rmm::device_uvector<size_type> scan_out(offset_size_at_level, stream);
     thrust::exclusive_scan(
       rmm::exec_policy(stream), scan_it, scan_it + offset_size_at_level, scan_out.begin());
@@ -1927,10 +1928,10 @@ dremel_data get_dremel_data(column_view h_col,
 
     // Scan to get distance by which each offset value is shifted due to the insertion of dremel
     // level value fof an empty list
-    auto scan_it =
-      thrust::make_transform_iterator(thrust::make_counting_iterator(column_offsets[level]),
-                                      [off = lcv.offsets().data<size_type>()] __device__(
-                                        auto i) -> int { return off[i] == off[i + 1]; });
+    auto scan_it = cudf::detail::make_counting_transform_iterator(
+      column_offsets[level], [off = lcv.offsets().data<size_type>()] __device__(auto i) -> int {
+        return off[i] == off[i + 1];
+      });
     rmm::device_uvector<size_type> scan_out(offset_size_at_level, stream);
     thrust::exclusive_scan(
       rmm::exec_policy(stream), scan_it, scan_it + offset_size_at_level, scan_out.begin());

diff --git a/cpp/src/lists/copying/segmented_gather.cu b/cpp/src/lists/copying/segmented_gather.cu
@@ -16,6 +16,7 @@
 #include <cudf/detail/gather.cuh>
 #include <cudf/detail/gather.hpp>
 #include <cudf/detail/indexalator.cuh>
+#include <cudf/detail/iterator.cuh>
 #include <cudf/detail/null_mask.hpp>
 #include <cudf/lists/detail/gather.cuh>
 
@@ -59,14 +60,12 @@ std::unique_ptr<column> segmented_gather(lists_column_view const& value_column,
     // Add sub_index to value_column offsets, to get gather indices of child of value_column
     return value_offsets[offset_idx] + wrapped_sub_index - value_offsets[0];
   };
-  auto child_gather_index_begin =
-    thrust::make_transform_iterator(thrust::make_counting_iterator<size_type>(0), transformer);
-  auto child_gather_index_end = child_gather_index_begin + gather_map_size;
+  auto child_gather_index_begin = cudf::detail::make_counting_transform_iterator(0, transformer);
 
   // Call gather on child of value_column
   auto child_table = cudf::detail::gather(table_view({value_column.get_sliced_child(stream)}),
                                           child_gather_index_begin,
-                                          child_gather_index_end,
+                                          child_gather_index_begin + gather_map_size,
                                           out_of_bounds_policy::DONT_CHECK,
                                           stream,
                                           mr);

diff --git a/cpp/src/reshape/tile.cu b/cpp/src/reshape/tile.cu
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,6 +16,7 @@
 
 #include <cudf/copying.hpp>
 #include <cudf/detail/gather.cuh>
+#include <cudf/detail/iterator.cuh>
 #include <cudf/detail/nvtx/ranges.hpp>
 #include <cudf/detail/reshape.hpp>
 #include <cudf/table/table.hpp>
@@ -46,13 +47,12 @@ std::unique_ptr<table> tile(const table_view &in,
 {
   CUDF_EXPECTS(count >= 0, "Count cannot be negative");
 
-  auto in_num_rows = in.num_rows();
+  auto const in_num_rows = in.num_rows();
 
   if (count == 0 or in_num_rows == 0) { return empty_like(in); }
 
   auto out_num_rows = in_num_rows * count;
-  auto counting_it  = thrust::make_counting_iterator<size_type>(0);
-  auto tiled_it     = thrust::make_transform_iterator(counting_it, tile_functor{in_num_rows});
+  auto tiled_it     = cudf::detail::make_counting_transform_iterator(0, tile_functor{in_num_rows});
 
   return detail::gather(
     in, tiled_it, tiled_it + out_num_rows, out_of_bounds_policy::DONT_CHECK, stream, mr);