Skip to content

Commit 7fe53da

Browse files
kevkristkevkrist
committed
Switch from cudf::size_type to int64_t for string offsets in the expression executor
1 parent d0e47d9 commit 7fe53da

20 files changed

Lines changed: 382 additions & 343 deletions

‎src/cuda/expression_executor/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,5 +2,6 @@ set(CUDA_SOURCES
22
${CUDA_SOURCES}
33
${CMAKE_CURRENT_SOURCE_DIR}/gpu_dispatch_string.cu
44
${CMAKE_CURRENT_SOURCE_DIR}/gpu_dispatch_materialize.cu
5+
${CMAKE_CURRENT_SOURCE_DIR}/gpu_dispatch_select.cu
56
PARENT_SCOPE
67
)

‎src/cuda/expression_executor/gpu_dispatch_materialize.cu‎

Lines changed: 35 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,10 @@ struct MaterializeNumeric
2121
static std::unique_ptr<cudf::column> Do(const T* input_data,
2222
const uint64_t* row_ids,
2323
uint64_t row_id_count,
24-
rmm::device_async_resource_ref mr)
24+
rmm::device_async_resource_ref mr,
25+
rmm::cuda_stream_view stream = rmm::cuda_stream_default)
2526
{
26-
auto stream = cudf::get_default_stream();
27+
// Define the thrust execution policy
2728
rmm::mr::thrust_allocator<uint8_t> thrust_allocator(stream, mr);
2829
auto exec = thrust::cuda::par(thrust_allocator).on(stream);
2930

@@ -55,9 +56,9 @@ struct MaterializeString
5556
{}
5657

5758
// Operator
58-
__device__ __forceinline__ cudf::size_type operator()(uint64_t row_id) const
59+
__device__ __forceinline__ int64_t operator()(uint64_t row_id) const
5960
{
60-
return static_cast<cudf::size_type>(input_offsets[row_id + 1] - input_offsets[row_id]);
61+
return static_cast<int64_t>(input_offsets[row_id + 1] - input_offsets[row_id]);
6162
}
6263
};
6364

@@ -68,14 +69,14 @@ struct MaterializeString
6869
const uint64_t* row_ids;
6970
const uint64_t* input_offsets;
7071
const uint8_t* input_data;
71-
cudf::size_type* output_offsets;
72+
int64_t* output_offsets;
7273
uint8_t* output_data;
7374

7475
// Constructor
7576
CopyData(const uint64_t* row_ids,
7677
const uint64_t* input_offsets,
7778
const uint8_t* input_data,
78-
cudf::size_type* output_offsets,
79+
int64_t* output_offsets,
7980
uint8_t* output_data)
8081
: row_ids(row_ids)
8182
, input_offsets(input_offsets)
@@ -101,21 +102,21 @@ struct MaterializeString
101102
const uint64_t* input_offsets,
102103
const uint64_t* row_ids,
103104
uint64_t row_id_count,
104-
rmm::device_async_resource_ref mr)
105+
rmm::device_async_resource_ref mr,
106+
rmm::cuda_stream_view stream = rmm::cuda_stream_default)
105107
{
106108
static_assert(std::is_same_v<int32_t, cudf::size_type>); // Sanity check
107109

108-
auto stream = cudf::get_default_stream();
110+
// Define the thrust execution policy
109111
rmm::mr::thrust_allocator<uint8_t> thrust_allocator(stream, mr);
110112
auto exec = thrust::cuda::par(thrust_allocator).on(stream);
111113
uint64_t offset_count = row_id_count + 1;
112114

113115
// Allocate temporary string length and output offsets buffer
114-
rmm::device_uvector<cudf::size_type> temp_string_lengths(offset_count, stream, mr);
115-
rmm::device_uvector<cudf::size_type> output_offsets(offset_count, stream, mr);
116+
rmm::device_uvector<int64_t> temp_string_lengths(offset_count, stream, mr);
117+
rmm::device_uvector<int64_t> output_offsets(offset_count, stream, mr);
116118
// Set the last string length to 0, so that the exclusive scan places the total sum at the end
117-
CUDF_CUDA_TRY(
118-
cudaMemset(temp_string_lengths.data() + row_id_count, 0, sizeof(cudf::size_type)));
119+
CUDF_CUDA_TRY(cudaMemset(temp_string_lengths.data() + row_id_count, 0, sizeof(int64_t)));
119120

120121
// Gather the string lengths
121122
thrust::transform(exec,
@@ -129,16 +130,9 @@ struct MaterializeString
129130
temp_string_lengths.begin(),
130131
temp_string_lengths.end(),
131132
output_offsets.begin(),
132-
static_cast<cudf::size_type>(0));
133+
static_cast<int64_t>(0));
133134
const auto output_bytes = output_offsets.back_element(stream);
134135

