Skip to content

Commit

Permalink
Track GPU memory
Browse files Browse the repository at this point in the history
  • Loading branch information
@JayjeetAtGithub
JayjeetAtGithub committed Jul 31, 2024
1 parent 4344fac commit 25173ee
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 31 deletions.
48 changes: 30 additions & 18 deletions cpp/benchmarks/common/cudf_datagen/dbgen.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -171,13 +171,14 @@ std::unique_ptr<cudf::table> generate_lineitem_partial(
auto const o_orderkey = orders_independent.column(1);
auto const o_orderdate_ts = orders_independent.column(3);

auto const left_table = cudf::table_view({l_pkey->view()});
auto const right_table = cudf::table_view({o_pkey, o_orderkey, o_orderdate_ts});
auto const l_base_unsorted = perform_left_join(left_table, right_table, {0}, {0});
auto const l_base = cudf::sort_by_key(l_base_unsorted->view(),
auto const left_table = cudf::table_view({l_pkey->view()});
auto const right_table = cudf::table_view({o_pkey, o_orderkey, o_orderdate_ts});
auto const l_base_unsorted =
perform_left_join(left_table, right_table, {0}, {0}, cudf::null_equality::EQUAL, stream, mr);
auto const l_base = cudf::sort_by_key(l_base_unsorted->view(),
cudf::table_view({l_base_unsorted->get_column(2).view()}),
{},
{},
{},
{},
stream,
mr);

Expand Down Expand Up @@ -569,17 +570,24 @@ auto generate_orders_lineitem_part(

// Join the `part` and partial `lineitem` tables, then calculate the `l_extendedprice` column,
// add the column to the `lineitem` table, and write the `lineitem` table to a parquet file
auto lineitem_joined_part = perform_left_join(lineitem_partial->view(), part->view(), {2}, {0});
auto const l_quantity = lineitem_joined_part->get_column(5);
auto const l_quantity_fp = cudf::cast(l_quantity.view(), cudf::data_type{cudf::type_id::FLOAT64});
auto const p_retailprice = lineitem_joined_part->get_column(23);

auto l_extendedprice = cudf::binary_operation(l_quantity_fp->view(),
p_retailprice.view(),
cudf::binary_operator::MUL,
cudf::data_type{cudf::type_id::FLOAT64},
stream,
mr);

auto l_extendedprice = [&]() {
auto const left = cudf::table_view(
{lineitem_partial->get_column(2).view(), lineitem_partial->get_column(5).view()});
auto const right = cudf::table_view({part->get_column(0).view(), part->get_column(7).view()});
auto joined_table =
perform_left_join(left, right, {0}, {0}, cudf::null_equality::EQUAL, stream, mr);
auto const l_quantity = joined_table->get_column(1);
auto const l_quantity_fp =
cudf::cast(l_quantity.view(), cudf::data_type{cudf::type_id::FLOAT64});
auto const p_retailprice = joined_table->get_column(3);
return cudf::binary_operation(l_quantity_fp->view(),
p_retailprice.view(),
cudf::binary_operator::MUL,
cudf::data_type{cudf::type_id::FLOAT64},
stream,
mr);
}();

auto lineitem_columns = lineitem_partial->release();
lineitem_columns.push_back(std::move(l_extendedprice));
Expand Down Expand Up @@ -830,6 +838,10 @@ int main(int argc, char** argv)
auto resource = create_memory_resource(memory_resource_type);
rmm::mr::set_current_device_resource(resource.get());

auto const [free, total] = rmm::available_device_memory();
std::cout << "Total GPU memory: " << total << std::endl;
std::cout << "Available GPU memory: " << free << std::endl;

auto const mem_stats_logger = memory_stats_logger();

auto [orders, lineitem, part] = generate_orders_lineitem_part(scale_factor);
Expand All @@ -852,7 +864,7 @@ int main(int argc, char** argv)
auto region = generate_region();
write_parquet(std::move(region), "region.parquet", schema_region);

std::cout << "Peak Memory Usage: " << mem_stats_logger.peak_memory_usage() << std::endl;
std::cout << "Peak GPU Memory Usage: " << mem_stats_logger.peak_memory_usage() << std::endl;

return 0;
}
27 changes: 14 additions & 13 deletions cpp/benchmarks/common/cudf_datagen/utils.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -121,28 +121,29 @@ void write_parquet(std::unique_ptr<cudf::table> tbl,
cudf::io::write_parquet(options);
}

std::unique_ptr<cudf::table> perform_left_join(
cudf::table_view const& left_input,
cudf::table_view const& right_input,
std::vector<cudf::size_type> const& left_on,
std::vector<cudf::size_type> const& right_on,
cudf::null_equality compare_nulls = cudf::null_equality::EQUAL)
std::unique_ptr<cudf::table> perform_left_join(cudf::table_view const& left_input,
cudf::table_view const& right_input,
std::vector<cudf::size_type> const& left_on,
std::vector<cudf::size_type> const& right_on,
cudf::null_equality compare_nulls,
rmm::cuda_stream_view stream,
rmm::device_async_resource_ref mr)
{
CUDF_FUNC_RANGE();
constexpr auto oob_policy = cudf::out_of_bounds_policy::NULLIFY;
auto const left_selected = left_input.select(left_on);
auto const right_selected = right_input.select(right_on);
auto const [left_join_indices, right_join_indices] = cudf::left_join(
left_selected, right_selected, compare_nulls, rmm::mr::get_current_device_resource());
constexpr auto oob_policy = cudf::out_of_bounds_policy::NULLIFY;
auto const left_selected = left_input.select(left_on);
auto const right_selected = right_input.select(right_on);
auto const [left_join_indices, right_join_indices] =
cudf::left_join(left_selected, right_selected, compare_nulls, mr);

auto const left_indices_span = cudf::device_span<cudf::size_type const>{*left_join_indices};
auto const right_indices_span = cudf::device_span<cudf::size_type const>{*right_join_indices};

auto const left_indices_col = cudf::column_view{left_indices_span};
auto const right_indices_col = cudf::column_view{right_indices_span};

auto const left_result = cudf::gather(left_input, left_indices_col, oob_policy);
auto const right_result = cudf::gather(right_input, right_indices_col, oob_policy);
auto const left_result = cudf::gather(left_input, left_indices_col, oob_policy, stream, mr);
auto const right_result = cudf::gather(right_input, right_indices_col, oob_policy, stream, mr);

auto joined_cols = left_result->release();
auto right_cols = right_result->release();
Expand Down

0 comments on commit 25173ee

Please sign in to comment.