[FEA] LTO IR Support (3) - Replace JITIFY usage with LIBRTCX

.devcontainer/Dockerfile

@@ -34,7 +34,7 @@ ENV PYTHONUNBUFFERED="1"
 ENV PYTHONDONTWRITEBYTECODE="1"
 
 ENV HISTFILE="/home/coder/.cache/._bash_history"
-ENV LIBCUDF_KERNEL_CACHE_PATH="/home/coder/cudf/cpp/build/${PYTHON_PACKAGE_MANAGER}/cuda-${CUDA_VERSION}/latest/jitify_cache"
+ENV LIBCUDF_KERNEL_CACHE_PATH="/home/coder/cudf/cpp/build/${PYTHON_PACKAGE_MANAGER}/cuda-${CUDA_VERSION}/latest/libcudf_kernel_cache"
 
 ###
 # sccache configuration

cpp/CMakeLists.txt

@@ -20,7 +20,7 @@ rapids_cuda_init_architectures(CUDF)
 project(
   CUDF
   VERSION "${RAPIDS_VERSION}"
-  LANGUAGES C CXX CUDA
+  LANGUAGES C CXX CUDA ASM
 )
 
 if(CMAKE_CUDA_COMPILER_ID STREQUAL "NVIDIA" AND CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 12.2)
@@ -38,7 +38,6 @@ option(USE_NVTX "Build with NVTX support" ON)
 option(BUILD_TESTS "Configure CMake to build tests" ON)
 option(BUILD_BENCHMARKS "Configure CMake to build benchmarks" OFF)
 option(BUILD_SHARED_LIBS "Build cuDF shared libraries" ON)
-option(JITIFY_USE_CACHE "Use a file cache for JIT compiled kernels" ON)
 option(CUDF_BUILD_TESTUTIL "Whether to build the test utilities contained in libcudf" ON)
 mark_as_advanced(CUDF_BUILD_TESTUTIL)
 option(CUDF_LARGE_STRINGS_DISABLED "Build with large string support disabled" OFF)
@@ -110,7 +109,6 @@ message(VERBOSE "CUDF: Build with NVTX support: ${USE_NVTX}")
 message(VERBOSE "CUDF: Configure CMake to build tests: ${BUILD_TESTS}")
 message(VERBOSE "CUDF: Configure CMake to build benchmarks: ${BUILD_BENCHMARKS}")
 message(VERBOSE "CUDF: Build cuDF shared libraries: ${BUILD_SHARED_LIBS}")
