Expose streams in public filling APIs (#13990)

Contributes to #925 Authors: - Vyas Ramasubramani (https://github.com/vyasr) Approvers: - Nghia Truong (https://github.com/ttnghia) - Bradley Dice (https://github.com/bdice) URL: #13990
rapidsai · Aug 29, 2023 · e2e92c4 · e2e92c4
1 parent cd56cc2
commit e2e92c4
Show file tree

Hide file tree

Showing 7 changed files with 109 additions and 10 deletions.
diff --git a/cpp/include/cudf/filling.hpp b/cpp/include/cudf/filling.hpp
@@ -17,6 +17,7 @@
 #pragma once
 
 #include <cudf/types.hpp>
+#include <cudf/utilities/default_stream.hpp>
 
 #include <rmm/mr/device/per_device_resource.hpp>
 
@@ -54,11 +55,13 @@ namespace cudf {
  * @param begin The starting index of the fill range (inclusive)
  * @param end The index of the last element in the fill range (exclusive)
  * @param value The scalar value to fill
+ * @param stream CUDA stream used for device memory operations and kernel launches
  */
 void fill_in_place(mutable_column_view& destination,
                    size_type begin,
                    size_type end,
-                   scalar const& value);
+                   scalar const& value,
+                   rmm::cuda_stream_view stream = cudf::get_default_stream());
 
 /**
  * @brief Fills a range of elements in a column out-of-place with a scalar
@@ -79,6 +82,7 @@ void fill_in_place(mutable_column_view& destination,
  * @param begin The starting index of the fill range (inclusive)
  * @param end The index of the last element in the fill range (exclusive)
  * @param value The scalar value to fill
+ * @param stream CUDA stream used for device memory operations and kernel launches
  * @param mr Device memory resource used to allocate the returned column's device memory
  * @return The result output column
  */
@@ -87,6 +91,7 @@ std::unique_ptr<column> fill(
   size_type begin,
   size_type end,
   scalar const& value,
+  rmm::cuda_stream_view stream        = cudf::get_default_stream(),
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
 /**
@@ -113,12 +118,14 @@ std::unique_ptr<column> fill(
  *
  * @param input_table Input table
  * @param count Non-nullable column of an integral type
+ * @param stream CUDA stream used for device memory operations and kernel launches
  * @param mr Device memory resource used to allocate the returned table's device memory
  * @return The result table containing the repetitions
  */
 std::unique_ptr<table> repeat(
   table_view const& input_table,
   column_view const& count,
+  rmm::cuda_stream_view stream        = cudf::get_default_stream(),
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
 /**
@@ -136,12 +143,14 @@ std::unique_ptr<table> repeat(
  *
  * @param input_table Input table
  * @param count Number of repetitions
+ * @param stream CUDA stream used for device memory operations and kernel launches
  * @param mr Device memory resource used to allocate the returned table's device memory
  * @return The result table containing the repetitions
  */
 std::unique_ptr<table> repeat(
   table_view const& input_table,
   size_type count,
+  rmm::cuda_stream_view stream        = cudf::get_default_stream(),
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
 /**
@@ -164,13 +173,15 @@ std::unique_ptr<table> repeat(
  * @param size Size of the output column
  * @param init First value in the sequence
  * @param step Increment value
+ * @param stream CUDA stream used for device memory operations and kernel launches
  * @param mr Device memory resource used to allocate the returned column's device memory
  * @return The result column containing the generated sequence
  */
 std::unique_ptr<column> sequence(
   size_type size,
   scalar const& init,
   scalar const& step,
+  rmm::cuda_stream_view stream        = cudf::get_default_stream(),
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
 /**
@@ -190,12 +201,14 @@ std::unique_ptr<column> sequence(
  *
  * @param size Size of the output column
  * @param init First value in the sequence
+ * @param stream CUDA stream used for device memory operations and kernel launches
  * @param mr Device memory resource used to allocate the returned column's device memory
  * @return The result column containing the generated sequence
  */
 std::unique_ptr<column> sequence(
   size_type size,
   scalar const& init,
+  rmm::cuda_stream_view stream        = cudf::get_default_stream(),
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
 /**
@@ -217,6 +230,7 @@ std::unique_ptr<column> sequence(
  * @param size Number of timestamps to generate
  * @param init The initial timestamp
  * @param months Months to increment
+ * @param stream CUDA stream used for device memory operations and kernel launches
  * @param mr Device memory resource used to allocate the returned column's device memory
  *
  * @return Timestamps column with sequences of months
@@ -225,6 +239,7 @@ std::unique_ptr<cudf::column> calendrical_month_sequence(
   size_type size,
   scalar const& init,
   size_type months,
+  rmm::cuda_stream_view stream        = cudf::get_default_stream(),
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
 /** @} */  // end of group

diff --git a/cpp/src/filling/calendrical_month_sequence.cu b/cpp/src/filling/calendrical_month_sequence.cu
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -40,10 +40,11 @@ std::unique_ptr<cudf::column> calendrical_month_sequence(size_type size,
 std::unique_ptr<cudf::column> calendrical_month_sequence(size_type size,
                                                          scalar const& init,
                                                          size_type months,
+                                                         rmm::cuda_stream_view stream,
                                                          rmm::mr::device_memory_resource* mr)
 {
   CUDF_FUNC_RANGE();
-  return detail::calendrical_month_sequence(size, init, months, cudf::get_default_stream(), mr);
+  return detail::calendrical_month_sequence(size, init, months, stream, mr);
 }
 
 }  // namespace cudf
diff --git a/cpp/src/filling/fill.cu b/cpp/src/filling/fill.cu
@@ -246,20 +246,22 @@ std::unique_ptr<column> fill(column_view const& input,
 void fill_in_place(mutable_column_view& destination,
                    size_type begin,
                    size_type end,
-                   scalar const& value)
+                   scalar const& value,
+                   rmm::cuda_stream_view stream)
 {
   CUDF_FUNC_RANGE();
-  return detail::fill_in_place(destination, begin, end, value, cudf::get_default_stream());
+  return detail::fill_in_place(destination, begin, end, value, stream);
 }
 
 std::unique_ptr<column> fill(column_view const& input,
                              size_type begin,
                              size_type end,
                              scalar const& value,
+                             rmm::cuda_stream_view stream,
                              rmm::mr::device_memory_resource* mr)
 {
   CUDF_FUNC_RANGE();
-  return detail::fill(input, begin, end, value, cudf::get_default_stream(), mr);
+  return detail::fill(input, begin, end, value, stream, mr);
 }
 
 }  // namespace cudf
diff --git a/cpp/src/filling/repeat.cu b/cpp/src/filling/repeat.cu
@@ -156,18 +156,20 @@ std::unique_ptr<table> repeat(table_view const& input_table,
 
 std::unique_ptr<table> repeat(table_view const& input_table,
                               column_view const& count,
+                              rmm::cuda_stream_view stream,
                               rmm::mr::device_memory_resource* mr)
 {
   CUDF_FUNC_RANGE();
-  return detail::repeat(input_table, count, cudf::get_default_stream(), mr);
+  return detail::repeat(input_table, count, stream, mr);
 }
 
 std::unique_ptr<table> repeat(table_view const& input_table,
                               size_type count,
+                              rmm::cuda_stream_view stream,
                               rmm::mr::device_memory_resource* mr)
 {
   CUDF_FUNC_RANGE();
-  return detail::repeat(input_table, count, cudf::get_default_stream(), mr);
+  return detail::repeat(input_table, count, stream, mr);
 }
 
 }  // namespace cudf
diff --git a/cpp/src/filling/sequence.cu b/cpp/src/filling/sequence.cu
@@ -150,18 +150,20 @@ std::unique_ptr<column> sequence(size_type size,
 std::unique_ptr<column> sequence(size_type size,
                                  scalar const& init,
                                  scalar const& step,
+                                 rmm::cuda_stream_view stream,
                                  rmm::mr::device_memory_resource* mr)
 {
   CUDF_FUNC_RANGE();
-  return detail::sequence(size, init, step, cudf::get_default_stream(), mr);
+  return detail::sequence(size, init, step, stream, mr);
 }
 
 std::unique_ptr<column> sequence(size_type size,
                                  scalar const& init,
+                                 rmm::cuda_stream_view stream,
                                  rmm::mr::device_memory_resource* mr)
 {
   CUDF_FUNC_RANGE();
-  return detail::sequence(size, init, cudf::get_default_stream(), mr);
+  return detail::sequence(size, init, stream, mr);
 }
 
 }  // namespace cudf
diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt
@@ -624,6 +624,7 @@ ConfigureTest(STREAM_HASHING_TEST streams/hash_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_COPYING_TEST streams/copying_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_GROUPBY_TEST streams/groupby_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_CONCATENATE_TEST streams/concatenate_test.cpp STREAM_MODE testing)
+ConfigureTest(STREAM_FILLING_TEST streams/filling_test.cpp STREAM_MODE testing)
 
 # ##################################################################################################
 # Install tests ####################################################################################

diff --git a/cpp/tests/streams/filling_test.cpp b/cpp/tests/streams/filling_test.cpp
@@ -0,0 +1,76 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <cudf/column/column_view.hpp>
+#include <cudf/filling.hpp>
+#include <cudf/scalar/scalar.hpp>
+
+#include <cudf_test/base_fixture.hpp>
+#include <cudf_test/column_wrapper.hpp>
+#include <cudf_test/default_stream.hpp>
+
+class FillingTest : public cudf::test::BaseFixture {};
+
+TEST_F(FillingTest, FillInPlace)
+{
+  cudf::test::fixed_width_column_wrapper<int> col({0, 0, 0, 0, 0});
+  auto scalar = cudf::numeric_scalar<int>(5, true, cudf::test::get_default_stream());
+  cudf::mutable_column_view mut_view = col;
+  cudf::fill_in_place(mut_view, 0, 4, scalar, cudf::test::get_default_stream());
+}
+
+TEST_F(FillingTest, Fill)
+{
+  cudf::test::fixed_width_column_wrapper<int> const col({0, 0, 0, 0, 0});
+  auto scalar = cudf::numeric_scalar<int>(5, true, cudf::test::get_default_stream());
+  cudf::fill(col, 0, 4, scalar, cudf::test::get_default_stream());
+}
+
+TEST_F(FillingTest, RepeatVariable)
+{
+  cudf::test::fixed_width_column_wrapper<int> const col({0, 0, 0, 0, 0});
+  cudf::table_view const table({col});
+  cudf::test::fixed_width_column_wrapper<int> const counts({1, 2, 3, 4, 5});
+  cudf::repeat(table, counts, cudf::test::get_default_stream());
+}
+
+TEST_F(FillingTest, RepeatConst)
+{
+  cudf::test::fixed_width_column_wrapper<int> const col({0, 0, 0, 0, 0});
+  cudf::table_view const table({col});
+  cudf::repeat(table, 5, cudf::test::get_default_stream());
+}
+
+TEST_F(FillingTest, SequenceStep)
+{
+  auto init = cudf::numeric_scalar<int>(5, true, cudf::test::get_default_stream());
+  auto step = cudf::numeric_scalar<int>(2, true, cudf::test::get_default_stream());
+  cudf::sequence(10, init, step, cudf::test::get_default_stream());
+}
+
+TEST_F(FillingTest, Sequence)
+{
+  auto init = cudf::numeric_scalar<int>(5, true, cudf::test::get_default_stream());
+  cudf::sequence(10, init, cudf::test::get_default_stream());
+}
+
+TEST_F(FillingTest, CalendricalMonthSequence)
+{
+  cudf::timestamp_scalar<cudf::timestamp_s> init(
+    1629852896L, true, cudf::test::get_default_stream());  // 2021-08-25 00:54:56 GMT
+
+  cudf::calendrical_month_sequence(10, init, 2, cudf::test::get_default_stream());
+}