[FEA] Implement Multi-output AST JIT & IR CSE

cpp/include/cudf/transform.hpp

@@ -11,8 +11,10 @@
 #include <cudf/utilities/export.hpp>
 #include <cudf/utilities/memory_resource.hpp>
 
+#include <functional>
 #include <memory>
 #include <optional>
+#include <span>
 #include <variant>
 #include <vector>
 
@@ -285,8 +287,6 @@ std::unique_ptr<column> 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<column> 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<table> compute_table_jit(
+  table_view const& table,
+  std::span<std::reference_wrapper<ast::expression const> 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.
  *

cpp/src/jit/row_ir.cpp

@@ -8,10 +8,12 @@
 #include "runtime/context.hpp"
 
 #include <cudf/column/column_factories.hpp>
+#include <cudf/hashing/detail/hashing.hpp>
 
 #include <cuda/std/inplace_vector>
 
 #include <algorithm>
+#include <array>
 #include <format>
 #include <iostream>
 #include <numeric>
@@ -437,6 +439,21 @@ int32_t instance_context::add_output()
 
 int32_t instance_context::add_input(input in)
 {
+  if (auto* column = std::get_if<column_input>(&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<column_input>(&inputs_[i]);
+      if (existing != nullptr && existing->table_source == column->table_source &&
+          existing->column_index == column->column_index) {
+        return static_cast<int32_t>(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<int32_t>(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<var_info const> instance_context::get_input_vars() const { return inpu
 
 std::span<untyped_var_info const> instance_context::get_output_vars() const { return output_vars_; }
 
+size_t node::compute_hash() const
+{
+  auto h = std::hash<int>{}(static_cast<int>(op_));
+  h = cudf::hashing::detail::hash_combine(h, std::hash<int>{}(static_cast<int>(error_policy_)));
+  h = cudf::hashing::detail::hash_combine(h, std::hash<size_t>{}(reference_.index()));
+  std::visit(
+    [&](auto& r) {
+      using reference_type = std::decay_t<decltype(r)>;
+      if constexpr (!std::is_same_v<reference_type, std::monostate>) {
+        h = cudf::hashing::detail::hash_combine(h, std::hash<int32_t>{}(r.index));
+      }
+    },
+    reference_);
+  h = cudf::hashing::detail::hash_combine(
+    h, target_scale_.has_value() ? std::hash<int32_t>{}(*target_scale_) : 0);
+  for (auto& arg : args_) {
+    h = cudf::hashing::detail::hash_combine(h, arg->hash());
+  }
+  return h;
+}
+
 node::node(opcode op,
            std::optional<int32_t> 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<node> 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<node>(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<int>(info.id)),
                 std::invalid_argument);
   }
+
+  emitted_ = true;
 }
 
 std::unique_ptr<row_ir::node> 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<std::string, null_aware, output_nullability> ast_converter::generate_code(
-  target target_id, ast::expression const& expr, std::string_view function_name)
+std::tuple<std::string, null_aware, std::vector<output_nullability>> ast_converter::generate_code(
+  target target_id,
+  std::span<std::reference_wrapper<ast::expression const> 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<row_ir::node>(output_reference{0}, expr.accept(*this)));
+  for (auto& expression : expressions) {
+    auto output_id = instance_.add_output();
+    output_irs_.emplace_back(
+      std::make_unique<row_ir::node>(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<std::string, null_aware, output_nullability> 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<output_nullability> 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<std::string, null_aware, output_nullability> 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<column_view, scalar_column_view> get_column_view(scalar_input const& in)
@@ -869,21 +973,19 @@ std::variant<column_view, scalar_column_view> 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<std::reference_wrapper<ast::expression const> 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<std::variant<column_view, scalar_column_view>> inputs;
   std::vector<std::unique_ptr<column>> scalar_columns;
   std::vector<std::optional<int32_t>> 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<transform_output> 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<std::reference_wrapper<ast::expression const>, 1> expressions{filter};
+  return compute_table(target_id, expressions, left_table, right_table, function_name, stream, mr);
 }
 
 }  // namespace cudf::detail::row_ir