diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile index 57ccf6302c5b..1b89799aa330 100644 --- a/.devcontainer/Dockerfile +++ b/.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 diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 63c292646773..ca95937ff1d4 100644 --- a/cpp/CMakeLists.txt +++ b/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 "$" - "$" - "$" + PUBLIC "$" "$" "$" PRIVATE "$" "$" "$" "$" + "$" + "$" INTERFACE "$" ) @@ -1056,20 +1092,6 @@ target_compile_definitions( "$:${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 $<$: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 diff --git a/cpp/cmake/Modules/JitifyPreprocessKernels.cmake b/cpp/cmake/Modules/JitifyPreprocessKernels.cmake deleted file mode 100644 index d035e1ea6ab4..000000000000 --- a/cpp/cmake/Modules/JitifyPreprocessKernels.cmake +++ /dev/null @@ -1,82 +0,0 @@ -# ============================================================================= -# cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 -# cmake-format: on -# ============================================================================= - -# Create `jitify_preprocess` executable -add_executable(jitify_preprocess "${JITIFY_INCLUDE_DIR}/jitify2_preprocess.cpp") - -target_compile_definitions(jitify_preprocess PRIVATE "_FILE_OFFSET_BITS=64") -rapids_cuda_set_runtime(jitify_preprocess USE_STATIC ON) -target_link_libraries(jitify_preprocess PUBLIC ${CMAKE_DL_LIBS}) - -# Take a list of files to JIT-compile and run them through jitify_preprocess. -function(jit_preprocess_files) - cmake_parse_arguments(ARG "" "SOURCE_DIRECTORY" "FILES" ${ARGN}) - - 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() - foreach(ARG_FILE ${ARG_FILES}) - set(ARG_OUTPUT ${CUDF_GENERATED_INCLUDE_DIR}/include/jit_preprocessed_files/${ARG_FILE}.jit.hpp) - get_filename_component(jit_output_directory "${ARG_OUTPUT}" DIRECTORY) - list(APPEND JIT_PREPROCESSED_FILES "${ARG_OUTPUT}") - - get_filename_component(ARG_OUTPUT_DIR "${ARG_OUTPUT}" DIRECTORY) - - # Note: need to pass _FILE_OFFSET_BITS=64 in COMMAND due to a limitation in how conda builds - # glibc - add_custom_command( - OUTPUT ${ARG_OUTPUT} - DEPENDS jitify_preprocess "${ARG_SOURCE_DIRECTORY}/${ARG_FILE}" - WORKING_DIRECTORY ${ARG_SOURCE_DIRECTORY} - VERBATIM - COMMAND ${CMAKE_COMMAND} -E make_directory "${jit_output_directory}" - COMMAND - "${CMAKE_COMMAND}" -E env LD_LIBRARY_PATH=${CUDAToolkit_LIBRARY_DIR} - $ ${ARG_FILE} -o ${ARG_OUTPUT_DIR} -i -std=c++20 - -remove-unused-globals -D_FILE_OFFSET_BITS=64 -D__CUDACC_RTC__ -DCUDF_RUNTIME_JIT - -I${CUDF_SOURCE_DIR}/include -I${CUDF_SOURCE_DIR}/src ${includes} - --no-preinclude-workarounds --no-replace-pragma-once --diag-suppress=47 --device-int128 - COMMENT "Custom command to JIT-compile files." - ) - endforeach() - set(JIT_PREPROCESSED_FILES - "${JIT_PREPROCESSED_FILES}" - PARENT_SCOPE - ) -endfunction() - -if(NOT (EXISTS "${CUDF_GENERATED_INCLUDE_DIR}/include")) - make_directory("${CUDF_GENERATED_INCLUDE_DIR}/include") -endif() - -jit_preprocess_files( - SOURCE_DIRECTORY ${CUDF_SOURCE_DIR}/src FILES binaryop/jit/kernel.cu rolling/jit/kernel.cu - transform/jit/kernel.cu join/jit/filter_join_kernel.cu -) - -add_custom_target( - jitify_preprocess_run - DEPENDS ${JIT_PREPROCESSED_FILES} - COMMENT "Target representing jitified files." -) - -# when a user requests CMake to clean the build directory -# -# * `cmake --build --target clean` -# * `cmake --build --clean-first` -# * ninja clean -# -# We also remove the jitify2 program cache as well. This ensures that we don't keep older versions -# of the programs in cache -set(cache_path "$ENV{HOME}/.cudf") -if(ENV{LIBCUDF_KERNEL_CACHE_PATH}) - set(cache_path "$ENV{LIBCUDF_KERNEL_CACHE_PATH}") -endif() -cmake_path(APPEND cache_path "${CUDF_VERSION}/") -set_target_properties(jitify_preprocess_run PROPERTIES ADDITIONAL_CLEAN_FILES "${cache_path}") diff --git a/cpp/cmake/thirdparty/get_jitify.cmake b/cpp/cmake/thirdparty/get_jitify.cmake deleted file mode 100644 index 4595b33c5160..000000000000 --- a/cpp/cmake/thirdparty/get_jitify.cmake +++ /dev/null @@ -1,25 +0,0 @@ -# ============================================================================= -# cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 -# cmake-format: on -# ============================================================================= - -# Jitify doesn't have a version :/ - -# This function finds Jitify and sets any additional necessary environment variables. -function(find_and_configure_jitify) - rapids_cpm_find( - jitify 2.0.0 - GIT_REPOSITORY https://github.com/NVIDIA/jitify.git - GIT_TAG 44e978b21fc8bdb6b2d7d8d179523c8350db72e5 # jitify2 branch as of 23rd Aug 2025 - GIT_SHALLOW FALSE - DOWNLOAD_ONLY TRUE - ) - set(JITIFY_INCLUDE_DIR - "${jitify_SOURCE_DIR}" - PARENT_SCOPE - ) -endfunction() - -find_and_configure_jitify() diff --git a/cpp/include/cudf/detail/kernel_instance.cuh b/cpp/include/cudf/detail/kernel_instance.cuh new file mode 100644 index 000000000000..2f6e002a543a --- /dev/null +++ b/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 diff --git a/cpp/include/cudf/detail/operation-udf.hpp b/cpp/include/cudf/detail/operation_udf.cuh similarity index 100% rename from cpp/include/cudf/detail/operation-udf.hpp rename to cpp/include/cudf/detail/operation_udf.cuh diff --git a/cpp/include/cudf/detail/utilities/getenv_or.hpp b/cpp/include/cudf/detail/utilities/getenv_or.hpp index 5003a11df084..1c282525b2dd 100644 --- a/cpp/include/cudf/detail/utilities/getenv_or.hpp +++ b/cpp/include/cudf/detail/utilities/getenv_or.hpp @@ -8,6 +8,7 @@ #include #include +#include #include #include @@ -43,6 +44,18 @@ T getenv_or(std::string_view env_var_name, T default_val) return converted_val; } +template +std::optional 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. */ diff --git a/cpp/librtcx/embed.hpp b/cpp/librtcx/embed.hpp index 30b45a19fd5b..d29c75bc6f88 100644 --- a/cpp/librtcx/embed.hpp +++ b/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 @@ -321,7 +321,6 @@ constexpr std::uint8_t hash[{}] = }} )***", namespace_decl, - id, include_dirs.size(), include_dirs_list, file_indices_list, diff --git a/cpp/librtcx/rtcx.hpp b/cpp/librtcx/rtcx.hpp index d0ea4bff235c..bc3ff8fd5569 100644 --- a/cpp/librtcx/rtcx.hpp +++ b/cpp/librtcx/rtcx.hpp @@ -809,7 +809,7 @@ void teardown(); * appropriate CUDA string representations. */ template -std::string reflect(T value); +std::string reflect(T value) = delete; /** * @brief Reflect a boolean value into its CUDA string representation ("true" or "false") diff --git a/cpp/librtcx/sha256.hpp b/cpp/librtcx/sha256.hpp index 14543fa684d8..1149477f4b38 100644 --- a/cpp/librtcx/sha256.hpp +++ b/cpp/librtcx/sha256.hpp @@ -13,26 +13,29 @@ namespace rtcx { -inline namespace detail { +namespace functions { -std::uint32_t ror(std::uint32_t x, std::uint32_t n) { return (x >> n) | (x << (32 - n)); } +inline std::uint32_t ror(std::uint32_t x, std::uint32_t n) { return (x >> n) | (x << (32 - n)); } -std::uint32_t ch(std::uint32_t x, std::uint32_t y, std::uint32_t z) { return z ^ (x & (y ^ z)); } +inline std::uint32_t ch(std::uint32_t x, std::uint32_t y, std::uint32_t z) +{ + return z ^ (x & (y ^ z)); +} -std::uint32_t maj(std::uint32_t x, std::uint32_t y, std::uint32_t z) +inline std::uint32_t maj(std::uint32_t x, std::uint32_t y, std::uint32_t z) { return ((x | y) & z) | (x & y); } -std::uint32_t sigma0(std::uint32_t x) { return ror(x, 2) ^ ror(x, 13) ^ ror(x, 22); } +inline std::uint32_t sigma0(std::uint32_t x) { return ror(x, 2) ^ ror(x, 13) ^ ror(x, 22); } -std::uint32_t sigma1(std::uint32_t x) { return ror(x, 6) ^ ror(x, 11) ^ ror(x, 25); } +inline std::uint32_t sigma1(std::uint32_t x) { return ror(x, 6) ^ ror(x, 11) ^ ror(x, 25); } -std::uint32_t gamma0(std::uint32_t x) { return ror(x, 7) ^ ror(x, 18) ^ (x >> 3); } +inline std::uint32_t gamma0(std::uint32_t x) { return ror(x, 7) ^ ror(x, 18) ^ (x >> 3); } -std::uint32_t gamma1(std::uint32_t x) { return ror(x, 17) ^ ror(x, 19) ^ (x >> 10); } +inline std::uint32_t gamma1(std::uint32_t x) { return ror(x, 17) ^ ror(x, 19) ^ (x >> 10); } -void put_be32(void* ptr, std::uint32_t value) +inline void put_be32(void* ptr, std::uint32_t value) { auto* p = (std::uint8_t*)ptr; p[0] = (value >> 24) & 0xff; @@ -41,14 +44,14 @@ void put_be32(void* ptr, std::uint32_t value) p[3] = (value >> 0) & 0xff; } -std::uint32_t get_be32(void const* ptr) +inline std::uint32_t get_be32(void const* ptr) { auto const* p = (std::uint8_t const*)ptr; return (std::uint32_t)p[0] << 24 | (std::uint32_t)p[1] << 16 | (std::uint32_t)p[2] << 8 | (std::uint32_t)p[3] << 0; } -} // namespace detail +} // namespace functions struct [[nodiscard]] sha256_hex_string { char data_[65]; // NOLINT(modernize-avoid-c-arrays) @@ -110,8 +113,8 @@ struct [[nodiscard]] sha256 { struct sha256_context { private: - static constexpr size_t BLOCK_SIZE = 64; - std::uint32_t state_[8] = // NOLINT(modernize-avoid-c-arrays) + static constexpr std::size_t BLOCK_SIZE = 64; + std::uint32_t state_[8] = // NOLINT(modernize-avoid-c-arrays) {0x6a09'e667ul, 0xbb67'ae85ul, 0x3c6e'f372ul, @@ -143,16 +146,16 @@ struct sha256_context { /* copy the state into 512-bits into W[0..15] */ for (i = 0; i < 16; i++, buf += sizeof(std::uint32_t)) - W[i] = get_be32(buf); + W[i] = functions::get_be32(buf); /* fill W[16..63] */ for (i = 16; i < 64; i++) - W[i] = gamma1(W[i - 2]) + W[i - 7] + gamma0(W[i - 15]) + W[i - 16]; + W[i] = functions::gamma1(W[i - 2]) + W[i - 7] + functions::gamma0(W[i - 15]) + W[i - 16]; -#define RND(a, b, c, d, e, f, g, h, i, ki) \ - t0 = h + sigma1(e) + ch(e, f, g) + ki + W[i]; \ - t1 = sigma0(a) + maj(a, b, c); \ - d += t0; \ +#define RND(a, b, c, d, e, f, g, h, i, ki) \ + t0 = h + functions::sigma1(e) + functions::ch(e, f, g) + ki + W[i]; \ + t1 = functions::sigma0(a) + functions::maj(a, b, c); \ + d += t0; \ h = t0 + t1; RND(S[0], S[1], S[2], S[3], S[4], S[5], S[6], S[7], 0, 0x428a'2f98); @@ -240,7 +243,7 @@ struct sha256_context { if (len_buf) { std::uint32_t left = 64 - len_buf; if (len < left) left = len; - memcpy(len_buf + buf_, data, left); + std::memcpy(len_buf + buf_, data, left); len_buf = (len_buf + left) & 63; len -= left; data = (data + left); @@ -254,7 +257,7 @@ struct sha256_context { len -= 64; } - if (len) memcpy(buf_, data, len); + if (len) std::memcpy(buf_, data, len); } sha256 finalize() @@ -276,7 +279,7 @@ struct sha256_context { /* copy output */ for (i = 0; i < 8; i++, digest += sizeof(std::uint32_t)) { - put_be32(digest, state_[i]); + functions::put_be32(digest, state_[i]); } return out; } diff --git a/cpp/src/binaryop/binaryop.cpp b/cpp/src/binaryop/binaryop.cpp index bb4a48629ec5..7a28e61319fb 100644 --- a/cpp/src/binaryop/binaryop.cpp +++ b/cpp/src/binaryop/binaryop.cpp @@ -43,8 +43,6 @@ #include -#include - #include namespace cudf { @@ -156,18 +154,24 @@ void binary_operation(mutable_column_view& out, {2, cudf::type_to_name(rhs.type())}, }); - std::string kernel_reflection = jitify2::reflection::Template("cudf::binops::jit::kernel_v_v") - .instantiate(output_type_name, // list of template arguments - cudf::type_to_name(lhs.type()), - cudf::type_to_name(rhs.type()), - std::string("cudf::binops::jit::UserDefinedOp")); - - cudf::jit::get_udf_kernel(*binaryop_jit_kernel_cu_jit, kernel_reflection, cuda_source) - ->configure_1d_max_occupancy(0, 0, nullptr, stream.value()) - ->launch(out.size(), - cudf::jit::get_data_ptr(out), - cudf::jit::get_data_ptr(lhs), - cudf::jit::get_data_ptr(rhs)); + 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(out.size()), + cudf::jit::get_data_ptr(out), + cudf::jit::get_data_ptr(lhs), + cudf::jit::get_data_ptr(rhs)); } } // namespace jit diff --git a/cpp/src/binaryop/jit/kernel.cu b/cpp/src/binaryop/jit/kernel.cu index 893c1866cb44..9f65e258e635 100644 --- a/cpp/src/binaryop/jit/kernel.cu +++ b/cpp/src/binaryop/jit/kernel.cu @@ -33,11 +33,12 @@ #include #pragma nv_hdrstop // The above headers are used by the kernel below and need to be included before - // it. Each UDF will have a different operation-udf.hpp generated for it, so we + // it. Each UDF will have a different operation_udf.cuh generated for it, so we // need to put this pragma before including it to avoid PCH mismatch. // clang-format off -#include +#include +#include // clang-format on namespace cudf { @@ -56,45 +57,31 @@ struct UserDefinedOp { }; template -CUDF_KERNEL void kernel_v_v(cudf::size_type size, - TypeOut* out_data, - TypeLhs* lhs_data, - TypeRhs* rhs_data) +__device__ void binary_op_kernel(cudf::size_type size, + void* p_out_data, + void* p_lhs_data, + void* p_rhs_data) { auto const start = cudf::detail::grid_1d::global_thread_id(); auto const step = cudf::detail::grid_1d::grid_stride(); - for (auto i = start; i < size; i += step) { - out_data[i] = TypeOpe::template operate(lhs_data[i], rhs_data[i]); - } -} - -template -CUDF_KERNEL void kernel_v_v_with_validity(cudf::size_type size, - TypeOut* out_data, - TypeLhs* lhs_data, - TypeRhs* rhs_data, - cudf::bitmask_type* output_mask, - cudf::bitmask_type const* lhs_mask, - cudf::size_type lhs_offset, - cudf::bitmask_type const* rhs_mask, - cudf::size_type rhs_offset) -{ - auto const start = cudf::detail::grid_1d::global_thread_id(); - auto const step = cudf::detail::grid_1d::grid_stride(); + auto* out_data = static_cast(p_out_data); + auto* lhs_data = static_cast(p_lhs_data); + auto* rhs_data = static_cast(p_rhs_data); for (auto i = start; i < size; i += step) { - bool output_valid = false; - out_data[i] = TypeOpe::template operate( - lhs_data[i], - rhs_data[i], - lhs_mask ? cudf::bit_is_set(lhs_mask, lhs_offset + i) : true, - rhs_mask ? cudf::bit_is_set(rhs_mask, rhs_offset + i) : true, - output_valid); - if (output_mask && !output_valid) cudf::clear_bit(output_mask, i); + out_data[i] = TypeOpe::template operate(lhs_data[i], rhs_data[i]); } } } // namespace jit } // namespace binops } // namespace cudf + +extern "C" __global__ void cudf_kernel_entry(cudf::size_type size, + void* out_data, + void* lhs_data, + void* rhs_data) +{ + CUDF_KERNEL_INSTANCE(size, out_data, lhs_data, rhs_data); +} diff --git a/cpp/src/jit/cache.cpp b/cpp/src/jit/cache.cpp index a2fb3d10fd51..2b063bceeee2 100644 --- a/cpp/src/jit/cache.cpp +++ b/cpp/src/jit/cache.cpp @@ -1,162 +1,344 @@ + /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ -#include "runtime/context.hpp" - -#include #include -#include #include +#include + +#include #include +#include +#include #include +#include +#include +#include + +namespace CUDF_EXPORT cudf { -namespace cudf { namespace { -// Get the directory in home to use for storing the cache -std::filesystem::path get_user_home_cache_dir() +rtcx::sha256 hash(std::span input) { - auto home_dir = std::getenv("HOME"); - if (home_dir != nullptr) { - return std::filesystem::path(home_dir) / ".cudf"; - } else { - return {}; + rtcx::sha256_context ctx; + ctx.update(std::span{reinterpret_cast(input.data()), input.size()}); + return ctx.finalize(); +} + +rtcx::sha256 hash(std::span inputs) +{ + rtcx::sha256_context ctx; + for (auto const* input : inputs) { + ctx.update(std::span{reinterpret_cast(input), std::strlen(input)}); } + return ctx.finalize(); } -// Default `LIBCUDF_KERNEL_CACHE_PATH` to `$HOME/.cudf/$CUDF_VERSION`. -// This definition can be overridden at compile time by specifying a -// `-DLIBCUDF_KERNEL_CACHE_PATH=/kernel/cache/path` CMake argument. -// Use `std::filesystem` for cross-platform path resolution and dir -// creation. This path is used in the `getCacheDir()` function below. -#if !defined(LIBCUDF_KERNEL_CACHE_PATH) -#define LIBCUDF_KERNEL_CACHE_PATH get_user_home_cache_dir() -#endif - -/** - * @brief Get the string path to the JITIFY kernel cache directory. - * - * This path can be overridden at runtime by defining an environment variable - * named `LIBCUDF_KERNEL_CACHE_PATH`. The value of this variable must be a path - * under which the process' user has read/write privileges. - * - * This function returns a path to the cache directory, creating it if it - * doesn't exist. - * - * The default cache directory is `$HOME/.cudf/$CUDF_VERSION`. If no overrides - * are used and if $HOME is not defined, returns an empty path and file - * caching is not used. - */ -std::filesystem::path get_cache_dir() +void install_file_set( + std::string_view target_dir, + std::span compressed_binary, + size_t uncompressed_size, + std::span file_ranges, // NOLINT(modernize-avoid-c-arrays) + std::span destinations, + std::string_view compression) { - // The environment variable always overrides the - // default/compile-time value of `LIBCUDF_KERNEL_CACHE_PATH` - auto kernel_cache_path_env = std::getenv("LIBCUDF_KERNEL_CACHE_PATH"); - auto kernel_cache_path = std::filesystem::path( - kernel_cache_path_env != nullptr ? kernel_cache_path_env : LIBCUDF_KERNEL_CACHE_PATH); - - // Cache path could be empty when env HOME is unset or LIBCUDF_KERNEL_CACHE_PATH is defined to be - // empty, to disallow use of file cache at runtime. - if (not kernel_cache_path.empty()) { - kernel_cache_path /= std::string{CUDF_STRINGIFY(CUDF_VERSION)}; - - // Make per device cache based on compute capability. This is to avoid multiple devices of - // different compute capability to access the same kernel cache. - int device = 0; - int cc_major = 0; - int cc_minor = 0; - CUDF_CUDA_TRY(cudaGetDevice(&device)); - CUDF_CUDA_TRY(cudaDeviceGetAttribute(&cc_major, cudaDevAttrComputeCapabilityMajor, device)); - CUDF_CUDA_TRY(cudaDeviceGetAttribute(&cc_minor, cudaDevAttrComputeCapabilityMinor, device)); - int const cc = cc_major * 10 + cc_minor; - - kernel_cache_path /= std::to_string(cc); - - try { - // `mkdir -p` the kernel cache path if it doesn't exist - std::filesystem::create_directories(kernel_cache_path); - } catch (std::exception const& e) { - // if directory creation fails for any reason, return empty path - return {}; + auto decompressed = rtcx::decompress_blob(compressed_binary, uncompressed_size, compression); + for (size_t i = 0; i < file_ranges.size(); ++i) { + auto file_data_range = file_ranges[i]; + auto file_data = std::span{decompressed.data() + file_data_range[0], file_data_range[1]}; + auto dst_path = destinations[i]; + auto target_path = std::format("{}/{}", target_dir, dst_path); + + std::filesystem::create_directories(std::filesystem::path{target_path}.parent_path()); + + std::ofstream file(target_path, std::ios::binary); + if (!file) { + throw std::runtime_error( + std::format("Failed to open file for writing at path: {}", target_path)); + } + + file.write(reinterpret_cast(file_data.data()), file_data.size()); + if (!file) { + throw std::runtime_error(std::format("Failed to write file at path: {}", target_path)); } } - return kernel_cache_path; } -std::string get_program_cache_dir() +std::string read_file_string(char const* path) { -#if defined(JITIFY_USE_CACHE) - return get_cache_dir().string(); -#else - return {}; -#endif + std::ifstream file(std::string{path}, std::ios::binary | std::ios::ate); + if (!file) { throw std::runtime_error(std::format("Failed to open file at path: {}", path)); } + + auto size = file.tellg(); + file.seekg(0, std::ios::beg); + + std::string contents(size, '\0'); + if (!file.read(contents.data(), size)) { + throw std::runtime_error(std::format("Failed to read file at path: {}", path)); + } + + return contents; +} + +void install_cudf_jit_files(std::string const& target_dir, std::string const& tmp_dir) +{ + // directory does not exist, so create it + auto tmp_dir_path_str = std::format("{}/cudf-jit-tmpdir_XXXXXX", tmp_dir); + char* tmp_dir_path = ::mkdtemp(tmp_dir_path_str.data()); + CUDF_EXPECTS( + tmp_dir_path != nullptr, + std::format("Failed to create temporary directory for JIT file installation in tmp dir: {}", + tmp_dir), + std::runtime_error); + + install_file_set(tmp_dir_path, + cudf_cuda_embed::files, + cudf_cuda_embed::files_uncompressed_size, + cudf_cuda_embed::file_ranges, + cudf_cuda_embed::file_destinations, + cudf_cuda_embed::files_compression); + + // rename the temporary directory to the target install directory + if (::rename(tmp_dir_path, target_dir.c_str()) == -1) { + auto errc = errno; + // another process created it + if (errc == ENOTEMPTY || errc == EEXIST) { + std::filesystem::remove_all(tmp_dir_path); + } else { + CUDF_FAIL( + std::format("Failed to install JIT files to target directory: {} with error ({}): {}", + target_dir, + errc, + std::strerror(errc)), + std::runtime_error); + } + } } } // namespace -jitify2::ProgramCache<>& jit::program_cache::get(jitify2::PreprocessedProgramData const& preprog) +jit_bundle_t::jit_bundle_t(std::string install_dir, rtcx::cache_t& cache) + : install_dir_{std::move(install_dir)}, cache_{&cache} +{ + ensure_installed(); +} + +void jit_bundle_t::ensure_installed() const { CUDF_FUNC_RANGE(); - std::lock_guard caches_lock(_caches_mutex); - auto existing_cache = _caches.find(preprog.name()); + auto expected_hash = get_hash(); + auto expected_path = std::format("{}/{}", install_dir_, expected_hash); - if (existing_cache == _caches.end() || _disabled.load(std::memory_order_seq_cst)) { - auto res = - _caches.emplace(preprog.name(), - std::make_unique>( - _kernel_limit_proc, preprog, nullptr, _cache_dir, _kernel_limit_disk)); - existing_cache = res.first; + if (!std::filesystem::exists(expected_path)) { + // ensure base install directory exists + std::filesystem::create_directories(install_dir_); + install_cudf_jit_files(expected_path.c_str(), cache_->get_tmp_dir()); + } else { + // directory exists, perform minor sanity check + CUDF_EXPECTS(std::filesystem::is_directory(expected_path), // throws if path does not exist + std::format("JIT install path ({}) exists but is not a directory", expected_path), + std::runtime_error); } +} - return *(existing_cache->second); +std::string jit_bundle_t::get_hash() const +{ + auto str = rtcx::sha256_hex_string::make(cudf_cuda_embed::hash); + return std::string{str.view()}; } -void jit::program_cache::clear() +std::string jit_bundle_t::get_directory() const { - CUDF_FUNC_RANGE(); - std::lock_guard caches_lock(_caches_mutex); + return std::format("{}/{}", install_dir_, get_hash()); +} - _caches.clear(); +std::vector jit_bundle_t::get_include_directories() const +{ + std::vector directories; + auto base_dir = get_directory(); + + for (auto dir : cudf_cuda_embed::include_directories) { + directories.emplace_back(std::format("{}/{}", base_dir, dir)); + } - // non-atomic - std::filesystem::remove_all(_cache_dir); + return directories; } -void jit::program_cache::enable(bool enable) +namespace { + +constexpr int MIN_CUDA_VERSION_PCH = 12800; // minimum CUDA version for the "--pch" NVRTC flag +constexpr int MIN_CUDA_VERSION_MINIMAL = + 12800; // minimum CUDA version for the "--minimal" NVRTC flag + +int32_t get_driver_version() { - _disabled.store(!enable, std::memory_order_seq_cst); + int32_t driver_version; + CUDF_CUDA_TRY(cudaDriverGetVersion(&driver_version)); + return driver_version; } -bool jit::program_cache::is_enabled() const { return !_disabled.load(std::memory_order_seq_cst); } +int32_t get_runtime_version() +{ + int32_t runtime_version; + CUDF_CUDA_TRY(cudaRuntimeGetVersion(&runtime_version)); + return runtime_version; +} + +int32_t get_current_device_compute_capability() +{ + int32_t device; + CUDF_CUDA_TRY(cudaGetDevice(&device)); + + cudaDeviceProp props; + CUDF_CUDA_TRY(cudaGetDeviceProperties(&props, device)); + + return props.major * 10 + props.minor; +} -std::unique_ptr jit::program_cache::create() +std::tuple compile_library( + char const* name, + char const* cuda_code, + std::span extra_header_include_names, + std::span extra_headers, + std::span name_expressions) { - auto const kernel_limit_proc = - cudf::detail::getenv_or("LIBCUDF_KERNEL_CACHE_LIMIT_PER_PROCESS", 10'000); - auto const kernel_limit_disk = - cudf::detail::getenv_or("LIBCUDF_KERNEL_CACHE_LIMIT_DISK", 100'000); - auto const disabled = cudf::detail::get_bool_env_or("LIBCUDF_KERNEL_CACHE_DISABLED", false); - auto const clear_cache = cudf::detail::get_bool_env_or("LIBCUDF_KERNEL_CACHE_CLEAR", false); + CUDF_FUNC_RANGE(); + + auto& ctx = cudf::get_context(); + auto& cfg = ctx.config(); + auto& bundle = ctx.jit_bundle(); + auto sm = get_current_device_compute_capability(); + auto runtime = get_runtime_version(); + + auto include_dirs = bundle.get_include_directories(); + auto pch_dir = ctx.get_jit_pch_dir(); + + auto use_pch = runtime >= MIN_CUDA_VERSION_PCH; + auto use_minimal = runtime >= MIN_CUDA_VERSION_MINIMAL; + + std::vector options; + + for (auto const& include_dir : include_dirs) { + options.emplace_back(std::format("-I{}", include_dir)); + } + + options.emplace_back(std::format("--gpu-architecture=sm_{}", sm)); + + options.emplace_back("--diag-suppress=47"); + options.emplace_back("--device-int128"); + + if (sm >= 100) { options.emplace_back("--device-float128"); } + + options.emplace_back("-std=c++20"); + options.emplace_back("--device-as-default-execution-space"); + options.emplace_back("--generate-line-info"); + options.emplace_back("--dopt=on"); + + if (use_minimal) { options.emplace_back("--minimal"); } + + if (use_pch) { + options.emplace_back("--pch"); + options.emplace_back(std::format("--pch-dir={}", pch_dir)); - // if kernel_limit_disk is zero, jitify will assign it the value of kernel_limit_proc. - // to avoid this, we treat zero as "disable disk caching" by not providing the cache dir. - auto cache_dir = kernel_limit_disk == 0 ? std::string{} : get_program_cache_dir(); + if (cfg.jit_verbose) { + options.emplace_back("--pch-verbose=true"); + options.emplace_back("--pch-messages=true"); + } else { + options.emplace_back("--pch-verbose=false"); + options.emplace_back("--pch-messages=false"); + } + } + + if (cfg.disable_cuda_cache) { options.emplace_back("--no-cache"); } + + if (cfg.dump_jit_trace) { options.emplace_back("--time=-"); } + + if (cfg.dump_jit_time_profile) { + options.emplace_back(std::format("--fdevice-time-trace=cudf_kernel_{}_trace", name)); + } + + std::vector options_cstr; + for (auto const& option : options) { + options_cstr.emplace_back(option.c_str()); + } - auto cache = - std::make_unique(kernel_limit_proc, kernel_limit_disk, cache_dir, disabled); + auto params = rtcx::compile_params{.name = name, + .source = cuda_code, + .header_include_names = extra_header_include_names, + .headers = extra_headers, + .options = options_cstr, + .name_expressions = name_expressions, + .target_type = rtcx::binary_type::CUBIN}; - if (clear_cache) { cache->clear(); } + auto cubin = rtcx::compile(params); + auto library = rtcx::load_library(cubin); + auto blob = rtcx::blob_t::from_buffer(std::move(cubin)); - return cache; + return std::make_tuple(library, std::make_shared(std::move(blob))); } -jitify2::ProgramCache<>& jit::get_program_cache(jitify2::PreprocessedProgramData const& preprog) +} // namespace + +kernel get_kernel(std::string const& name, + std::string const& source_file_id, + std::span header_include_names, + std::span headers, + std::string const& kernel_instance) { - return cudf::get_context().program_cache().get(preprog); + CUDF_FUNC_RANGE(); + + auto& cache = cudf::get_context().rtcx_cache(); + auto& bundle = cudf::get_context().jit_bundle(); + + auto runtime = get_runtime_version(); + auto driver = get_driver_version(); + auto sm = get_current_device_compute_capability(); + auto header_include_names_hash = hash(header_include_names).to_hex_string(); + auto headers_hash = hash(headers).to_hex_string(); + auto bundle_hash = bundle.get_hash(); + auto source_file = std::format("{}/{}", bundle.get_directory(), source_file_id); + + auto cache_key = std::format(R"***(cuLibrary +name={} +binary_type=CUBIN +cuda_runtime={} +cuda_driver={} +arch={} +bundle={} +source_file={} +header_include_names={} +headers={} +kernel_instance={} +)***", + name, + runtime, + driver, + sm, + bundle_hash, + source_file, + header_include_names_hash.view(), + headers_hash.view(), + kernel_instance); + + auto cache_key_sha256 = hash(cache_key); + + auto compile = [&] { + auto bundle_dir = cudf::get_context().jit_bundle().get_directory(); + auto source = read_file_string(source_file.c_str()); + return compile_library(name.c_str(), source.c_str(), header_include_names, headers, {}); + }; + + auto fut = + cache.get_or_add_library(cache_key_sha256, rtcx::library_compile_func::from_functor(compile)); + + auto lib = fut.get(); + return kernel{lib, lib->get_kernel("cudf_kernel_entry")}; } -} // namespace cudf + +} // namespace CUDF_EXPORT cudf diff --git a/cpp/src/jit/cache.hpp b/cpp/src/jit/cache.hpp index 0e6738fed176..e9b68f9e933b 100644 --- a/cpp/src/jit/cache.hpp +++ b/cpp/src/jit/cache.hpp @@ -1,63 +1,92 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once -#pragma GCC diagnostic ignored "-Wignored-attributes" // Work-around for JITIFY2's false-positive - // warnings when compiled with GCC13 - #include -#include +#include + +#include -#include -#include -#include -#include -#include +namespace CUDF_EXPORT cudf { -namespace cudf { -namespace jit { +struct [[nodiscard]] jit_bundle_t { + private: + std::string install_dir_; + rtcx::cache_t* cache_; -class program_cache { - std::mutex _caches_mutex; - std::unordered_map>> _caches; - int32_t _kernel_limit_proc; - int32_t _kernel_limit_disk; - std::filesystem::path _cache_dir; - std::atomic _disabled; + void ensure_installed() const; public: - program_cache(int32_t kernel_limit_proc, - int32_t kernel_limit_disk, - std::filesystem::path cache_dir, - bool disabled) - : _kernel_limit_proc{kernel_limit_proc}, - _kernel_limit_disk{kernel_limit_disk}, - _cache_dir{std::move(cache_dir)}, - _disabled{disabled} - { - } + jit_bundle_t(std::string install_dir, rtcx::cache_t& cache); + + [[nodiscard]] std::string get_hash() const; + + [[nodiscard]] std::string get_directory() const; + + [[nodiscard]] std::vector get_include_directories() const; +}; - program_cache(program_cache const&) = delete; - program_cache(program_cache&&) = delete; - program_cache& operator=(program_cache const&) = delete; - program_cache& operator=(program_cache&&) = delete; - ~program_cache() = default; +struct [[nodiscard]] kernel { + private: + rtcx::library _library; + rtcx::kernel_ref _kernel; - jitify2::ProgramCache<>& get(jitify2::PreprocessedProgramData const& preprog); + public: + kernel(rtcx::library lib, rtcx::kernel_ref kernel) : _library(std::move(lib)), _kernel(kernel) {} + kernel(kernel const&) = default; + kernel(kernel&&) = default; + kernel& operator=(kernel const&) = default; + kernel& operator=(kernel&&) = default; + ~kernel() = default; - void clear(); + rtcx::kernel_ref get() const { return _kernel; } - void enable(bool enable); + rtcx::kernel_occupancy_config max_occupancy_config(size_t dynamic_shared_memory_bytes, + int32_t block_size_limit) const + { + return _kernel.max_occupancy_config(dynamic_shared_memory_bytes, block_size_limit); + } - bool is_enabled() const; + void launch(rtcx::cuda_dim3 grid_dim, + rtcx::cuda_dim3 block_dim, + uint32_t shared_mem_bytes, + rmm::cuda_stream_view stream, + void** kernel_params) const + { + return _kernel.launch(grid_dim, block_dim, shared_mem_bytes, stream.value(), kernel_params); + } - static std::unique_ptr create(); + template + void launch_with(rtcx::cuda_dim3 grid_dim, + rtcx::cuda_dim3 block_dim, + uint32_t shared_mem_bytes, + rmm::cuda_stream_view stream, + Args&&... args) + requires(sizeof...(Args) > 0) + { + void const* params[] = {&args...}; // NOLINT(modernize-avoid-c-arrays) + launch(grid_dim, block_dim, shared_mem_bytes, stream, const_cast(params)); + } }; -jitify2::ProgramCache<>& get_program_cache(jitify2::PreprocessedProgramData const& preprog); +/** + * @brief Gets a kernel from an embedded CUDA source file + * @param name Debug name for the kernel (used for caching and logging) + * @param source_file_id Identifier for the embedded source file (used to locate the source and for + * caching) + * @param header_include_names Names of any additional embedded header files to include during + * compilation + * @param headers Contents of any additional embedded header files to include during compilation + * @param kernel_instance String identifier for the specific kernel instance being requested (used + * for caching) + */ +kernel get_kernel(std::string const& name, + std::string const& source_file_id, + std::span header_include_names, + std::span headers, + std::string const& kernel_instance); -} // namespace jit -} // namespace cudf +} // namespace CUDF_EXPORT cudf diff --git a/cpp/src/jit/helpers.cpp b/cpp/src/jit/helpers.cpp index d035f43efe28..f7d78c08360f 100644 --- a/cpp/src/jit/helpers.cpp +++ b/cpp/src/jit/helpers.cpp @@ -8,6 +8,8 @@ #include #include +#include +#include namespace cudf { namespace jit { @@ -61,7 +63,7 @@ std::map build_ptx_params(std::span ou if (has_user_data) { params.emplace(index++, "void *"); - params.emplace(index++, jitify2::reflection::reflect()); + params.emplace(index++, "cudf::size_type"); } for (auto& name : output_typenames) { @@ -87,32 +89,22 @@ std::vector input_type_names( return names; } -jitify2::Kernel get_udf_kernel(jitify2::PreprocessedProgramData const& preprocessed_program_data, - std::string const& kernel_name, - std::string const& cuda_source, - std::vector const& extra_options) +kernel get_udf_kernel(std::string const& source_file, + std::string const& kernel_name, + std::string const& cuda_source) { CUDF_FUNC_RANGE(); - int runtime_version; - CUDF_CUDA_TRY(cudaRuntimeGetVersion(&runtime_version)); + auto kernel_instance_source = std::format(R"***( + #define CUDF_KERNEL_INSTANCE {} + )***", + kernel_name); + char const* include_names[] = // NOLINT(modernize-avoid-c-arrays) + {"cudf/detail/operation_udf.cuh", "cudf/detail/kernel_instance.cuh"}; + char const* include_headers[] = // NOLINT(modernize-avoid-c-arrays) + {cuda_source.c_str(), kernel_instance_source.c_str()}; - constexpr int min_pch_cuda_version = 12800; // CUDA 12.8 - constexpr int min_minimal_cuda_version = 12800; // CUDA 12.8 - - std::vector options; - options.emplace_back("-arch=sm_."); - - if (runtime_version >= min_minimal_cuda_version) { options.emplace_back("-minimal"); } - - if (runtime_version >= min_pch_cuda_version) { options.emplace_back("-pch"); } - - for (auto& opt : extra_options) { - options.push_back(opt); - } - - return cudf::jit::get_program_cache(preprocessed_program_data) - .get_kernel(kernel_name, {}, {{"cudf/detail/operation-udf.hpp", cuda_source}}, options); + return get_kernel(source_file, source_file, include_names, include_headers, kernel_name); } } // namespace jit diff --git a/cpp/src/jit/helpers.hpp b/cpp/src/jit/helpers.hpp index 4a3f27b048c1..8a7685a63ebd 100644 --- a/cpp/src/jit/helpers.hpp +++ b/cpp/src/jit/helpers.hpp @@ -75,10 +75,9 @@ column_views_to_device(std::span views, std::vector input_type_names( std::span const> views); -jitify2::Kernel get_udf_kernel(jitify2::PreprocessedProgramData const& preprocessed_program_data, - std::string const& kernel_name, - std::string const& cuda_source, - std::vector const& extra_options = {}); +kernel get_udf_kernel(std::string const& source_file, + std::string const& kernel_name, + std::string const& cuda_source); } // namespace jit } // namespace cudf diff --git a/cpp/src/join/filter_join_indices_jit.cu b/cpp/src/join/filter_join_indices_jit.cu index 15b532cac030..ae57a4263058 100644 --- a/cpp/src/join/filter_join_indices_jit.cu +++ b/cpp/src/join/filter_join_indices_jit.cu @@ -38,7 +38,6 @@ #include #include #include -#include #include #include @@ -48,55 +47,54 @@ namespace detail { namespace { -jitify2::StringVec build_join_filter_template_params( +std::vector build_join_filter_template_params( std::span inputs, std::span> table_sources, bool has_user_data, - null_aware is_null_aware) + bool is_null_aware) { - jitify2::StringVec template_params; - template_params.emplace_back(jitify2::reflection::reflect(has_user_data)); - template_params.emplace_back(jitify2::reflection::reflect(is_null_aware)); + std::vector template_params; + template_params.emplace_back(rtcx::reflect(has_user_data)); + template_params.emplace_back(rtcx::reflect(is_null_aware)); - jitify2::StringVec accessors; + std::vector accessors; for (size_t i = 0; i < inputs.size(); ++i) { auto const& input = inputs[i]; if (auto* col = std::get_if(&input)) { - auto element = cudf::type_to_name(col->type()); - accessors.emplace_back( - jitify2::reflection::Template("cudf::jit::column_accessor") - .instantiate( - i, "cudf::column_device_view_core", element, false, table_sources[i].value())); + accessors.emplace_back(rtcx::reflect_template("cudf::jit::column_accessor", + rtcx::reflect(i), + "cudf::column_device_view_core", + cudf::type_to_name(col->type()), + rtcx::reflect(false), + rtcx::reflect(table_sources[i].value()))); } else { auto& scalar = std::get(input); - auto element = cudf::type_to_name(scalar.as_column_view().type()); - accessors.emplace_back(jitify2::reflection::Template("cudf::jit::column_accessor") - .instantiate( - i, - "cudf::column_device_view_core", - element, - true, - 0 // scalars don't belong to a table, so just use 0 as placeholder - )); + accessors.emplace_back(rtcx::reflect_template( + "cudf::jit::column_accessor", + rtcx::reflect(i), + "cudf::column_device_view_core", + cudf::type_to_name(scalar.as_column_view().type()), + rtcx::reflect(true), + rtcx::reflect(0) // scalars don't belong to a table, so just use 0 as placeholder + )); } } - template_params.push_back( - jitify2::reflection::Template("cudf::jit::type_list").instantiate(accessors)); + template_params.push_back(rtcx::reflect_template("cudf::jit::type_list", accessors)); return template_params; } // Build the JIT kernel for join filtering -jitify2::ConfiguredKernel build_join_filter_kernel(std::string const& predicate_code, - std::span inputs, - std::span> table_sources, - bool is_ptx, - bool has_user_data, - null_aware is_null_aware, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) +kernel build_join_filter_kernel(std::string const& predicate_code, + std::span inputs, + std::span> table_sources, + bool is_ptx, + bool has_user_data, + bool is_null_aware, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { CUDF_FUNC_RANGE(); @@ -123,18 +121,15 @@ jitify2::ConfiguredKernel build_join_filter_kernel(std::string const& predicate_ // Build template parameters and kernel name auto template_args = build_join_filter_template_params(inputs, table_sources, has_user_data, is_null_aware); - auto kernel_name = - jitify2::reflection::Template("cudf::join::jit::filter_join_kernel").instantiate(template_args); + auto kernel_name = rtcx::reflect_template("cudf::join::jit::filter_join_kernel", template_args); // Get compiled kernel - auto kernel = - cudf::jit::get_udf_kernel(*join_jit_filter_join_kernel_cu_jit, kernel_name, cuda_source); - - return kernel->configure_1d_max_occupancy(0, 0, nullptr, stream.value()); + return cudf::jit::get_udf_kernel( + "cudf/cpp/src/join/jit/filter_join_kernel.cu", kernel_name, cuda_source); } // Launch the JIT kernel for join filtering -void launch_join_filter_kernel(jitify2::ConfiguredKernel& kernel, +void launch_join_filter_kernel(kernel const& kernel, cudf::device_span left_indices, cudf::device_span right_indices, std::span inputs, @@ -173,7 +168,9 @@ void launch_join_filter_kernel(jitify2::ConfiguredKernel& kernel, &predicate_results, &user_data_ptr}; - kernel->launch_raw(args); + auto cfg = kernel.max_occupancy_config(0, 0); + + kernel.launch({cfg.min_grid_size}, {cfg.block_size}, 0, stream, args); } // Same join semantics handling as the AST version @@ -426,7 +423,7 @@ filter_join_indices_jit(cudf::table_view const& left, table_sources, is_ptx, false, // has_user_data = false for now - null_aware::NO, + false, stream, mr); @@ -482,23 +479,22 @@ filter_join_indices_jit(cudf::table_view const& left, auto filter_result = row_ir::ast_converter::filter( row_ir::target::CUDA, predicate, left, right, "filter_operation", stream, mr); - auto template_args = build_join_filter_template_params(filter_result.inputs, - filter_result.input_table_sources, - filter_result.user_data.has_value(), - filter_result.is_null_aware); + auto template_args = + build_join_filter_template_params(filter_result.inputs, + filter_result.input_table_sources, + filter_result.user_data.has_value(), + filter_result.is_null_aware == null_aware::YES); auto const cuda_source = cudf::jit::parse_single_function_cuda(filter_result.udf, "GENERIC_JOIN_FILTER_OP"); - auto kernel_name = - jitify2::reflection::Template("cudf::join::jit::filter_join_kernel").instantiate(template_args); - auto kernel = - cudf::jit::get_udf_kernel(*join_jit_filter_join_kernel_cu_jit, kernel_name, cuda_source); - auto configured_kernel = kernel->configure_1d_max_occupancy(0, 0, nullptr, stream.value()); + auto kernel_name = rtcx::reflect_template("cudf::join::jit::filter_join_kernel", template_args); + auto kernel = cudf::jit::get_udf_kernel( + "cudf/cpp/src/join/jit/filter_join_kernel.cu", kernel_name, cuda_source); // Allocate and compute predicate results auto predicate_results = rmm::device_uvector(left_indices.size(), stream); - launch_join_filter_kernel(configured_kernel, + launch_join_filter_kernel(kernel, left_indices, right_indices, filter_result.inputs, diff --git a/cpp/src/join/jit/filter_join_kernel.cu b/cpp/src/join/jit/filter_join_kernel.cu index 61f26d261067..c760df20af78 100644 --- a/cpp/src/join/jit/filter_join_kernel.cu +++ b/cpp/src/join/jit/filter_join_kernel.cu @@ -17,13 +17,14 @@ #include #pragma nv_hdrstop // The above headers are used by the kernel below and need to be included before - // it. Each UDF will have a different operation-udf.hpp generated for it, so we + // it. Each UDF will have a different operation_udf.cuh generated for it, so we // need to put this pragma before including it to avoid PCH mismatch. // clang-format off // This header is an inlined header that defines the GENERIC_JOIN_FILTER_OP function. It is placed here // so the symbols in the headers above can be used by it. -#include +#include +#include // clang-format on namespace cudf::join::jit { @@ -45,13 +46,13 @@ __device__ void execute_predicate_op(void* user_data, } } -template -CUDF_KERNEL void filter_join_kernel(cudf::size_type num_rows, - cudf::size_type const* __restrict__ left_indices, - cudf::size_type const* __restrict__ right_indices, - cudf::column_device_view_core const* __restrict__ columns, - bool* __restrict__ predicate_results, - void* __restrict__ user_data) +template +__device__ void filter_join_kernel(cudf::size_type num_rows, + cudf::size_type const* __restrict__ left_indices, + cudf::size_type const* __restrict__ right_indices, + cudf::column_device_view_core const* __restrict__ columns, + bool* __restrict__ predicate_results, + void* __restrict__ user_data) { auto const start = cudf::detail::grid_1d::global_thread_id(); auto const stride = cudf::detail::grid_1d::grid_stride(); @@ -67,7 +68,7 @@ CUDF_KERNEL void filter_join_kernel(cudf::size_type num_rows, // Each accessor receives both tables and both indices, and internally selects // the appropriate table based on whether it's a left or right accessor. - if constexpr (is_null_aware == null_aware::YES) { + if constexpr (is_null_aware) { // Null-aware path: pass optional inputs, get optional result cuda::std::optional result{false}; auto inputs = Accessors::map([&]() { @@ -96,3 +97,15 @@ CUDF_KERNEL void filter_join_kernel(cudf::size_type num_rows, } } // namespace cudf::join::jit + +extern "C" __global__ void cudf_kernel_entry( + cudf::size_type num_rows, + cudf::size_type const* __restrict__ left_indices, + cudf::size_type const* __restrict__ right_indices, + cudf::column_device_view_core const* __restrict__ columns, + bool* __restrict__ predicate_results, + void* __restrict__ user_data) +{ + CUDF_KERNEL_INSTANCE( + num_rows, left_indices, right_indices, columns, predicate_results, user_data); +} diff --git a/cpp/src/rolling/detail/rolling_udf.cuh b/cpp/src/rolling/detail/rolling_udf.cuh index 2981a14e64d8..d0f2641e3d78 100644 --- a/cpp/src/rolling/detail/rolling_udf.cuh +++ b/cpp/src/rolling/detail/rolling_udf.cuh @@ -21,8 +21,6 @@ #include -#include - #include namespace cudf { @@ -90,24 +88,29 @@ inline std::unique_ptr rolling_window_udf_impl( 0, stream, cudf::get_current_device_resource_ref()}; std::string kernel_reflection = - jitify2::reflection::Template("cudf::rolling::jit::rolling_window_kernel") // - .instantiate(cudf::type_to_name(input.type()), // list of template arguments - cudf::type_to_name(output->type()), - udf_agg._operator_name, - preceding_window_str, - following_window_str); - - cudf::jit::get_udf_kernel(*rolling_jit_kernel_cu_jit, kernel_reflection, cuda_source) - ->configure_1d_max_occupancy(0, 0, nullptr, stream.value()) - ->launch(input.size(), - cudf::jit::get_data_ptr(input), - input.null_mask(), - cudf::jit::get_data_ptr(output_view), - output_view.null_mask(), - device_valid_count.data(), - preceding_window, - following_window, - min_periods); + rtcx::reflect_template("cudf::rolling::jit::rolling_window_kernel", + cudf::type_to_name(input.type()), // list of template arguments + cudf::type_to_name(output->type()), + udf_agg._operator_name, + preceding_window_str, + following_window_str); + + auto kernel = + cudf::jit::get_udf_kernel("cudf/cpp/src/rolling/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, + input.size(), + cudf::jit::get_data_ptr(input), + input.null_mask(), + cudf::jit::get_data_ptr(output->mutable_view()), + output_view.null_mask(), + device_valid_count.data(), + preceding_window, + following_window, + min_periods); output->set_null_count(output->size() - device_valid_count.value(stream)); diff --git a/cpp/src/rolling/jit/kernel.cu b/cpp/src/rolling/jit/kernel.cu index e13add4e326a..141758a71666 100644 --- a/cpp/src/rolling/jit/kernel.cu +++ b/cpp/src/rolling/jit/kernel.cu @@ -10,10 +10,11 @@ #include #pragma nv_hdrstop // The above headers are used by the kernel below and need to be included before - // it. Each UDF will have a different operation-udf.hpp generated for it, so we + // it. Each UDF will have a different operation_udf.cuh generated for it, so we // need to put this pragma before including it to avoid PCH mismatch. -#include +#include +#include struct rolling_udf_ptx { template @@ -44,20 +45,22 @@ template -CUDF_KERNEL void rolling_window_kernel(cudf::size_type nrows, - InType const* const __restrict__ in_col, - cudf::bitmask_type const* const __restrict__ in_col_valid, - OutType* __restrict__ out_col, - cudf::bitmask_type* __restrict__ out_col_valid, - cudf::size_type* __restrict__ output_valid_count, - detail::window_wrapper_base b_preceding_window_begin, - detail::window_wrapper_base b_following_window_begin, - cudf::size_type min_periods) +__device__ void rolling_window_kernel(cudf::size_type nrows, + void const* __restrict__ p_in_col, + cudf::bitmask_type const* const __restrict__ in_col_valid, + void* __restrict__ p_out_col, + cudf::bitmask_type* __restrict__ out_col_valid, + cudf::size_type* __restrict__ output_valid_count, + detail::window_wrapper_base b_preceding_window_begin, + detail::window_wrapper_base b_following_window_begin, + cudf::size_type min_periods) { auto i = cudf::detail::grid_1d::global_thread_id(); auto const stride = cudf::detail::grid_1d::grid_stride(); PrecedingWindowType const preceding_window_begin = b_preceding_window_begin; FollowingWindowType const following_window_begin = b_following_window_begin; + auto const* const in_col = static_cast(p_in_col); + auto* const out_col = static_cast(p_out_col); cudf::size_type warp_valid_count{0}; @@ -109,3 +112,25 @@ CUDF_KERNEL void rolling_window_kernel(cudf::size_type nrows, } // namespace jit } // namespace rolling } // namespace cudf + +extern "C" __global__ void cudf_kernel_entry( + cudf::size_type nrows, + void const* __restrict__ in_col, + cudf::bitmask_type const* __restrict__ in_col_valid, + void* __restrict__ out_col, + cudf::bitmask_type* __restrict__ out_col_valid, + cudf::size_type* __restrict__ output_valid_count, + cudf::detail::window_wrapper_base preceding_window_begin, + cudf::detail::window_wrapper_base following_window_begin, + cudf::size_type min_periods) +{ + CUDF_KERNEL_INSTANCE(nrows, + in_col, + in_col_valid, + out_col, + out_col_valid, + output_valid_count, + preceding_window_begin, + following_window_begin, + min_periods); +} diff --git a/cpp/src/runtime/context.cpp b/cpp/src/runtime/context.cpp index 9e945c6a2446..2a15321862c0 100644 --- a/cpp/src/runtime/context.cpp +++ b/cpp/src/runtime/context.cpp @@ -12,12 +12,13 @@ #include #include +#include #include namespace cudf { -context::context(context_config const& cfg, init_flags flags) - : _config{cfg}, _program_cache_init_flag{}, _program_cache{nullptr} +context::context(context_config cfg, init_flags flags) + : _config{std::move(cfg)}, _jit_cache_init_flag{} { initialize_components(flags); } @@ -26,20 +27,54 @@ void context::ensure_nvcomp_loaded() { io::detail::nvcomp::load_nvcomp_library() void context::ensure_jit_cache_initialized() { - std::call_once(_program_cache_init_flag, - [&]() { _program_cache = jit::program_cache::create(); }); + std::call_once(_jit_cache_init_flag, [&]() { + // make sure the required directories exist + std::filesystem::create_directories(_config.rtcx_cache_dir); + std::filesystem::create_directories(_config.jit_bundle_dir); + std::filesystem::create_directories(_config.jit_pch_dir); + std::filesystem::create_directories(_config.jit_tmp_dir); + + rtcx::initialize(); + + auto limits = rtcx::cache_limits{.num_mem_blobs = _config.kernel_cache_limit_process, + .num_mem_libraries = _config.kernel_cache_limit_process}; + + _rtcx_cache = std::make_unique(_config.rtcx_cache_dir, + _config.jit_tmp_dir, + limits, + bool{_config.preload_jit_cache}, + bool{_config.disable_jit_cache}); + + if (_config.clear_jit_cache) { + _rtcx_cache->clear_memory_store(); + _rtcx_cache->clear_disk_store(); + } + + // note that jit_bundle depends on rtcx_cache, so we ensure rtcx_cache is initialized first. + _jit_bundle = std::make_unique(_config.jit_bundle_dir, *_rtcx_cache); + }); } -jit::program_cache& context::program_cache() +context::~context() { rtcx::teardown(); } + +rtcx::cache_t& context::rtcx_cache() { ensure_jit_cache_initialized(); - return *_program_cache; + return *_rtcx_cache; +} + +jit_bundle_t& context::jit_bundle() +{ + ensure_jit_cache_initialized(); + return *_jit_bundle; } bool context::dump_codegen() const { return _config.dump_codegen; } bool context::use_jit() const { return _config.use_jit; } +std::string const& context::get_jit_pch_dir() const { return _config.jit_pch_dir; } + void context::initialize_components(init_flags flags) { if (has_flag(flags, init_flags::INIT_JIT_CACHE)) { ensure_jit_cache_initialized(); } @@ -47,6 +82,26 @@ void context::initialize_components(init_flags flags) if (has_flag(flags, init_flags::LOAD_NVCOMP)) { io::detail::nvcomp::load_nvcomp_library(); } } +std::filesystem::path get_cudf_kernel_cache_dir() +{ + if (auto libcudf_kernel_cache_path = + detail::getenv_optional("LIBCUDF_KERNEL_CACHE_PATH"); + libcudf_kernel_cache_path.has_value()) { + return std::filesystem::path(*libcudf_kernel_cache_path); + } + + if (auto home = detail::getenv_optional("HOME"); home.has_value()) { + return std::filesystem::path(*home) / ".libcudf"; + } + + CUDF_FAIL( + "Unable to determine the CUDF root directory. Please set the `LIBCUDF_KERNEL_CACHE_PATH` or " + "`HOME` " + "environment variables to allow automatic resolution of the root " + "directory.", + std::runtime_error); +} + static std::optional _context{std::nullopt}; static std::optional _context_init_flag{std::in_place}; static std::optional _context_deinit_flag{std::in_place}; @@ -58,15 +113,43 @@ namespace CUDF_EXPORT cudf { void initialize(init_flags flags) { std::call_once(*_context_init_flag, [&]() { - bool dump_codegen = cudf::detail::get_bool_env_or("LIBCUDF_JIT_DUMP_CODEGEN", false); - bool use_jit = cudf::detail::get_bool_env_or("LIBCUDF_JIT_ENABLED", false); + auto const dump_codegen = detail::get_bool_env_or("LIBCUDF_JIT_DUMP_CODEGEN", false); + auto const use_jit = detail::get_bool_env_or("LIBCUDF_JIT_ENABLED", false); + auto const preload_jit_cache = detail::get_bool_env_or("LIBCUDF_KERNEL_CACHE_PRELOAD", false); + auto const disable_jit_cache = detail::get_bool_env_or("LIBCUDF_KERNEL_CACHE_DISABLED", false); + auto const clear_jit_cache = detail::get_bool_env_or("LIBCUDF_KERNEL_CACHE_CLEAR", false); + auto const disable_cuda_cache = + detail::get_bool_env_or("LIBCUDF_JIT_DISABLE_CUDA_CACHE", false); + auto const jit_verbose = detail::get_bool_env_or("LIBCUDF_JIT_VERBOSE", false); + auto const dump_jit_trace = detail::get_bool_env_or("LIBCUDF_JIT_DUMP_TRACE", false); + auto const dump_jit_time_profile = + detail::get_bool_env_or("LIBCUDF_JIT_DUMP_TIME_PROFILE", false); + + auto const kernel_cache_limit_process = + detail::getenv_or("LIBCUDF_KERNEL_CACHE_LIMIT_PER_PROCESS", 16'384U); flags = flags | (use_jit ? init_flags::INIT_JIT_CACHE : init_flags::NONE); - context_config cfg{ - .dump_codegen = dump_codegen, - .use_jit = use_jit, - }; + auto const cache_dir = get_cudf_kernel_cache_dir(); + auto const jit_bundle_dir = cache_dir / "bundle"; + auto const rtcx_cache_dir = cache_dir / "rtcx_cache"; + auto const jit_pch_dir = cache_dir / "pch"; + auto const jit_tmp_dir = cache_dir / "tmp"; + + context_config cfg{.dump_codegen = dump_codegen, + .use_jit = use_jit, + .preload_jit_cache = preload_jit_cache, + .disable_jit_cache = disable_jit_cache, + .clear_jit_cache = clear_jit_cache, + .disable_cuda_cache = disable_cuda_cache, + .jit_verbose = jit_verbose, + .dump_jit_trace = dump_jit_trace, + .dump_jit_time_profile = dump_jit_time_profile, + .rtcx_cache_dir = rtcx_cache_dir, + .jit_bundle_dir = jit_bundle_dir, + .jit_pch_dir = jit_pch_dir, + .jit_tmp_dir = jit_tmp_dir, + .kernel_cache_limit_process = kernel_cache_limit_process}; _context.emplace(cfg, flags); }); @@ -85,9 +168,18 @@ void teardown() }); } -void enable_jit_cache(bool enable) { get_context().program_cache().enable(enable); } +void enable_jit_cache(bool enabled) +{ + auto& cache = get_context().rtcx_cache(); + cache.enable(enabled); +} -void clear_jit_cache() { get_context().program_cache().clear(); } +void clear_jit_cache() +{ + auto& cache = get_context().rtcx_cache(); + cache.clear_memory_store(); + cache.clear_disk_store(); +} context& get_context() { diff --git a/cpp/src/runtime/context.hpp b/cpp/src/runtime/context.hpp index 39d259dd2fb2..8477fcc89f7b 100644 --- a/cpp/src/runtime/context.hpp +++ b/cpp/src/runtime/context.hpp @@ -11,15 +11,33 @@ #include #include +namespace rtcx { +struct cache_t; +} // namespace rtcx + namespace cudf { namespace jit { class program_cache; } +struct jit_bundle_t; + struct [[nodiscard]] context_config { - bool dump_codegen = false; - bool use_jit = false; + bool dump_codegen : 1 = false; + bool use_jit : 1 = false; + bool preload_jit_cache : 1 = false; + bool disable_jit_cache : 1 = false; + bool clear_jit_cache : 1 = false; + bool disable_cuda_cache : 1 = false; + bool jit_verbose : 1 = false; + bool dump_jit_trace : 1 = false; + bool dump_jit_time_profile : 1 = false; + std::string rtcx_cache_dir = {}; + std::string jit_bundle_dir = {}; + std::string jit_pch_dir = {}; + std::string jit_tmp_dir = {}; + uint32_t kernel_cache_limit_process = 0; }; /// @brief The context object contains global state internal to CUDF. @@ -29,8 +47,9 @@ class context { public: private: context_config _config; - std::once_flag _program_cache_init_flag; - std::unique_ptr _program_cache; + std::once_flag _jit_cache_init_flag; + std::unique_ptr _rtcx_cache; + std::unique_ptr _jit_bundle; private: void ensure_nvcomp_loaded(); @@ -38,19 +57,25 @@ class context { void ensure_jit_cache_initialized(); public: - context(context_config const& cfg = {}, init_flags flags = init_flags::DEFAULT); + context(context_config cfg = {}, init_flags flags = init_flags::DEFAULT); context(context const&) = delete; context& operator=(context const&) = delete; context(context&&) = delete; context& operator=(context&&) = delete; - ~context() = default; + ~context(); + + rtcx::cache_t& rtcx_cache(); - jit::program_cache& program_cache(); + jit_bundle_t& jit_bundle(); [[nodiscard]] bool dump_codegen() const; [[nodiscard]] bool use_jit() const; + [[nodiscard]] context_config const& config() const { return _config; } + + [[nodiscard]] std::string const& get_jit_pch_dir() const; + /// @brief Initialize additional components based on the provided flags /// @param flags The initialization flags to process void initialize_components(init_flags flags); diff --git a/cpp/src/transform/jit/kernel.cu b/cpp/src/transform/jit/kernel.cu index 2456433d0bc6..f8887b93c2d3 100644 --- a/cpp/src/transform/jit/kernel.cu +++ b/cpp/src/transform/jit/kernel.cu @@ -22,13 +22,14 @@ #include #pragma nv_hdrstop // The above headers are used by the kernel below and need to be included before - // it. Each UDF will have a different operation-udf.hpp generated for it, so we + // it. Each UDF will have a different operation_udf.cuh generated for it, so we // need to put this pragma before including it to avoid PCH mismatch. // clang-format off // This header is an inlined header that defines the GENERIC_TRANSFORM_OP function. It is placed here // so the symbols in the headers above can be used by it. -#include +#include +#include // clang-format on namespace cudf { @@ -47,22 +48,18 @@ __device__ void execute_transform_op(void* user_data, size_type element_idx, Arg } /// @brief The generic transform kernel. Supports all types and nullability combinations. -template -CUDF_KERNEL void transform_kernel(size_type row_size, - bitmask_type const* __restrict__ stencil, - void* __restrict__ user_data, - column_device_view_core const* __restrict__ input_cols, - mutable_column_device_view_core const* __restrict__ output_cols) +template +__device__ void transform_kernel(size_type row_size, + bitmask_type const* __restrict__ stencil, + void* __restrict__ user_data, + column_device_view_core const* __restrict__ input_cols, + mutable_column_device_view_core const* __restrict__ output_cols) { - // TODO: ensure block size is a multiple of warp size for correct warp-synchronous behavior auto start = detail::grid_1d::global_thread_id(); auto stride = detail::grid_1d::grid_stride(); for (auto element_idx = start; element_idx < row_size; element_idx += stride) { - if constexpr (is_null_aware == null_aware::NO) { + if constexpr (!is_null_aware) { if (stencil != nullptr && !bit_is_set(stencil, element_idx)) { continue; } auto ins = InputAccessors::map( @@ -111,3 +108,22 @@ CUDF_KERNEL void transform_kernel(size_type row_size, } // namespace jit } // namespace cudf + +// The entry point for the JIT compiled kernel. This is the C-ABI function that will be used to +// retrieve the `CuFunction` for the kernel from the compiled module. This is because we don't want +// to track the scope-dependent C++ mangled name of the kernel, and can just use a fixed name to +// retrieve the `CuFunction` of the kernel. +// A C++-mangled symbol has ambiguous and complex resolution rules, and can change based on the +// scope of the function, the types of the arguments, and other factors that will not be known until +// after compilation. By using a fixed C-ABI symbol name for the kernel entry point, we can avoid +// these issues and ensure that we can always retrieve the correct `CuFunction` for the kernel +// regardless of the context in which it was compiled or used. +extern "C" __global__ void cudf_kernel_entry( + cudf::size_type row_size, + cudf::bitmask_type const* __restrict__ stencil, + void* __restrict__ user_data, + cudf::column_device_view_core const* __restrict__ input_cols, + cudf::mutable_column_device_view_core const* __restrict__ output_cols) +{ + CUDF_KERNEL_INSTANCE(row_size, stencil, user_data, input_cols, output_cols); +} diff --git a/cpp/src/transform/transform.cu b/cpp/src/transform/transform.cu index b635c1c42d50..80b0509cfddf 100644 --- a/cpp/src/transform/transform.cu +++ b/cpp/src/transform/transform.cu @@ -29,7 +29,6 @@ #include #include #include -#include #include #include @@ -158,14 +157,14 @@ using handle = std::variant< namespace jit_transform { -jitify2::Kernel instantiate(null_aware is_null_aware, - bool has_user_data, - std::string const& ins, - std::string const& outs, - std::vector const& ptx_input_types, - std::vector const& ptx_output_types, - std::string const& udf, - udf_source_type source_type) +kernel instantiate(bool is_null_aware, + bool has_user_data, + std::string const& ins, + std::string const& outs, + std::vector const& ptx_input_types, + std::vector const& ptx_output_types, + std::string const& udf, + udf_source_type source_type) { CUDF_FUNC_RANGE(); auto cuda_source = (source_type == udf_source_type::PTX) @@ -175,14 +174,16 @@ jitify2::Kernel instantiate(null_aware is_null_aware, jit::build_ptx_params(ptx_output_types, ptx_input_types, has_user_data)) : jit::parse_single_function_cuda(udf, "GENERIC_TRANSFORM_OP"); - auto kernel = jitify2::reflection::Template("cudf::jit::transform_kernel") - .instantiate(is_null_aware, has_user_data, ins, outs); + auto kernel = rtcx::reflect_template("cudf::jit::transform_kernel", + rtcx::reflect(is_null_aware), + rtcx::reflect(has_user_data), + ins, + outs); - return jit::get_udf_kernel( - *transform_jit_kernel_cu_jit, kernel, cuda_source, {"-restrict", "--dopt=on"}); + return jit::get_udf_kernel("cudf/cpp/src/transform/jit/kernel.cu", kernel, cuda_source); } -void launch(jitify2::Kernel const& kernel, +void launch(cudf::kernel const& kernel, size_type row_size, bitmask_type const* stencil, void* user_data, @@ -192,7 +193,11 @@ void launch(jitify2::Kernel const& kernel, { CUDF_FUNC_RANGE(); void* args[] = {&row_size, &stencil, &user_data, &input_cols, &output_cols}; - kernel->configure_1d_max_occupancy(0, 0, nullptr, stream.value())->launch_raw(args); + auto cfg = kernel.max_occupancy_config(0, 0); + CUDF_EXPECTS(cfg.block_size % cudf::detail::warp_size == 0, + "Expected block size to be a multiple of warp size", + std::runtime_error); + kernel.launch({cfg.min_grid_size}, {cfg.block_size}, 0, stream, args); } std::string reflect_input_element(column_view const& c) { return type_to_name(c.type()); } @@ -244,8 +249,12 @@ auto reflect(udf_source_type source_type, auto column = std::visit([](auto& c) { return reflect_input_column(c); }, in); auto element = std::visit([](auto& c) { return reflect_input_element(c); }, in); bool as_scalar = std::holds_alternative(in); - auto accessor = jitify2::reflection::Template("cudf::jit::column_accessor") - .instantiate(i, column, element, as_scalar, 0); + auto accessor = rtcx::reflect_template("cudf::jit::column_accessor", + rtcx::reflect(i), + column, + element, + rtcx::reflect(as_scalar), + rtcx::reflect(0)); in_types.push_back(accessor); } @@ -256,14 +265,18 @@ auto reflect(udf_source_type source_type, auto column = std::visit([](auto& c) { return reflect_output_column(c); }, out); auto element = std::visit([](auto& c) { return reflect_output_element(c); }, out); bool as_scalar = false; // never scalar - auto accessor = jitify2::reflection::Template("cudf::jit::column_accessor") - .instantiate(i, column, element, as_scalar, 0); + auto accessor = rtcx::reflect_template("cudf::jit::column_accessor", + rtcx::reflect(i), + column, + element, + rtcx::reflect(as_scalar), + rtcx::reflect(0)); out_types.push_back(accessor); } - auto ins = jitify2::reflection::Template("cudf::jit::type_list").instantiate(in_types); - auto outs = jitify2::reflection::Template("cudf::jit::type_list").instantiate(out_types); + auto ins = rtcx::reflect_template("cudf::jit::type_list", in_types); + auto outs = rtcx::reflect_template("cudf::jit::type_list", out_types); std::vector ptx_in_types; std::vector ptx_out_types; @@ -319,7 +332,7 @@ auto to_args(std::span inputs, return std::make_tuple(std::move(d_args), std::move(handles)); } -void run(null_aware is_null_aware, +void run(bool is_null_aware, bool has_user_data, size_type row_size, bitmask_type const* d_stencil, @@ -819,7 +832,7 @@ std::unique_ptr execute_transform(std::string const& udf, auto stencil_arg = stencil.has_value() ? stencil->first : nullptr; auto stencil_has_nulls = stencil.has_value() ? (stencil->second > 0) : false; - jit_transform::run(is_null_aware, + jit_transform::run(is_null_aware == null_aware::YES, user_data.has_value(), row_size, stencil_has_nulls ? stencil_arg : nullptr,