Skip to content

Commit

Permalink
Refactor from_arrow_device/host to use resource_ref (#16160)
Browse files Browse the repository at this point in the history 
Fixes #16159

Also fixes typos / leftovers in  dictionary `add_keys` copydocs.

Authors:
  - Mark Harris (https://github.com/harrism)

Approvers:
  - Paul Mattione (https://github.com/pmattione-nvidia)
  - David Wendt (https://github.com/davidwendt)
  - Vyas Ramasubramani (https://github.com/vyasr)

URL: #16160
  • Loading branch information
@harrism
harrism authored Jul 3, 2024
1 parent 3aedeea commit 39de5a2
Show file tree
Hide file tree
Showing 5 changed files with 47 additions and 44 deletions.
10 changes: 5 additions & 5 deletions cpp/include/cudf/dictionary/detail/update_keys.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ namespace dictionary {
namespace detail {
/**
* @copydoc cudf::dictionary::add_keys(dictionary_column_view const&,column_view
* const&,mm::mr::device_memory_resource*)
* const&,rmm::device_async_resource_ref)
*
* @param stream CUDA stream used for device memory operations and kernel launches.
*/
Expand All @@ -40,7 +40,7 @@ std::unique_ptr<column> add_keys(dictionary_column_view const& dictionary_column

/**
* @copydoc cudf::dictionary::remove_keys(dictionary_column_view const&,column_view
* const&,mm::mr::device_memory_resource*)
* const&,rmm::device_async_resource_ref)
*
* @param stream CUDA stream used for device memory operations and kernel launches.
*/
Expand All @@ -51,7 +51,7 @@ std::unique_ptr<column> remove_keys(dictionary_column_view const& dictionary_col

/**
* @copydoc cudf::dictionary::remove_unused_keys(dictionary_column_view
* const&,mm::mr::device_memory_resource*)
* const&,rmm::device_async_resource_ref)
*
* @param stream CUDA stream used for device memory operations and kernel launches.
*/
Expand All @@ -61,7 +61,7 @@ std::unique_ptr<column> remove_unused_keys(dictionary_column_view const& diction

/**
* @copydoc cudf::dictionary::set_keys(dictionary_column_view
* const&,mm::mr::device_memory_resource*)
* const&,rmm::device_async_resource_ref)
*
* @param stream CUDA stream used for device memory operations and kernel launches.
*/
Expand All @@ -72,7 +72,7 @@ std::unique_ptr<column> set_keys(dictionary_column_view const& dictionary_column

/**
* @copydoc
* cudf::dictionary::match_dictionaries(std::vector<cudf::dictionary_column_view>,mm::mr::device_memory_resource*)
* cudf::dictionary::match_dictionaries(std::vector<cudf::dictionary_column_view>,rmm::device_async_resource_ref)
*
* @param stream CUDA stream used for device memory operations and kernel launches.
*/
Expand Down
29 changes: 15 additions & 14 deletions cpp/include/cudf/interop.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@
#include <cudf/utilities/span.hpp>

#include <rmm/mr/device/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>

#include <utility>

Expand Down Expand Up @@ -372,8 +373,8 @@ std::unique_ptr<cudf::scalar> from_arrow(
std::unique_ptr<cudf::table> from_arrow(
ArrowSchema const* schema,
ArrowArray const* input,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());

/**
* @brief Create `cudf::column` from a given ArrowArray and ArrowSchema input
Expand All @@ -391,8 +392,8 @@ std::unique_ptr<cudf::table> from_arrow(
std::unique_ptr<cudf::column> from_arrow_column(
ArrowSchema const* schema,
ArrowArray const* input,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());

/**
* @brief Create `cudf::table` from given ArrowDeviceArray input
Expand All @@ -415,8 +416,8 @@ std::unique_ptr<cudf::column> from_arrow_column(
std::unique_ptr<table> from_arrow_host(
ArrowSchema const* schema,
ArrowDeviceArray const* input,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());

/**
* @brief Create `cudf::table` from given ArrowArrayStream input
Expand All @@ -433,8 +434,8 @@ std::unique_ptr<table> from_arrow_host(
*/
std::unique_ptr<table> from_arrow_stream(
ArrowArrayStream* input,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());

/**
* @brief Create `cudf::column` from given ArrowDeviceArray input
Expand All @@ -456,8 +457,8 @@ std::unique_ptr<table> from_arrow_stream(
std::unique_ptr<column> from_arrow_host_column(
ArrowSchema const* schema,
ArrowDeviceArray const* input,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());

/**
* @brief typedef for a vector of owning columns, used for conversion from ArrowDeviceArray
Expand Down Expand Up @@ -537,8 +538,8 @@ using unique_table_view_t =
unique_table_view_t from_arrow_device(
ArrowSchema const* schema,
ArrowDeviceArray const* input,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());

/**
* @brief typedef for a unique_ptr to a `cudf::column_view` with custom deleter
Expand Down Expand Up @@ -580,8 +581,8 @@ using unique_column_view_t =
unique_column_view_t from_arrow_device_column(
ArrowSchema const* schema,
ArrowDeviceArray const* input,
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
rmm::cuda_stream_view stream = cudf::get_default_stream(),
rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());

/** @} */ // end of group
} // namespace cudf
27 changes: 14 additions & 13 deletions cpp/src/interop/from_arrow_device.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
#include <rmm/cuda_device.hpp>
#include <rmm/cuda_stream_view.hpp>
#include <rmm/device_buffer.hpp>
#include <rmm/resource_ref.hpp>

#include <nanoarrow/nanoarrow.h>
#include <nanoarrow/nanoarrow.hpp>
Expand All @@ -56,7 +57,7 @@ struct dispatch_from_arrow_device {
data_type,
bool,
rmm::cuda_stream_view,
rmm::mr::device_memory_resource*)
rmm::device_async_resource_ref)
{
CUDF_FAIL("Unsupported type in from_arrow_device", cudf::data_type_error);
}
Expand All @@ -68,7 +69,7 @@ struct dispatch_from_arrow_device {
data_type type,
bool skip_mask,
rmm::cuda_stream_view,
rmm::mr::device_memory_resource*)
rmm::device_async_resource_ref mr)
{
size_type const num_rows = input->length;
size_type const offset = input->offset;
Expand All @@ -90,15 +91,15 @@ dispatch_tuple_t get_column(ArrowSchemaView* schema,
data_type type,
bool skip_mask,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr);
rmm::device_async_resource_ref mr);

template <>
dispatch_tuple_t dispatch_from_arrow_device::operator()<bool>(ArrowSchemaView* schema,
ArrowArray const* input,
data_type type,
bool skip_mask,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
if (input->length == 0) {
return std::make_tuple<column_view, owned_columns_t>(
Expand Down Expand Up @@ -141,7 +142,7 @@ dispatch_tuple_t dispatch_from_arrow_device::operator()<cudf::string_view>(
data_type type,
bool skip_mask,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
CUDF_EXPECTS(schema->type != NANOARROW_TYPE_LARGE_STRING,
"Large strings are not yet supported in from_arrow_device",
Expand Down Expand Up @@ -182,7 +183,7 @@ dispatch_tuple_t dispatch_from_arrow_device::operator()<cudf::dictionary32>(
data_type type,
bool skip_mask,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
ArrowSchemaView keys_schema_view;
NANOARROW_THROW_NOT_OK(
Expand Down Expand Up @@ -238,7 +239,7 @@ dispatch_tuple_t dispatch_from_arrow_device::operator()<cudf::struct_view>(
data_type type,
bool skip_mask,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
std::vector<column_view> children;
owned_columns_t out_owned_cols;
Expand Down Expand Up @@ -283,7 +284,7 @@ dispatch_tuple_t dispatch_from_arrow_device::operator()<cudf::list_view>(
data_type type,
bool skip_mask,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
size_type const num_rows = input->length;
size_type const offset = input->offset;
Expand Down Expand Up @@ -324,7 +325,7 @@ dispatch_tuple_t get_column(ArrowSchemaView* schema,
data_type type,
bool skip_mask,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
return type.id() != type_id::EMPTY
? std::move(type_dispatcher(
Expand All @@ -342,7 +343,7 @@ dispatch_tuple_t get_column(ArrowSchemaView* schema,
unique_table_view_t from_arrow_device(ArrowSchema const* schema,
ArrowDeviceArray const* input,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
CUDF_EXPECTS(schema != nullptr && input != nullptr,
"input ArrowSchema and ArrowDeviceArray must not be NULL",
Expand Down Expand Up @@ -397,7 +398,7 @@ unique_table_view_t from_arrow_device(ArrowSchema const* schema,
unique_column_view_t from_arrow_device_column(ArrowSchema const* schema,
ArrowDeviceArray const* input,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
CUDF_EXPECTS(schema != nullptr && input != nullptr,
"input ArrowSchema and ArrowDeviceArray must not be NULL",
Expand Down Expand Up @@ -429,7 +430,7 @@ unique_column_view_t from_arrow_device_column(ArrowSchema const* schema,
unique_table_view_t from_arrow_device(ArrowSchema const* schema,
ArrowDeviceArray const* input,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
CUDF_FUNC_RANGE();

Expand All @@ -439,7 +440,7 @@ unique_table_view_t from_arrow_device(ArrowSchema const* schema,
unique_column_view_t from_arrow_device_column(ArrowSchema const* schema,
ArrowDeviceArray const* input,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
CUDF_FUNC_RANGE();

Expand Down
19 changes: 10 additions & 9 deletions cpp/src/interop/from_arrow_host.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
#include <rmm/cuda_device.hpp>
#include <rmm/cuda_stream_view.hpp>
#include <rmm/device_buffer.hpp>
#include <rmm/resource_ref.hpp>

#include <nanoarrow/nanoarrow.h>
#include <nanoarrow/nanoarrow.hpp>
Expand All @@ -49,7 +50,7 @@ namespace {

struct dispatch_copy_from_arrow_host {
rmm::cuda_stream_view stream;
rmm::mr::device_memory_resource* mr;
rmm::device_async_resource_ref mr;

std::unique_ptr<rmm::device_buffer> get_mask_buffer(ArrowArray const* array)
{
Expand Down Expand Up @@ -131,7 +132,7 @@ std::unique_ptr<column> get_column_copy(ArrowSchemaView* schema,
data_type type,
bool skip_mask,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr);
rmm::device_async_resource_ref mr);

template <>
std::unique_ptr<column> dispatch_copy_from_arrow_host::operator()<bool>(ArrowSchemaView* schema,
Expand Down Expand Up @@ -388,7 +389,7 @@ std::unique_ptr<column> get_column_copy(ArrowSchemaView* schema,
data_type type,
bool skip_mask,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
return type.id() != type_id::EMPTY
? std::move(type_dispatcher(
Expand All @@ -405,7 +406,7 @@ std::unique_ptr<column> get_column_copy(ArrowSchemaView* schema,
std::unique_ptr<table> from_arrow_host(ArrowSchema const* schema,
ArrowDeviceArray const* input,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
CUDF_EXPECTS(schema != nullptr && input != nullptr,
"input ArrowSchema and ArrowDeviceArray must not be NULL",
Expand Down Expand Up @@ -441,7 +442,7 @@ std::unique_ptr<table> from_arrow_host(ArrowSchema const* schema,
std::unique_ptr<column> from_arrow_host_column(ArrowSchema const* schema,
ArrowDeviceArray const* input,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
CUDF_EXPECTS(schema != nullptr && input != nullptr,
"input ArrowSchema and ArrowDeviceArray must not be NULL",
Expand All @@ -462,7 +463,7 @@ std::unique_ptr<column> from_arrow_host_column(ArrowSchema const* schema,
std::unique_ptr<table> from_arrow_host(ArrowSchema const* schema,
ArrowDeviceArray const* input,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
CUDF_FUNC_RANGE();

Expand All @@ -472,7 +473,7 @@ std::unique_ptr<table> from_arrow_host(ArrowSchema const* schema,
std::unique_ptr<column> from_arrow_host_column(ArrowSchema const* schema,
ArrowDeviceArray const* input,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
CUDF_FUNC_RANGE();

Expand All @@ -482,7 +483,7 @@ std::unique_ptr<column> from_arrow_host_column(ArrowSchema const* schema,
std::unique_ptr<table> from_arrow(ArrowSchema const* schema,
ArrowArray const* input,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
CUDF_FUNC_RANGE();

Expand All @@ -497,7 +498,7 @@ std::unique_ptr<table> from_arrow(ArrowSchema const* schema,
std::unique_ptr<column> from_arrow_column(ArrowSchema const* schema,
ArrowArray const* input,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
CUDF_FUNC_RANGE();

Expand Down
6 changes: 3 additions & 3 deletions cpp/src/interop/from_arrow_stream.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ namespace {

std::unique_ptr<column> make_empty_column_from_schema(ArrowSchema const* schema,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
ArrowSchemaView schema_view;
NANOARROW_THROW_NOT_OK(ArrowSchemaViewInit(&schema_view, schema, nullptr));
Expand Down Expand Up @@ -81,7 +81,7 @@ std::unique_ptr<column> make_empty_column_from_schema(ArrowSchema const* schema,

std::unique_ptr<table> from_arrow_stream(ArrowArrayStream* input,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
CUDF_EXPECTS(input != nullptr, "input ArrowArrayStream must not be NULL", std::invalid_argument);

Expand Down Expand Up @@ -135,7 +135,7 @@ std::unique_ptr<table> from_arrow_stream(ArrowArrayStream* input,

std::unique_ptr<table> from_arrow_stream(ArrowArrayStream* input,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
rmm::device_async_resource_ref mr)
{
CUDF_FUNC_RANGE();
return detail::from_arrow_stream(input, stream, mr);
Expand Down

0 comments on commit 39de5a2

Please sign in to comment.