Skip to content

[DRAFT] Add AOT JIT+LTO capability #22390

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
divyegala wants to merge 16 commits into from
Closed

[DRAFT] Add AOT JIT+LTO capability #22390

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
82c4bf1
First pass
vyasr Apr 23, 2026
9832a01
Merge remote-tracking branch 'origin/main' into feat/nvjitlink_kernels
divyegala May 5, 2026
77c483c
add jit+lto, tests pass
divyegala May 6, 2026
1881bfa
remove udf reference
divyegala May 6, 2026
5bfce79
run pre-commit
divyegala May 6, 2026
eb63608
Merge branch 'main' into feat/nvjitlink_kernels
divyegala May 6, 2026
7a0c3f1
fix style check
divyegala May 7, 2026
0292e26
Merge remote-tracking branch 'divye/feat/nvjitlink_kernels' into feat…
divyegala May 7, 2026
dcf0f20
Merge remote-tracking branch 'origin/main' into feat/nvjitlink_kernels
divyegala May 7, 2026
22f1761
exclude nvjitlink from wheel install path
divyegala May 7, 2026
551693f
remove double compilation
divyegala May 7, 2026
de922ce
Merge remote-tracking branch 'origin/main' into feat/nvjitlink_kernels
divyegala May 7, 2026
bab9125
make dispatcher its own fragment
divyegala May 7, 2026
0caa15d
specialized dispatch and lto measurement
divyegala May 8, 2026
c9c5de0
bench options
divyegala May 11, 2026
fa91814
bench scripts
divyegala May 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions ci/build_wheel_libcudf.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ export SKBUILD_CMAKE_ARGS="-DUSE_NVCOMP_RUNTIME_WHEEL=ON"
python -m auditwheel repair \
--exclude libkvikio.so \
--exclude libnvcomp.so.5 \
--exclude libnvJitLink.so.* \
--exclude librapids_logger.so \
--exclude librmm.so \
-w "${RAPIDS_WHEEL_BLD_OUTPUT_DIR}" \
Expand Down
108 changes: 107 additions & 1 deletion cpp/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,6 +236,7 @@ rapids_find_package(
BUILD_EXPORT_SET cudf-exports
INSTALL_EXPORT_SET cudf-exports
)
find_package(CUDAToolkit REQUIRED COMPONENTS nvJitLink)
include(cmake/Modules/ConfigureCUDA.cmake) # set other CUDA compilation flags

# ##################################################################################################
Expand Down Expand Up @@ -324,6 +325,106 @@ if(NOT BUILD_SHARED_LIBS)
endif()
endif()

# ##################################################################################################
# * JIT+LTO (MurmurHash3 x86 32 device fragments + nvjitlink) --------------------------------------
include(cmake/Modules/generate_jit_lto_kernels.cmake)

add_library(cudf_jit_lto_kernel_usage_requirements INTERFACE)
target_include_directories(
cudf_jit_lto_kernel_usage_requirements
INTERFACE "${CMAKE_CURRENT_SOURCE_DIR}/include" "${CMAKE_CURRENT_SOURCE_DIR}/src"
"${CMAKE_CURRENT_SOURCE_DIR}/src/hash/jit_lto_kernels"
)
target_compile_options(
cudf_jit_lto_kernel_usage_requirements INTERFACE "$<$<COMPILE_LANGUAGE:CXX>:${CUDF_CXX_FLAGS}>"
"$<$<COMPILE_LANGUAGE:CUDA>:${CUDF_CUDA_FLAGS}>"
)
target_compile_features(cudf_jit_lto_kernel_usage_requirements INTERFACE cuda_std_20)
target_link_libraries(
cudf_jit_lto_kernel_usage_requirements INTERFACE CCCL::CCCL rmm::rmm
$<BUILD_LOCAL_INTERFACE:cuco::cuco>
)

set(JIT_LTO_TARGET_ARCHITECTURE "70-real")
if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
set(JIT_LTO_TARGET_ARCHITECTURE "75-real")
endif()

block(PROPAGATE cudf_jit_lto_generated_sources)
set(cudf_jit_lto_generated_sources)
set(CMAKE_CUDA_ARCHITECTURES ${JIT_LTO_TARGET_ARCHITECTURE})

