Skip to content

Commit 77e0486

Browse files
Yifei-yang7Yifei-yang7
committed
Merge remote-tracking branch 'origin/main' into issue/11_28
2 parents 4ed0f71 + ae6afcb commit 77e0486

10 files changed

Lines changed: 937 additions & 326 deletions

File tree

‎src/cuda/expression_executor/CMakeLists.txt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
set(CUDA_SOURCES
22
${CUDA_SOURCES}
3-
${CMAKE_CURRENT_SOURCE_DIR}/gpu_dispatcher.cu
3+
${CMAKE_CURRENT_SOURCE_DIR}/gpu_dispatch_string.cu
4+
${CMAKE_CURRENT_SOURCE_DIR}/gpu_dispatch_materialize.cu
45
PARENT_SCOPE
56
)

‎src/cuda/expression_executor/gpu_dispatch_materialize.cu‎

Lines changed: 215 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
#include "duckdb/common/exception.hpp"
2+
#include "expression_executor/gpu_dispatcher.hpp"
3+
#include "gpu_columns.hpp"
4+
#include <rmm/device_uvector.hpp>
5+
#include <thrust/for_each.h>
6+
#include <thrust/gather.h>
7+
#include <thrust/iterator/counting_iterator.h>
8+
#include <thrust/scan.h>
9+
#include <thrust/transform.h>
10+
11+
namespace duckdb
12+
{
13+
namespace sirius
14+
{
15+
16+
// Helper (template) functors to reduce bloat
17+
//----------Numerics----------//
18+
template <typename T>
19+
struct MaterializeNumeric
20+
{
21+
static std::unique_ptr<cudf::column> Do(const T* input_data,
22+
const uint64_t* row_ids,
23+
uint64_t row_id_count,
24+
rmm::device_async_resource_ref mr)
25+
{
26+
auto stream = cudf::get_default_stream();
27+
rmm::mr::thrust_allocator<uint8_t> thrust_allocator(stream, mr);
28+
auto exec = thrust::cuda::par(thrust_allocator).on(stream);
29+
30+
// Allocate output buffer
31+
rmm::device_uvector<T> output_data(row_id_count, stream, mr);
32+
33+
// Gather the input data
34+
thrust::gather(exec, row_ids, row_ids + row_id_count, input_data, output_data.begin());
35+
36+
// Construct and return a cudf::column
37+
return std::make_unique<cudf::column>(std::move(output_data),
38+
rmm::device_buffer(0, stream, mr),
39+
0);
40+
}
41+
};
42+
43+
//----------Strings----------//
44+
struct MaterializeString
45+
{
46+
// Functor for calculating the string lengths
47+
struct GatherLengths
48+
{
49+
// Fields
50+
const uint64_t* input_offsets;
51+
52+
// Constructor
53+
explicit GatherLengths(const uint64_t* input_offset)
54+
: input_offsets(input_offset)
55+
{}
56+
57+
// Operator
58+
__device__ __forceinline__ cudf::size_type operator()(uint64_t row_id) const
59+
{
60+
return static_cast<cudf::size_type>(input_offsets[row_id + 1] - input_offsets[row_id]);
61+
}
62+
};
63+
64+
// Functor for copying string data
65+
struct CopyData
66+
{
67+
// Fields
68+
const uint64_t* row_ids;
69+
const uint64_t* input_offsets;
70+
const uint8_t* input_data;
71+
cudf::size_type* output_offsets;
72+
uint8_t* output_data;
73+
74+
// Constructor
75+
CopyData(const uint64_t* row_ids,
76+
const uint64_t* input_offsets,
77+
const uint8_t* input_data,
78+
cudf::size_type* output_offsets,
79+
uint8_t* output_data)
80+
: row_ids(row_ids)
81+
, input_offsets(input_offsets)
82+
, input_data(input_data)
83+
, output_offsets(output_offsets)
84+
, output_data(output_data)
85+
{}
86+
87+
// Operator
88+
__device__ __forceinline__ void operator()(cudf::size_type output_idx)
89+
{
90+
const auto row_id = row_ids[output_idx];
91+
const auto output_offset = output_offsets[output_idx];
92+
const auto input_offset = input_offsets[row_id];
93+
const auto next_offset = input_offsets[row_id + 1];
94+
const auto length = next_offset - input_offset;
95+
96+
memcpy(output_data + output_offset, input_data + input_offset, length * sizeof(uint8_t));
97+
}
98+
};
99+
100+
static std::unique_ptr<cudf::column> Do(const uint8_t* input_data,
101+
const uint64_t* input_offsets,
102+
const uint64_t* row_ids,
103+
uint64_t row_id_count,
104+
rmm::device_async_resource_ref mr)
105+
{
106+
static_assert(std::is_same_v<int32_t, cudf::size_type>); // Sanity check
107+
108+
auto stream = cudf::get_default_stream();
109+
rmm::mr::thrust_allocator<uint8_t> thrust_allocator(stream, mr);
110+
auto exec = thrust::cuda::par(thrust_allocator).on(stream);
111+
uint64_t offset_count = row_id_count + 1;
112+
113+
// 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+
// 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+
120+
// Gather the string lengths
121+
thrust::transform(exec,
122+
row_ids,
123+
row_ids + row_id_count,
124+
temp_string_lengths.begin(),
125+
GatherLengths{input_offsets});
126+
127+
// Calculate the output offsets
128+
thrust::exclusive_scan(exec,
129+
temp_string_lengths.begin(),
130+
temp_string_lengths.end(),
131+
output_offsets.begin(),
132+
static_cast<cudf::size_type>(0));
133+
const auto output_bytes = output_offsets.back_element(stream);
134+
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+
142+
// Allocate output data buffer
143+
rmm::device_uvector<uint8_t> output_data(output_bytes, stream, mr);
144+
145+
// Gather the string data
146+
thrust::for_each_n(
147+
exec,
148+
thrust::counting_iterator<cudf::size_type>(0),
149+
row_id_count,
150+
CopyData(row_ids, input_offsets, input_data, output_offsets.data(), output_data.data()));
151+
152+
// Return a cudf::column
153+
auto offsets_col = std::make_unique<cudf::column>(cudf::data_type{cudf::type_id::INT32},
154+
static_cast<cudf::size_type>(offset_count),
155+
output_offsets.release(),
156+
rmm::device_buffer{0, stream, mr},
157+
0);
158+
return cudf::make_strings_column(static_cast<cudf::size_type>(row_id_count),
159+
std::move(offsets_col),
160+
output_data.release(),
161+
0,
162+
rmm::device_buffer{0, stream, mr});
163+
}
164+
};
165+
166+
// Materialize a GPUColumn directly into a cudf::column
167+
std::unique_ptr<cudf::column> GpuDispatcher::DispatchMaterialize(const GPUColumn* input,
168+
rmm::device_async_resource_ref mr)
169+
{
170+
D_ASSERT(input->row_ids != nullptr);
171+
D_ASSERT(input->row_ids->size() > 0);
172+
173+
const auto* input_data = input->data_wrapper.data;
174+
const auto* input_offsets = input->data_wrapper.offset; // Maybe unused
175+
176+
switch (input->data_wrapper.type)
177+
{
178+
case ColumnType::INT32:
179+
return MaterializeNumeric<int32_t>::Do(reinterpret_cast<const int32_t*>(input_data),
180+
input->row_ids,
181+
input->row_id_count,
182+
mr);
183+
case ColumnType::INT64:
184+
return MaterializeNumeric<uint64_t>::Do(reinterpret_cast<const uint64_t*>(input_data),
185+
input->row_ids,
186+
input->row_id_count,
187+
mr);
188+
case ColumnType::FLOAT32:
189+
return MaterializeNumeric<float_t>::Do(reinterpret_cast<const float_t*>(input_data),
190+
input->row_ids,
191+
input->row_id_count,
192+
mr);
193+
case ColumnType::FLOAT64:
194+
return MaterializeNumeric<double_t>::Do(reinterpret_cast<const double_t*>(input_data),
195+
input->row_ids,
196+
input->row_id_count,
197+
mr);
198+
case ColumnType::BOOLEAN:
199+
return MaterializeNumeric<bool>::Do(reinterpret_cast<const bool*>(input_data),
200+
input->row_ids,
201+
input->row_id_count,
202+
mr);
203+
case ColumnType::VARCHAR:
204+
return MaterializeString::Do(input_data,
205+
input_offsets,
206+
input->row_ids,
207+
input->row_id_count,
208+
mr);
209+
default:
210+
throw InternalException("Dispatch[Materialize]: Unsupported column type!");
211+
}
212+
}
213+
214+
} // namespace sirius
215+
} // namespace duckdb