135-
// Check for overflow
136-
if (output_bytes < 0)
137-
{
138-
throw InternalException("Dispatch[Materialize]: String data greater than INT32_MAX "
139-
"detedted!");
140-
}
141-
142136
// Allocate output data buffer
143137
rmm::device_uvector<uint8_t> output_data(output_bytes, stream, mr);
144138

@@ -150,7 +144,7 @@ struct MaterializeString
150144
CopyData(row_ids, input_offsets, input_data, output_offsets.data(), output_data.data()));
151145

152146
// Return a cudf::column
153-
auto offsets_col = std::make_unique<cudf::column>(cudf::data_type{cudf::type_id::INT32},
147+
auto offsets_col = std::make_unique<cudf::column>(cudf::data_type{cudf::type_id::INT64},
154148
static_cast<cudf::size_type>(offset_count),
155149
output_offsets.release(),
156150
rmm::device_buffer{0, stream, mr},
@@ -165,7 +159,8 @@ struct MaterializeString
165159

166160
// Materialize a GPUColumn directly into a cudf::column
167161
std::unique_ptr<cudf::column> GpuDispatcher::DispatchMaterialize(const GPUColumn* input,
168-
rmm::device_async_resource_ref mr)
162+
rmm::device_async_resource_ref mr,
163+
rmm::cuda_stream_view stream)
169164
{
170165
D_ASSERT(input->row_ids != nullptr);
171166
D_ASSERT(input->row_ids->size() > 0);
@@ -179,38 +174,46 @@ std::unique_ptr<cudf::column> GpuDispatcher::DispatchMaterialize(const GPUColumn
179174
return MaterializeNumeric<int32_t>::Do(reinterpret_cast<const int32_t*>(input_data),
180175
input->row_ids,
181176
input->row_id_count,
182-
mr);
177+
mr,
178+
stream);
183179
case GPUColumnTypeId::INT64:
184180
return MaterializeNumeric<uint64_t>::Do(reinterpret_cast<const uint64_t*>(input_data),
185181
input->row_ids,
186182
input->row_id_count,
187-
mr);
183+
mr,
184+
stream);
188185
case GPUColumnTypeId::FLOAT32:
189186
return MaterializeNumeric<float_t>::Do(reinterpret_cast<const float_t*>(input_data),
190187
input->row_ids,
191188
input->row_id_count,
192-
mr);
189+
mr,
190+
stream);
193191
case GPUColumnTypeId::FLOAT64:
194192
return MaterializeNumeric<double_t>::Do(reinterpret_cast<const double_t*>(input_data),
195193
input->row_ids,
196194
input->row_id_count,
197-
mr);
195+
mr,
196+
stream);
198197
case GPUColumnTypeId::BOOLEAN:
199198
return MaterializeNumeric<bool>::Do(reinterpret_cast<const bool*>(input_data),
200199
input->row_ids,
201200
input->row_id_count,
202-
mr);
201+
mr,
202+
stream);
203203
case GPUColumnTypeId::DATE:
204-
return MaterializeNumeric<cudf::timestamp_D>::Do(reinterpret_cast<const cudf::timestamp_D*>(input_data),
205-
input->row_ids,
206-
input->row_id_count,
207-
mr);
204+
return MaterializeNumeric<cudf::timestamp_D>::Do(
205+
reinterpret_cast<const cudf::timestamp_D*>(input_data),
206+
input->row_ids,
207+
input->row_id_count,
208+
mr,
209+
stream);
208210
case GPUColumnTypeId::VARCHAR:
209211
return MaterializeString::Do(input_data,
210212
input_offsets,
211213
input->row_ids,
212214
input->row_id_count,
213-
mr);
215+
mr,
216+
stream);
214217
default:
215218
throw InternalException("Unsupported sirius column type in `Dispatch[Materialize]`: %d",
216219
static_cast<int>(input->data_wrapper.type.id()));

