diff --git a/cpp/include/cudf/transform.hpp b/cpp/include/cudf/transform.hpp index 2be1a88a9020..805c3d844e07 100644 --- a/cpp/include/cudf/transform.hpp +++ b/cpp/include/cudf/transform.hpp @@ -11,8 +11,10 @@ #include #include +#include #include #include +#include #include #include @@ -285,8 +287,6 @@ std::unique_ptr compute_column( * transform. * * @throws cudf::logic_error if passed an expression operating on table_reference::RIGHT. - * @throws cudf::data_type_error if the expression applies a non-comparison binary operator to - * decimal128 operands. * @throws cudf::evaluation_error if the evaluation of the expression results in an error during * execution. * @@ -302,6 +302,34 @@ std::unique_ptr compute_column_jit( rmm::cuda_stream_view stream = cudf::get_default_stream(), rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); +/** + * @brief Compute a new table by evaluating expression trees on a table using a JIT-compiled + * kernel. + * + * This evaluates expressions over a table to produce a new table. Also called an n-ary + * transform. Expressions are evaluated in the order supplied, and output column `i` contains the + * result of `expressions[i]`. Common subexpressions shared by multiple outputs are evaluated once + * by the generated function. + * + * @pre `expressions` must not be empty. + * + * @throws cudf::logic_error if passed an empty collection of expressions. + * @throws cudf::logic_error if passed an expression operating on table_reference::RIGHT. + * @throws cudf::evaluation_error if the evaluation of the expression results in an error during + * execution. + * + * @param table The table used for expression evaluation + * @param expressions Non-empty collection of expression-tree roots, one per output column + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource + * @return Table containing one output column per expression, in the same order as `expressions` + */ +std::unique_ptr compute_table_jit( + table_view const& table, + std::span const> expressions, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + /** * @brief Creates a bitmask from a column of boolean elements. * diff --git a/cpp/src/jit/row_ir.cpp b/cpp/src/jit/row_ir.cpp index 74d4f581a243..7b888feb3a38 100644 --- a/cpp/src/jit/row_ir.cpp +++ b/cpp/src/jit/row_ir.cpp @@ -8,10 +8,12 @@ #include "runtime/context.hpp" #include +#include #include #include +#include #include #include #include @@ -437,6 +439,21 @@ int32_t instance_context::add_output() int32_t instance_context::add_input(input in) { + if (auto* column = std::get_if(&in); + column != nullptr && column->table_source.has_value() && column->column_index.has_value()) { + for (size_t i = 0; i < inputs_.size(); ++i) { + auto* existing = std::get_if(&inputs_[i]); + if (existing != nullptr && existing->table_source == column->table_source && + existing->column_index == column->column_index) { + return static_cast(i); + } + } + } + + // TODO: deduplication for scalar inputs once they use scalar_column_view instead of a column type + // Alternatively, use an inlined literal value type that is host-device accessible and hashable + // + auto id = static_cast(inputs_.size()); auto id_str = std::format("in_{}", id); @@ -453,6 +470,19 @@ int32_t instance_context::add_input(input in) return id; } +node const* instance_context::find_equivalent(node const& candidate) const +{ + auto [first, last] = cse_nodes_.equal_range(candidate.hash()); + auto found = + std::find_if(first, last, [&](auto& entry) { return candidate.is_equivalent(*entry.second); }); + return found == last ? nullptr : found->second; +} + +void instance_context::add_cse_node(node const& candidate) +{ + cse_nodes_.emplace(candidate.hash(), &candidate); +} + std::string instance_context::make_tmp_id() { return std::format("{}{}", tmp_prefix_, num_tmp_vars_++); @@ -468,6 +498,27 @@ std::span instance_context::get_input_vars() const { return inpu std::span instance_context::get_output_vars() const { return output_vars_; } +size_t node::compute_hash() const +{ + auto h = std::hash{}(static_cast(op_)); + h = cudf::hashing::detail::hash_combine(h, std::hash{}(static_cast(error_policy_))); + h = cudf::hashing::detail::hash_combine(h, std::hash{}(reference_.index())); + std::visit( + [&](auto& r) { + using reference_type = std::decay_t; + if constexpr (!std::is_same_v) { + h = cudf::hashing::detail::hash_combine(h, std::hash{}(r.index)); + } + }, + reference_); + h = cudf::hashing::detail::hash_combine( + h, target_scale_.has_value() ? std::hash{}(*target_scale_) : 0); + for (auto& arg : args_) { + h = cudf::hashing::detail::hash_combine(h, arg->hash()); + } + return h; +} + node::node(opcode op, std::optional target_scale, error_policy error_policy, @@ -505,22 +556,40 @@ node::node(opcode op, std::format("Target scale must be provided for RESCALE operator and must be nullopt " "for other operators.")); } + + hash_ = compute_hash(); } node::node(input_reference input) : reference_{input}, op_{opcode::GET_INPUT} // NOLINT(modernize-use-default-member-init) { + hash_ = compute_hash(); } node::node(output_reference reference, std::unique_ptr arg) : reference_{reference}, op_{opcode::SET_OUTPUT} { args_.emplace_back(std::move(arg)); + hash_ = compute_hash(); } node::node(output_reference reference, node arg) : node{reference, std::make_unique(std::move(arg))} { + hash_ = compute_hash(); +} + +size_t node::hash() const { return hash_; } + +bool node::is_equivalent(node const& other) const +{ + return hash_ == other.hash_ && op_ == other.op_ && target_scale_ == other.target_scale_ && + error_policy_ == other.error_policy_ && reference_ == other.reference_ && + args_.size() == other.args_.size() && + std::equal( + args_.begin(), args_.end(), other.args_.begin(), [](auto const& lhs, auto const& rhs) { + return lhs->is_equivalent(*rhs); + }); } std::string_view node::get_id() const { return id_; } @@ -589,6 +658,16 @@ void node::instantiate(instance_context& ctx) arg->instantiate(ctx); } + // check if an equivalent node has already been instantiated in this context. If so, alias the + // node. + if (auto equivalent = ctx.find_equivalent(*this)) { + id_ = equivalent->id_; + type_ = equivalent->type_; + scale_reference_ = equivalent->scale_reference_; + alias_ = equivalent; + return; + } + id_ = ctx.make_tmp_id(); switch (op_) { @@ -613,14 +692,24 @@ void node::instantiate(instance_context& ctx) type_ = get_return_type(op_, arg_types, target_scale_); } break; } + + // add this node to the context's CSE map so that future equivalent nodes can reuse it. + ctx.add_cse_node(*this); } -void node::emit_code(instance_context& instance, target_info const& info, code_sink& sink) const +void node::emit_code(instance_context& instance, target_info const& info, code_sink& sink) { for (auto& arg : args_) { arg->emit_code(instance, info, sink); } + if (alias_ != nullptr) { + CUDF_EXPECTS(alias_->emitted_, + "Alias node has not been emitted yet. This should never happen.", + std::runtime_error); + return; + } + switch (info.id) { case target::CUDA: { auto type = to_cuda_type(type_, instance.has_nulls()); @@ -717,6 +806,8 @@ if(expected__{1}.has_value()) {{ CUDF_FAIL(std::format("Unsupported target: {}", static_cast(info.id)), std::invalid_argument); } + + emitted_ = true; } std::unique_ptr ast_converter::add_ir_node(ast::literal const& expr) @@ -776,13 +867,18 @@ bool is_nullable(scalar_input const& in) { return in.scalar_column->view().nulla bool is_nullable(column_input const& in) { return in.column.nullable(); } -std::tuple ast_converter::generate_code( - target target_id, ast::expression const& expr, std::string_view function_name) +std::tuple> ast_converter::generate_code( + target target_id, + std::span const> expressions, + std::string_view function_name) { - // add 1 auto-deduced output variable - [[maybe_unused]] auto output_id = instance_.add_output(); + CUDF_EXPECTS(!expressions.empty(), "At least one output expression is required"); - output_irs_.emplace_back(std::make_unique(output_reference{0}, expr.accept(*this))); + for (auto& expression : expressions) { + auto output_id = instance_.add_output(); + output_irs_.emplace_back( + std::make_unique(output_reference{output_id}, expression.get().accept(*this))); + } bool has_nullable_inputs = std::any_of(instance_.inputs_.begin(), instance_.inputs_.end(), [&](auto& in) { @@ -792,15 +888,22 @@ std::tuple ast_converter::generate_ bool is_null_aware = std::any_of( output_irs_.cbegin(), output_irs_.cend(), [](auto& ir) { return ir->is_null_aware(); }); - bool output_is_always_valid = std::all_of( - output_irs_.cbegin(), output_irs_.cend(), [](auto& ir) { return ir->is_always_valid(); }); - - bool may_evaluate_null = output_is_always_valid ? false : (has_nullable_inputs || is_null_aware); + std::vector null_policies; + std::transform( + output_irs_.cbegin(), output_irs_.cend(), std::back_inserter(null_policies), [&](auto& ir) { + auto may_evaluate_null = + !ir->is_always_valid() && (has_nullable_inputs || ir->is_null_aware()); + return may_evaluate_null ? output_nullability::PRESERVE : output_nullability::ALL_VALID; + }); - auto null_policy = - may_evaluate_null ? output_nullability::PRESERVE : output_nullability::ALL_VALID; + // In a multi-output UDF, if any input is nullable, we need to generate a null mask for each + // output. + // Instead of generating a single null mask or multiple null masks, we make the UDF + // null-aware and let the UDF handle null propagation for each output. + auto needs_per_output_nullmask = output_irs_.size() > 1 && has_nullable_inputs; + auto generate_null_aware_udf = is_null_aware || needs_per_output_nullmask; - instance_.set_has_nulls(is_null_aware); + instance_.set_has_nulls(generate_null_aware_udf); // instantiate the IR nodes for (auto& ir : output_irs_) { @@ -854,7 +957,8 @@ std::tuple ast_converter::generate_ ir->emit_code(instance_, target, sink); } sink.emit("return cudf::errc::SUCCESS;\n}"); - return {sink.get_code(), is_null_aware ? null_aware::YES : null_aware::NO, null_policy}; + return { + sink.get_code(), generate_null_aware_udf ? null_aware::YES : null_aware::NO, null_policies}; } std::variant get_column_view(scalar_input const& in) @@ -869,21 +973,19 @@ std::variant get_column_view(column_input const // Due to the AST expression tree structure, we can't generate the IR without the target // tables -transform_args ast_converter::compute_column(target target_id, - ast::expression const& expr, - table_view const& left_table, - table_view const& right_table, - std::string_view function_name, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) +transform_args ast_converter::compute_table( + target target_id, + std::span const> expressions, + table_view const& left_table, + table_view const& right_table, + std::string_view function_name, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { ast_converter converter{stream, mr, left_table, right_table}; - // TODO(lamarrr): consider deduplicating ast expression's input column references. See - // TransformTest/1.DeeplyNestedArithmeticLogicalExpression for reference - - auto [code, is_null_aware, output_nullability] = - converter.generate_code(target_id, expr, function_name); + auto [code, is_null_aware, output_nullabilities] = + converter.generate_code(target_id, expressions, function_name); std::vector> inputs; std::vector> scalar_columns; std::vector> table_sources; @@ -908,9 +1010,11 @@ transform_args ast_converter::compute_column(target target_id, } } - auto& out = converter.output_irs_[0]; - auto output_column_type = out->get_type(); - auto output = transform_output{.type = output_column_type, .nullability = output_nullability}; + std::vector outputs; + for (size_t i = 0; i < converter.output_irs_.size(); ++i) { + outputs.push_back(transform_output{.type = converter.output_irs_[i]->get_type(), + .nullability = output_nullabilities[i]}); + } auto row_size = std::max({left_table.num_rows(), right_table.num_rows()}); auto result = transform_args{.scalar_columns = std::move(scalar_columns), .input_table_sources = std::move(table_sources), @@ -920,7 +1024,7 @@ transform_args ast_converter::compute_column(target target_id, .is_null_aware = is_null_aware, .user_data = std::nullopt, .inputs = inputs, - .outputs{output}, + .outputs = std::move(outputs), .string_offsets{}, .row_size = row_size}; if (get_context().dump_codegen()) { @@ -939,7 +1043,8 @@ transform_args ast_converter::filter(target target_id, rmm::device_async_resource_ref mr) { auto filter = ast::detail::predicate{expr}; - return compute_column(target_id, filter, left_table, right_table, function_name, stream, mr); + std::array, 1> expressions{filter}; + return compute_table(target_id, expressions, left_table, right_table, function_name, stream, mr); } } // namespace cudf::detail::row_ir diff --git a/cpp/src/jit/row_ir.hpp b/cpp/src/jit/row_ir.hpp index d54ac08b42cd..96e338fa9105 100644 --- a/cpp/src/jit/row_ir.hpp +++ b/cpp/src/jit/row_ir.hpp @@ -21,12 +21,15 @@ #include #include +#include #include #include #include #include #include #include +#include +#include #include #include @@ -98,6 +101,8 @@ struct [[nodiscard]] transform_args { std::optional row_size = std::nullopt; }; +struct node; + /** * @brief The context within which the IR is instantiated. * This context is used to generate temporary variable identifiers and any state setup needed for @@ -111,6 +116,7 @@ struct [[nodiscard]] instance_context { std::vector inputs_; ///< The inputs for the IR std::vector input_vars_; ///< The input variables for the IR std::vector output_vars_; ///< The output variables for the IR + std::unordered_multimap cse_nodes_; ///< multimap of IR nodes rmm::cuda_stream_view stream_; ///< The CUDA stream for any device operations during IR generation rmm::device_async_resource_ref @@ -135,8 +141,21 @@ struct [[nodiscard]] instance_context { ~instance_context() = default; ///< Destructor + /** + * @brief Adds an output variable to the generated transform. + * + * @return Index of the newly added output + */ [[nodiscard]] int32_t add_output(); + /** + * @brief Adds an input to the generated transform. + * + * Column inputs that identify the same source table and column reuse the existing input. + * + * @param in Input to add + * @return Index of the new or reused input + */ [[nodiscard]] int32_t add_input(input in); [[nodiscard]] int32_t add_input(scalar const& scalar) @@ -150,6 +169,21 @@ struct [[nodiscard]] instance_context { return add_input(column_input{.column = column}); } + /** + * @brief Finds a structurally equivalent node belonging to a previously completed output. + * + * @param candidate Node for which to find an equivalent common subexpression + * @return Equivalent node, or `nullptr` if none exists + */ + [[nodiscard]] node const* find_equivalent(node const& candidate) const; + + /** + * @brief Stages a newly instantiated node for registration after its output is completed. + * + * @param candidate Newly instantiated node + */ + void add_cse_node(node const& candidate); + /** * @brief Generate a globally unique temporary variable identifier * @return A unique temporary variable identifier @@ -210,27 +244,54 @@ struct [[nodiscard]] code_sink { struct [[nodiscard]] input_reference { int32_t index = 0; ///< The index of the input variable + + constexpr bool operator==(input_reference const& other) const { return index == other.index; } + + constexpr bool operator!=(input_reference const& other) const { return index != other.index; } }; struct [[nodiscard]] output_reference { int32_t index = 0; ///< The index of the output variable + + constexpr bool operator==(output_reference const& other) const { return index == other.index; } + + constexpr bool operator!=(output_reference const& other) const { return index != other.index; } }; struct [[nodiscard]] node { private: std::variant reference_ = std::monostate{}; ///< The index of the input/output variable - opcode op_ = opcode::GET_INPUT; ///< The operation code - std::optional target_scale_ = std::nullopt; ///< The target scale for decimal - input_reference scale_reference_ = {}; ///< The index of the target scale as an IR input + + opcode op_ = opcode::GET_INPUT; ///< The operation code + + std::optional target_scale_ = std::nullopt; ///< The target scale for decimal + error_policy error_policy_ = - cudf::error_policy::PROPAGATE; ///< The error policy for the operation + cudf::error_policy::PROPAGATE; ///< The error policy for the operation + std::vector> args_ = {}; ///< The arguments of the operation - data_type type_ = {}; ///< The resolved type information of the IR node + size_t hash_ = 0; ///< The structural hash of the IR node std::string id_ = {}; ///< The identifier of the IR node + data_type type_ = {}; ///< The resolved type information of the IR node + + input_reference scale_reference_ = {}; ///< The index of the target scale as an IR input + + bool emitted_ = false; ///< Whether the IR node has been emitted into the generated function + + node const* alias_ = nullptr; ///< The equivalent IR node that this IR aliases, if any. This is + ///< used to avoid emitting duplicate code for equivalent IR nodes. + + /** + * @brief Computes the structural hash of this node and its arguments. + * + * @return The structural hash + */ + [[nodiscard]] size_t compute_hash() const; + /** * @brief Create a set of argument IR nodes */ @@ -307,14 +368,14 @@ struct [[nodiscard]] node { /** * @brief Construct a new output reference IR node - * @param output The index of the output variable + * @param reference The output variable reference * @param arg The argument node that produces the value to be set to the output variable */ node(output_reference reference, std::unique_ptr arg); /** * @brief Construct a new output reference IR node - * @param output The index of the output variable + * @param reference The output variable reference * @param arg The argument node that produces the value to be set to the output variable */ node(output_reference reference, node arg); @@ -325,6 +386,21 @@ struct [[nodiscard]] node { node& operator=(node&& other) = default; ///< Move assignment operator ~node() = default; ///< Destructor + /** + * @brief Gets the cached structural hash of the IR node. + * + * @return The structural hash + */ + [[nodiscard]] size_t hash() const; + + /** + * @brief Checks whether another node represents the same expression. + * + * @param other Node to compare with this node + * @return `true` if the nodes are structurally equivalent + */ + [[nodiscard]] bool is_equivalent(node const& other) const; + /** * @brief Get the identifier of the IR node * @return The identifier of the IR node @@ -373,7 +449,6 @@ struct [[nodiscard]] node { * @brief Instantiate the IR node with the given context and instance information, setting up any * necessary state and preprocessing needed for code generation. * @param ctx The context within which the IR is instantiated - * @param info The instance information */ void instantiate(instance_context& ctx); @@ -381,10 +456,9 @@ struct [[nodiscard]] node { * @brief Generate the code for the IR node based on the instance context and target information. * @param ctx The context within which the IR is instantiated * @param info The target information - * @param instance The instance information * @param sink The code sink to which the generated code is emitted */ - void emit_code(instance_context& ctx, target_info const& info, code_sink& sink) const; + void emit_code(instance_context& ctx, target_info const& info, code_sink& sink); }; /** @@ -406,6 +480,8 @@ struct [[nodiscard]] ast_converter { * @brief Construct a new AST Converter object * @param stream CUDA stream used for device memory operations and kernel launches. * @param mr Device memory resource used to allocate the returned table's device memory + * @param left_table Left input table referenced by expressions + * @param right_table Right input table referenced by expressions */ ast_converter(rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr, @@ -439,28 +515,40 @@ struct [[nodiscard]] ast_converter { [[nodiscard]] std::unique_ptr add_ir_node(ast::jit::detail::operation const& expr); - [[nodiscard]] std::tuple generate_code( - target target, ast::expression const& expr, std::string_view function_name); + /** + * @brief Converts multiple AST expressions into one generated transform function. + * + * @pre `expressions` is not empty + * + * @param target Code generation target + * @param expressions AST expressions, one for each output column + * @param function_name Name of the generated transform function + * @return Generated source, function null-awareness, and nullability policy for each output + */ + [[nodiscard]] std::tuple> generate_code( + target target, + std::span const> expressions, + std::string_view function_name); /** - * @brief Convert an AST `compute_column` expression to a `cudf::transform` + * @brief Converts AST expressions to arguments for a multi-output `cudf::transform`. * @param target The target for which the IR is generated - * @param expr The AST expression to convert + * @param expressions AST expressions, one for each output column * @param left_table The left input table for the expression * @param right_table The right input table for the expression - * @param table The input table for the expression * @param function_name The name of the generated function * @param stream CUDA stream used for device memory operations and kernel launches. * @param mr Device memory resource used to allocate the returned table's device memory * @return The result of the conversion, containing the transform arguments and scalar columns */ - static transform_args compute_column(target target, - ast::expression const& expr, - table_view const& left_table, - table_view const& right_table, - std::string_view function_name, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr); + static transform_args compute_table( + target target, + std::span const> expressions, + table_view const& left_table, + table_view const& right_table, + std::string_view function_name, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); /** * @brief Convert an AST `filter` expression to a `cudf::filter` @@ -468,7 +556,6 @@ struct [[nodiscard]] ast_converter { * @param expr The AST expression to convert * @param left_table The left input table for the expression * @param right_table The right input table for the expression - * @param table The input table for the expression * @param function_name The name of the generated function * @param stream CUDA stream used for device memory operations and kernel launches. * @param mr Device memory resource used to allocate the returned table's device memory diff --git a/cpp/src/transform/transform.cu b/cpp/src/transform/transform.cu index 9e08a0950557..12fb36540a3a 100644 --- a/cpp/src/transform/transform.cu +++ b/cpp/src/transform/transform.cu @@ -33,6 +33,7 @@ #include #include +#include #include #include #include @@ -1164,8 +1165,9 @@ std::unique_ptr compute_column_jit(table_view const& table, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { - auto args = detail::row_ir::ast_converter::compute_column( - detail::row_ir::target::CUDA, expr, table, {}, "compute_operation", stream, mr); + std::array, 1> expressions{expr}; + auto args = detail::row_ir::ast_converter::compute_table( + detail::row_ir::target::CUDA, expressions, table, {}, "compute_operation", stream, mr); auto result = transform(args.udf, args.source_type, args.is_null_aware, @@ -1180,6 +1182,26 @@ std::unique_ptr compute_column_jit(table_view const& table, return std::move(cols[0]); } +std::unique_ptr
compute_table_jit( + table_view const& table, + std::span const> expressions, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + auto args = detail::row_ir::ast_converter::compute_table( + detail::row_ir::target::CUDA, expressions, table, {}, "compute_operation", stream, mr); + return transform(args.udf, + args.source_type, + args.is_null_aware, + args.user_data, + args.inputs, + args.outputs, + std::move(args.string_offsets), + args.row_size, + stream, + mr); +} + // if we have a matching pre-compiled kernel fragment for the given transform configuration, return // it to use for LTO linking instead of compiling a new one std::optional, lto_binary_type, char const*>> diff --git a/cpp/tests/ast/transform_tests.cpp b/cpp/tests/ast/transform_tests.cpp index fe4960b9dfec..f2a591062e3a 100644 --- a/cpp/tests/ast/transform_tests.cpp +++ b/cpp/tests/ast/transform_tests.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include @@ -25,8 +26,11 @@ #include #include +#include +#include #include #include +#include #include // NOTE: each test in this file must be run twice - once with the AST Interpreter executor @@ -1558,4 +1562,107 @@ TYPED_TEST(TransformTest, Decimal128IdentityOutput) } } +struct ComputeTableJitTest : public cudf::test::BaseFixture {}; + +TEST_F(ComputeTableJitTest, CommonSubexpression) +{ + auto c0 = column_wrapper{1, 2, 3, 4}; + auto c1 = column_wrapper{10, 20, 30, 40}; + auto c2 = column_wrapper{2, 3, 4, 5}; + auto table = cudf::table_view{{c0, c1, c2}}; + + auto ref0 = cudf::ast::column_reference{0}; + auto ref1 = cudf::ast::column_reference{1}; + auto ref2 = cudf::ast::column_reference{2}; + auto sum = cudf::ast::operation{cudf::ast::ast_operator::ADD, ref0, ref1}; + auto mul = cudf::ast::operation{cudf::ast::ast_operator::MUL, sum, ref2}; + auto sub = cudf::ast::operation{cudf::ast::ast_operator::SUB, sum, ref2}; + + std::array, 3> expressions{mul, sub, sum}; + auto result = cudf::compute_table_jit(table, expressions); + + auto expected_mul = column_wrapper{22, 66, 132, 220}; + auto expected_sub = column_wrapper{9, 19, 29, 39}; + auto expected_sum = column_wrapper{11, 22, 33, 44}; + auto expected = cudf::table_view{{expected_mul, expected_sub, expected_sum}}; + + CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result->view()); +} + +TEST_F(ComputeTableJitTest, PerOutputNullability) +{ + auto c0 = column_wrapper{{1, 0, 3, 0}, {1, 0, 1, 0}}; + auto c1 = column_wrapper{10, 20, 30, 40}; + auto table = cudf::table_view{{c0, c1}}; + + auto ref0 = cudf::ast::column_reference{0}; + auto ref1 = cudf::ast::column_reference{1}; + auto is_null = cudf::ast::operation{cudf::ast::ast_operator::IS_NULL, ref0}; + auto sum = cudf::ast::operation{cudf::ast::ast_operator::ADD, ref0, ref1}; + + std::array, 2> expressions{is_null, sum}; + auto result = cudf::compute_table_jit(table, expressions); + + auto expected_is_null = column_wrapper{false, true, false, true}; + auto expected_sum = column_wrapper{{11, 0, 33, 0}, {1, 0, 1, 0}}; + auto expected = cudf::table_view{{expected_is_null, expected_sum}}; + + CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result->view()); + EXPECT_FALSE(result->view().column(0).nullable()); + EXPECT_TRUE(result->view().column(1).nullable()); +} + +TEST_F(ComputeTableJitTest, IndependentNullMasks) +{ + auto c0 = column_wrapper{{1, 2, 3, 4}, {1, 0, 1, 0}}; + auto c1 = column_wrapper{{10, 20, 30, 40}, {1, 1, 0, 0}}; + auto table = cudf::table_view{{c0, c1}}; + + auto ref0 = cudf::ast::column_reference{0}; + auto ref1 = cudf::ast::column_reference{1}; + auto out0 = cudf::ast::operation{cudf::ast::ast_operator::IDENTITY, ref0}; + auto out1 = cudf::ast::operation{cudf::ast::ast_operator::IDENTITY, ref1}; + + std::array, 2> expressions{out0, out1}; + auto result = cudf::compute_table_jit(table, expressions); + + auto expected0 = column_wrapper{{1, 2, 3, 4}, {1, 0, 1, 0}}; + auto expected1 = column_wrapper{{10, 20, 30, 40}, {1, 1, 0, 0}}; + auto expected = cudf::table_view{{expected0, expected1}}; + + CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result->view()); +} + +TEST_F(ComputeTableJitTest, UnrelatedInputNulls) +{ + auto nullable_input = column_wrapper{{1, 2, 3, 4}, {1, 0, 1, 0}}; + auto valid_input = column_wrapper{10, 20, 30, 40}; + auto table = cudf::table_view{{nullable_input, valid_input}}; + + auto nullable_ref = cudf::ast::column_reference{0}; + auto valid_ref = cudf::ast::column_reference{1}; + auto nullable_out = cudf::ast::operation{cudf::ast::ast_operator::IDENTITY, nullable_ref}; + auto valid_out = cudf::ast::operation{cudf::ast::ast_operator::IDENTITY, valid_ref}; + + std::array, 2> expressions{nullable_out, + valid_out}; + auto result = cudf::compute_table_jit(table, expressions); + + auto expected_nullable = column_wrapper{{1, 2, 3, 4}, {1, 0, 1, 0}}; + auto expected_valid = column_wrapper{10, 20, 30, 40}; + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_nullable, result->view().column(0)); + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected_valid, result->view().column(1)); + EXPECT_EQ(result->view().column(1).null_count(), 0); +} + +TEST_F(ComputeTableJitTest, EmptyExpressions) +{ + auto input = column_wrapper{1, 2, 3}; + auto table = cudf::table_view{{input}}; + auto expressions = std::span const>{}; + + EXPECT_THROW(cudf::compute_table_jit(table, expressions), cudf::logic_error); +} + CUDF_TEST_PROGRAM_MAIN() diff --git a/cpp/tests/jit/row_ir.cpp b/cpp/tests/jit/row_ir.cpp index f4cbde494515..e6ee4c7b12f6 100644 --- a/cpp/tests/jit/row_ir.cpp +++ b/cpp/tests/jit/row_ir.cpp @@ -16,7 +16,9 @@ #include #include +#include #include +#include namespace row_ir = cudf::detail::row_ir; @@ -186,8 +188,7 @@ TEST_F(RowIRCudaCodeGenTest, BinaryOperation) auto expected_code = R"***(int32_t tmp_0 = in_0; -int32_t tmp_1 = in_0; -int32_t tmp_2 = cudf::detail::row_ir::evaluate(tmp_0, tmp_1); +int32_t tmp_1 = cudf::detail::row_ir::evaluate(tmp_0, tmp_0); )***"; EXPECT_EQ(sink.get_code(), expected_code); @@ -209,8 +210,7 @@ int32_t tmp_2 = cudf::detail::row_ir::evaluate(tmp_0, tmp_1); +numeric::decimal32 tmp_1 = cudf::detail::row_ir::evaluate(tmp_0, tmp_0); )***"; EXPECT_EQ(sink.get_code(), expected_null_code); @@ -237,12 +237,11 @@ TEST_F(RowIRCudaCodeGenTest, BinaryOperationOverflow) auto expected_code = R"***(int32_t tmp_0 = in_0; -int32_t tmp_1 = in_0; -auto expected__tmp_2 = cudf::detail::row_ir::evaluate(tmp_0, tmp_1); -if(!expected__tmp_2.has_value()) { - return expected__tmp_2.error(); +auto expected__tmp_1 = cudf::detail::row_ir::evaluate(tmp_0, tmp_0); +if(!expected__tmp_1.has_value()) { + return expected__tmp_1.error(); } -int32_t tmp_2 = expected__tmp_2.value(); +int32_t tmp_1 = expected__tmp_1.value(); )***"; EXPECT_EQ(sink.get_code(), expected_code); @@ -265,8 +264,7 @@ int32_t tmp_2 = expected__tmp_2.value(); auto expected_code = R"***(cuda::std::optional tmp_0 = in_0; -cuda::std::optional tmp_1 = in_0; -cuda::std::optional tmp_2 = cudf::detail::row_ir::evaluate(tmp_0, tmp_1); +cuda::std::optional tmp_1 = cudf::detail::row_ir::evaluate(tmp_0, tmp_0); )***"; EXPECT_EQ(sink.get_code(), expected_code); @@ -319,15 +317,13 @@ TEST_F(RowIRCudaCodeGenTest, VectorLengthOperation) auto expected_code = R"***(double tmp_0 = in_0; -double tmp_1 = in_0; -double tmp_2 = cudf::detail::row_ir::evaluate(tmp_0, tmp_1); -double tmp_3 = in_1; -double tmp_4 = in_1; -double tmp_5 = cudf::detail::row_ir::evaluate(tmp_3, tmp_4); -double tmp_6 = cudf::detail::row_ir::evaluate(tmp_2, tmp_5); -double tmp_7 = cudf::detail::row_ir::evaluate(tmp_6); -double tmp_8 = tmp_7; -*out_0 = tmp_8; +double tmp_1 = cudf::detail::row_ir::evaluate(tmp_0, tmp_0); +double tmp_2 = in_1; +double tmp_3 = cudf::detail::row_ir::evaluate(tmp_2, tmp_2); +double tmp_4 = cudf::detail::row_ir::evaluate(tmp_1, tmp_3); +double tmp_5 = cudf::detail::row_ir::evaluate(tmp_4); +double tmp_6 = tmp_5; +*out_0 = tmp_6; )***"; EXPECT_EQ(sink.get_code(), expected_code); @@ -350,14 +346,15 @@ TEST_F(RowIRCudaCodeGenTest, AstConversionBasic) auto expected = cudf::test::fixed_width_column_wrapper(expected_iter, expected_iter + column->size()); + std::array, 1> output_expressions{add_op}; auto transform_args = - row_ir::ast_converter::compute_column(row_ir::target::CUDA, - add_op, - cudf::table_view{{*column}}, - cudf::table_view{}, - "expression", - cudf::get_default_stream(), - cudf::get_current_device_resource_ref()); + row_ir::ast_converter::compute_table(row_ir::target::CUDA, + output_expressions, + cudf::table_view{{*column}}, + cudf::table_view{}, + "expression", + cudf::get_default_stream(), + cudf::get_current_device_resource_ref()); ASSERT_EQ(transform_args.scalar_columns.size(), 1); ASSERT_EQ(transform_args.scalar_columns[0]->view().size(), 1); @@ -450,4 +447,36 @@ cuda::std::optional tmp_1 = cudf::detail::row_ir::evaluate, 2> expressions{mul, sub}; + + row_ir::ast_converter converter{ + cudf::get_default_stream(), cudf::get_current_device_resource_ref(), table, {}}; + auto const& [code, null_aware, nullability] = + converter.generate_code(row_ir::target::CUDA, expressions, "compute_operation"); + + auto expected_code = + R"***(__device__ cudf::errc compute_operation(int32_t* out_0, int32_t* out_1, int32_t in_0) +{ +int32_t tmp_0 = in_0; +int32_t tmp_1 = cudf::detail::row_ir::evaluate(tmp_0, tmp_0); +int32_t tmp_2 = cudf::detail::row_ir::evaluate(tmp_1, tmp_0); +int32_t tmp_3 = tmp_2; +*out_0 = tmp_3; +int32_t tmp_4 = cudf::detail::row_ir::evaluate(tmp_1, tmp_0); +int32_t tmp_5 = tmp_4; +*out_1 = tmp_5; +return cudf::errc::SUCCESS; +})***"; + + EXPECT_EQ(code, expected_code); + EXPECT_EQ(null_aware, cudf::null_aware::NO); + EXPECT_EQ(nullability.size(), 2); +} + CUDF_TEST_PROGRAM_MAIN()