Adopt LTO-IR For JIT Compilation

cpp/CMakeLists.txt

@@ -329,6 +329,31 @@ if(NOT BUILD_SHARED_LIBS)
   endif()
 endif()
 
+set(CUDF_LTO_FLAGS -fatbin --relocatable-device-code=true --gen-opt-lto --ltoir -dlto)
+set(CUDF_LTO_CXX_STANDARD 20)
+
+function(add_lto_ir_target target_name source_file)
+  get_target_property(libcudacxx_raw_includes CCCL::libcudacxx INTERFACE_INCLUDE_DIRECTORIES)
+  set(includes)
+  foreach(inc IN LISTS libcudacxx_raw_includes CUDAToolkit_INCLUDE_DIRS)
+    list(APPEND includes "-I${inc}")
+  endforeach()
+  set(SOURCE_FILE ${CMAKE_CURRENT_SOURCE_DIR}/${source_file})
+  add_custom_target(
+    ${target_name}
+    COMMAND
+      ${CMAKE_CUDA_COMPILER} -std=c++${CUDF_LTO_CXX_STANDARD} ${CUDF_LTO_FLAGS}
+      -I${CUDF_SOURCE_DIR}/include -I${CUDF_SOURCE_DIR}/src -I${CMAKE_CURRENT_SOURCE_DIR} ${includes} -D__CUDACC_RTC_ -DCUDF_RUNTIME_JIT  -x cu -c
+      ${SOURCE_FILE} -o ${target_name}
+    COMMENT "Generating LTO fatbin ${target_name}: ${SOURCE_FILE}"
+    WORKING_DIRECTORY ${CMAKE_BINARY_DIR}
+    DEPENDS ${SOURCE_FILE}
+    VERBATIM
+  )
+endfunction()
+
+add_lto_ir_target(cudf_lto_kernel.fatbin src/transform/jit/lto_def.cu)
+
 # ##################################################################################################
 # * library targets -------------------------------------------------------------------------------
 add_library(
@@ -814,6 +839,7 @@ add_library(
   src/transform/nans_to_nulls.cu
   src/transform/one_hot_encode.cu
   src/transform/row_bit_count.cu
+  src/transform/jit/launch.cpp
   src/transform/transform.cpp
   src/transpose/transpose.cu
   src/unary/cast_ops.cu
@@ -835,6 +861,8 @@ add_library(
   src/utilities/type_dispatcher.cpp
 )
 
+add_dependencies(cudf cudf_lto_kernel.fatbin)
+
 # 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(

cpp/src/transform/jit/launch.cpp

@@ -0,0 +1,28 @@
+
+
+#include <nvJitLink.h>
+
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+
+#define CHECK_JITLINK_ERROR(...)                          \
+  if (auto err = __VA_ARGS__; err != NVJITLINK_SUCCESS) { \
+    fprintf(stderr, "JITLink error: %d\n", err);          \
+    exit(EXIT_FAILURE);                                   \
+  }
+
+void load_jit_kernel()
+{
+  const char* smbuf = "-arch=sm_86";
+  nvJitLinkHandle handle;
+  const char* lopts[] = {"-lto", smbuf};
+  CHECK_JITLINK_ERROR(nvJitLinkCreate(&handle, 2, lopts));
+  CHECK_JITLINK_ERROR(nvJitLinkAddFile(
+    handle, NVJITLINK_INPUT_FATBIN, "/home/coder/cudf/cpp/build/cudf_lto_operation.fatbin"));
+  CHECK_JITLINK_ERROR(nvJitLinkAddFile(
+    handle, NVJITLINK_INPUT_FATBIN, "/home/coder/cudf/cpp/build/cudf_lto_operators.fatbin"));
+  CHECK_JITLINK_ERROR(nvJitLinkAddFile(
+    handle, NVJITLINK_INPUT_FATBIN, "/home/coder/cudf/cpp/build/cudf_lto_kernel.fatbin"));
+  CHECK_JITLINK_ERROR(nvJitLinkComplete(handle));
+}