[BUG] Hoist cudaGetDeviceProperties to cudf::context

cpp/CMakeLists.txt

@@ -413,6 +413,9 @@ include(cmake/thirdparty/get_nanoarrow.cmake)
 # find thread_pool
 include(cmake/thirdparty/get_thread_pool.cmake)
 
+# find xxhash
+include(cmake/thirdparty/get_xxhash.cmake)
+
 # find zstd
 include(cmake/thirdparty/get_zstd.cmake)
 
@@ -422,6 +425,9 @@ add_subdirectory(librtcx)
 # JIT Embedding helper functions
 include(librtcx/embed.cmake)
 
+# Pre-compiled fragment management helper functions
+include(cmake/Modules/AddFragment.cmake)
+
 # Workaround until https://github.com/rapidsai/rapids-cmake/issues/176 is resolved
 if(NOT BUILD_SHARED_LIBS)
   include("${rapids-cmake-dir}/export/find_package_file.cmake")
@@ -444,39 +450,39 @@ if(NOT BUILD_SHARED_LIBS)
   )
 endif()
 
-add_embed(cudf_cuda_embed)
+rtcx_add_embed(cudf_cuda_embed)
 
-embed_includes(
+rtcx_embed_includes(
   cudf_cuda_embed SOURCE_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/librtcx/libcxx DEST_DIRECTORY
   librtcx/libcxx INCLUDE_DIRECTORIES librtcx/libcxx
 )
 
-embed_includes(
+rtcx_embed_includes(
   cudf_cuda_embed SOURCE_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/include/cudf DEST_DIRECTORY
   cudf/cpp/include/cudf INCLUDE_DIRECTORIES cudf/cpp/include
 )
 
-embed_includes(
+rtcx_embed_includes(
   cudf_cuda_embed SOURCE_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/src/jit DEST_DIRECTORY
   cudf/cpp/src/jit INCLUDE_DIRECTORIES cudf/cpp/src
 )
 
-embed_includes(
+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
 )
 
-embed_includes(
+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
 )
 
-embed_includes(
+rtcx_embed_includes(
   cudf_cuda_embed SOURCE_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/src/rolling DEST_DIRECTORY
   cudf/cpp/src/rolling INCLUDE_DIRECTORIES cudf/cpp/src
 )
 
-embed_includes(
+rtcx_embed_includes(
   cudf_cuda_embed SOURCE_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/src/transform/jit DEST_DIRECTORY
   cudf/cpp/src/transform/jit INCLUDE_DIRECTORIES cudf/cpp/src
 )
@@ -486,13 +492,115 @@ get_target_property(LIBCUDACXX_RAW_INCLUDE_DIRS CCCL::libcudacxx INTERFACE_INCLU
 foreach(INC_DIR IN LISTS LIBCUDACXX_RAW_INCLUDE_DIRS)
   cmake_path(GET INC_DIR FILENAME INC_DIR_NAME)
 
-  embed_includes(
+  rtcx_embed_includes(
     cudf_cuda_embed SOURCE_DIRECTORY ${INC_DIR} DEST_DIRECTORY CCCL/libcudacxx/${INC_DIR_NAME}
     INCLUDE_DIRECTORIES CCCL/libcudacxx/${INC_DIR_NAME}
   )
 endforeach()
 
-embed(cudf_cuda_embed COMPRESSION zstd OUTPUT_DIRECTORY "${CUDF_GENERATED_INCLUDE_DIR}/rtcx_embed")
+rtcx_embed(
+  cudf_cuda_embed COMPRESSION zstd OUTPUT_DIRECTORY "${CUDF_GENERATED_INCLUDE_DIR}/rtcx_embed"
+)
+
+rtcx_add_embed(cudf_fragments)
+
+list(
+  APPEND
+  CUDF_PRECOMPILE_TYPES
+  uint8_t
+  uint16_t
+  uint32_t
+  uint64_t
+  int8_t
+  int16_t
+  int32_t
+  int64_t
+  float
+  double
+  numeric::decimal32
+  numeric::decimal64
+  numeric::decimal128
+  cudf::timestamp_D
+  cudf::timestamp_s
+  cudf::timestamp_ms
+  cudf::timestamp_us
+  cudf::timestamp_ns
+  cudf::duration_D
+  cudf::duration_s
+  cudf::duration_ms
+  cudf::duration_us
+  cudf::duration_ns
+)
+
+foreach(TYPE IN ITEMS ${CUDF_PRECOMPILE_TYPES})
+  set(FRAGMENT_NAME transform_kernel)
+  get_property(
+    FILE_INDEX
+    TARGET cudf_fragments__embed_props
+    PROPERTY EMBED_FILE_INDEX
+  )
+  set(VARIANT_NAME transform_kernel_${FILE_INDEX})
+  set(INSTANCE
+      "cudf::jit::transform_kernel<false, false, cudf::jit::type_list<cudf::jit::column_accessor<0ULL, cudf::column_device_view_core, ${TYPE}, false, 0>>, cudf::jit::type_list<cudf::jit::column_accessor<0ULL, cudf::mutable_column_device_view_core, ${TYPE}, false, 0>>>"
+  )
+  add_fragment(
+    cudf_fragments
+    FRAGMENT
+    ${VARIANT_NAME}
+    SOURCE
+    src/transform/jit/kernel.cu
+    KERNEL_INSTANCE
+    ${INSTANCE}
+    UDF_TYPE
+    "int(${TYPE} *, ${TYPE})"
+    DEFINITIONS
+    CUDF_LTO_MODE
+    ARRAY_IDS
+    ${FRAGMENT_NAME}_FILE_INDEX
+    ${FRAGMENT_NAME}_INSTANCE
+    ARRAY_VALUES
+    ${FILE_INDEX}
+    "${INSTANCE}"
+  )
+endforeach()
+
+foreach(TYPE IN ITEMS ${CUDF_PRECOMPILE_TYPES})
+  foreach(RHS_IS_SCALAR IN ITEMS "false" "true")
+    set(FRAGMENT_NAME transform_kernel)
+    get_property(
+      FILE_INDEX
+      TARGET cudf_fragments__embed_props
+      PROPERTY EMBED_FILE_INDEX
+    )
+    set(VARIANT_NAME transform_kernel_${FILE_INDEX})
+    set(INSTANCE
+        "cudf::jit::transform_kernel<false, false, cudf::jit::type_list<cudf::jit::column_accessor<0ULL, cudf::column_device_view_core, ${TYPE}, false, 0>, cudf::jit::column_accessor<1ULL, cudf::column_device_view_core, ${TYPE}, ${RHS_IS_SCALAR}, 0>>, cudf::jit::type_list<cudf::jit::column_accessor<0ULL, cudf::mutable_column_device_view_core, ${TYPE}, false, 0>>>"
+    )
+    add_fragment(
+      cudf_fragments
+      FRAGMENT
+      ${VARIANT_NAME}
+      SOURCE
+      src/transform/jit/kernel.cu
+      KERNEL_INSTANCE
+      ${INSTANCE}
+      UDF_TYPE
+      "int(${TYPE} *, ${TYPE}, ${TYPE})"
+      DEFINITIONS
+      CUDF_LTO_MODE
+      ARRAY_IDS
+      ${FRAGMENT_NAME}_FILE_INDEX
+      ${FRAGMENT_NAME}_INSTANCE
+      ARRAY_VALUES
+      ${FILE_INDEX}
+      "${INSTANCE}"
+    )
+  endforeach()
+endforeach()
+
+rtcx_embed(
+  cudf_fragments COMPRESSION none OUTPUT_DIRECTORY "${CUDF_GENERATED_INCLUDE_DIR}/rtcx_embed"
+)
 
 # ##################################################################################################
 # * library targets -------------------------------------------------------------------------------
@@ -1056,9 +1164,10 @@ add_library(
   src/utilities/type_checks.cpp
   src/utilities/type_dispatcher.cpp
   ${cudf_cuda_embed_SOURCE_DIR}/cudf_cuda_embed.s
+  ${cudf_fragments_SOURCE_DIR}/cudf_fragments.s
 )
 
-add_dependencies(cudf cudf_cuda_embed)
+add_dependencies(cudf cudf_cuda_embed cudf_fragments)
 
 set_property(
   SOURCE src/io/parquet/writer_impl.cu
@@ -1128,7 +1237,9 @@ target_include_directories(
           "$<BUILD_INTERFACE:${nanoarrow_SOURCE_DIR}/src>"
           "$<BUILD_INTERFACE:${FlatBuffers_SOURCE_DIR}/include>"
           "$<BUILD_INTERFACE:${ZSTD_INCLUDE_DIR}>"
+          "$<BUILD_INTERFACE:${CMAKE_CURRENT_LIST_DIR}/librtcx>"
           "$<BUILD_INTERFACE:${cudf_cuda_embed_INCLUDE_DIRS}>"
+          "$<BUILD_INTERFACE:${cudf_fragments_INCLUDE_DIRS}>"
   INTERFACE "$<INSTALL_INTERFACE:include>"
 )
 
@@ -1170,8 +1281,15 @@ target_compile_definitions(cudf PRIVATE THRUST_FORCE_32_BIT_OFFSET_TYPE=1 CCCL_A
 target_link_libraries(
   cudf
   PUBLIC CCCL::CCCL $<BUILD_LOCAL_INTERFACE:BS::thread_pool>
-  PRIVATE $<BUILD_LOCAL_INTERFACE:nvtx3::nvtx3-cpp> $<BUILD_LOCAL_INTERFACE:cuco::cuco> ZLIB::ZLIB
-          ${CUDF_nvcomp_TARGET} kvikio::kvikio ${CUDF_nanoarrow_TARGET} zstd rtcx::rtcx
+  PRIVATE $<BUILD_LOCAL_INTERFACE:nvtx3::nvtx3-cpp>
+          $<BUILD_LOCAL_INTERFACE:cuco::cuco>
+          ZLIB::ZLIB
+          ${CUDF_nvcomp_TARGET}
+          kvikio::kvikio
+          ${CUDF_nanoarrow_TARGET}
+          zstd
+          $<BUILD_LOCAL_INTERFACE:xxhash>
+          rtcx::rtcx
 )
 
 # When rmm is a static library being absorbed via whole-archive, strip nvtx3 from its public

cpp/benchmarks/CMakeLists.txt

@@ -1,6 +1,6 @@
 # =============================================================================
 # cmake-format: off
-# SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 # cmake-format: on
 # =============================================================================
@@ -386,11 +386,24 @@ ConfigureNVBench(CSV_WRITER_NVBENCH io/csv/csv_writer.cpp)
 # * ast benchmark ---------------------------------------------------------------------------------
 ConfigureNVBench(AST_NVBENCH ast/polynomials.cpp ast/transform.cpp)
 
+# ##################################################################################################
+# * LTO Fragments ----------------------------------------------------------------------------
+rtcx_add_embed(cudf_benchmark_fragments)
+add_fragment(cudf_benchmark_fragments FRAGMENT add_f32 SOURCE binaryop/fragments/add_f32.cu)
+add_fragment(cudf_benchmark_fragments FRAGMENT mul_f32 SOURCE binaryop/fragments/mul_f32.cu)
+rtcx_embed(
+  cudf_benchmark_fragments COMPRESSION none OUTPUT_DIRECTORY
+  "${CUDF_GENERATED_INCLUDE_DIR}/rtcx_embed"
+)
+
 # ##################################################################################################
 # * binaryop benchmark ----------------------------------------------------------------------------
 ConfigureNVBench(
   BINARYOP_NVBENCH binaryop/binaryop.cpp binaryop/compiled_binaryop.cpp binaryop/polynomials.cpp
+  ${cudf_benchmark_fragments_SOURCE_DIR}/cudf_benchmark_fragments.s
 )
+target_include_directories(BINARYOP_NVBENCH PRIVATE ${cudf_benchmark_fragments_SOURCE_DIR})
+add_dependencies(BINARYOP_NVBENCH cudf_benchmark_fragments)
 
 # ##################################################################################################
 # * transform benchmark

cpp/benchmarks/binaryop/compiled_binaryop.cpp

@@ -1,13 +1,15 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #include <benchmarks/common/generate_input.hpp>
 #include <benchmarks/common/memory_stats.hpp>
 
 #include <cudf/binaryop.hpp>
+#include <cudf/transform.hpp>
 
+#include <cudf_benchmark_fragments.hpp>
 #include <nvbench/nvbench.cuh>
 
 template <typename TypeLhs, typename TypeRhs, typename TypeOut>
@@ -72,6 +74,7 @@ BINARYOP_BENCHMARK_DEFINE(timestamp_s,  duration_s,   ADD,                  time
 BINARYOP_BENCHMARK_DEFINE(duration_s,   duration_D,   SUB,                  duration_ms);
 BINARYOP_BENCHMARK_DEFINE(int64_t,      int64_t,      SUB,                  int64_t);
 BINARYOP_BENCHMARK_DEFINE(float,        float,        MUL,                  int64_t);
+BINARYOP_BENCHMARK_DEFINE(float,        float,        MUL,                  float);
 BINARYOP_BENCHMARK_DEFINE(duration_s,   int64_t,      MUL,                  duration_s);
 BINARYOP_BENCHMARK_DEFINE(int64_t,      int64_t,      DIV,                  int64_t);
 BINARYOP_BENCHMARK_DEFINE(duration_ms,  int32_t,      DIV,                  duration_ms);
@@ -101,3 +104,115 @@ BINARYOP_BENCHMARK_DEFINE(duration_ms,  duration_ns,  NULL_EQUALS,          bool
 BINARYOP_BENCHMARK_DEFINE(duration_ms,  duration_ns,  NULL_NOT_EQUALS,      bool);
 BINARYOP_BENCHMARK_DEFINE(decimal32,    decimal32,    NULL_MAX,             decimal32);
 BINARYOP_BENCHMARK_DEFINE(timestamp_D,  timestamp_s,  NULL_MIN,             timestamp_s);
+// clang-format on
+
+template <typename TypeLhs, typename TypeRhs, typename TypeOut>
+void BM_jit_binaryop(nvbench::state& state, cudf::binary_operator binop)
+{
+  constexpr auto const jit_mul_cuda = R"***(
+__device__ void transform(float* out, float a, float b) {
+  *out = a * b;
+}
+)***";
+
+  constexpr auto const jit_add_cuda = R"***(
+__device__ void transform(float* out, float a, float b) {
+  *out = a + b;
+}
+)***";
+
+  auto const num_rows = static_cast<cudf::size_type>(state.get_int64("num_rows"));
+  auto const use_lto  = state.get_string("use_lto") == "true";
+  static_assert(std::is_same_v<TypeLhs, TypeRhs> && std::is_same_v<TypeRhs, TypeOut>);
+  static_assert(std::is_same_v<TypeLhs, float>);
+
+  auto const source_table = create_random_table(
+    {cudf::type_to_id<TypeLhs>(), cudf::type_to_id<TypeRhs>()}, row_count{num_rows});
+
+  auto lhs = cudf::column_view(source_table->get_column(0));
+  auto rhs = cudf::column_view(source_table->get_column(1));
+
+  size_t fragment_id = 0;
+  char const* cuda   = nullptr;
+
+  switch (binop) {
+    case cudf::binary_operator::ADD: {
+      fragment_id = cudf_benchmark_fragments::add_f32;
+      cuda        = jit_add_cuda;
+    } break;
+    case cudf::binary_operator::MUL: {
+      fragment_id = cudf_benchmark_fragments::mul_f32;
+      cuda        = jit_mul_cuda;
+    } break;
+    default: throw std::runtime_error("Unsupported binary operator for JIT benchmark");
+  }
+
+  // Call once for hot cache.
+  cudf::transform_input inputs[]   = {lhs, rhs};
+  cudf::transform_output outputs[] = {
+    {cudf::data_type{cudf::type_to_id<TypeOut>()}, cudf::output_nullability::ALL_VALID}};
+
+  auto const range = cudf_benchmark_fragments::file_ranges[fragment_id];
+  std::span<uint8_t const> udf{cudf_benchmark_fragments::files.subspan(range[0], range[1])};
+
+  auto result = use_lto ? cudf::transform_lto(udf,
+                                              cudf::lto_binary_type::FATBIN,
+                                              cudf::null_aware::NO,
+                                              std::nullopt,
+                                              inputs,
+                                              outputs,
+                                              {},
+                                              std::nullopt)
+                        : cudf::multi_transform(cuda,
+                                                cudf::udf_source_type::CUDA,
+                                                cudf::null_aware::NO,
+                                                std::nullopt,
+                                                inputs,
+                                                outputs,
+                                                {},
+                                                std::nullopt);
+
+  // use number of bytes read and written to global memory
+  state.add_global_memory_reads<TypeLhs>(num_rows);
+  state.add_global_memory_reads<TypeRhs>(num_rows);
+  state.add_global_memory_writes<TypeOut>(num_rows);
+
+  state.exec(nvbench::exec_tag::sync, [&](nvbench::launch&) {
+    [[maybe_unused]] auto result = use_lto ? cudf::transform_lto(udf,
+                                                                 cudf::lto_binary_type::FATBIN,
+                                                                 cudf::null_aware::NO,
+                                                                 std::nullopt,
+                                                                 inputs,
+                                                                 outputs,
+                                                                 {},
+                                                                 std::nullopt)
+                                           : cudf::multi_transform(cuda,
+                                                                   cudf::udf_source_type::CUDA,
+                                                                   cudf::null_aware::NO,
+                                                                   std::nullopt,
+                                                                   inputs,
+                                                                   outputs,
+                                                                   {},
+                                                                   std::nullopt);
+  });
+}
+
+#define BM_JIT_BINARYOP_BENCHMARK_DEFINE(name, lhs, rhs, bop, tout)                    \
+  static void name(::nvbench::state& st)                                               \
+  {                                                                                    \
+    ::BM_jit_binaryop<lhs, rhs, tout>(st, ::cudf::binary_operator::bop);               \
+  }                                                                                    \
+  NVBENCH_BENCH(name)                                                                  \
+    .set_name("jit_binary_op_" BM_STRINGIFY(name))                                     \
+    .add_int64_axis("num_rows", {10'000, 100'000, 1'000'000, 10'000'000, 100'000'000}) \
+    .add_string_axis("use_lto", {"true", "false"})
+
+#define build_name_jit(a, b, c, d) a##_##b##_##c##_##d##_jit
+
+#define JIT_BINARYOP_BENCHMARK_DEFINE(lhs, rhs, bop, tout) \
+  BM_JIT_BINARYOP_BENCHMARK_DEFINE(build_name_jit(bop, lhs, rhs, tout), lhs, rhs, bop, tout)
+
+// clang-format off
+JIT_BINARYOP_BENCHMARK_DEFINE(float,        float,      ADD,                  float);
+JIT_BINARYOP_BENCHMARK_DEFINE(float,        float,      MUL,                  float);
+// clang-format on