Skip to content

Commit

Permalink
Migrated Dynamic AST Expression Trees in Benchmarks and Tests to use …
Browse files Browse the repository at this point in the history 
…AST Tree (#17697)

Follows-up and closes #17400.
This merge cleans up some of the ast tree expressions in CUDF's benchmarks and tests, specifically the dynamically constructed ones.
The other existing AST trees/expressions in the tests and benchmarks are linear and without complicated lifetimes.

Authors:
  - Basit Ayantunde (https://github.com/lamarrr)

Approvers:
  - Vyas Ramasubramani (https://github.com/vyasr)
  - Yunsong Wang (https://github.com/PointKernel)
  - Bradley Dice (https://github.com/bdice)

URL: #17697
  • Loading branch information
@lamarrr
lamarrr authored Jan 24, 2025
1 parent 0d2b29d commit c57cb6e
Show file tree
Hide file tree
Showing 3 changed files with 69 additions and 82 deletions.
78 changes: 34 additions & 44 deletions cpp/benchmarks/ast/transform.cpp
View file
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.
Expand Down Expand Up @@ -56,54 +56,46 @@ static void BM_ast_transform(nvbench::state& state)
auto const tree_levels = static_cast<cudf::size_type>(state.get_int64("tree_levels"));

// Create table data
auto const n_cols = reuse_columns ? 1 : tree_levels + 1;
auto const num_columns = reuse_columns ? 1 : tree_levels + 1;
auto const source_table =
create_sequence_table(cycle_dtypes({cudf::type_to_id<key_type>()}, n_cols),
create_sequence_table(cycle_dtypes({cudf::type_to_id<key_type>()}, num_columns),
row_count{num_rows},
Nullable ? std::optional<double>{0.5} : std::nullopt);
auto table = source_table->view();

cudf::ast::tree tree;

// Create column references
auto column_refs = std::vector<cudf::ast::column_reference>();
std::transform(thrust::make_counting_iterator(0),
thrust::make_counting_iterator(n_cols),
std::back_inserter(column_refs),
[](auto const& column_id) {
return cudf::ast::column_reference(reuse_columns ? 0 : column_id);
});
std::for_each(
thrust::make_counting_iterator(0),
thrust::make_counting_iterator(num_columns),
[&](int column_id) { tree.push(cudf::ast::column_reference(reuse_columns ? 0 : column_id)); });

// Create expression trees

// Note that a std::list is required here because of its guarantees against reference invalidation
// when items are added or removed. References to items in a std::vector are not safe if the
// vector must re-allocate.
auto expressions = std::list<cudf::ast::operation>();

// Construct tree that chains additions like (((a + b) + c) + d)
auto const op = cudf::ast::ast_operator::ADD;
if (reuse_columns) {
expressions.push_back(cudf::ast::operation(op, column_refs.at(0), column_refs.at(0)));
tree.push(cudf::ast::operation(op, tree.at(0), tree.at(0)));
for (cudf::size_type i = 0; i < tree_levels - 1; i++) {
expressions.push_back(cudf::ast::operation(op, expressions.back(), column_refs.at(0)));
tree.push(cudf::ast::operation(op, tree.back(), tree.at(0)));
}
} else {
expressions.push_back(cudf::ast::operation(op, column_refs.at(0), column_refs.at(1)));
std::transform(std::next(column_refs.cbegin(), 2),
column_refs.cend(),
std::back_inserter(expressions),
[&](auto const& column_ref) {
return cudf::ast::operation(op, expressions.back(), column_ref);
});
tree.push(cudf::ast::operation(op, tree.at(0), tree.at(1)));
std::for_each(
thrust::make_counting_iterator(2),
thrust::make_counting_iterator(num_columns),
[&](int col_id) { tree.push(cudf::ast::operation(op, tree.back(), tree.at(col_id))); });
}

auto const& expression_tree_root = expressions.back();
auto const& root_expression = tree.back();

// Use the number of bytes read from global memory
state.add_global_memory_reads<key_type>(static_cast<size_t>(num_rows) * (tree_levels + 1));
state.add_global_memory_writes<key_type>(num_rows);

state.exec(nvbench::exec_tag::sync,
[&](nvbench::launch&) { cudf::compute_column(table, expression_tree_root); });
[&](nvbench::launch&) { cudf::compute_column(table, root_expression); });
}

template <cudf::ast::ast_operator cmp_op, cudf::ast::ast_operator reduce_op>
Expand All @@ -117,10 +109,10 @@ static void BM_string_compare_ast_transform(nvbench::state& state)
CUDF_EXPECTS(tree_levels > 0, "benchmarks require 1 or more comparisons");

// Create table data
auto const num_cols = tree_levels * 2;
auto const num_columns = tree_levels * 2;
std::vector<std::unique_ptr<cudf::column>> columns;
std::for_each(
thrust::make_counting_iterator(0), thrust::make_counting_iterator(num_cols), [&](size_t) {
thrust::make_counting_iterator(0), thrust::make_counting_iterator(num_columns), [&](size_t) {
columns.emplace_back(create_string_column(num_rows, string_width, hit_rate));
});

Expand All @@ -135,38 +127,36 @@ static void BM_string_compare_ast_transform(nvbench::state& state)
return size + cudf::strings_column_view{column}.chars_size(cudf::get_default_stream());
});

// Create column references
auto column_refs = std::vector<cudf::ast::column_reference>();
std::transform(thrust::make_counting_iterator(0),
thrust::make_counting_iterator(num_cols),
std::back_inserter(column_refs),
[](auto const& column_id) { return cudf::ast::column_reference(column_id); });
// Create expression tree
cudf::ast::tree tree;

// Create expression trees
std::list<cudf::ast::operation> expressions;
// Create column references
std::for_each(thrust::make_counting_iterator(0),
thrust::make_counting_iterator(num_columns),
[&](int column_id) { tree.push(cudf::ast::column_reference{column_id}); });

// Construct AST tree (a == b && c == d && e == f && ...)

expressions.emplace_back(cudf::ast::operation(cmp_op, column_refs[0], column_refs[1]));
tree.push(cudf::ast::operation(cmp_op, tree[0], tree[1]));

std::for_each(thrust::make_counting_iterator(1),
thrust::make_counting_iterator(tree_levels),
[&](size_t idx) {
auto const& lhs = expressions.back();
auto const& rhs = expressions.emplace_back(
cudf::ast::operation(cmp_op, column_refs[idx * 2], column_refs[idx * 2 + 1]));
expressions.emplace_back(cudf::ast::operation(reduce_op, lhs, rhs));
auto const& lhs = tree.back();
auto const& rhs =
tree.push(cudf::ast::operation(cmp_op, tree[idx * 2], tree[idx * 2 + 1]));
tree.push(cudf::ast::operation(reduce_op, lhs, rhs));
});

auto const& expression_tree_root = expressions.back();

// Use the number of bytes read from global memory
state.add_element_count(chars_size, "chars_size");
state.add_global_memory_reads<nvbench::uint8_t>(chars_size);
state.add_global_memory_writes<nvbench::int32_t>(num_rows);

auto const& expression = tree.back();

state.exec(nvbench::exec_tag::sync,
[&](nvbench::launch&) { cudf::compute_column(table, expression_tree_root); });
[&](nvbench::launch&) { cudf::compute_column(table, expression); });
}

#define AST_TRANSFORM_BENCHMARK_DEFINE(name, key_type, tree_type, reuse_columns, nullable) \
Expand Down
17 changes: 8 additions & 9 deletions cpp/include/cudf/ast/expressions.hpp
View file
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.
Expand Down Expand Up @@ -556,7 +556,7 @@ class column_name_reference : public expression {

/**
* @brief An AST expression tree. It owns and contains multiple dependent expressions. All the
* expressions are destroyed when the tree is destructed.
* expressions are destroyed when the tree is destroyed.
*/
class tree {
public:
Expand Down Expand Up @@ -588,12 +588,11 @@ class tree {
* @returns a reference to the added expression
*/
template <typename Expr, typename... Args>
Expr const& emplace(Args&&... args)
std::enable_if_t<std::is_base_of_v<expression, Expr>, Expr const&> emplace(Args&&... args)
{
static_assert(std::is_base_of_v<expression, Expr>);
auto expr = std::make_shared<Expr>(std::forward<Args>(args)...);
auto expr = std::make_unique<Expr>(std::forward<Args>(args)...);
Expr const& expr_ref = *expr;
expressions.emplace_back(std::static_pointer_cast<expression>(std::move(expr)));
expressions.emplace_back(std::move(expr));
return expr_ref;
}

Expand All @@ -603,7 +602,7 @@ class tree {
* @returns a reference to the added expression
*/
template <typename Expr>
Expr const& push(Expr expr)
decltype(auto) push(Expr expr)
{
return emplace<Expr>(std::move(expr));
}
Expand Down Expand Up @@ -641,9 +640,9 @@ class tree {
expression const& operator[](size_t index) const { return *expressions[index]; }

private:
// TODO: use better ownership semantics, the shared_ptr here is redundant. Consider using a bump
// TODO: use better ownership semantics, the unique_ptr here is redundant. Consider using a bump
// allocator with type-erased deleters.
std::vector<std::shared_ptr<expression>> expressions;
std::vector<std::unique_ptr<expression>> expressions;
};

/** @} */ // end of group
Expand Down
56 changes: 27 additions & 29 deletions cpp/tests/ast/transform_tests.cpp
View file
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.
Expand Down Expand Up @@ -359,18 +359,15 @@ TEST_F(TransformTest, DeeplyNestedArithmeticLogicalExpression)
constexpr int64_t left_depth_level = 100;
constexpr int64_t right_depth_level = 75;

auto generate_ast_expr = [](int64_t depth_level,
cudf::ast::column_reference const& col_ref,
cudf::ast::ast_operator root_operator,
cudf::ast::ast_operator arithmetic_operator,
bool nested_left_tree) {
// Note that a std::list is required here because of its guarantees against reference
// invalidation when items are added or removed. References to items in a std::vector are not
// safe if the vector must re-allocate.
auto expressions = std::list<cudf::ast::operation>();
cudf::ast::tree tree;

auto generate_ast_expr = [&](int64_t depth_level,
cudf::ast::column_reference const& col_ref,
cudf::ast::ast_operator root_operator,
cudf::ast::ast_operator arithmetic_operator,
bool nested_left_tree) -> cudf::ast::expression const& {
auto op = arithmetic_operator;
expressions.emplace_back(op, col_ref, col_ref);
tree.push(cudf::ast::operation{op, col_ref, col_ref});

for (int64_t i = 0; i < depth_level - 1; i++) {
if (i == depth_level - 2) {
Expand All @@ -379,34 +376,35 @@ TEST_F(TransformTest, DeeplyNestedArithmeticLogicalExpression)
op = arithmetic_operator;
}
if (nested_left_tree) {
expressions.emplace_back(op, expressions.back(), col_ref);
tree.push(cudf::ast::operation{op, tree.back(), col_ref});
} else {
expressions.emplace_back(op, col_ref, expressions.back());
tree.push(cudf::ast::operation{op, col_ref, tree.back()});
}
}
return expressions;

return tree.back();
};

auto c_0 = column_wrapper<int64_t>{0, 0, 0};
auto c_1 = column_wrapper<int32_t>{0, 0, 0};
auto table = cudf::table_view{{c_0, c_1}};

auto col_ref_0 = cudf::ast::column_reference(0);
auto col_ref_1 = cudf::ast::column_reference(1);
auto const& col_ref_0 = tree.push(cudf::ast::column_reference(0));
auto const& col_ref_1 = tree.push(cudf::ast::column_reference(1));

auto left_expression = generate_ast_expr(left_depth_level,
col_ref_0,
cudf::ast::ast_operator::LESS,
cudf::ast::ast_operator::ADD,
false);
auto right_expression = generate_ast_expr(right_depth_level,
col_ref_1,
cudf::ast::ast_operator::EQUAL,
cudf::ast::ast_operator::SUB,
true);
auto const& left_expression = generate_ast_expr(left_depth_level,
col_ref_0,
cudf::ast::ast_operator::LESS,
cudf::ast::ast_operator::ADD,
false);
auto const& right_expression = generate_ast_expr(right_depth_level,
col_ref_1,
cudf::ast::ast_operator::EQUAL,
cudf::ast::ast_operator::SUB,
true);

auto expression_tree = cudf::ast::operation(
cudf::ast::ast_operator::LOGICAL_OR, left_expression.back(), right_expression.back());
auto const& expression = tree.push(
cudf::ast::operation(cudf::ast::ast_operator::LOGICAL_OR, left_expression, right_expression));

// Expression:
// OR(<(+(+(+(+($0, $0), $0), $0), $0), $0), ==($1, -($1, -($1, -($1, -($1, $1))))))
Expand All @@ -421,7 +419,7 @@ TEST_F(TransformTest, DeeplyNestedArithmeticLogicalExpression)
// If all $1 values and $R values are zeros, the result is true because of the equality check
// combined with the OR operator in OR(<($L, $0), ==($1, $R)).

auto result = cudf::compute_column(table, expression_tree);
auto result = cudf::compute_column(table, expression);
auto expected = column_wrapper<bool>{true, true, true};

CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), verbosity);
Expand Down

0 comments on commit c57cb6e

Please sign in to comment.