set(murmur_jit_ns "cudf::hashing::detail::jit_lto")
generate_jit_lto_kernels(
cudf_jit_lto_generated_sources
NAME_FORMAT
"murmurhash_jit_hasher_@abbrev@"
MATRIX_JSON_FILE
"${CMAKE_CURRENT_SOURCE_DIR}/src/hash/jit_lto_kernels/murmurhash_hasher_matrix.json"
KERNEL_INPUT_FILE
"${CMAKE_CURRENT_SOURCE_DIR}/src/hash/jit_lto_kernels/murmurhash_hasher_fragment.cu.in"
FRAGMENT_TAG_FORMAT
"${murmur_jit_ns}::fragment_tag_murmur_hasher<${murmur_jit_ns}::tag_@abbrev@>"
FRAGMENT_TAG_HEADER_FILES
"<cudf/hashing/detail/murmurhash3_x86_32_jit_tags.hpp>"
OUTPUT_DIRECTORY
"${CMAKE_CURRENT_BINARY_DIR}/generated_jit_lto/murmurhash/hasher"
KERNEL_LINK_LIBRARIES
cudf_jit_lto_kernel_usage_requirements
)
generate_jit_lto_kernels(
cudf_jit_lto_generated_sources
NAME_FORMAT
"murmurhash_jit_hasher_noop_@abbrev@"
MATRIX_JSON_FILE
"${CMAKE_CURRENT_SOURCE_DIR}/src/hash/jit_lto_kernels/murmurhash_hasher_matrix.json"
KERNEL_INPUT_FILE
"${CMAKE_CURRENT_SOURCE_DIR}/src/hash/jit_lto_kernels/murmurhash_hasher_noop_fragment.cu.in"
FRAGMENT_TAG_FORMAT
"${murmur_jit_ns}::fragment_tag_murmur_hasher_noop<${murmur_jit_ns}::tag_@abbrev@>"
FRAGMENT_TAG_HEADER_FILES
"<cudf/hashing/detail/murmurhash3_x86_32_jit_tags.hpp>"
OUTPUT_DIRECTORY
"${CMAKE_CURRENT_BINARY_DIR}/generated_jit_lto/murmurhash/hasher_noop"
KERNEL_LINK_LIBRARIES
cudf_jit_lto_kernel_usage_requirements
)
generate_jit_lto_kernels(
cudf_jit_lto_generated_sources
NAME_FORMAT
"murmurhash_jit_@k@"
MATRIX_JSON_FILE
"${CMAKE_CURRENT_SOURCE_DIR}/src/hash/jit_lto_kernels/murmurhash_entry_matrix.json"
KERNEL_INPUT_FILE
"${CMAKE_CURRENT_SOURCE_DIR}/src/hash/jit_lto_kernels/murmurhash_entry_kernel.cu.in"
FRAGMENT_TAG_FORMAT
"${murmur_jit_ns}::fragment_tag_murmur_entry"
FRAGMENT_TAG_HEADER_FILES
"<cudf/hashing/detail/murmurhash3_x86_32_jit_tags.hpp>"
OUTPUT_DIRECTORY
"${CMAKE_CURRENT_BINARY_DIR}/generated_jit_lto/murmurhash/entry"
KERNEL_LINK_LIBRARIES
cudf_jit_lto_kernel_usage_requirements
)
generate_jit_lto_kernels(
cudf_jit_lto_generated_sources
NAME_FORMAT
"murmurhash_jit_dispatch_@suffix@"
MATRIX_JSON_FILE
"${CMAKE_CURRENT_SOURCE_DIR}/src/hash/jit_lto_kernels/murmurhash_dispatch_matrix.json"
KERNEL_INPUT_FILE
"${CMAKE_CURRENT_SOURCE_DIR}/src/hash/jit_lto_kernels/murmurhash_dispatch_fragment.cu.in"
FRAGMENT_TAG_FORMAT
"${murmur_jit_ns}::fragment_tag_murmur_dispatch_@suffix@"
FRAGMENT_TAG_HEADER_FILES
"<cudf/hashing/detail/murmurhash3_x86_32_jit_tags.hpp>"
OUTPUT_DIRECTORY
"${CMAKE_CURRENT_BINARY_DIR}/generated_jit_lto/murmurhash/dispatch"
KERNEL_LINK_LIBRARIES
cudf_jit_lto_kernel_usage_requirements
)
endblock()

# ##################################################################################################
# * library targets -------------------------------------------------------------------------------
add_library(
Expand Down Expand Up @@ -456,6 +557,11 @@ add_library(
src/groupby/sort/sort_helper.cu
src/hash/md5_hash.cu
src/hash/murmurhash3_x86_32.cu
src/jit_lto/AlgorithmLauncher.cpp
src/jit_lto/AlgorithmPlanner.cpp
src/jit_lto/FragmentEntry.cpp
src/jit_lto/nvjitlink_checker.cpp
${cudf_jit_lto_generated_sources}
src/hash/murmurhash3_x64_128.cu
src/hash/sha1_hash.cu
src/hash/sha224_hash.cu
Expand Down Expand Up @@ -1000,7 +1106,7 @@ target_link_libraries(
cudf
PUBLIC CCCL::CCCL rapids_logger::rapids_logger rmm::rmm $<BUILD_LOCAL_INTERFACE:BS::thread_pool>
PRIVATE $<BUILD_LOCAL_INTERFACE:nvtx3::nvtx3-cpp> $<BUILD_LOCAL_INTERFACE:cuco::cuco> ZLIB::ZLIB
nvcomp::nvcomp kvikio::kvikio nanoarrow::nanoarrow zstd
nvcomp::nvcomp kvikio::kvikio nanoarrow::nanoarrow zstd CUDA::nvJitLink
)

# Add Conda library, and include paths if specified
Expand Down
55 changes: 55 additions & 0 deletions cpp/cmake/Modules/compute_matrix_product.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
# =============================================================================
# cmake-format: off
# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
# SPDX-License-Identifier: Apache-2.0
# cmake-format: on
# =============================================================================

include_guard(GLOBAL)

# This function runs compute_matrix_product.py and stores the JSON matrix product in output_var.
function(compute_matrix_product output_var)
set(options)
set(one_value MATRIX_JSON_FILE MATRIX_JSON_STRING)
set(multi_value)

cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN})