‎src/cuda/expression_executor/gpu_dispatch_select.cu‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
#include "expression_executor/gpu_dispatcher.hpp"
2+
#include "gpu_buffer_manager.hpp"
3+
#include <cub/cub.cuh>
4+
#include <tuple>
5+
6+
namespace duckdb
7+
{
8+
namespace sirius
9+
{
10+
11+
//----------Select----------//
12+
std::tuple<uint64_t*, uint64_t> GpuDispatcher::DispatchSelect(const cudf::column_view& bitmap,
13+
rmm::device_async_resource_ref mr,
14+
rmm::cuda_stream_view stream)
15+
{
16+
// The row ids are owned by the query executor and so must be managed by the buffer manager
17+
auto* gpu_buffer_manager = &GPUBufferManager::GetInstance();
18+
uint64_t* row_ids = gpu_buffer_manager->customCudaMalloc<uint64_t>(bitmap.size(), 0, false);
19+
rmm::device_scalar<uint64_t> d_num_selected(0, stream, mr);
20+
21+
size_t temp_storage_bytes = 0;
22+
uint64_t num_selected = 0;
23+
cub::DeviceSelect::Flagged(nullptr,
24+
temp_storage_bytes,
25+
thrust::make_counting_iterator<uint64_t>(0),
26+
bitmap.data<bool>(),
27+
row_ids,
28+
d_num_selected.data(),
29+
bitmap.size(),
30+
stream);
31+
rmm::device_buffer temp_storage(temp_storage_bytes, stream, mr);
32+
cub::DeviceSelect::Flagged(temp_storage.data(),
33+
temp_storage_bytes,
34+
thrust::make_counting_iterator<uint64_t>(0),
35+
bitmap.data<bool>(),
36+
row_ids,
37+
d_num_selected.data(),
38+
bitmap.size(),
39+
stream);
40+
num_selected = d_num_selected.value(stream);
41+
return std::make_tuple(row_ids, num_selected);
42+
}
43+
44+
} // namespace sirius
45+
} // namespace duckdb

‎src/cuda/expression_executor/gpu_dispatch_string.cu‎

Lines changed: 22 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
#include "expression_executor/gpu_dispatcher.hpp"
2-
#include "gpu_buffer_manager.hpp"
32
#include "gpu_physical_strings_matching.hpp"
43
#include "gpu_physical_substring.hpp"
54
#include <cub/cub.cuh>
@@ -8,7 +7,6 @@
87
#include <cudf/unary.hpp>
98
#include <memory>
109
#include <thrust/iterator/counting_iterator.h>
11-
#include <tuple>
1210

1311
namespace duckdb
1412
{
@@ -21,36 +19,39 @@ namespace sirius
2119
std::unique_ptr<cudf::column> GpuDispatcher::DispatchSubstring(const cudf::column_view& input,
2220
uint64_t start_idx,
2321
uint64_t len,
24-
rmm::device_async_resource_ref mr)
22+
rmm::device_async_resource_ref mr,
23+
rmm::cuda_stream_view stream)
2524
{
2625
return DoSubstring(input.data<char>(),
2726
input.size(),
28-
input.child(0).data<cudf::size_type>(),
29-
static_cast<cudf::size_type>(start_idx),
30-
static_cast<cudf::size_type>(len),
31-
mr);
27+
input.child(0).data<int64_t>(),
28+
static_cast<int64_t>(start_idx),
29+
static_cast<int64_t>(len),
30+
mr,
31+
stream);
3232
};
3333