‎src/cuda/expression_executor/gpu_dispatch_string.cu‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
#include "expression_executor/gpu_dispatcher.hpp"
2+
#include "gpu_buffer_manager.hpp"
3+
#include "gpu_physical_strings_matching.hpp"
4+
#include "gpu_physical_substring.hpp"
5+
#include <cub/cub.cuh>
6+
#include <cudf/column/column.hpp>
7+
#include <cudf/strings/strings_column_view.hpp>
8+
#include <cudf/unary.hpp>
9+
#include <memory>
10+
#include <thrust/iterator/counting_iterator.h>
11+
#include <tuple>
12+
13+
namespace duckdb
14+
{
15+
namespace sirius
16+
{
17+
18+
#define SPLIT_DELIMITER "%"
19+
20+
//----------Substring----------//
21+
std::unique_ptr<cudf::column> GpuDispatcher::DispatchSubstring(const cudf::column_view& input,
22+
uint64_t start_idx,
23+
uint64_t len,
24+
rmm::device_async_resource_ref mr)
25+
{
26+
return DoSubstring(input.data<char>(),
27+
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);
32+
};
33+
34+
//----------String Matching----------//
35+
template <StringMatchingType MatchType>
36+
std::unique_ptr<cudf::column>
37+
GpuDispatcher::DispatchStringMatching(const cudf::column_view& input,
38+
const std::string& match_str,
39+
rmm::device_async_resource_ref mr)
40+
{
41+
auto stream = cudf::get_default_stream();
42+
cudf::strings_column_view input_view(input);
43+
const auto byte_count = static_cast<cudf::size_type>(input_view.chars_size(stream));
44+
45+
if constexpr (MatchType == StringMatchingType::LIKE || MatchType == StringMatchingType::NOT_LIKE)
46+
{
47+
std::vector<std::string> match_terms = string_split(match_str, SPLIT_DELIMITER);
48+
auto result = DoMultiStringMatching(input.data<char>(),
49+
input.size(),
50+
input.child(0).data<cudf::size_type>(),
51+
byte_count,
52+
match_terms,
53+
mr);
54+
if constexpr (MatchType == StringMatchingType::LIKE)
55+
{
56+
return std::move(result);
57+
}
58+
59+
// Otherwise, we need to invert the result
60+
return cudf::unary_operation(result->view(), cudf::unary_operator::NOT, stream, mr);
61+
}
62+
else if constexpr (MatchType == StringMatchingType::CONTAINS)
63+
{
64+
return DoStringMatching(input.data<char>(),
65+
input.size(),
66+
input.child(0).data<cudf::size_type>(),
67+
byte_count,
68+
match_str,
69+
mr);
70+
}
71+
else if constexpr (MatchType == StringMatchingType::PREFIX)
72+
{
73+
return DoPrefixMatching(input.data<char>(),
74+
input.size(),
75+
input.child(0).data<cudf::size_type>(),
76+
byte_count,
77+
match_str,
78+
mr);
79+
}
80+
}
81+
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+
115+
//----------Instantiations----------//
116+
#define INSTANTIATE_STR_MATCHING(T) \
117+
template std::unique_ptr<cudf::column> \
118+
GpuDispatcher::DispatchStringMatching<StringMatchingType::T>(const cudf::column_view& input, \
119+
const std::string& match_str, \
120+
rmm::device_async_resource_ref mr);
121+
122+
INSTANTIATE_STR_MATCHING(CONTAINS)
123+
INSTANTIATE_STR_MATCHING(LIKE)
124+
INSTANTIATE_STR_MATCHING(NOT_LIKE)
125+
INSTANTIATE_STR_MATCHING(PREFIX)
126+
127+
#undef SPLIT_DELIMITER
128+
#undef INSTANTIATE_STR_MATCHING
129+
130+
} // namespace sirius
131+
} // namespace duckdb

0 commit comments

Comments
 (0)