-message(VERBOSE "CUDF: Use a file cache for JIT compiled kernels: ${JITIFY_USE_CACHE}")
 message(VERBOSE "CUDF: Build with per-thread default stream: ${CUDF_USE_PER_THREAD_DEFAULT_STREAM}")
 message(
   VERBOSE
@@ -345,9 +343,6 @@ endif()
 
 create_logger_macros(CUDF "cudf::default_logger()" include/cudf)
 
-# find jitify
-include(cmake/thirdparty/get_jitify.cmake)
-
 # find NVTX
 include(cmake/thirdparty/get_nvtx.cmake)
 
@@ -377,9 +372,6 @@ if(CUDF_BUILD_TESTUTIL)
   include(cmake/thirdparty/get_gtest.cmake)
 endif()
 
-# preprocess jitify-able kernels
-include(cmake/Modules/JitifyPreprocessKernels.cmake)
-
 # find KvikIO
 include(cmake/thirdparty/get_kvikio.cmake)
 
@@ -392,6 +384,9 @@ include(cmake/thirdparty/get_thread_pool.cmake)
 # find zstd
 include(cmake/thirdparty/get_zstd.cmake)
 
+# JIT Embedding helper functions
+include(librtcx/embed.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")
@@ -409,6 +404,56 @@ if(NOT BUILD_SHARED_LIBS)
   endif()
 endif()
 
+add_embed(cudf_cuda_embed)
+
+embed_includes(
+  cudf_cuda_embed SOURCE_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/librtcx/libcxx DEST_DIRECTORY
+  librtcx/libcxx INCLUDE_DIRECTORIES librtcx/libcxx
+)
+
+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(
+  cudf_cuda_embed SOURCE_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/src/jit DEST_DIRECTORY
+  cudf/cpp/src/jit INCLUDE_DIRECTORIES cudf/cpp/src
+)
+
+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(
+  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(
+  cudf_cuda_embed SOURCE_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/src/rolling DEST_DIRECTORY
+  cudf/cpp/src/rolling INCLUDE_DIRECTORIES cudf/cpp/src
+)
+
+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
+)
+
+get_target_property(LIBCUDACXX_RAW_INCLUDE_DIRS CCCL::libcudacxx INTERFACE_INCLUDE_DIRECTORIES)
+
+foreach(INC_DIR IN LISTS LIBCUDACXX_RAW_INCLUDE_DIRS)
+  cmake_path(GET INC_DIR FILENAME INC_DIR_NAME)
+
+  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)
+
 # ##################################################################################################
 # * library targets -------------------------------------------------------------------------------
 add_library(
@@ -963,20 +1008,11 @@ add_library(
   src/utilities/traits.cpp
   src/utilities/type_checks.cpp
   src/utilities/type_dispatcher.cpp
+  librtcx/rtcx.cpp
+  ${cudf_cuda_embed_SOURCE_DIR}/cudf_cuda_embed.s
 )
 
-# Anything that includes jitify needs to be compiled with _FILE_OFFSET_BITS=64 due to a limitation
-# in how conda builds glibc
-set_source_files_properties(
-  src/binaryop/binaryop.cpp
-  src/jit/cache.cpp
-  src/rolling/detail/rolling_fixed_window.cu
-  src/rolling/detail/rolling_variable_window.cu
-  src/rolling/grouped_rolling.cu
-  src/rolling/rolling.cu
-  src/transform/transform.cu
-  PROPERTIES COMPILE_DEFINITIONS "_FILE_OFFSET_BITS=64"
-)
+add_dependencies(cudf cudf_cuda_embed)
 
 set_property(
   SOURCE src/io/parquet/writer_impl.cu
@@ -1040,14 +1076,14 @@ target_compile_options(
 # Specify include paths for the current target and dependents
 target_include_directories(
   cudf
-  PUBLIC "$<BUILD_INTERFACE:${DLPACK_INCLUDE_DIR}>"
-         "$<BUILD_INTERFACE:${JITIFY_INCLUDE_DIR}>"
-         "$<BUILD_INTERFACE:${CUDF_SOURCE_DIR}/include>"
+  PUBLIC "$<BUILD_INTERFACE:${DLPACK_INCLUDE_DIR}>" "$<BUILD_INTERFACE:${CUDF_SOURCE_DIR}/include>"
          "$<BUILD_INTERFACE:${CUDF_GENERATED_INCLUDE_DIR}/include>"
   PRIVATE "$<BUILD_INTERFACE:${CUDF_SOURCE_DIR}/src>"
           "$<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}>"
   INTERFACE "$<INSTALL_INTERFACE:include>"
 )
 
@@ -1056,20 +1092,6 @@ target_compile_definitions(
               "$<BUILD_INTERFACE:$<$<COMPILE_LANGUAGE:CUDA>:${CUDF_CUDA_DEFINITIONS}>>"
 )
 
-# Disable Jitify log printing. See https://github.com/NVIDIA/jitify/issues/79
-target_compile_definitions(cudf PRIVATE "JITIFY_PRINT_LOG=0")
-
-if(JITIFY_USE_CACHE)
-  # Instruct src/jit/cache what version of cudf we are building so it can compute a cal-ver cache
-  # directory. We isolate this definition to the single source so it doesn't effect compiling
-  # caching for all of libcudf
-  set_property(
-    SOURCE src/jit/cache.cpp
-    APPEND
-    PROPERTY COMPILE_DEFINITIONS "JITIFY_USE_CACHE" "CUDF_VERSION=${PROJECT_VERSION}"
-  )
-endif()
-
 # Per-thread default stream
 if(CUDF_USE_PER_THREAD_DEFAULT_STREAM)
   target_compile_definitions(
@@ -1099,9 +1121,6 @@ target_compile_definitions(cudf PRIVATE $<$<BOOL:${CUDF_KVIKIO_REMOTE_IO}>:CUDF_
 # https://github.com/NVIDIA/cccl/pull/2844
 target_compile_definitions(cudf PRIVATE THRUST_FORCE_32_BIT_OFFSET_TYPE=1 CCCL_AVOID_SORT_UNROLL=1)
 
-# Compile stringified JIT sources first
-add_dependencies(cudf jitify_preprocess_run)
-
 # Specify the target module library dependencies
 target_link_libraries(
   cudf

cpp/include/cudf/detail/kernel_instance.cuh

@@ -0,0 +1,8 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+// This file serves as a placeholder for kernel entry instantiations, so NVRTC/NVCC can override it

cpp/include/cudf/detail/utilities/getenv_or.hpp

@@ -8,6 +8,7 @@
 #include <cudf/logger.hpp>
 
 #include <cstdlib>
+#include <optional>
 #include <sstream>
 #include <string>
 
@@ -43,6 +44,18 @@ T getenv_or(std::string_view env_var_name, T default_val)
   return converted_val;
 }
 
+template <typename T>
+std::optional<T> getenv_optional(std::string_view env_var_name)
+{
+  auto const env_val = std::getenv(env_var_name.data());
+  if (env_val == nullptr) { return std::nullopt; }
+
+  std::stringstream sstream(env_val);
+  T converted_val;
+  sstream >> converted_val;
+  return converted_val;
+}
+
 /**
  * @brief Specialization of getenv_or for bool, to allow common "ON"/"OFF" string values.
  */

cpp/librtcx/embed.hpp

@@ -266,7 +266,7 @@ embed_output generate_cxx_source_files_data(std::string_view id,
 
   auto cxx_header = std::format(
     R"***(
-// Auto-generated header for embedded files with ID: {}
+// Auto-generated header for embedded files
 #pragma once
 
 #include <cstdint>
@@ -321,7 +321,6 @@ constexpr std::uint8_t hash[{}] =
 }}
 )***",
     namespace_decl,
-    id,
     include_dirs.size(),
     include_dirs_list,
     file_indices_list,

cpp/librtcx/rtcx.hpp

@@ -809,7 +809,7 @@ void teardown();
  * appropriate CUDA string representations.
  */
 template <typename T>
-std::string reflect(T value);
+std::string reflect(T value) = delete;
 
 /**
  * @brief Reflect a boolean value into its CUDA string representation ("true" or "false")