3434
//----------String Matching----------//
3535
template <StringMatchingType MatchType>
3636
std::unique_ptr<cudf::column>
3737
GpuDispatcher::DispatchStringMatching(const cudf::column_view& input,
3838
const std::string& match_str,
39-
rmm::device_async_resource_ref mr)
39+
rmm::device_async_resource_ref mr,
40+
rmm::cuda_stream_view stream)
4041
{
41-
auto stream = cudf::get_default_stream();
4242
cudf::strings_column_view input_view(input);
43-
const auto byte_count = static_cast<cudf::size_type>(input_view.chars_size(stream));
43+
const auto byte_count = input_view.chars_size(stream);
4444

4545
if constexpr (MatchType == StringMatchingType::LIKE || MatchType == StringMatchingType::NOT_LIKE)
4646
{
4747
std::vector<std::string> match_terms = string_split(match_str, SPLIT_DELIMITER);
48-
auto result = DoMultiStringMatching(input.data<char>(),
48+
auto result = DoMultiStringMatching(input.data<char>(),
4949
input.size(),
50-
input.child(0).data<cudf::size_type>(),
50+
input.child(0).data<int64_t>(),
5151
byte_count,
5252
match_terms,
53-
mr);
53+
mr,
54+
stream);
5455
if constexpr (MatchType == StringMatchingType::LIKE)
5556
{
5657
return std::move(result);
@@ -63,61 +64,31 @@ GpuDispatcher::DispatchStringMatching(const cudf::column_view& input,
6364
{
6465
return DoStringMatching(input.data<char>(),
6566
input.size(),
66-
input.child(0).data<cudf::size_type>(),
67+
input.child(0).data<int64_t>(),
6768
byte_count,
6869
match_str,
69-
mr);
70+
mr,
71+
stream);
7072
}
7173
else if constexpr (MatchType == StringMatchingType::PREFIX)
7274
{
7375
return DoPrefixMatching(input.data<char>(),
7476
input.size(),
75-
input.child(0).data<cudf::size_type>(),
77+
input.child(0).data<int64_t>(),
7678
byte_count,
7779
match_str,
78-
mr);
80+
mr,
81+
stream);
7982
}
8083
}
8184

82-
std::tuple<uint64_t*, uint64_t> GpuDispatcher::DispatchSelect(const cudf::column_view& bitmap,
83-
rmm::device_async_resource_ref mr)
84-
{
85-
auto stream = cudf::get_default_stream();
86-
auto* gpu_buffer_manager = &GPUBufferManager::GetInstance();
87-
88-
// The row ids are owned by the query executor and so must be managed by the buffer manager
89-
uint64_t* row_ids = gpu_buffer_manager->customCudaMalloc<uint64_t>(bitmap.size(), 0, false);
90-
rmm::device_scalar<uint64_t> d_num_selected(0, stream, mr);
91-
92-
size_t temp_storage_bytes = 0;
93-
uint64_t num_selected = 0;
94-
cub::DeviceSelect::Flagged(nullptr,
95-
temp_storage_bytes,
96-
thrust::make_counting_iterator<uint64_t>(0),
97-
bitmap.data<bool>(),
98-
row_ids,
99-
d_num_selected.data(),
100-
bitmap.size(),
101-
cudf::get_default_stream());
102-
rmm::device_buffer temp_storage(temp_storage_bytes, stream, mr);
103-
cub::DeviceSelect::Flagged(temp_storage.data(),
104-
temp_storage_bytes,
105-
thrust::make_counting_iterator<uint64_t>(0),
106-
bitmap.data<bool>(),
107-
row_ids,
108-
d_num_selected.data(),
109-
bitmap.size(),
110-
cudf::get_default_stream());
111-
num_selected = d_num_selected.value(stream);
112-
return std::make_tuple(row_ids, num_selected);
113-
}
114-
11585
//----------Instantiations----------//
11686
#define INSTANTIATE_STR_MATCHING(T) \
11787
template std::unique_ptr<cudf::column> \
11888
GpuDispatcher::DispatchStringMatching<StringMatchingType::T>(const cudf::column_view& input, \
11989
const std::string& match_str, \
120-
rmm::device_async_resource_ref mr);
90+
rmm::device_async_resource_ref mr, \
91+
rmm::cuda_stream_view stream);
12192

12293
INSTANTIATE_STR_MATCHING(CONTAINS)
12394
INSTANTIATE_STR_MATCHING(LIKE)

0 commit comments

Comments
 (0)