jitify 2 support

cpp/CMakeLists.txt

@@ -137,8 +137,8 @@ include(cmake/thirdparty/CUDF_GetDLPack.cmake)
 include(cmake/thirdparty/CUDF_GetLibcudacxx.cmake)
 # find or install GoogleTest
 include(cmake/thirdparty/CUDF_GetGTest.cmake)
-# Stringify libcudf and libcudacxx headers used in JIT operations
-include(cmake/Modules/StringifyJITHeaders.cmake)
+# preprocess jitify-able kernels
+include(cmake/Modules/JitifyPreprocessKernels.cmake)
 # find cuFile
 include(cmake/Modules/FindcuFile.cmake)
 
@@ -153,9 +153,6 @@ add_library(cudf
     src/ast/transform.cu
     src/binaryop/binaryop.cpp
     src/binaryop/compiled/binary_ops.cu
-    src/binaryop/jit/code/kernel.cpp
-    src/binaryop/jit/code/operation.cpp
-    src/binaryop/jit/code/traits.cpp
     src/labeling/label_bins.cu
     src/bitmask/null_mask.cu
     src/column/column.cu
@@ -256,7 +253,6 @@ add_library(cudf
     src/io/utilities/parsing_utils.cu
     src/io/utilities/type_conversion.cpp
     src/jit/cache.cpp
-    src/jit/launcher.cpp
     src/jit/parser.cpp
     src/jit/type.cpp
     src/join/cross_join.cu
@@ -302,8 +298,6 @@ add_library(cudf
     src/reshape/interleave_columns.cu
     src/reshape/tile.cu
     src/rolling/grouped_rolling.cu
-    src/rolling/jit/code/kernel.cpp
-    src/rolling/jit/code/operation.cpp
     src/rolling/rolling.cu
     src/round/round.cu
     src/scalar/scalar.cpp
@@ -387,7 +381,6 @@ add_library(cudf
     src/text/tokenize.cu
     src/transform/bools_to_mask.cu
     src/transform/encode.cu
-    src/transform/jit/code/kernel.cpp
     src/transform/mask_to_bools.cu
     src/transform/nans_to_nulls.cu
     src/transform/row_bit_count.cu
@@ -465,7 +458,7 @@ endif()
 target_compile_definitions(cudf PUBLIC "SPDLOG_ACTIVE_LEVEL=SPDLOG_LEVEL_${RMM_LOGGING_LEVEL}")
 
 # Compile stringified JIT sources first
-add_dependencies(cudf stringify_run)
+add_dependencies(cudf jitify_preprocess_run)
 
 # Specify the target module library dependencies
 target_link_libraries(cudf
@@ -476,9 +469,9 @@ target_link_libraries(cudf
                   rmm::rmm)
 
 if(CUDA_STATIC_RUNTIME)
-    target_link_libraries(cudf PUBLIC CUDA::nvrtc CUDA::cudart_static CUDA::cuda_driver)
+    target_link_libraries(cudf PUBLIC CUDA::cudart_static CUDA::cuda_driver)
 else()
-    target_link_libraries(cudf PUBLIC CUDA::nvrtc CUDA::cudart CUDA::cuda_driver)
+    target_link_libraries(cudf PUBLIC CUDA::cudart CUDA::cuda_driver)
 endif()
 
 # Add cuFile interface if available

cpp/cmake/Modules/JitifyPreprocessKernels.cmake

@@ -0,0 +1,70 @@
+#=============================================================================
+# Copyright (c) 2021, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#=============================================================================
+
+cmake_minimum_required(VERSION 3.18)
+
+file(MAKE_DIRECTORY "${CUDF_GENERATED_INCLUDE_DIR}/include/jit_preprocessed_files")
+
+# Create `jitify_preprocess` executable
+project(jitify_preprocess VERSION 2.0 LANGUAGES CXX CUDA)
+add_executable(jitify_preprocess "${JITIFY_INCLUDE_DIR}/jitify2_preprocess.cpp")
+
+target_link_libraries(jitify_preprocess CUDA::cudart ${CMAKE_DL_LIBS})
+
+function(jit_preprocess_files)
+    cmake_parse_arguments(ARG
+                          ""
+                          "SOURCE_DIRECTORY"
+                          "FILES"
+                          ${ARGN}
+                          )
+
+    foreach(ARG_FILE ${ARG_FILES})
+        set(ARG_OUTPUT ${CUDF_GENERATED_INCLUDE_DIR}/include/jit_preprocessed_files/${ARG_FILE}.jit)
+        list(APPEND JIT_PREPROCESSED_FILES "${ARG_OUTPUT}")
+        add_custom_command(WORKING_DIRECTORY ${ARG_SOURCE_DIRECTORY}
+                           DEPENDS jitify_preprocess
+                           OUTPUT ${ARG_OUTPUT}
+                           VERBATIM
+                           COMMAND jitify_preprocess ${ARG_FILE}
+                                    -o ${CUDF_GENERATED_INCLUDE_DIR}/include/jit_preprocessed_files
+                                    -v
+                                    -i
+                                    -m
+                                    -std=c++14
+                                    -remove-unused-globals
+                                    -D__CUDACC_RTC__
+                                    -I${CUDF_SOURCE_DIR}/include
+                                    -I${CUDF_SOURCE_DIR}/src
+                                    -I${LIBCUDACXX_INCLUDE_DIR}
+                                    -I${CUDAToolkit_INCLUDE_DIRS}
+                                    --no-preinclude-workarounds
+                                    --no-replace-pragma-once
+                           )
+    endforeach()
+    set(JIT_PREPROCESSED_FILES "${JIT_PREPROCESSED_FILES}" PARENT_SCOPE)
+endfunction()
+
+jit_preprocess_files(SOURCE_DIRECTORY      ${CUDF_SOURCE_DIR}/src
+                     FILES                 binaryop/jit/kernel.cu
+                                           transform/jit/kernel.cu
+                                           rolling/jit/kernel.cu
+                     )
+
+add_custom_target(jitify_preprocess_run DEPENDS ${JIT_PREPROCESSED_FILES})
+
+file(COPY "${LIBCUDACXX_INCLUDE_DIR}/" DESTINATION "${CUDF_GENERATED_INCLUDE_DIR}/include/libcudacxx")
+file(COPY "${LIBCXX_INCLUDE_DIR}"      DESTINATION "${CUDF_GENERATED_INCLUDE_DIR}/include/libcxx")