[FEA] Persistent JIT Program References

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<cudf::size_type>(state.get_int64("num_rows"));
+  auto const expression_depth = static_cast<cudf::size_type>(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<cudf::transform_program> program;
+  if (api == "transform_program") {
+    program = std::make_unique<cudf::transform_program>(table, expression);
+  } else {
+    // Populate the JIT cache before timing, matching transform_program construction.
+    cudf::compute_column_jit(table, expression);
+  }
+
+  state.add_global_memory_reads<int32_t>(static_cast<std::size_t>(num_rows) *
+                                         (expression_depth + 1));
+  state.add_global_memory_writes<int32_t>(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 <cudf::ast::ast_operator cmp_op, cudf::ast::ast_operator reduce_op>
 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});

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<cudf::size_type>(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<cudf::transform_input, 1> 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<cudf::transform_program> program;
+  if (api == "transform_program") {
+    program =
+      std::make_unique<cudf::transform_program>(udf,
+                                                cudf::udf_source_type::CUDA,
+                                                cudf::null_aware::NO,
+                                                std::nullopt,
+                                                inputs,
+                                                outputs,
+                                                std::span<std::unique_ptr<cudf::column> 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<int32_t>(num_rows);
+  state.add_global_memory_writes<int32_t>(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});

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<transform_input_spec> 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<transform_output_spec> 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> 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<void*> user_data,
+                    std::span<transform_input const> inputs,
+                    std::span<transform_output const> outputs,
+                    std::span<std::unique_ptr<column> 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<void*> user_data,
+                    std::span<transform_input_spec const> inputs,
+                    std::span<transform_output_spec const> 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<table> run(
+    std::span<transform_input const> inputs,
+    std::span<transform_output const> outputs,
+    std::vector<std::unique_ptr<column>>&& string_offsets,
+    std::optional<size_type> 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<column> 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.