Skip to content
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.

Sign up for GitHub

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

Avoid overflow in fused_concatenate_kernel output_index #10344

Merged
merged 6 commits into from
Feb 28, 2022
Merged

Avoid overflow in fused_concatenate_kernel output_index #10344

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
bdd0317
Avoid overflow in fused_concatenate_kernel output_index
abellina Feb 22, 2022
94d5808
Fix styles
abellina Feb 22, 2022
b3f9c7a
Fix copyrights
abellina Feb 22, 2022
2527ab1
Use int64_t instead of std::size_t
abellina Feb 23, 2022
f135b2b
Allow up to INT_MAX elements as a result of fused_concatenate
abellina Feb 24, 2022
3adaa31
Fix comment in test
abellina Feb 24, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions cpp/src/copying/concatenate.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2020-2021, NVIDIA CORPORATION.
* Copyright (c) 2020-2022, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -166,7 +166,7 @@ __global__ void fused_concatenate_kernel(column_device_view const* input_views,
auto const output_size = output_view.size();
auto* output_data = output_view.data<T>();

size_type output_index = threadIdx.x + blockIdx.x * blockDim.x;
int64_t output_index = threadIdx.x + blockIdx.x * blockDim.x;
size_type warp_valid_count = 0;

unsigned active_mask;
Expand Down Expand Up @@ -222,7 +222,7 @@ std::unique_ptr<column> fused_concatenate(host_span<column_view const> views,
auto const& d_offsets = std::get<2>(device_views);
auto const output_size = std::get<3>(device_views);

CUDF_EXPECTS(output_size < static_cast<std::size_t>(std::numeric_limits<size_type>::max()),
CUDF_EXPECTS(output_size <= static_cast<std::size_t>(std::numeric_limits<size_type>::max()),
"Total number of concatenated rows exceeds size_type range");

// Allocate output
Expand Down
18 changes: 17 additions & 1 deletion cpp/tests/copying/concatenate_tests.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2020-2021, NVIDIA CORPORATION.
* Copyright (c) 2020-2022, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -340,6 +340,22 @@ struct OverflowTest : public cudf::test::BaseFixture {
TEST_F(OverflowTest, OverflowTest)
{
using namespace cudf;
// should concatenate up to size_type::max rows.
{
// 5 x size + size_last adds to size_type::max
constexpr auto size = static_cast<size_type>(static_cast<uint32_t>(250) * 1024 * 1024);
constexpr auto size_last = static_cast<size_type>(836763647);

auto many_chars = cudf::make_fixed_width_column(data_type{type_id::INT8}, size);
auto many_chars_last = cudf::make_fixed_width_column(data_type{type_id::INT8}, size_last);

table_view tbl({*many_chars});
table_view tbl_last({*many_chars_last});
std::vector<cudf::table_view> table_views_to_concat({tbl, tbl, tbl, tbl, tbl, tbl_last});
std::unique_ptr<cudf::table> concatenated_tables = cudf::concatenate(table_views_to_concat);
EXPECT_NO_THROW(rmm::cuda_stream_default.synchronize());
ASSERT_EQ(concatenated_tables->num_rows(), std::numeric_limits<size_type>::max());
}

// primitive column
{
Expand Down