Merge branch 'main' into disable-large-extent-join-tests

Author: PointKernel

cpp/CMakeLists.txt

@@ -466,11 +466,6 @@ rtcx_embed_includes(
   cudf/cpp/src/jit INCLUDE_DIRECTORIES cudf/cpp/src
 )
 
-rtcx_embed_includes(
-  cudf_cuda_embed SOURCE_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/src/binaryop/jit DEST_DIRECTORY
-  cudf/cpp/src/binaryop/jit INCLUDE_DIRECTORIES cudf/cpp/src
-)
-
 rtcx_embed_includes(
   cudf_cuda_embed SOURCE_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/src/join/jit DEST_DIRECTORY
   cudf/cpp/src/join/jit INCLUDE_DIRECTORIES cudf/cpp/src

cpp/include/cudf/binaryop.hpp

@@ -7,10 +7,20 @@
 
 #include <cudf/column/column.hpp>
 #include <cudf/scalar/scalar.hpp>
+#include <cudf/types.hpp>
+#include <cudf/utilities/default_stream.hpp>
 #include <cudf/utilities/export.hpp>
 #include <cudf/utilities/memory_resource.hpp>
+#include <cudf/utilities/traits.hpp>
+
+#include <cuda/buffer>
+#include <cuda/std/type_traits>
+#include <cuda/stream>
 
 #include <memory>
+#include <string>
+#include <type_traits>
+#include <utility>
 
 /**
  * @file
@@ -333,6 +343,9 @@ std::unique_ptr<column> binary_operation(
  * @brief Performs a binary operation between two columns using a
  * user-defined PTX function.
  *
+ * @deprecated in release 26.12 and will be removed in a future release. Use `cudf::transform`
+ * instead.
+ *
  * The output contains the result of `op(lhs[i], rhs[i])` for all `0 <= i < lhs.size()`
  *
  * Regardless of the operator, the validity of the output value is the logical
@@ -352,7 +365,7 @@ std::unique_ptr<column> binary_operation(
  * @throw cudf::logic_error if @p lhs and @p rhs dtypes aren't numeric
  * @throw cudf::logic_error if @p output_type dtype isn't numeric
  */
-std::unique_ptr<column> binary_operation(
+[[deprecated("Use cudf::transform instead")]] std::unique_ptr<column> binary_operation(
   column_view const& lhs,
   column_view const& rhs,
   std::string const& ptx,

cpp/include/cudf/join/join.hpp

@@ -23,6 +23,7 @@
 #include <cstdint>
 #include <memory>
 #include <optional>
+#include <string>
 #include <utility>
 
 /**
@@ -522,8 +523,8 @@ filter_join_indices_output_size(
  * For both null-aware and non-null-aware predicates, the output value is uninitialized and must
  * be set on every execution path.
  * The remaining parameters correspond to columns in the left table followed by the right table.
- * The CUDA function may return `cudf::errc` instead of `void` to report evaluation errors. PTX UDFs
- * do not support returning errors.
+ * The CUDA function may return `cudf::errc` instead of `void` to report evaluation errors.
+ * PTX modules functions must return `cudf::errc` as a 32-bit integer status.
  *
  *
  * @throw std::invalid_argument if join_kind is not INNER_JOIN, LEFT_JOIN, or FULL_JOIN.

cpp/include/cudf/transform.hpp

@@ -9,15 +9,20 @@
 #include <cudf/column/scalar_column_view.hpp>
 #include <cudf/types.hpp>
 #include <cudf/udf.hpp>
+#include <cudf/utilities/default_stream.hpp>
 #include <cudf/utilities/export.hpp>
 #include <cudf/utilities/memory_resource.hpp>
 
 #include <cuda/buffer>
+#include <cuda/stream>
 
+#include <cstdint>
 #include <functional>
 #include <memory>
 #include <optional>
 #include <span>
+#include <string>
+#include <utility>
 #include <variant>
 #include <vector>
 
@@ -237,6 +242,8 @@ struct transform_program {
  * Computes:
  * `(outputs[i]...) =  UDF(inputs[i]...)`.
  *
+ * PTX UDFs are linked as complete modules. The function must follow the CUDF UDF ABI.
+ *
  * For both null-aware and non-null-aware UDFs, output values have the following initialization:
  * - Fixed-width outputs have uninitialized values.
  * - Mutable string outputs (`cuda::std::span<char>`, used with preallocated string offsets) contain
@@ -310,8 +317,9 @@ std::unique_ptr<table> transform(
  * The size of the resulting column is the `row_size` if provided, otherwise it is inferred from
  * the input and pre-allocated output columns.
  *
- * @param udf           The LTO-IR fragment containing the transform function to apply. The UDF must
- * be named `transform` and follow the CUDF UDF ABI
+ * @param udf           The LTO-IR, FATBIN, or PTX fragment containing the transform function to
+ * apply. The UDF must follow the CUDF UDF ABI and be named `transform`. For PTX, the first function
+ * and references to it are renamed to `transform` before linking. PTX input must be null-terminated
  * @param binary_type   The type of the LTO binary provided in `udf`
  * @param is_null_aware Signifies the UDF will receive row inputs as optional values
  * @param user_data     User-defined device data to pass to the UDF

cpp/include/cudf/udf.hpp

@@ -21,7 +21,8 @@ namespace CUDF_EXPORT cudf {
  */
 enum class lto_binary_type : uint8_t {
   LTO_IR,  //< LTO-IR binary
-  FATBIN   //< FATBIN binary
+  FATBIN,  //< FATBIN binary
+  PTX      //< Null-terminated PTX code string
 };
 
 /**

cpp/src/binaryop/binaryop.cpp

@@ -22,30 +22,34 @@
  */
 
 #include "compiled/binary_ops.hpp"
-#include "jit/cache.hpp"
-#include "jit/helpers.hpp"
-#include "jit/parser.hpp"
-#include "jit/util.hpp"
 
 #include <cudf/binaryop.hpp>
 #include <cudf/column/column_factories.hpp>
 #include <cudf/detail/binaryop.hpp>
 #include <cudf/detail/null_mask.hpp>
 #include <cudf/detail/nvtx/ranges.hpp>
-#include <cudf/replace.hpp>
 #include <cudf/scalar/scalar.hpp>
+#include <cudf/table/table.hpp>
 #include <cudf/table/table_view.hpp>
+#include <cudf/transform.hpp>
 #include <cudf/types.hpp>
 #include <cudf/utilities/error.hpp>
 #include <cudf/utilities/traits.hpp>
-#include <cudf/utilities/type_dispatcher.hpp>
 
 #include <cuda/std/optional>
 #include <cuda/stream>
 
 #include <transform/checked_arithmetic.hpp>
 
+#include <algorithm>
+#include <array>
+#include <cstdint>
+#include <optional>
+#include <span>
+#include <stdexcept>
 #include <string>
+#include <type_traits>
+#include <utility>
 
 namespace cudf {
 namespace binops {
@@ -149,45 +153,6 @@ bool is_same_scale_necessary(binary_operator op)
          op != binary_operator::MUL_OVERFLOW && op != binary_operator::DIV_OVERFLOW;
 }
 
-namespace jit {
-void binary_operation(mutable_column_view& out,
-                      column_view const& lhs,
-                      column_view const& rhs,
-                      std::string const& ptx,
-                      cuda::stream_ref stream)
-{
-  std::string const output_type_name = cudf::type_to_name(out.type());
-
-  std::string cuda_source =
-    cudf::jit::parse_single_function_ptx(ptx,
-                                         "GENERIC_BINARY_OP",
-                                         {
-                                           {0, output_type_name + " *"},
-                                           {1, cudf::type_to_name(lhs.type())},
-                                           {2, cudf::type_to_name(rhs.type())},
-                                         });
-
-  std::string kernel_reflection = rtcx::reflect_template("cudf::binops::jit::binary_op_kernel",
-                                                         output_type_name,
-                                                         cudf::type_to_name(lhs.type()),
-                                                         cudf::type_to_name(rhs.type()),
-                                                         "cudf::binops::jit::UserDefinedOp");
-
-  auto kernel = cudf::jit::get_udf_kernel(
-    "cudf/cpp/src/binaryop/jit/kernel.cu", kernel_reflection, cuda_source, {});
-  auto cfg = kernel.max_occupancy_config(0, 0);
-
-  kernel.launch_with({cfg.min_grid_size},
-                     {cfg.block_size},
-                     0,
-                     stream,
-                     static_cast<size_type>(out.size()),
-                     cudf::jit::get_data_ptr(out),
-                     cudf::jit::get_data_ptr(lhs),
-                     cudf::jit::get_data_ptr(rhs));
-}
-}  // namespace jit
-
 // Compiled Binary operation
 namespace compiled {
 
@@ -409,17 +374,25 @@ std::unique_ptr<column> binary_operation(column_view const& lhs,
 
   CUDF_EXPECTS((lhs.size() == rhs.size()), "Column sizes don't match");
 
-  auto [new_mask, null_count] = cudf::detail::bitmask_and(table_view({lhs, rhs}), stream, mr);
-  auto out =
-    make_fixed_width_column(output_type, lhs.size(), std::move(new_mask), null_count, stream, mr);
-
-  // Check for 0 sized data
-  if (lhs.is_empty() or rhs.is_empty()) return out;
-
-  auto out_view = out->mutable_view();
-  binops::jit::binary_operation(out_view, lhs, rhs, ptx, stream);
-  out->set_null_count(cudf::detail::null_count(out_view.null_mask(), 0, out->size(), stream));
-  return out;
+  if (lhs.is_empty()) { return make_empty_column(output_type); }
+
+  std::array<transform_input, 2> inputs{lhs, rhs};
+  std::array<transform_output, 1> outputs{transform_output{.type = output_type}};
+
+  auto result =
+    cudf::transform_lto(std::span{reinterpret_cast<uint8_t const*>(ptx.data()), ptx.size()},
+                        lto_binary_type::PTX,
+                        null_aware::NO,
+                        std::nullopt,
+                        inputs,
+                        outputs,
+                        {},
+                        std::nullopt,
+                        stream,
+                        mr.get_output_mr());
+
+  auto columns = result->release();
+  return std::move(columns.front());
 }
 }  // namespace detail
 

cpp/src/binaryop/jit/kernel.cu

@@ -371,6 +371,19 @@ rtcx::blob compile_fragment(char const* name,
 
 }  // namespace
 
+rtcx::binary_type as_rtcx_binary_type(lto_binary_type type)
+{
+  switch (type) {
+    case lto_binary_type::LTO_IR: return rtcx::binary_type::LTO_IR;
+    case lto_binary_type::FATBIN: return rtcx::binary_type::FATBIN;
+    case lto_binary_type::PTX: return rtcx::binary_type::PTX;
+    default:
+      CUDF_FAIL(
+        std::format("Unrecognized LTO binary type {} for LTO transform", static_cast<int>(type)),
+        std::invalid_argument);
+  }
+}
+
 kernel get_kernel(std::string const& name,
                   std::string const& source_file_id,
                   std::span<char const* const> header_include_names,

cpp/src/jit/cache.cpp

@@ -4,6 +4,7 @@
  */
 
 #pragma once
+#include <cudf/udf.hpp>
 #include <cudf/utilities/export.hpp>
 
 #include <cuda/stream>
@@ -72,6 +73,13 @@ struct [[nodiscard]] kernel {
   }
 };
 
+/**
+ * @brief Converts an LTO binary type to the corresponding RTCX binary type
+ * @param type The LTO binary type to convert
+ * @return The corresponding RTCX binary type
+ */
+rtcx::binary_type as_rtcx_binary_type(lto_binary_type type);
+
 /**
  * @brief Gets a kernel from an embedded CUDA source file
  * @param name Debug name for the kernel (used for caching and logging)