Expose streams in public concatenate APIs (#13987)

Contributes to #925

Authors:
  - Vyas Ramasubramani (https://github.com/vyasr)

Approvers:
  - Mark Harris (https://github.com/harrism)
  - Divye Gala (https://github.com/divyegala)

URL: #13987

cpp/include/cudf/concatenate.hpp

@@ -17,6 +17,7 @@
 
 #include <cudf/column/column_view.hpp>
 #include <cudf/table/table_view.hpp>
+#include <cudf/utilities/default_stream.hpp>
 #include <cudf/utilities/span.hpp>
 
 #include <rmm/mr/device/per_device_resource.hpp>
@@ -40,10 +41,12 @@ namespace cudf {
  *
  * @param views Column views whose bitmasks will be concatenated
  * @param mr Device memory resource used for allocating the returned memory
+ * @param stream CUDA stream used for device memory operations and kernel launches
  * @return Bitmasks of all the column views in the views vector
  */
 rmm::device_buffer concatenate_masks(
   host_span<column_view const> views,
+  rmm::cuda_stream_view stream        = cudf::get_default_stream(),
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
 /**
@@ -53,12 +56,14 @@ rmm::device_buffer concatenate_masks(
  * @throws std::overflow_error If the total number of output rows exceeds cudf::size_type
  *
  * @param columns_to_concat Column views to be concatenated into a single column
+ * @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 A single column having all the rows from the elements of `columns_to_concat` respectively
  * in the same order.
  */
 std::unique_ptr<column> concatenate(
   host_span<column_view const> columns_to_concat,
+  rmm::cuda_stream_view stream        = cudf::get_default_stream(),
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
 /**
@@ -80,12 +85,14 @@ std::unique_ptr<column> concatenate(
  * @throws std::overflow_error If the total number of output rows exceeds cudf::size_type
  *
  * @param tables_to_concat Table views to be concatenated into a single table
+ * @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 A single table having all the rows from the elements of
  * `tables_to_concat` respectively in the same order.
  */
 std::unique_ptr<table> concatenate(
   host_span<table_view const> tables_to_concat,
+  rmm::cuda_stream_view stream        = cudf::get_default_stream(),
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
 /** @} */  // end of group

cpp/include/cudf_test/column_wrapper.hpp

@@ -1596,12 +1596,10 @@ class lists_column_wrapper : public detail::column_wrapper {
     thrust::copy_if(
       std::cbegin(cols), std::cend(cols), valids, std::back_inserter(children), thrust::identity{});
 
-    // TODO: Once the public concatenate API exposes streams, use that instead.
-    auto data =
-      children.empty()
-        ? cudf::empty_like(expected_hierarchy)
-        : cudf::detail::concatenate(
-            children, cudf::test::get_default_stream(), rmm::mr::get_current_device_resource());
+    auto data = children.empty() ? cudf::empty_like(expected_hierarchy)
+                                 : cudf::concatenate(children,
+                                                     cudf::test::get_default_stream(),
+                                                     rmm::mr::get_current_device_resource());
 
     // increment depth
     depth = expected_depth + 1;

cpp/src/copying/concatenate.cu

@@ -574,25 +574,28 @@ rmm::device_buffer concatenate_masks(host_span<column_view const> views,
 }  // namespace detail
 
 rmm::device_buffer concatenate_masks(host_span<column_view const> views,
+                                     rmm::cuda_stream_view stream,
                                      rmm::mr::device_memory_resource* mr)
 {
   CUDF_FUNC_RANGE();
-  return detail::concatenate_masks(views, cudf::get_default_stream(), mr);
+  return detail::concatenate_masks(views, stream, mr);
 }
 
 // Concatenates the elements from a vector of column_views
 std::unique_ptr<column> concatenate(host_span<column_view const> columns_to_concat,
+                                    rmm::cuda_stream_view stream,
                                     rmm::mr::device_memory_resource* mr)
 {
   CUDF_FUNC_RANGE();
-  return detail::concatenate(columns_to_concat, cudf::get_default_stream(), mr);
+  return detail::concatenate(columns_to_concat, stream, mr);
 }
 
 std::unique_ptr<table> concatenate(host_span<table_view const> tables_to_concat,
+                                   rmm::cuda_stream_view stream,
                                    rmm::mr::device_memory_resource* mr)
 {
   CUDF_FUNC_RANGE();
-  return detail::concatenate(tables_to_concat, cudf::get_default_stream(), mr);
+  return detail::concatenate(tables_to_concat, stream, mr);
 }
 
 }  // namespace cudf

cpp/tests/CMakeLists.txt

@@ -623,6 +623,7 @@ ConfigureTest(
 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)
 
 # ##################################################################################################
 # Install tests ####################################################################################

cpp/tests/streams/concatenate_test.cpp

@@ -0,0 +1,51 @@
+/*
+ * 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/concatenate.hpp>
+
+#include <cudf_test/base_fixture.hpp>
+#include <cudf_test/column_wrapper.hpp>
+#include <cudf_test/default_stream.hpp>
+
+class ConcatenateTest : public cudf::test::BaseFixture {};
+
+TEST_F(ConcatenateTest, Column)
+{
+  cudf::test::fixed_width_column_wrapper<int> const input1({0, 0, 0, 0, 0});
+  cudf::test::fixed_width_column_wrapper<int> const input2({1, 1, 1, 1, 1});
+  std::vector<cudf::column_view> views{input1, input2};
+  auto result = cudf::concatenate(views, cudf::test::get_default_stream());
+}
+
+TEST_F(ConcatenateTest, Table)
+{
+  cudf::test::fixed_width_column_wrapper<int> const input1({0, 0, 0, 0, 0});
+  cudf::test::fixed_width_column_wrapper<int> const input2({1, 1, 1, 1, 1});
+  cudf::table_view tbl1({input1, input2});
+  cudf::table_view tbl2({input2, input1});
+  std::vector<cudf::table_view> views{tbl1, tbl2};
+  auto result = cudf::concatenate(views, cudf::test::get_default_stream());
+}
+
+TEST_F(ConcatenateTest, Masks)
+{
+  cudf::test::fixed_width_column_wrapper<int> const input1(
+    {{0, 0, 0, 0, 0}, {false, false, false, false, false}});
+  cudf::test::fixed_width_column_wrapper<int> const input2(
+    {{0, 0, 0, 0, 0}, {true, true, true, true, true}});
+  std::vector<cudf::column_view> views{input1, input2};
+  auto result = cudf::concatenate_masks(views, cudf::test::get_default_stream());
+}