From 70fbec809a45fb4d462d7f3ef22464d00d2640e0 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Mon, 28 Aug 2023 20:32:38 -0500 Subject: [PATCH] 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: https://github.com/rapidsai/cudf/pull/13987 --- cpp/include/cudf/concatenate.hpp | 7 ++++ cpp/include/cudf_test/column_wrapper.hpp | 10 ++--- cpp/src/copying/concatenate.cu | 9 +++-- cpp/tests/CMakeLists.txt | 1 + cpp/tests/streams/concatenate_test.cpp | 51 ++++++++++++++++++++++++ 5 files changed, 69 insertions(+), 9 deletions(-) create mode 100644 cpp/tests/streams/concatenate_test.cpp diff --git a/cpp/include/cudf/concatenate.hpp b/cpp/include/cudf/concatenate.hpp index 11c6a02c225..9ee55275a5e 100644 --- a/cpp/include/cudf/concatenate.hpp +++ b/cpp/include/cudf/concatenate.hpp @@ -17,6 +17,7 @@ #include #include +#include #include #include @@ -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 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 concatenate( host_span 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 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 concatenate( host_span 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 diff --git a/cpp/include/cudf_test/column_wrapper.hpp b/cpp/include/cudf_test/column_wrapper.hpp index 1e311322de1..cc8cac35ef4 100644 --- a/cpp/include/cudf_test/column_wrapper.hpp +++ b/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; diff --git a/cpp/src/copying/concatenate.cu b/cpp/src/copying/concatenate.cu index a53ec295512..35f06e47436 100644 --- a/cpp/src/copying/concatenate.cu +++ b/cpp/src/copying/concatenate.cu @@ -574,25 +574,28 @@ rmm::device_buffer concatenate_masks(host_span views, } // namespace detail rmm::device_buffer concatenate_masks(host_span 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 concatenate(host_span 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
concatenate(host_span 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 diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index d9e34c739ea..c97e2a58ca4 100644 --- a/cpp/tests/CMakeLists.txt +++ b/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 #################################################################################### diff --git a/cpp/tests/streams/concatenate_test.cpp b/cpp/tests/streams/concatenate_test.cpp new file mode 100644 index 00000000000..6e6ff58686f --- /dev/null +++ b/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 + +#include +#include +#include + +class ConcatenateTest : public cudf::test::BaseFixture {}; + +TEST_F(ConcatenateTest, Column) +{ + cudf::test::fixed_width_column_wrapper const input1({0, 0, 0, 0, 0}); + cudf::test::fixed_width_column_wrapper const input2({1, 1, 1, 1, 1}); + std::vector views{input1, input2}; + auto result = cudf::concatenate(views, cudf::test::get_default_stream()); +} + +TEST_F(ConcatenateTest, Table) +{ + cudf::test::fixed_width_column_wrapper const input1({0, 0, 0, 0, 0}); + cudf::test::fixed_width_column_wrapper const input2({1, 1, 1, 1, 1}); + cudf::table_view tbl1({input1, input2}); + cudf::table_view tbl2({input2, input1}); + std::vector views{tbl1, tbl2}; + auto result = cudf::concatenate(views, cudf::test::get_default_stream()); +} + +TEST_F(ConcatenateTest, Masks) +{ + cudf::test::fixed_width_column_wrapper const input1( + {{0, 0, 0, 0, 0}, {false, false, false, false, false}}); + cudf::test::fixed_width_column_wrapper const input2( + {{0, 0, 0, 0, 0}, {true, true, true, true, true}}); + std::vector views{input1, input2}; + auto result = cudf::concatenate_masks(views, cudf::test::get_default_stream()); +}