Skip to content

Commit

Permalink
Add stream param to stream compaction APIs (#16295)
Browse files Browse the repository at this point in the history 
Add `stream` param to a bunch of stream compaction APIs.

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

Approvers:
  - Nghia Truong (https://github.com/ttnghia)
  - Mark Harris (https://github.com/harrism)
  - Karthikeyan (https://github.com/karthikeyann)
  - Mike Wilson (https://github.com/hyperbolic2346)

URL: #16295
  • Loading branch information
@JayjeetAtGithub
JayjeetAtGithub authored Aug 7, 2024
1 parent d11d2cf commit 3fd8783
Show file tree
Hide file tree
Showing 13 changed files with 362 additions and 121 deletions.
30 changes: 6 additions & 24 deletions cpp/include/cudf/detail/stream_compaction.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,9 +29,7 @@ namespace CUDF_EXPORT cudf {
namespace detail {
/**
* @copydoc cudf::drop_nulls(table_view const&, std::vector<size_type> const&,
* cudf::size_type, rmm::device_async_resource_ref)
*
* @param[in] stream CUDA stream used for device memory operations and kernel launches.
* cudf::size_type, rmm::cuda_stream_view, rmm::device_async_resource_ref)
*/
std::unique_ptr<table> drop_nulls(table_view const& input,
std::vector<size_type> const& keys,
Expand All @@ -41,9 +39,7 @@ std::unique_ptr<table> drop_nulls(table_view const& input,

/**
* @copydoc cudf::drop_nans(table_view const&, std::vector<size_type> const&,
* cudf::size_type, rmm::device_async_resource_ref)
*
* @param[in] stream CUDA stream used for device memory operations and kernel launches.
* cudf::size_type, rmm::cuda_stream_view, rmm::device_async_resource_ref)
*/
std::unique_ptr<table> drop_nans(table_view const& input,
std::vector<size_type> const& keys,
Expand All @@ -53,8 +49,6 @@ std::unique_ptr<table> drop_nans(table_view const& input,

/**
* @copydoc cudf::apply_boolean_mask
*
* @param[in] stream CUDA stream used for device memory operations and kernel launches.
*/
std::unique_ptr<table> apply_boolean_mask(table_view const& input,
column_view const& boolean_mask,
Expand All @@ -63,8 +57,6 @@ std::unique_ptr<table> apply_boolean_mask(table_view const& input,

/**
* @copydoc cudf::unique
*
* @param[in] stream CUDA stream used for device memory operations and kernel launches.
*/
std::unique_ptr<table> unique(table_view const& input,
std::vector<size_type> const& keys,
Expand All @@ -75,8 +67,6 @@ std::unique_ptr<table> unique(table_view const& input,

/**
* @copydoc cudf::distinct
*
* @param[in] stream CUDA stream used for device memory operations and kernel launches.
*/
std::unique_ptr<table> distinct(table_view const& input,
std::vector<size_type> const& keys,
Expand Down Expand Up @@ -110,38 +100,30 @@ rmm::device_uvector<size_type> distinct_indices(table_view const& input,
rmm::device_async_resource_ref mr);

/**
* @copydoc cudf::unique_count(column_view const&, null_policy, nan_policy)
*
* @param[in] stream CUDA stream used for device memory operations and kernel launches.
* @copydoc cudf::unique_count(column_view const&, null_policy, nan_policy, rmm::cuda_stream_view)
*/
cudf::size_type unique_count(column_view const& input,
null_policy null_handling,
nan_policy nan_handling,
rmm::cuda_stream_view stream);

/**
* @copydoc cudf::unique_count(table_view const&, null_equality)
*
* @param[in] stream CUDA stream used for device memory operations and kernel launches.
* @copydoc cudf::unique_count(table_view const&, null_equality, rmm::cuda_stream_view)
*/
cudf::size_type unique_count(table_view const& input,
null_equality nulls_equal,
rmm::cuda_stream_view stream);

/**
* @copydoc cudf::distinct_count(column_view const&, null_policy, nan_policy)
*
* @param[in] stream CUDA stream used for device memory operations and kernel launches.
* @copydoc cudf::distinct_count(column_view const&, null_policy, nan_policy, rmm::cuda_stream_view)
*/
cudf::size_type distinct_count(column_view const& input,
null_policy null_handling,
nan_policy nan_handling,
rmm::cuda_stream_view stream);

/**
* @copydoc cudf::distinct_count(table_view const&, null_equality)
*
* @param[in] stream CUDA stream used for device memory operations and kernel launches.
* @copydoc cudf::distinct_count(table_view const&, null_equality, rmm::cuda_stream_view)
*/
cudf::size_type distinct_count(table_view const& input,
null_equality nulls_equal,
Expand Down
9 changes: 2 additions & 7 deletions cpp/include/cudf/lists/detail/stream_compaction.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,20 +26,15 @@ namespace CUDF_EXPORT cudf {
namespace lists::detail {

/**
* @copydoc cudf::lists::apply_boolean_mask(lists_column_view const&, lists_column_view const&,
* rmm::device_async_resource_ref)
*
* @param stream CUDA stream used for device memory operations and kernel launches
* @copydoc cudf::lists::apply_boolean_mask
*/
std::unique_ptr<column> apply_boolean_mask(lists_column_view const& input,
lists_column_view const& boolean_mask,
rmm::cuda_stream_view stream,
rmm::device_async_resource_ref mr);

/**
* @copydoc cudf::list::distinct
*
* @param stream CUDA stream used for device memory operations and kernel launches.
* @copydoc cudf::lists::distinct
*/
std::unique_ptr<column> distinct(lists_column_view const& input,
null_equality nulls_equal,
Expand Down
30 changes: 26 additions & 4 deletions cpp/include/cudf/stream_compaction.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@ namespace CUDF_EXPORT cudf {
* @param[in] keys vector of indices representing key columns from `input`
* @param[in] keep_threshold The minimum number of non-null fields in a row
* required to keep the row.
* @param[in] stream CUDA stream used for device memory operations and kernel launches
* @param[in] mr Device memory resource used to allocate the returned table's device memory
* @return Table containing all rows of the `input` with at least @p
* keep_threshold non-null fields in @p keys.
Expand All @@ -75,6 +76,7 @@ std::unique_ptr<table> drop_nulls(
table_view const& input,
std::vector<size_type> const& keys,
cudf::size_type keep_threshold,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());

/**
Expand All @@ -99,13 +101,15 @@ std::unique_ptr<table> drop_nulls(
*
* @param[in] input The input `table_view` to filter
* @param[in] keys vector of indices representing key columns from `input`
* @param[in] stream CUDA stream used for device memory operations and kernel launches
* @param[in] mr Device memory resource used to allocate the returned table's device memory
* @return Table containing all rows of the `input` without nulls in the columns
* of @p keys.
*/
std::unique_ptr<table> drop_nulls(
table_view const& input,
std::vector<size_type> const& keys,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());

/**
Expand Down Expand Up @@ -141,6 +145,7 @@ std::unique_ptr<table> drop_nulls(
* @param[in] keys vector of indices representing key columns from `input`
* @param[in] keep_threshold The minimum number of non-NAN elements in a row
* required to keep the row.
* @param[in] stream CUDA stream used for device memory operations and kernel launches
* @param[in] mr Device memory resource used to allocate the returned table's device memory
* @return Table containing all rows of the `input` with at least @p
* keep_threshold non-NAN elements in @p keys.
Expand All @@ -149,6 +154,7 @@ std::unique_ptr<table> drop_nans(
table_view const& input,
std::vector<size_type> const& keys,
cudf::size_type keep_threshold,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());

/**
Expand All @@ -174,13 +180,15 @@ std::unique_ptr<table> drop_nans(
*
* @param[in] input The input `table_view` to filter
* @param[in] keys vector of indices representing key columns from `input`
* @param[in] stream CUDA stream used for device memory operations and kernel launches
* @param[in] mr Device memory resource used to allocate the returned table's device memory
* @return Table containing all rows of the `input` without NANs in the columns
* of @p keys.
*/
std::unique_ptr<table> drop_nans(
table_view const& input,
std::vector<size_type> const& keys,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());

/**
Expand All @@ -200,13 +208,15 @@ std::unique_ptr<table> drop_nans(
* @param[in] input The input table_view to filter
* @param[in] boolean_mask A nullable column_view of type type_id::BOOL8 used
* as a mask to filter the `input`.
* @param[in] stream CUDA stream used for device memory operations and kernel launches
* @param[in] mr Device memory resource used to allocate the returned table's device memory
* @return Table containing copy of all rows of @p input passing
* the filter defined by @p boolean_mask.
*/
std::unique_ptr<table> apply_boolean_mask(
table_view const& input,
column_view const& boolean_mask,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());

/**
Expand Down Expand Up @@ -241,6 +251,7 @@ enum class duplicate_keep_option {
* @param[in] keep keep any, first, last, or none of the found duplicates
* @param[in] nulls_equal flag to denote nulls are equal if null_equality::EQUAL, nulls are not
* equal if null_equality::UNEQUAL
* @param[in] stream CUDA stream used for device memory operations and kernel launches
* @param[in] mr Device memory resource used to allocate the returned table's device
* memory
*
Expand All @@ -251,6 +262,7 @@ std::unique_ptr<table> unique(
std::vector<size_type> const& keys,
duplicate_keep_option keep,
null_equality nulls_equal = null_equality::EQUAL,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());

/**
Expand All @@ -269,6 +281,7 @@ std::unique_ptr<table> unique(
* @param keep Copy any, first, last, or none of the found duplicates
* @param nulls_equal Flag to specify whether null elements should be considered as equal
* @param nans_equal Flag to specify whether NaN elements should be considered as equal
* @param stream CUDA stream used for device memory operations and kernel launches
* @param mr Device memory resource used to allocate the returned table
* @return Table with distinct rows in an unspecified order
*/
Expand All @@ -278,6 +291,7 @@ std::unique_ptr<table> distinct(
duplicate_keep_option keep = duplicate_keep_option::KEEP_ANY,
null_equality nulls_equal = null_equality::EQUAL,
nan_equality nans_equal = nan_equality::ALL_EQUAL,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());

/**
Expand Down Expand Up @@ -346,24 +360,28 @@ std::unique_ptr<table> stable_distinct(
* @param[in] input The column_view whose consecutive groups of equivalent rows will be counted
* @param[in] null_handling flag to include or ignore `null` while counting
* @param[in] nan_handling flag to consider `NaN==null` or not
* @param[in] stream CUDA stream used for device memory operations and kernel launches
*
* @return number of consecutive groups of equivalent rows in the column
*/
cudf::size_type unique_count(column_view const& input,
null_policy null_handling,
nan_policy nan_handling);
nan_policy nan_handling,
rmm::cuda_stream_view stream = cudf::get_default_stream());

/**
* @brief Count the number of consecutive groups of equivalent rows in a table.
*
* @param[in] input Table whose consecutive groups of equivalent rows will be counted
* @param[in] nulls_equal flag to denote if null elements should be considered equal
* nulls are not equal if null_equality::UNEQUAL.
* @param[in] stream CUDA stream used for device memory operations and kernel launches
*
* @return number of consecutive groups of equivalent rows in the column
*/
cudf::size_type unique_count(table_view const& input,
null_equality nulls_equal = null_equality::EQUAL);
null_equality nulls_equal = null_equality::EQUAL,
rmm::cuda_stream_view stream = cudf::get_default_stream());

/**
* @brief Count the distinct elements in the column_view.
Expand All @@ -382,24 +400,28 @@ cudf::size_type unique_count(table_view const& input,
* @param[in] input The column_view whose distinct elements will be counted
* @param[in] null_handling flag to include or ignore `null` while counting
* @param[in] nan_handling flag to consider `NaN==null` or not
* @param[in] stream CUDA stream used for device memory operations and kernel launches
*
* @return number of distinct rows in the table
*/
cudf::size_type distinct_count(column_view const& input,
null_policy null_handling,
nan_policy nan_handling);
nan_policy nan_handling,
rmm::cuda_stream_view stream = cudf::get_default_stream());

/**
* @brief Count the distinct rows in a table.
*
* @param[in] input Table whose distinct rows will be counted
* @param[in] nulls_equal flag to denote if null elements should be considered equal.
* nulls are not equal if null_equality::UNEQUAL.
* @param[in] stream CUDA stream used for device memory operations and kernel launches
*
* @return number of distinct rows in the table
*/
cudf::size_type distinct_count(table_view const& input,
null_equality nulls_equal = null_equality::EQUAL);
null_equality nulls_equal = null_equality::EQUAL,
rmm::cuda_stream_view stream = cudf::get_default_stream());

/** @} */
} // namespace CUDF_EXPORT cudf
3 changes: 2 additions & 1 deletion cpp/src/stream_compaction/apply_boolean_mask.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,9 +91,10 @@ std::unique_ptr<table> apply_boolean_mask(table_view const& input,
*/
std::unique_ptr<table> apply_boolean_mask(table_view const& input,
column_view const& boolean_mask,
rmm::cuda_stream_view stream,
rmm::device_async_resource_ref mr)
{
CUDF_FUNC_RANGE();
return detail::apply_boolean_mask(input, boolean_mask, cudf::get_default_stream(), mr);
return detail::apply_boolean_mask(input, boolean_mask, stream, mr);
}
} // namespace cudf
4 changes: 2 additions & 2 deletions cpp/src/stream_compaction/distinct.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -150,11 +150,11 @@ std::unique_ptr<table> distinct(table_view const& input,
duplicate_keep_option keep,
null_equality nulls_equal,
nan_equality nans_equal,
rmm::cuda_stream_view stream,
rmm::device_async_resource_ref mr)
{
CUDF_FUNC_RANGE();
return detail::distinct(
input, keys, keep, nulls_equal, nans_equal, cudf::get_default_stream(), mr);
return detail::distinct(input, keys, keep, nulls_equal, nans_equal, stream, mr);
}

std::unique_ptr<column> distinct_indices(table_view const& input,
Expand Down
11 changes: 7 additions & 4 deletions cpp/src/stream_compaction/distinct_count.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -218,15 +218,18 @@ cudf::size_type distinct_count(column_view const& input,

cudf::size_type distinct_count(column_view const& input,
null_policy null_handling,
nan_policy nan_handling)
nan_policy nan_handling,
rmm::cuda_stream_view stream)
{
CUDF_FUNC_RANGE();
return detail::distinct_count(input, null_handling, nan_handling, cudf::get_default_stream());
return detail::distinct_count(input, null_handling, nan_handling, stream);
}

cudf::size_type distinct_count(table_view const& input, null_equality nulls_equal)
cudf::size_type distinct_count(table_view const& input,
null_equality nulls_equal,
rmm::cuda_stream_view stream)
{
CUDF_FUNC_RANGE();
return detail::distinct_count(input, nulls_equal, cudf::get_default_stream());
return detail::distinct_count(input, nulls_equal, stream);
}
} // namespace cudf
6 changes: 4 additions & 2 deletions cpp/src/stream_compaction/drop_nans.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -117,20 +117,22 @@ std::unique_ptr<table> drop_nans(table_view const& input,
std::unique_ptr<table> drop_nans(table_view const& input,
std::vector<size_type> const& keys,
cudf::size_type keep_threshold,
rmm::cuda_stream_view stream,
rmm::device_async_resource_ref mr)
{
CUDF_FUNC_RANGE();
return detail::drop_nans(input, keys, keep_threshold, cudf::get_default_stream(), mr);
return detail::drop_nans(input, keys, keep_threshold, stream, mr);
}
/*
* Filters a table to remove nan elements.
*/
std::unique_ptr<table> drop_nans(table_view const& input,
std::vector<size_type> const& keys,
rmm::cuda_stream_view stream,
rmm::device_async_resource_ref mr)
{
CUDF_FUNC_RANGE();
return detail::drop_nans(input, keys, keys.size(), cudf::get_default_stream(), mr);
return detail::drop_nans(input, keys, keys.size(), stream, mr);
}

} // namespace cudf
6 changes: 4 additions & 2 deletions cpp/src/stream_compaction/drop_nulls.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -90,20 +90,22 @@ std::unique_ptr<table> drop_nulls(table_view const& input,
std::unique_ptr<table> drop_nulls(table_view const& input,
std::vector<size_type> const& keys,
cudf::size_type keep_threshold,
rmm::cuda_stream_view stream,
rmm::device_async_resource_ref mr)
{
CUDF_FUNC_RANGE();
return detail::drop_nulls(input, keys, keep_threshold, cudf::get_default_stream(), mr);
return detail::drop_nulls(input, keys, keep_threshold, stream, mr);
}
/*
* Filters a table to remove null elements.
*/
std::unique_ptr<table> drop_nulls(table_view const& input,
std::vector<size_type> const& keys,
rmm::cuda_stream_view stream,
rmm::device_async_resource_ref mr)
{
CUDF_FUNC_RANGE();
return detail::drop_nulls(input, keys, keys.size(), cudf::get_default_stream(), mr);
return detail::drop_nulls(input, keys, keys.size(), stream, mr);
}

} // namespace cudf
Loading

0 comments on commit 3fd8783

Please sign in to comment.