diff --git a/cpp/benchmarks/ast/transform.cpp b/cpp/benchmarks/ast/transform.cpp index bb5d1bfaedaa..4c753a4e1966 100644 --- a/cpp/benchmarks/ast/transform.cpp +++ b/cpp/benchmarks/ast/transform.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -118,6 +118,44 @@ static void BM_ast_transform(nvbench::state& state) mem_stats_logger.peak_memory_usage(), "peak_memory_usage", "peak_memory_usage"); } +static void BM_ast_transform_dispatch(nvbench::state& state) +{ + auto const num_rows = static_cast(state.get_int64("num_rows")); + auto const expression_depth = static_cast(state.get_int64("expression_depth")); + auto const api = state.get_string("api"); + auto source_table = + create_sequence_table({cudf::type_id::INT32}, row_count{num_rows}, std::nullopt); + auto const table = source_table->view(); + + cudf::ast::tree tree; + tree.push(cudf::ast::column_reference{0}); + tree.push(cudf::ast::operation{cudf::ast::ast_operator::ADD, tree.at(0), tree.at(0)}); + for (cudf::size_type level = 1; level < expression_depth; ++level) { + tree.push(cudf::ast::operation{cudf::ast::ast_operator::ADD, tree.back(), tree.at(0)}); + } + auto const& expression = tree.back(); + + std::unique_ptr program; + if (api == "transform_program") { + program = std::make_unique(table, expression); + } else { + // Populate the JIT cache before timing, matching transform_program construction. + cudf::compute_column_jit(table, expression); + } + + state.add_global_memory_reads(static_cast(num_rows) * + (expression_depth + 1)); + state.add_global_memory_writes(num_rows); + + state.exec(nvbench::exec_tag::sync, [&](nvbench::launch& launch) { + if (program) { + program->run(table, launch.get_stream().get_stream()); + } else { + cudf::compute_column_jit(table, expression, launch.get_stream().get_stream()); + } + }); +} + template static void BM_string_compare_ast_transform(nvbench::state& state) { @@ -238,3 +276,9 @@ AST_TRANSFORM_BENCHMARK_DEFINE( AST_STRING_COMPARE_TRANSFORM_BENCHMARK_DEFINE(ast_string_equal_logical_and, cudf::ast::ast_operator::EQUAL, cudf::ast::ast_operator::LOGICAL_AND); + +NVBENCH_BENCH(BM_ast_transform_dispatch) + .set_name("ast_transform_dispatch") + .add_string_axis("api", {"compute_column_jit", "transform_program"}) + .add_int64_axis("expression_depth", {1, 5, 10}) + .add_int64_axis("num_rows", {1, 1'000, 100'000, 1'000'000}); diff --git a/cpp/benchmarks/transform/transform.cpp b/cpp/benchmarks/transform/transform.cpp index 059dbffc802a..2b9e9975a71b 100644 --- a/cpp/benchmarks/transform/transform.cpp +++ b/cpp/benchmarks/transform/transform.cpp @@ -109,6 +109,60 @@ static void BM_transform(nvbench::state& state) mem_stats_logger.peak_memory_usage(), "peak_memory_usage", "peak_memory_usage"); } +static void BM_transform_dispatch(nvbench::state& state) +{ + auto const num_rows = static_cast(state.get_int64("num_rows")); + auto const api = state.get_string("api"); + auto source_table = + create_sequence_table({cudf::type_id::INT32}, row_count{num_rows}, std::nullopt); + + std::array inputs{source_table->get_column(0).view()}; + std::array outputs{cudf::transform_output{cudf::data_type{cudf::type_id::INT32}, + cudf::output_nullability::ALL_VALID}}; + std::string const udf = "__device__ void transform(int32_t* out, int32_t in) { *out = in + 1; }"; + + std::unique_ptr program; + if (api == "transform_program") { + program = + std::make_unique(udf, + cudf::udf_source_type::CUDA, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + std::span const>{}); + } else { + // Populate the JIT cache before timing, matching transform_program construction. + cudf::transform(udf, + cudf::udf_source_type::CUDA, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + {}, + std::nullopt); + } + + state.add_global_memory_reads(num_rows); + state.add_global_memory_writes(num_rows); + + state.exec(nvbench::exec_tag::sync, [&](nvbench::launch& launch) { + if (program) { + program->run(inputs, outputs, {}, std::nullopt, launch.get_stream().get_stream()); + } else { + cudf::transform(udf, + cudf::udf_source_type::CUDA, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + {}, + std::nullopt, + launch.get_stream().get_stream()); + } + }); +} + #define AST_TRANSFORM_BENCHMARK_DEFINE(name, key_type, tree_type, reuse_columns, nullable) \ static void name(::nvbench::state& st) \ { \ @@ -125,3 +179,8 @@ AST_TRANSFORM_BENCHMARK_DEFINE( transform_int32_imbalanced_reuse, int32_t, TreeType::IMBALANCED_LEFT, true, false); AST_TRANSFORM_BENCHMARK_DEFINE( transform_double_imbalanced_unique, double, TreeType::IMBALANCED_LEFT, false, false); + +NVBENCH_BENCH(BM_transform_dispatch) + .set_name("transform_dispatch") + .add_string_axis("api", {"transform", "transform_program"}) + .add_int64_axis("num_rows", {1, 1'000, 100'000, 1'000'000}); diff --git a/cpp/include/cudf/transform.hpp b/cpp/include/cudf/transform.hpp index 2be1a88a9020..1d8fe4ec1a24 100644 --- a/cpp/include/cudf/transform.hpp +++ b/cpp/include/cudf/transform.hpp @@ -77,6 +77,168 @@ struct transform_output { rmm::cuda_stream_view stream = cudf::get_default_stream(), rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); +/** + * @brief Describes a transform input independently of a particular column. + * + * An input specification contains the type information needed to reflect and retrieve a transform + * kernel. Dictionary specifications recursively describe their indices and keys through `children`. + * String specifications retain their offsets child type so `INT32` and `INT64` layouts can be + * distinguished. + */ +struct transform_input_spec { + type_id type = type_id::EMPTY; ///< Logical type of the input + bool is_scalar = false; ///< Whether the input is presented to the UDF as a scalar + std::vector children = + {}; ///< Specifications of dictionary children or string offsets +}; + +/** + * @brief Describes a transform output independently of a particular output column. + * + * The string-offset setting identifies the device-view representation required by the kernel. The + * nullability setting is retained so inputs to `transform_program::run` can be validated against + * the output policy used to construct the program. + */ +struct transform_output_spec { + type_id type = type_id::EMPTY; ///< Logical type of the output + output_nullability nullability = + output_nullability::PRESERVE; ///< Null-mask policy for the output + bool has_string_offsets = false; ///< Whether a string output uses preallocated offsets + std::vector children = + {}; ///< Specifications of string offsets or nested child columns +}; + +/** + * @brief A reusable transform program that retains a JIT-compiled kernel. + * + * Construction retrieves the kernel for the UDF and the supplied input and output specifications. + * Subsequent calls to `run` reuse that kernel and otherwise follow the regular `transform` + * execution path. Runtime inputs and outputs must match the specifications used at construction. + */ +struct transform_program { + private: + struct impl; + + std::unique_ptr impl_; ///< The implementation of the transform program + + public: + /** + * @brief Constructs a reusable program by deriving specifications from transform arguments. + * + * The UDF kernel is retrieved during construction and retained for subsequent calls to `run`. + * The input and output objects are inspected only to derive their specifications and are not + * retained. + * + * @param udf The PTX or CUDA source for the transform UDF + * @param source_type The source type of `udf` + * @param is_null_aware Whether the UDF receives row inputs as optional values + * @param user_data User-defined device data, not owned by the program, retained and passed to the + * UDF by `run` + * @param inputs Inputs from which to derive the input specifications + * @param outputs Outputs from which to derive the output type and nullability specifications + * @param string_offsets Optional string offsets used to determine each string output + * representation + */ + transform_program(std::string const& udf, + udf_source_type source_type, + null_aware is_null_aware, + std::optional user_data, + std::span inputs, + std::span outputs, + std::span const> string_offsets); + + /** + * @brief Constructs a reusable program from explicit input and output specifications. + * + * This overload enables composition without requiring concrete columns when the program is + * created. The UDF kernel is retrieved during construction and retained for subsequent calls to + * `run`. + * + * @param udf The PTX or CUDA source for the transform UDF + * @param source_type The source type of `udf` + * @param is_null_aware Whether the UDF receives row inputs as optional values + * @param user_data User-defined device data, not owned by the program, retained and passed to the + * UDF by `run` + * @param inputs Specifications of the transform inputs + * @param outputs Specifications of the transform outputs + */ + transform_program(std::string const& udf, + udf_source_type source_type, + null_aware is_null_aware, + std::optional user_data, + std::span inputs, + std::span outputs); + + /** + * @brief Constructs a reusable program for an AST expression. + * + * The expression is lowered and its kernel is retrieved during construction. Literal values are + * retained by the program, while column inputs are rebound to the table passed to `run`. + * + * @param table A table whose schema is used to lower the expression and retrieve its kernel + * @param expr The root of the expression tree + * @param stream CUDA stream used for device memory operations during construction + * @param mr Device memory resource used for device memory allocations during construction + */ + transform_program(table_view const& table, + ast::expression const& expr, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + + transform_program(transform_program const&) = delete; ///< Deleted copy constructor + transform_program(transform_program&&); ///< Move constructor + transform_program& operator=(transform_program const&) = delete; ///< Deleted copy assignment + transform_program& operator=(transform_program&&); ///< Move assignment operator + ~transform_program(); ///< Destructor + + /** + * @brief Runs the transform program on the given inputs and outputs. + * + * The transform program must have a matching set of input and output specifications as the inputs + * and outputs provided to this function. + * + * @throws std::invalid_argument if the inputs, outputs, or string offsets do not match the + * specifications used to construct the program + * @throws std::logic_error if this is a moved-from program + * + * @param inputs The inputs to the transform program + * @param outputs The outputs of the transform program + * @param string_offsets For string output columns, the offsets can be pre-allocated and passed in + * to prevent overhead of compacting string views into run-end strings column. + * @param row_size The row size of the transform operation. If not provided, it will be inferred + * from the inputs. + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used to allocate the returned column's device memory + * @return A table containing the columns resulting from applying the transform function to every + * element of the input according to the output specifications + */ + std::unique_ptr run( + std::span inputs, + std::span outputs, + std::vector>&& string_offsets, + std::optional row_size, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + + /** + * @brief Evaluates the AST expression used to construct this program on a table. + * + * The input table must have types compatible with the table used during construction. + * + * @throws std::invalid_argument if this program was not constructed from an AST expression or + * if the referenced input columns are incompatible with the program + * + * @param table The table used for expression evaluation + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used to allocate the returned column device memory + * @return The column resulting from evaluating the expression + */ + std::unique_ptr run( + table_view const& table, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); +}; + /** * @brief Creates a new table by applying a transform function against every * element of the input columns. diff --git a/cpp/src/transform/transform.cu b/cpp/src/transform/transform.cu index 9e08a0950557..928a8bc6a804 100644 --- a/cpp/src/transform/transform.cu +++ b/cpp/src/transform/transform.cu @@ -209,129 +209,99 @@ void launch(cudf::kernel const& kernel, kernel.launch({cfg.min_grid_size}, {cfg.block_size}, 0, stream, args); } -std::string get_element_type_name(column_view const& view, bool use_physical_type); +std::string get_element_type_name(transform_input_spec const& spec, bool use_physical_type); struct element_type_name_fn { template - std::string operator()(column_view const& view, bool use_physical_type) const + std::string operator()(transform_input_spec const& spec, bool use_physical_type) const requires(is_fixed_width() || std::same_as) { - return type_to_name(use_physical_type ? jit::physical_type_of(view.type()) : view.type()); + auto type = data_type{spec.type}; + return type_to_name(use_physical_type ? jit::physical_type_of(type) : type); } template - std::string operator()(column_view const& view, bool use_physical_type) const + std::string operator()(transform_input_spec const& spec, bool use_physical_type) const requires(std::same_as) { + CUDF_EXPECTS(spec.children.size() > static_cast(dictionary_keys_column_index), + "Dictionary transform input specifications must include indices and keys", + std::invalid_argument); return std::format( "cudf::dictionary_element<{}, {}>", - get_element_type_name(view.child(cudf::dictionary_indices_column_index), use_physical_type), - get_element_type_name(view.child(cudf::dictionary_keys_column_index), use_physical_type)); + get_element_type_name(spec.children[dictionary_indices_column_index], use_physical_type), + get_element_type_name(spec.children[dictionary_keys_column_index], use_physical_type)); } template - std::string operator()(column_view const& view, bool use_physical_type) const + std::string operator()(transform_input_spec const& spec, bool) const requires(!is_fixed_width() && !std::same_as && !std::same_as) { - CUDF_FAIL("Unsupported type for JIT compilation: " + type_to_name(view.type())); + CUDF_FAIL("Unsupported type for JIT compilation: " + type_to_name(data_type{spec.type})); } }; -std::string get_element_type_name(column_view const& view, bool use_physical_type) -{ - return cudf::type_dispatcher(view.type(), element_type_name_fn{}, view, use_physical_type); -} - -std::string reflect_input_element(column_view const& c, bool use_physical_type) +std::string get_element_type_name(transform_input_spec const& spec, bool use_physical_type) { - return get_element_type_name(c, use_physical_type); + return cudf::type_dispatcher( + data_type{spec.type}, element_type_name_fn{}, spec, use_physical_type); } -std::string reflect_input_element(scalar_column_view const& c, bool use_physical_type) +std::string reflect_input_element(transform_input_spec const& spec, bool use_physical_type) { - return get_element_type_name(c.as_column_view(), use_physical_type); + return get_element_type_name(spec, use_physical_type); } -std::string reflect_output_element(fixed_width_column const& c, bool use_physical_type) +std::string reflect_output_element(transform_output_spec const& spec, bool use_physical_type) { - return get_element_type_name(c._col->view(), use_physical_type); -} - -std::string reflect_output_element(string_views_column const&, - [[maybe_unused]] bool use_physical_type) -{ - return "cudf::string_view"; -} - -std::string reflect_output_element(mutable_strings_column const&, - [[maybe_unused]] bool use_physical_type) -{ - return "cuda::std::span"; -} - -std::string reflect_input_value_type(column_view const& c, bool use_physical_type) -{ - return is_dictionary(c.type()) ? reflect_input_value_type( - c.child(cudf::dictionary_keys_column_index), use_physical_type) - : reflect_input_element(c, use_physical_type); -} - -std::string reflect_input_value_type(scalar_column_view const& c, bool use_physical_type) -{ - return reflect_input_value_type(c.as_column_view(), use_physical_type); -} - -std::string reflect_output_value_type(fixed_width_column const& c, bool use_physical_type) -{ - return reflect_output_element(c, use_physical_type); + if (spec.type == type_id::STRING) { + return spec.has_string_offsets ? "cuda::std::span" : "cudf::string_view"; + } + return get_element_type_name(transform_input_spec{.type = spec.type}, use_physical_type); } -std::string reflect_output_value_type(string_views_column const& c, bool use_physical_type) +std::string reflect_input_value_type(transform_input_spec const& spec, bool use_physical_type) { - return reflect_output_element(c, use_physical_type); + if (spec.type == type_id::DICTIONARY32) { + CUDF_EXPECTS(spec.children.size() > static_cast(dictionary_keys_column_index), + "Dictionary transform input specifications must include indices and keys", + std::invalid_argument); + return reflect_input_value_type(spec.children[dictionary_keys_column_index], use_physical_type); + } + return reflect_input_element(spec, use_physical_type); } -std::string reflect_output_value_type(mutable_strings_column const& c, bool use_physical_type) +std::string reflect_output_value_type(transform_output_spec const& spec, bool use_physical_type) { - return reflect_output_element(c, use_physical_type); + return reflect_output_element(spec, use_physical_type); } -std::string reflect_input_column(column_view const&) { return "cudf::column_device_view_core"; } - -std::string reflect_input_column(scalar_column_view const&) +std::string reflect_input_column(transform_input_spec const&) { return "cudf::column_device_view_core"; } -std::string reflect_output_column(fixed_width_column const&) +std::string reflect_output_column(transform_output_spec const& spec) { + if (spec.type == type_id::STRING) { + return spec.has_string_offsets ? "cudf::jit::mutable_strings_column_device_view" + : "cudf::jit::mutable_vector_device_view"; + } return "cudf::mutable_column_device_view_core"; } -std::string reflect_output_column(string_views_column const&) -{ - return "cudf::jit::mutable_vector_device_view"; -} - -std::string reflect_output_column(mutable_strings_column const&) -{ - return "cudf::jit::mutable_strings_column_device_view"; -} - auto reflect(std::variant source_type, - std::span inputs, - std::span outputs) + std::span inputs, + std::span outputs) { std::vector in_types; bool use_physical_types = std::holds_alternative(source_type); - for (size_t i = 0; i < inputs.size(); i++) { - auto& in = inputs[i]; - auto column = std::visit([&](auto& c) { return reflect_input_column(c); }, in); - auto element = - std::visit([&](auto& c) { return reflect_input_element(c, use_physical_types); }, in); - bool as_scalar = std::holds_alternative(in); + auto& in = inputs[i]; + auto column = reflect_input_column(in); + auto element = reflect_input_element(in, use_physical_types); + bool as_scalar = in.is_scalar; auto accessor = rtcx::reflect_template("cudf::jit::column_accessor", rtcx::reflect(i), column, @@ -342,12 +312,10 @@ auto reflect(std::variant source_type, } std::vector out_types; - for (size_t i = 0; i < outputs.size(); i++) { - auto& out = outputs[i]; - auto column = std::visit([&](auto& c) { return reflect_output_column(c); }, out); - auto element = - std::visit([&](auto& c) { return reflect_output_element(c, use_physical_types); }, out); + auto& out = outputs[i]; + auto column = reflect_output_column(out); + auto element = reflect_output_element(out, use_physical_types); bool as_scalar = false; // never scalar auto accessor = rtcx::reflect_template("cudf::jit::column_accessor", rtcx::reflect(i), @@ -364,44 +332,126 @@ auto reflect(std::variant source_type, std::vector ptx_in_types; std::vector ptx_out_types; - if (std::holds_alternative(source_type) && std::get(source_type) == udf_source_type::PTX) { for (auto& in : inputs) { - ptx_in_types.push_back( - std::visit([&](auto& c) { return reflect_input_value_type(c, use_physical_types); }, in)); + ptx_in_types.push_back(reflect_input_value_type(in, use_physical_types)); } for (auto& out : outputs) { - ptx_out_types.push_back( - std::visit([&](auto& c) { return reflect_output_value_type(c, use_physical_types); }, out)); + ptx_out_types.push_back(reflect_output_value_type(out, use_physical_types)); } } return std::make_tuple(ins, outs, ptx_in_types, ptx_out_types); } +transform_input_spec make_input_spec(column_view const& column, bool is_scalar) +{ + transform_input_spec result{.type = column.type().id(), .is_scalar = is_scalar}; + if (is_dictionary(column.type())) { + for (size_type i = 0; i < column.num_children(); ++i) { + result.children.push_back(make_input_spec(column.child(i), false)); + } + } else if (column.type().id() == type_id::STRING && + column.num_children() > strings_column_view::offsets_column_index) { + result.children.push_back( + make_input_spec(column.child(strings_column_view::offsets_column_index), false)); + } + return result; +} + +transform_input_spec make_input_spec(input_column_view const& input) +{ + return std::visit( + [](auto& value) { + return make_input_spec(as_column_view(value), + std::is_same_v, scalar_column_view>); + }, + input); +} + +std::vector make_input_specs(std::span inputs) +{ + std::vector result; + for (auto& input : inputs) { + result.push_back(make_input_spec(input)); + } + return result; +} + +transform_output_spec make_output_spec(fixed_width_column const& output) +{ + return {.type = output._col->type().id()}; +} + +transform_output_spec make_output_spec(string_views_column const&) +{ + return {.type = type_id::STRING}; +} + +transform_output_spec make_output_spec(mutable_strings_column const& output) +{ + auto offsets = output._col->view().child(strings_column_view::offsets_column_index); + return {.type = type_id::STRING, + .has_string_offsets = true, + .children = {{.type = offsets.type().id()}}}; +} + +std::vector make_output_specs(std::span outputs) +{ + std::vector result; + for (auto& output : outputs) { + result.push_back(std::visit([](auto& value) { return make_output_spec(value); }, output)); + } + return result; +} + +std::vector make_output_specs( + std::span outputs, + std::span const> string_offsets) +{ + CUDF_EXPECTS(string_offsets.empty() || string_offsets.size() == outputs.size(), + "Number of string offsets must be empty or match the number of outputs", + std::invalid_argument); + std::vector result; + for (size_t i = 0; i < outputs.size(); ++i) { + auto has_string_offsets = !string_offsets.empty() && string_offsets[i] != nullptr; + transform_output_spec spec{.type = outputs[i].type.id(), + .nullability = outputs[i].nullability, + .has_string_offsets = has_string_offsets}; + if (has_string_offsets) { spec.children.push_back({.type = string_offsets[i]->type().id()}); } + result.push_back(std::move(spec)); + } + return result; +} + +auto reflect(std::variant source_type, + std::span inputs, + std::span outputs) +{ + auto input_specs = make_input_specs(inputs); + auto output_specs = make_output_specs(outputs); + return reflect(source_type, input_specs, output_specs); +} + std::string reflect_udf_signature(bool is_null_aware, bool has_user_data, - std::span inputs, - std::span outputs, + std::span inputs, + std::span outputs, bool use_physical_types) { std::vector in_types; for (size_t i = 0; i < inputs.size(); i++) { - auto& in = inputs[i]; - auto element = - std::visit([&](auto& c) { return reflect_input_element(c, use_physical_types); }, in); + auto element = reflect_input_element(inputs[i], use_physical_types); in_types.push_back(is_null_aware ? std::format("cuda::std::optional<{}>", element) : element); } std::vector out_types; for (size_t i = 0; i < outputs.size(); i++) { - auto& out = outputs[i]; - auto element = - std::visit([&](auto& c) { return reflect_output_element(c, use_physical_types); }, out); + auto element = reflect_output_element(outputs[i], use_physical_types); out_types.push_back(is_null_aware ? std::format("cuda::std::optional<{}> *", element) : std::format("{} *", element)); } @@ -420,6 +470,17 @@ std::string reflect_udf_signature(bool is_null_aware, return std::format("int({})", joined); } +std::string reflect_udf_signature(bool is_null_aware, + bool has_user_data, + std::span inputs, + std::span outputs, + bool use_physical_types) +{ + auto input_specs = make_input_specs(inputs); + auto output_specs = make_output_specs(outputs); + return reflect_udf_signature( + is_null_aware, has_user_data, input_specs, output_specs, use_physical_types); +} std::tuple instantiate_fragment( bool is_null_aware, @@ -484,6 +545,42 @@ auto to_args(std::span inputs, return std::make_tuple(std::move(d_args), std::move(handles)); } +kernel get_kernel(bool is_null_aware, + bool has_user_data, + std::span inputs, + std::span outputs, + std::string const& udf, + udf_source_type source_type) +{ + auto [in_types, out_types, ptx_in_types, ptx_out_types] = reflect(source_type, inputs, outputs); + return instantiate(is_null_aware, + has_user_data, + in_types, + out_types, + ptx_in_types, + ptx_out_types, + udf, + source_type); +} + +kernel get_kernel(bool is_null_aware, + bool has_user_data, + std::span inputs, + std::span outputs, + std::string const& udf, + udf_source_type source_type) +{ + auto [in_types, out_types, ptx_in_types, ptx_out_types] = reflect(source_type, inputs, outputs); + return instantiate(is_null_aware, + has_user_data, + in_types, + out_types, + ptx_in_types, + ptx_out_types, + udf, + source_type); +} + void run(bool is_null_aware, bool has_user_data, size_type row_size, @@ -497,17 +594,27 @@ void run(bool is_null_aware, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { - auto [in_types, out_types, ptx_in_types, ptx_out_types] = reflect(source_type, inputs, outputs); - auto kernel = instantiate(is_null_aware, - has_user_data, - in_types, - out_types, - ptx_in_types, - ptx_out_types, - udf, - source_type); - auto [cols, handles] = to_args(inputs, outputs, stream, mr); - auto* input_cols = reinterpret_cast(cols.data()); + auto kernel = get_kernel(is_null_aware, has_user_data, inputs, outputs, udf, source_type); + auto [cols, handles] = to_args(inputs, outputs, stream, mr); + auto* input_cols = reinterpret_cast(cols.data()); + auto* output_cols = + reinterpret_cast(input_cols + inputs.size()); + return launch( + kernel, row_size, d_stencil, user_data, input_cols, output_cols, d_max_error, stream); +} + +void run(kernel const& kernel, + size_type row_size, + bitmask_type const* d_stencil, + void* user_data, + std::span inputs, + std::span outputs, + int32_t* d_max_error, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + auto [cols, handles] = to_args(inputs, outputs, stream, mr); + auto* input_cols = reinterpret_cast(cols.data()); auto* output_cols = reinterpret_cast(input_cols + inputs.size()); return launch( @@ -1045,6 +1152,7 @@ std::unique_ptr
execute_transform(std::string const& udf, std::span inputs, std::span outputs, std::vector> string_offsets, + kernel const* compiled_kernel, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { @@ -1064,18 +1172,30 @@ std::unique_ptr
execute_transform(std::string const& udf, rmm::device_scalar d_max_error(static_cast(errc::SUCCESS), stream, mr); - jit_transform::run(is_null_aware == null_aware::YES, - user_data.has_value(), - row_size, - stencil_has_nulls ? stencil_arg : nullptr, - user_data.value_or(nullptr), - inputs, - output_columns, - d_max_error.data(), - udf, - source_type, - stream, - mr); + if (compiled_kernel == nullptr) { + jit_transform::run(is_null_aware == null_aware::YES, + user_data.has_value(), + row_size, + stencil_has_nulls ? stencil_arg : nullptr, + user_data.value_or(nullptr), + inputs, + output_columns, + d_max_error.data(), + udf, + source_type, + stream, + mr); + } else { + jit_transform::run(*compiled_kernel, + row_size, + stencil_has_nulls ? stencil_arg : nullptr, + user_data.value_or(nullptr), + inputs, + output_columns, + d_max_error.data(), + stream, + mr); + } auto error = static_cast(d_max_error.value(stream)); @@ -1113,6 +1233,7 @@ std::unique_ptr
transform(std::string const& udf, inputs, outputs, std::move(string_offsets), + nullptr, stream, mr); } @@ -1281,4 +1402,163 @@ std::unique_ptr
transform_lto(std::span udf, return std::make_unique
(std::move(finalized)); } +struct transform_program::impl { + void validate(udf_source_type source_type, + std::span actual_inputs, + std::span actual_outputs, + std::span const> string_offsets) const + { + auto actual_input_specs = jit_transform::make_input_specs(actual_inputs); + auto actual_output_specs = jit_transform::make_output_specs(actual_outputs, string_offsets); + auto actual_reflection = + jit_transform::reflect(source_type, actual_input_specs, actual_output_specs); + CUDF_EXPECTS(reflection_ == actual_reflection, + "Transform program specifications do not match the provided inputs and outputs", + std::invalid_argument); + } + + impl(std::string const& udf, + udf_source_type source_type, + null_aware is_null_aware, + std::optional user_data, + std::vector inputs, + std::vector outputs) + : reflection_{jit_transform::reflect(source_type, inputs, outputs)}, + source_type_{source_type}, + is_null_aware_{is_null_aware}, + user_data_{user_data}, + kernel_{jit_transform::get_kernel(is_null_aware_ == null_aware::YES, + user_data_.has_value(), + inputs, + outputs, + udf, + source_type_)} + { + } + + std::tuple, std::vector> + reflection_; + udf_source_type source_type_; + null_aware is_null_aware_; + std::optional user_data_; + kernel kernel_; + std::vector> ast_scalar_columns_; + std::optional>> ast_input_column_indices_; + std::vector ast_input_types_; + std::vector ast_input_nullable_; + std::vector ast_outputs_; +}; + +transform_program::transform_program(std::string const& udf, + udf_source_type source_type, + null_aware is_null_aware, + std::optional user_data, + std::span inputs, + std::span outputs, + std::span const> string_offsets) + : transform_program(udf, + source_type, + is_null_aware, + user_data, + jit_transform::make_input_specs(inputs), + jit_transform::make_output_specs(outputs, string_offsets)) +{ +} + +transform_program::transform_program(std::string const& udf, + udf_source_type source_type, + null_aware is_null_aware, + std::optional user_data, + std::span inputs, + std::span outputs) +{ + impl_ = + std::make_unique(udf, + source_type, + is_null_aware, + user_data, + std::vector{inputs.begin(), inputs.end()}, + std::vector{outputs.begin(), outputs.end()}); +} + +transform_program::transform_program(table_view const& table, + ast::expression const& expr, + 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); + impl_ = + std::make_unique(args.udf, + args.source_type, + args.is_null_aware, + args.user_data, + jit_transform::make_input_specs(args.inputs), + jit_transform::make_output_specs(args.outputs, args.string_offsets)); + for (auto& input : args.inputs) { + impl_->ast_input_types_.push_back(std::visit([](auto& view) { return view.type(); }, input)); + impl_->ast_input_nullable_.push_back( + std::visit([](auto& view) { return view.nullable(); }, input)); + } + impl_->ast_scalar_columns_ = std::move(args.scalar_columns); + impl_->ast_input_column_indices_ = std::move(args.input_column_indices); + impl_->ast_outputs_ = std::move(args.outputs); +} + +transform_program::transform_program(transform_program&&) = default; +transform_program& transform_program::operator=(transform_program&&) = default; +transform_program::~transform_program() = default; + +std::unique_ptr
transform_program::run(std::span inputs, + std::span outputs, + std::vector>&& string_offsets, + std::optional row_size, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + CUDF_FUNC_RANGE(); + impl_->validate(impl_->source_type_, inputs, outputs, string_offsets); + perform_checks( + impl_->source_type_, impl_->is_null_aware_, row_size, inputs, outputs, string_offsets); + return execute_transform({}, + impl_->source_type_, + impl_->is_null_aware_, + row_size, + impl_->user_data_, + inputs, + outputs, + std::move(string_offsets), + &impl_->kernel_, + stream, + mr); +} + +std::unique_ptr transform_program::run(table_view const& table, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + CUDF_EXPECTS(impl_->ast_input_column_indices_.has_value(), + "Transform program was not constructed from an AST expression", + std::invalid_argument); + + std::vector inputs; + auto scalar_index = std::size_t{0}; + for (auto i = std::size_t{0}; i < impl_->ast_input_column_indices_->size(); ++i) { + auto& column_index = (*impl_->ast_input_column_indices_)[i]; + if (column_index.has_value()) { + auto input = table.column(*column_index); + CUDF_EXPECTS(input.type() == impl_->ast_input_types_[i] && + input.nullable() == impl_->ast_input_nullable_[i], + "AST transform program input schema does not match the provided table", + std::invalid_argument); + inputs.emplace_back(input); + } else { + inputs.emplace_back(scalar_column_view{impl_->ast_scalar_columns_[scalar_index++]->view()}); + } + } + auto result = run(inputs, impl_->ast_outputs_, {}, table.num_rows(), stream, mr); + auto columns = result->release(); + return std::move(columns.front()); +} + } // namespace cudf diff --git a/cpp/tests/ast/transform_tests.cpp b/cpp/tests/ast/transform_tests.cpp index fe4960b9dfec..72068f559433 100644 --- a/cpp/tests/ast/transform_tests.cpp +++ b/cpp/tests/ast/transform_tests.cpp @@ -67,6 +67,48 @@ using Executors = cudf::test::Types; TYPED_TEST_SUITE(TransformTest, Executors); +struct TransformProgramTest : public cudf::test::BaseFixture {}; + +TEST_F(TransformProgramTest, ReusesAstWithCompatibleTable) +{ + auto construction_input = column_wrapper{3, 20, 1, 50}; + auto construction_table = cudf::table_view{{construction_input}}; + auto column_ref = cudf::ast::column_reference{0}; + auto literal_value = cudf::numeric_scalar{2}; + auto literal = cudf::ast::literal{literal_value}; + auto expression = cudf::ast::operation{cudf::ast::ast_operator::ADD, column_ref, literal}; + + auto program = cudf::transform_program{construction_table, expression}; + + auto construction_expected = column_wrapper{5, 22, 3, 52}; + auto construction_result = program.run(construction_table); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(construction_expected, construction_result->view(), verbosity); + + auto input = column_wrapper{10, 20, 30}; + auto table = cudf::table_view{{input}}; + auto expected = column_wrapper{12, 22, 32}; + auto result = program.run(table); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), verbosity); +} + +TEST_F(TransformProgramTest, RejectsIncompatibleTable) +{ + auto construction_input = column_wrapper{3, 20, 1, 50}; + auto construction_table = cudf::table_view{{construction_input}}; + auto column_ref = cudf::ast::column_reference{0}; + auto program = cudf::transform_program{construction_table, column_ref}; + + auto input = column_wrapper{10, 20, 30}; + auto table = cudf::table_view{{input}}; + + EXPECT_THROW((void)program.run(table), std::invalid_argument); + + auto nullable_input = column_wrapper{{10, 20, 30}, {1, 1, 1}}; + auto nullable_table = cudf::table_view{{nullable_input}}; + EXPECT_THROW((void)program.run(nullable_table), std::invalid_argument); +} + TYPED_TEST(TransformTest, ColumnReference) { using Executor = TypeParam; diff --git a/cpp/tests/transform/integration/unary_transform_test.cpp b/cpp/tests/transform/integration/unary_transform_test.cpp index 528ef661cda6..0e867a454010 100644 --- a/cpp/tests/transform/integration/unary_transform_test.cpp +++ b/cpp/tests/transform/integration/unary_transform_test.cpp @@ -148,6 +148,76 @@ TEST_F(AssertsTest, NullSupport) .front())); } +TEST_F(RuntimeSupportTest, TransformProgram) +{ + cudf::transform_input inputs[] = {a, b, cudf::scalar_column_view(t)}; + std::array outputs{cudf::transform_output{cudf::data_type{cudf::type_id::FLOAT32}, + cudf::output_nullability::ALL_VALID}}; + + cudf::transform_program program{ + udf, cudf::udf_source_type::CUDA, cudf::null_aware::NO, std::nullopt, inputs, outputs, {}}; + + std::array input_specs{ + cudf::transform_input_spec{.type = cudf::type_id::FLOAT32}, + cudf::transform_input_spec{.type = cudf::type_id::FLOAT32}, + cudf::transform_input_spec{.type = cudf::type_id::FLOAT32, .is_scalar = true}}; + std::array output_specs{cudf::transform_output_spec{ + .type = cudf::type_id::FLOAT32, .nullability = cudf::output_nullability::ALL_VALID}}; + cudf::transform_program spec_program{udf, + cudf::udf_source_type::CUDA, + cudf::null_aware::NO, + std::nullopt, + input_specs, + output_specs}; + + auto expected = cudf::transform(udf, + cudf::udf_source_type::CUDA, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + {}, + std::nullopt); + auto result = program.run(inputs, outputs, {}, std::nullopt); + CUDF_TEST_EXPECT_TABLES_EQUAL(expected->view(), result->view()); + + result = spec_program.run(inputs, outputs, {}, std::nullopt); + CUDF_TEST_EXPECT_TABLES_EQUAL(expected->view(), result->view()); + + std::array compatible_outputs{cudf::transform_output{cudf::data_type{cudf::type_id::FLOAT32}, + cudf::output_nullability::PRESERVE}}; + result = spec_program.run(inputs, compatible_outputs, {}, std::nullopt); + CUDF_TEST_EXPECT_TABLES_EQUAL(expected->view(), result->view()); +} + +TEST_F(RuntimeSupportTest, TransformProgramIncompatibleSpecifications) +{ + cudf::transform_input inputs[] = {a, b, cudf::scalar_column_view(t)}; + std::array outputs{cudf::transform_output{cudf::data_type{cudf::type_id::FLOAT32}, + cudf::output_nullability::ALL_VALID}}; + cudf::transform_program program{ + udf, cudf::udf_source_type::CUDA, cudf::null_aware::NO, std::nullopt, inputs, outputs, {}}; + + auto integers = cudf::test::fixed_width_column_wrapper{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}; + cudf::transform_input different_type[] = {a, integers, cudf::scalar_column_view(t)}; + EXPECT_THROW((void)program.run(different_type, outputs, {}, std::nullopt), std::invalid_argument); + + cudf::transform_input different_input_kind[] = {a, b, a}; + EXPECT_THROW((void)program.run(different_input_kind, outputs, {}, std::nullopt), + std::invalid_argument); + + std::array fewer_inputs{a, b}; + EXPECT_THROW((void)program.run(fewer_inputs, outputs, {}, std::nullopt), std::invalid_argument); + + std::array different_output_type{cudf::transform_output{cudf::data_type{cudf::type_id::INT32}, + cudf::output_nullability::ALL_VALID}}; + EXPECT_THROW((void)program.run(inputs, different_output_type, {}, std::nullopt), + std::invalid_argument); + + std::array fewer_outputs{}; + EXPECT_THROW((void)program.run(inputs, fewer_outputs, {}, std::nullopt), std::invalid_argument); +} + struct UnaryOperationIntegrationTest : public cudf::test::BaseFixture {}; template @@ -1140,6 +1210,28 @@ TEST_F(StringOperationTest, OutputOffsetted) std::nullopt); CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->get_column(0)); + + auto program_offsets = + cudf::test::fixed_width_column_wrapper{0, 4, 9, 15, 22, 30, 39}.release(); + std::vector> program_string_offsets; + program_string_offsets.push_back(std::move(program_offsets)); + cudf::transform_program program{cuda, + cudf::udf_source_type::CUDA, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + program_string_offsets}; + + auto int64_offsets = + cudf::test::fixed_width_column_wrapper{0, 4, 9, 15, 22, 30, 39}.release(); + std::vector> int64_string_offsets; + int64_string_offsets.push_back(std::move(int64_offsets)); + result = program.run(inputs, outputs, std::move(int64_string_offsets), std::nullopt); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->get_column(0)); + + EXPECT_THROW((void)program.run(inputs, outputs, {}, std::nullopt), std::invalid_argument); } TEST_F(StringOperationTest, OutputOffsettedMixed)