diff --git a/cpp/benchmarks/CMakeLists.txt b/cpp/benchmarks/CMakeLists.txt index 1b57e3b23666..2a26e5d3220c 100644 --- a/cpp/benchmarks/CMakeLists.txt +++ b/cpp/benchmarks/CMakeLists.txt @@ -415,7 +415,7 @@ add_dependencies(BINARYOP_NVBENCH cudf_benchmark_fragments) # --------------------------------------------------------------------------------- ConfigureNVBench( TRANSFORM_NVBENCH transform/encode.cpp transform/polynomials.cpp - transform/polynomials_concurrent.cpp transform/transform.cpp + transform/polynomials_concurrent.cpp transform/transform.cpp transform/transform_widetable.cpp ) # ################################################################################################## diff --git a/cpp/benchmarks/transform/transform_widetable.cpp b/cpp/benchmarks/transform/transform_widetable.cpp new file mode 100644 index 000000000000..e9ec1d988bcc --- /dev/null +++ b/cpp/benchmarks/transform/transform_widetable.cpp @@ -0,0 +1,155 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +enum class executor_type : uint8_t { AST, JIT, JIT_OPT }; + +executor_type executor_from_string(std::string_view executor) +{ + if (executor == "ast") { return executor_type::AST; } + if (executor == "jit") { return executor_type::JIT; } + if (executor == "jit-opt") { return executor_type::JIT_OPT; } + CUDF_FAIL("unrecognized executor: " + std::string{executor}); +} + +template +std::vector make_expression_trees(cudf::size_type table_width, + cudf::size_type expression_depth, + LiteralFactory make_literal) +{ + std::vector trees; + trees.reserve(table_width); + + for (cudf::size_type column_index = 0; column_index < table_width; ++column_index) { + cudf::ast::tree tree; + cudf::ast::expression const* expression = &tree.push(cudf::ast::column_reference{column_index}); + + for (cudf::size_type level = 0; level < expression_depth; ++level) { + auto& literal = tree.push(make_literal(level)); + expression = + &tree.push(cudf::ast::operation{cudf::ast::ast_operator::ADD, *expression, literal}); + } + trees.push_back(std::move(tree)); + } + + return trees; +} + +void BM_ast_jit_wide_table(nvbench::state& state) +{ + auto table_width = static_cast(state.get_int64("table_width")); + auto rows_per_batch = static_cast(state.get_int64("rows_per_batch")); + auto total_rows = state.get_int64("total_rows"); + auto expression_depth = static_cast(state.get_int64("expression_depth")); + auto executor = executor_from_string(state.get_string("executor")); + + if (rows_per_batch > total_rows || total_rows % rows_per_batch != 0) { + state.skip("rows_per_batch must evenly divide total_rows"); + return; + } + + auto input = create_sequence_table(cycle_dtypes({cudf::type_id::INT32}, table_width), + row_count{rows_per_batch}); + + auto input_view = input->view(); + + auto num_batches = total_rows / rows_per_batch; + + std::vector> scalars; + scalars.reserve(expression_depth); + for (cudf::size_type level = 0; level < expression_depth; ++level) { + scalars.emplace_back(level + 1); + } + + auto scalar_trees = + make_expression_trees(table_width, expression_depth, [&scalars](cudf::size_type level) { + return cudf::ast::literal{scalars[level]}; + }); + + std::vector> scalar_columns; + scalar_columns.reserve(expression_depth); + for (auto& scalar : scalars) { + scalar_columns.push_back(cudf::make_column_from_scalar(scalar, 1)); + } + + auto scalar_column_view_trees = + make_expression_trees(table_width, expression_depth, [&scalar_columns](cudf::size_type level) { + return cudf::ast::literal{cudf::scalar_column_view{scalar_columns[level]->view()}}; + }); + + auto elements = static_cast(total_rows) * static_cast(table_width); + state.add_element_count(elements); + state.add_global_memory_reads(elements * (expression_depth + 1)); + state.add_global_memory_writes(elements); + + auto mem_stats_logger = cudf::memory_stats_logger(); + + state.exec(nvbench::exec_tag::sync, [&](nvbench::launch& launch) { + cudf::benchmark::scoped_range range{"benchmark_iteration"}; + auto stream = launch.get_stream().get_stream(); + std::vector> outputs; + outputs.reserve(table_width); + + for (int64_t batch = 0; batch < num_batches; ++batch) { + outputs.clear(); + + auto& trees = (executor == executor_type::JIT_OPT) ? scalar_column_view_trees : scalar_trees; + + for (auto& tree : trees) { + switch (executor) { + case executor_type::AST: { + outputs.push_back(cudf::compute_column(input_view, tree.back(), stream)); + break; + } + case executor_type::JIT: { + outputs.push_back(cudf::compute_column_jit(input_view, tree.back(), stream)); + break; + } + case executor_type::JIT_OPT: { + outputs.push_back(cudf::compute_column_jit(input_view, tree.back(), stream)); + break; + } + } + } + } + }); + + state.add_buffer_size( + mem_stats_logger.peak_memory_usage(), "peak_memory_usage", "peak_memory_usage"); +} + +} // namespace + +NVBENCH_BENCH(BM_ast_jit_wide_table) + .set_name("ast_jit_wide_table") + .add_int64_axis("table_width", {1, 16, 64}) + .add_int64_axis("rows_per_batch", {1'024, 16'384, 262'144}) + .add_int64_axis("total_rows", {262'144, 1'048'576}) + .add_int64_axis("expression_depth", {1, 4, 16}) + .add_string_axis("executor", {"ast", "jit", "jit-opt"}); diff --git a/cpp/include/cudf/ast/expressions.hpp b/cpp/include/cudf/ast/expressions.hpp index 9212400b37fc..4329f6904300 100644 --- a/cpp/include/cudf/ast/expressions.hpp +++ b/cpp/include/cudf/ast/expressions.hpp @@ -5,6 +5,7 @@ #pragma once #include +#include #include #include #include @@ -249,7 +250,8 @@ class literal : public expression { * @param value A numeric scalar value */ template - literal(cudf::numeric_scalar& value) : scalar(value), value(value) + literal(cudf::numeric_scalar& value) + : scalar{ast_scalar{std::ref(value), generic_scalar_device_view(value)}} { } @@ -260,7 +262,8 @@ class literal : public expression { * @param value A timestamp scalar value */ template - literal(cudf::timestamp_scalar& value) : scalar(value), value(value) + literal(cudf::timestamp_scalar& value) + : scalar{ast_scalar{std::ref(value), generic_scalar_device_view(value)}} { } @@ -271,7 +274,8 @@ class literal : public expression { * @param value A duration scalar value */ template - literal(cudf::duration_scalar& value) : scalar(value), value(value) + literal(cudf::duration_scalar& value) + : scalar{ast_scalar{std::ref(value), generic_scalar_device_view(value)}} { } @@ -280,7 +284,10 @@ class literal : public expression { * * @param value A string scalar value */ - literal(cudf::string_scalar& value) : scalar(value), value(value) {} + literal(cudf::string_scalar& value) + : scalar{ast_scalar{std::ref(value), generic_scalar_device_view(value)}} + { + } /** * @brief Construct a new literal object. @@ -288,30 +295,75 @@ class literal : public expression { * @param value A fixed-point scalar value */ template - literal(cudf::fixed_point_scalar& value) : scalar(value), value(value) + literal(cudf::fixed_point_scalar& value) + : scalar{ast_scalar{std::ref(value), generic_scalar_device_view(value)}} { } + /** + * @brief Construct a new literal object. + * + * @param value A scalar column view value + */ + literal(scalar_column_view value) : scalar{std::move(value)} {} + /** * @brief Get the data type. * * @return The data type of the literal */ - [[nodiscard]] cudf::data_type get_data_type() const { return get_value().type(); } + [[nodiscard]] cudf::data_type get_data_type() const + { + return std::visit( + [](auto const& value) { + if constexpr (std::is_same_v, ast_scalar>) { + return value.value.type(); + } else { + return value.type(); + } + }, + scalar); + } + + /** + * @brief Check whether the literal is backed by a scalar column view. + * + * @return true if the literal is backed by a scalar column view + */ + [[nodiscard]] bool is_scalar_column_view() const noexcept + { + return std::holds_alternative(scalar); + } /** * @brief Get the value object. * * @return The device scalar object */ - [[nodiscard]] generic_scalar_device_view get_value() const { return value; } + [[nodiscard]] generic_scalar_device_view get_value() const + { + return std::get(scalar).value; + } /** * @brief Get the scalar. * * @return The scalar object */ - [[nodiscard]] cudf::scalar const& get_scalar() const { return scalar; } + [[nodiscard]] cudf::scalar const& get_scalar() const + { + return std::get(scalar).scalar.get(); + } + + /** + * @brief Get scalar column view. + * + * @return The scalar column view object + */ + [[nodiscard]] scalar_column_view const& get_scalar_column_view() const + { + return std::get(scalar); + } /** * @copydoc expression::accept @@ -345,12 +397,21 @@ class literal : public expression { */ [[nodiscard]] bool is_valid(rmm::cuda_stream_view stream) const { - return scalar.is_valid(stream); + if (auto* s = std::get_if(&scalar)) { + return s->scalar.get().is_valid(stream); + } else { + auto& c = std::get(scalar); + return c.null_count() == 0; + } } private: - cudf::scalar const& scalar; - generic_scalar_device_view const value; + struct ast_scalar { + std::reference_wrapper scalar; + generic_scalar_device_view value; + }; + + std::variant scalar; }; /** diff --git a/cpp/src/jit/row_ir.cpp b/cpp/src/jit/row_ir.cpp index 74d4f581a243..3e31f78fb400 100644 --- a/cpp/src/jit/row_ir.cpp +++ b/cpp/src/jit/row_ir.cpp @@ -445,7 +445,11 @@ int32_t instance_context::add_input(input in) return col->column.type(); } else { auto& scalar = std::get(in); - return scalar.scalar_column->type(); + if (auto* s = std::get_if>(&scalar)) { + return (*s)->type(); + } else { + return std::get(scalar).type(); + } } }(); inputs_.emplace_back(std::move(in)); @@ -721,7 +725,8 @@ if(expected__{1}.has_value()) {{ std::unique_ptr ast_converter::add_ir_node(ast::literal const& expr) { - auto id = instance_.add_input(expr.get_scalar()); + auto id = expr.is_scalar_column_view() ? instance_.add_input(expr.get_scalar_column_view()) + : instance_.add_input(expr.get_scalar()); return std::make_unique(input_reference{id}); } @@ -772,7 +777,14 @@ std::unique_ptr ast_converter::add_ir_node(ast::jit::detail::opera expr.get_opcode(), expr.get_target_scale(), expr.get_error_policy(), std::move(args)); } -bool is_nullable(scalar_input const& in) { return in.scalar_column->view().nullable(); } +bool is_nullable(scalar_input const& in) +{ + if (auto* s = std::get_if>(&in)) { + return (*s)->nullable(); + } else { + return std::get(in).nullable(); + } +} bool is_nullable(column_input const& in) { return in.column.nullable(); } @@ -859,7 +871,11 @@ std::tuple ast_converter::generate_ std::variant get_column_view(scalar_input const& in) { - return scalar_column_view{in.scalar_column->view()}; + if (auto* s = std::get_if>(&in)) { + return scalar_column_view{**s}; + } else { + return std::get(in); + } } std::variant get_column_view(column_input const& in) @@ -884,16 +900,15 @@ transform_args ast_converter::compute_column(target target_id, auto [code, is_null_aware, output_nullability] = converter.generate_code(target_id, expr, function_name); - std::vector> inputs; + std::vector inputs; std::vector> scalar_columns; std::vector> table_sources; std::vector> column_indices; for (auto& input : converter.instance_.inputs_) { - if (std::holds_alternative(input)) { - auto& col = std::get(input); - table_sources.emplace_back(col.table_source); - column_indices.emplace_back(col.column_index); + if (auto* col = std::get_if(&input)) { + table_sources.emplace_back(col->table_source); + column_indices.emplace_back(col->column_index); } else { table_sources.emplace_back(std::nullopt); column_indices.emplace_back(std::nullopt); @@ -902,9 +917,10 @@ transform_args ast_converter::compute_column(target target_id, auto view = std::visit([](auto& in) { return get_column_view(in); }, input); inputs.emplace_back(view); - if (std::holds_alternative(input)) { - auto& scalar = std::get(input); - scalar_columns.emplace_back(std::move(scalar.scalar_column)); + if (auto* scalar = std::get_if(&input)) { + if (auto* c = std::get_if>(scalar)) { + scalar_columns.emplace_back(std::move(*c)); + } } } diff --git a/cpp/src/jit/row_ir.hpp b/cpp/src/jit/row_ir.hpp index d54ac08b42cd..810581523da9 100644 --- a/cpp/src/jit/row_ir.hpp +++ b/cpp/src/jit/row_ir.hpp @@ -68,10 +68,9 @@ struct target_info { target id = target::CUDA; ///< The target identifier }; -struct scalar_input { - std::unique_ptr scalar_column = - nullptr; ///< The scalar value represented as a column with a single element -}; +///< The scalar value represented as a column with a single element or a prepared scalar column +///< view. +using scalar_input = std::variant, scalar_column_view>; struct column_input { column_view column = {}; ///< The column input @@ -141,8 +140,12 @@ struct [[nodiscard]] instance_context { [[nodiscard]] int32_t add_input(scalar const& scalar) { - return add_input( - scalar_input{.scalar_column = make_column_from_scalar(scalar, 1, stream_, mr_)}); + return add_input(scalar_input{make_column_from_scalar(scalar, 1, stream_, mr_)}); + } + + [[nodiscard]] int32_t add_input(scalar_column_view const& column) + { + return add_input(scalar_input{column}); } [[nodiscard]] int32_t add_input(column_view const& column)