find_package(Python3 REQUIRED COMPONENTS Interpreter)

if(_JIT_LTO_MATRIX_JSON_FILE)
execute_process(
COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py"
"${_JIT_LTO_MATRIX_JSON_FILE}" #
OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY
)
else()
execute_process(
COMMAND ${CMAKE_COMMAND} -E echo "${_JIT_LTO_MATRIX_JSON_STRING}"
COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py"
-
OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY
)
endif()
Comment on lines +20 to +33

@coderabbitai coderabbitai Bot May 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Require exactly one matrix input source.

compute_matrix_product(...) silently prefers MATRIX_JSON_FILE when both inputs are set, and it falls through to the stdin path when neither is set. Both cases make configuration failures harder to diagnose and can generate the wrong matrix product.

Suggested change 
+  if((_JIT_LTO_MATRIX_JSON_FILE AND _JIT_LTO_MATRIX_JSON_STRING) OR
+     (NOT _JIT_LTO_MATRIX_JSON_FILE AND NOT _JIT_LTO_MATRIX_JSON_STRING))
+    message(FATAL_ERROR
+            "compute_matrix_product requires exactly one of MATRIX_JSON_FILE or MATRIX_JSON_STRING")
+  endif()
+
   if(_JIT_LTO_MATRIX_JSON_FILE)
     execute_process(
       COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py"
               "${_JIT_LTO_MATRIX_JSON_FILE}" #
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
if(_JIT_LTO_MATRIX_JSON_FILE)
execute_process(
COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py"
"${_JIT_LTO_MATRIX_JSON_FILE}" #
OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY
)
else()
execute_process(
COMMAND ${CMAKE_COMMAND} -E echo "${_JIT_LTO_MATRIX_JSON_STRING}"
COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py"
-
OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY
)
endif()
if((_JIT_LTO_MATRIX_JSON_FILE AND _JIT_LTO_MATRIX_JSON_STRING) OR
(NOT _JIT_LTO_MATRIX_JSON_FILE AND NOT _JIT_LTO_MATRIX_JSON_STRING))
message(FATAL_ERROR
"compute_matrix_product requires exactly one of MATRIX_JSON_FILE or MATRIX_JSON_STRING")
endif()
if(_JIT_LTO_MATRIX_JSON_FILE)
execute_process(
COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py"
"${_JIT_LTO_MATRIX_JSON_FILE}" #
OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY
)
else()
execute_process(
COMMAND ${CMAKE_COMMAND} -E echo "${_JIT_LTO_MATRIX_JSON_STRING}"
COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py"
-
OUTPUT_VARIABLE output COMMAND_ERROR_IS_FATAL ANY
)
endif()
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@cpp/cmake/Modules/compute_matrix_product.cmake` around lines 20 - 33, The
current logic in compute_matrix_product.cmake prefers _JIT_LTO_MATRIX_JSON_FILE
when both inputs are present and falls back to stdin when neither is provided;
change this to require exactly one input by adding a validation block that
checks the presence of _JIT_LTO_MATRIX_JSON_FILE and _JIT_LTO_MATRIX_JSON_STRING
(use the same variable names) and calls message(FATAL_ERROR ...) if both are set
or both are unset, otherwise proceed with the existing execute_process branches
(the branches that call "${Python3_EXECUTABLE}"
"${CMAKE_CURRENT_FUNCTION_LIST_DIR}/compute_matrix_product.py" with either the
file path or "-" for stdin).


set(${output_var}
"${output}"
PARENT_SCOPE
)
endfunction()

# This function unpacks a JSON object into CMake variables in the caller scope.
function(populate_matrix_variables matrix_json_entry)
string(JSON len LENGTH "${matrix_json_entry}")
math(EXPR last "${len} - 1")

# cmake-lint: disable=C0103,E1120
foreach(i RANGE "${last}")
string(JSON key MEMBER "${matrix_json_entry}" "${i}")
string(JSON value GET "${matrix_json_entry}" "${key}")
set(${key}
"${value}"
PARENT_SCOPE
)
endforeach()
endfunction()
Loading
Loading