-
Notifications
You must be signed in to change notification settings - Fork 1k
Add a public API for copying a table_view to device array #18450
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
rapids-bot
merged 27 commits into
rapidsai:branch-25.06
from
Matt711:fea/cpp/table-to-device-array
May 7, 2025
Merged
Changes from 25 commits
Commits
Show all changes
27 commits
Select commit
Hold shift + click to select a range
e201564
Add a public API for converting a table_view to device array
Matt711 6e9289e
support decimals and add more tests
Matt711 21a9201
Merge branch 'branch-25.06' into fea/cpp/table-to-device-array
Matt711 7eb3690
fallback if cuda version < 12.8
Matt711 854cbbf
Merge branch 'branch-25.06' into fea/cpp/table-to-device-array
Matt711 5a9d195
clean up
Matt711 e5e65cc
address reviews
Matt711 0343bb3
address review
Matt711 7530ecf
use snake case
Matt711 c7ab103
Merge branch 'branch-25.06' into fea/cpp/table-to-device-array
Matt711 eec23b5
Merge branch 'branch-25.06' into fea/cpp/table-to-device-array
Matt711 f027616
address reviews
Matt711 b3251fe
clean up
Matt711 f5a0096
Merge branch 'branch-25.06' into fea/cpp/table-to-device-array
Matt711 96ef619
address review
Matt711 82c5b22
pass a device_span instead of a raw pointer
Matt711 4399938
sort file names
Matt711 737fa22
Merge branch 'branch-25.06' into fea/cpp/table-to-device-array
Matt711 36abd46
Merge branch 'branch-25.06' into fea/cpp/table-to-device-array
Matt711 1fb9563
Merge branch 'branch-25.06' into fea/cpp/table-to-device-array
Matt711 c6785c7
add other impl for benchmarking purposes
Matt711 04c837f
clean up
Matt711 1530b5f
Merge branch 'branch-25.06' into fea/cpp/table-to-device-array
Matt711 fd0506b
address reviews
Matt711 db78162
address review
Matt711 f5bf21c
Merge branch 'branch-25.06' into fea/cpp/table-to-device-array
Matt711 b71ec3a
address review
Matt711 File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,58 @@ | ||
| /* | ||
| * Copyright (c) 2025, 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 <benchmarks/common/generate_input.hpp> | ||
|
|
||
| #include <cudf/reshape.hpp> | ||
| #include <cudf/utilities/default_stream.hpp> | ||
| #include <cudf/utilities/span.hpp> | ||
|
|
||
| #include <cuda/functional> | ||
|
|
||
| #include <nvbench/nvbench.cuh> | ||
|
|
||
| static void bench_table_to_array(nvbench::state& state) | ||
| { | ||
| auto const num_rows = static_cast<cudf::size_type>(state.get_int64("num_rows")); | ||
| auto const num_cols = static_cast<cudf::size_type>(state.get_int64("columns")); | ||
|
|
||
| data_profile profile = data_profile_builder() | ||
| .distribution(cudf::type_id::INT32, distribution_id::UNIFORM, 0, 1000) | ||
| .no_validity(); | ||
| std::vector<cudf::type_id> types(num_cols, cudf::type_id::INT32); | ||
| auto input_table = create_random_table(types, row_count{num_rows}, profile); | ||
|
|
||
| auto input_view = input_table->view(); | ||
| auto stream = cudf::get_default_stream(); | ||
| auto dtype = cudf::data_type{cudf::type_id::INT32}; | ||
|
|
||
| rmm::device_buffer output(num_rows * num_cols * sizeof(int32_t), stream); | ||
| auto span = cudf::device_span<cuda::std::byte>(reinterpret_cast<cuda::std::byte*>(output.data()), | ||
| output.size()); | ||
|
|
||
| state.set_cuda_stream(nvbench::make_cuda_stream_view(stream.value())); | ||
| state.add_global_memory_reads<int32_t>(num_rows * num_cols); // all bytes are read | ||
| state.add_global_memory_writes<int32_t>(num_rows * num_cols); // all bytes are written | ||
|
|
||
| state.exec(nvbench::exec_tag::sync, [&](nvbench::launch& launch) { | ||
| cudf::table_to_array(input_view, span, dtype, stream); | ||
| }); | ||
| } | ||
|
|
||
| NVBENCH_BENCH(bench_table_to_array) | ||
| .set_name("table_to_array") | ||
| .add_int64_axis("num_rows", {32768, 262144, 2097152, 16777216}) | ||
| .add_int64_axis("columns", {2, 10, 100}); |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,5 +1,5 @@ | ||
| /* | ||
| * Copyright (c) 2020-2024, NVIDIA CORPORATION. | ||
| * Copyright (c) 2020-2025, NVIDIA CORPORATION. | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
|
|
@@ -21,6 +21,9 @@ | |
| #include <cudf/types.hpp> | ||
| #include <cudf/utilities/export.hpp> | ||
| #include <cudf/utilities/memory_resource.hpp> | ||
| #include <cudf/utilities/span.hpp> | ||
|
|
||
| #include <cuda/functional> | ||
|
|
||
| #include <memory> | ||
|
|
||
|
|
@@ -107,6 +110,30 @@ std::unique_ptr<column> byte_cast( | |
| rmm::cuda_stream_view stream = cudf::get_default_stream(), | ||
| rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); | ||
|
|
||
| /** | ||
| * @brief Copies a table into a contiguous column-major device array. | ||
| * | ||
| * This function copies a `table_view` with columns of the same fixed-width type | ||
| * into a 2D device array stored in column-major order. | ||
| * | ||
| * The output buffer must be preallocated and passed as a `device_span` using | ||
| * a `device_span<cuda::std::byte>`. It must be large enough to hold | ||
| * `num_rows * num_columns * sizeof(output_dtype)` bytes. | ||
| * | ||
| * @throws cudf::logic_error if columns do not all have the same type as `output_dtype` | ||
| * @throws cudf::logic_error if `output_dtype` is not a fixed-width type | ||
| * @throws std::invalid_argument if the output span is too small | ||
| * | ||
| * @param input A table with fixed-width, non-nullable columns of the same type | ||
| * @param output A span representing preallocated device memory for the output | ||
| * @param output_dtype The data type of the output elements | ||
| * @param stream CUDA stream used for memory operations | ||
| */ | ||
| void table_to_array(table_view const& input, | ||
| device_span<cuda::std::byte> output, | ||
| cudf::data_type output_dtype, | ||
|
||
| rmm::cuda_stream_view stream = cudf::get_default_stream()); | ||
|
|
||
| /** @} */ // end of group | ||
|
|
||
| } // namespace CUDF_EXPORT cudf | ||
Matt711 marked this conversation as resolved.
Show resolved
Hide resolved
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,139 @@ | ||
| /* | ||
| * Copyright (c) 2025, 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/detail/nvtx/ranges.hpp> | ||
| #include <cudf/detail/reshape.hpp> | ||
| #include <cudf/detail/utilities/vector_factories.hpp> | ||
| #include <cudf/reshape.hpp> | ||
| #include <cudf/types.hpp> | ||
| #include <cudf/utilities/default_stream.hpp> | ||
| #include <cudf/utilities/error.hpp> | ||
| #include <cudf/utilities/span.hpp> | ||
| #include <cudf/utilities/type_checks.hpp> | ||
| #include <cudf/utilities/type_dispatcher.hpp> | ||
|
|
||
| #include <rmm/cuda_stream_view.hpp> | ||
| #include <rmm/device_uvector.hpp> | ||
|
|
||
| #include <cub/device/device_memcpy.cuh> | ||
| #include <cuda/functional> | ||
| #include <cuda_runtime.h> | ||
| #include <thrust/device_vector.h> | ||
| #include <thrust/iterator/constant_iterator.h> | ||
| #include <thrust/iterator/counting_iterator.h> | ||
| #include <thrust/iterator/transform_iterator.h> | ||
|
|
||
| namespace cudf { | ||
| namespace detail { | ||
| namespace { | ||
|
|
||
| template <typename T> | ||
| void table_to_array_impl(table_view const& input, | ||
| device_span<cuda::std::byte> output, | ||
| rmm::cuda_stream_view stream) | ||
| { | ||
| auto const num_columns = input.num_columns(); | ||
| auto const num_rows = input.num_rows(); | ||
| auto const item_size = sizeof(T); | ||
| auto const total_bytes = static_cast<size_t>(num_columns) * num_rows * item_size; | ||
|
|
||
| CUDF_EXPECTS(output.size() >= total_bytes, "Output span is too small", std::invalid_argument); | ||
| CUDF_EXPECTS(cudf::all_have_same_types(input.begin(), input.end()), | ||
| "All columns must have the same data type", | ||
| cudf::data_type_error); | ||
| CUDF_EXPECTS(!cudf::has_nulls(input), "All columns must contain no nulls", std::invalid_argument); | ||
|
|
||
| auto* base_ptr = output.data(); | ||
|
|
||
| auto h_srcs = make_host_vector<T const*>(num_columns, stream); | ||
| auto h_dsts = make_host_vector<T*>(num_columns, stream); | ||
|
|
||
| std::transform(input.begin(), input.end(), h_srcs.begin(), [](auto& col) { | ||
| return const_cast<T*>(col.template data<T>()); | ||
| }); | ||
|
|
||
| for (int i = 0; i < num_columns; ++i) { | ||
| h_dsts[i] = reinterpret_cast<T*>(base_ptr + i * item_size * num_rows); | ||
| } | ||
|
|
||
| auto const mr = cudf::get_current_device_resource_ref(); | ||
|
|
||
| auto d_srcs = cudf::detail::make_device_uvector_async(h_srcs, stream, mr); | ||
| auto d_dsts = cudf::detail::make_device_uvector_async(h_dsts, stream, mr); | ||
|
|
||
| thrust::constant_iterator<size_t> sizes(static_cast<size_t>(item_size * num_rows)); | ||
|
|
||
| size_t temp_storage_bytes = 0; | ||
| cub::DeviceMemcpy::Batched(nullptr, | ||
| temp_storage_bytes, | ||
| d_srcs.begin(), | ||
| d_dsts.begin(), | ||
| sizes, | ||
| num_columns, | ||
| stream.value()); | ||
| rmm::device_buffer temp_storage(temp_storage_bytes, stream); | ||
| cub::DeviceMemcpy::Batched(temp_storage.data(), | ||
| temp_storage_bytes, | ||
| d_srcs.begin(), | ||
| d_dsts.begin(), | ||
| sizes, | ||
| num_columns, | ||
| stream.value()); | ||
Matt711 marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
| } | ||
|
|
||
| struct table_to_array_dispatcher { | ||
| table_view const& input; | ||
| device_span<cuda::std::byte> output; | ||
| rmm::cuda_stream_view stream; | ||
|
|
||
| template <typename T, CUDF_ENABLE_IF(is_fixed_width<T>())> | ||
| void operator()() const | ||
| { | ||
| table_to_array_impl<T>(input, output, stream); | ||
| } | ||
|
|
||
| template <typename T, CUDF_ENABLE_IF(!is_fixed_width<T>())> | ||
| void operator()() const | ||
| { | ||
| CUDF_FAIL("Unsupported dtype"); | ||
| } | ||
| }; | ||
|
|
||
| } // namespace | ||
|
|
||
| void table_to_array(table_view const& input, | ||
| device_span<cuda::std::byte> output, | ||
| data_type output_dtype, | ||
| rmm::cuda_stream_view stream) | ||
| { | ||
| if (input.num_columns() == 0) { return; } | ||
|
|
||
| cudf::type_dispatcher<cudf::dispatch_storage_type>( | ||
| output_dtype, table_to_array_dispatcher{input, output, stream}); | ||
| } | ||
|
|
||
| } // namespace detail | ||
|
|
||
| void table_to_array(table_view const& input, | ||
| device_span<cuda::std::byte> output, | ||
| data_type output_dtype, | ||
| rmm::cuda_stream_view stream) | ||
| { | ||
| CUDF_FUNC_RANGE(); | ||
| cudf::detail::table_to_array(input, output, output_dtype, stream); | ||
| } | ||
|
|
||
| } // namespace cudf | ||
Matt711 marked this conversation as resolved.
Show resolved
Hide resolved
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.