diff --git a/ci/build_wheel_libcudf.sh b/ci/build_wheel_libcudf.sh index 3680f1ebe10a..b468400aa28b 100755 --- a/ci/build_wheel_libcudf.sh +++ b/ci/build_wheel_libcudf.sh @@ -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}" \ diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 6d684af8d990..0e26f419f7c3 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -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 # ################################################################################################## @@ -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 "$<$:${CUDF_CXX_FLAGS}>" + "$<$:${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 + $ +) + +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 + "" + 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 + "" + 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 + "" + 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 + "" + 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( @@ -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 @@ -1000,7 +1106,7 @@ target_link_libraries( cudf PUBLIC CCCL::CCCL rapids_logger::rapids_logger rmm::rmm $ PRIVATE $ $ 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 diff --git a/cpp/cmake/Modules/compute_matrix_product.cmake b/cpp/cmake/Modules/compute_matrix_product.cmake new file mode 100644 index 000000000000..a00070a6018d --- /dev/null +++ b/cpp/cmake/Modules/compute_matrix_product.cmake @@ -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() + + 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() diff --git a/cpp/cmake/Modules/compute_matrix_product.py b/cpp/cmake/Modules/compute_matrix_product.py new file mode 100644 index 000000000000..1d88942b2216 --- /dev/null +++ b/cpp/cmake/Modules/compute_matrix_product.py @@ -0,0 +1,245 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# This algorithm takes a JSON dictionary and computes a matrix product of all +# of its arrays. We use this to compute all matrix combinations for kernel +# generation. We *could* write this in CMake with `string(JSON)`, but writing +# it in Python is much easier. Once we have a version of CMake that has +# https://gitlab.kitware.com/cmake/cmake/-/merge_requests/11516, we may be able +# to port the algorithm to CMake script and use it in other RAPIDS projects. + +from __future__ import annotations + +import argparse +import json +import re +import sys +import warnings +from itertools import chain +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Generator, Iterator + + MatrixValue = None | bool | int | float | str + Matrix = MatrixValue | list["Matrix"] | dict[str, "Matrix"] + + +class NoKeyError(ValueError): + pass + + +class UnusedKeyWarning(UserWarning): + pass + + +class UsedKeyWarning(UserWarning): + pass + + +IDENTIFIER_RE: re.Pattern = re.compile(r"^(?P_*)(?P.*)$") + + +def iterate_matrix_product( + *, + matrix: "Matrix", + warn_unused=True, + warn_used=True, +) -> "Generator[dict[str, MatrixValue]]": + """Computes a matrix product of a JSON document + + This algorithm computes the product of a matrix in a more sophisticated + way than can be done with itertools.product(). Multiple related values + can be grouped together, and a dimension can have sub-dimensions. Given + the following JSON document: + + .. code-block:: json + + { + "value": ["one", "two"], + "_group": [ + { + "subgroup_value": "three" + "subgroup_subdim": ["four", "five"] + }, + { + "subgroup_value": "six", + "subgroup_subdim": ["seven", "eight"] + } + ] + } + + The following matrix product entries will be produced: + + .. code-block:: json + + {"subgroup_subdim": "four", "subgroup_value": "three", "value": "one"} + {"subgroup_subdim": "five", "subgroup_value": "three", "value": "one"} + {"subgroup_subdim": "seven", "subgroup_value": "six", "value": "one"} + {"subgroup_subdim": "eight", "subgroup_value": "six", "value": "one"} + {"subgroup_subdim": "four", "subgroup_value": "three", "value": "two"} + {"subgroup_subdim": "five", "subgroup_value": "three", "value": "two"} + {"subgroup_subdim": "seven", "subgroup_value": "six", "value": "two"} + {"subgroup_subdim": "eight", "subgroup_value": "six", "value": "two"} + + Notice that the name ``_group`` does not appear in any of the matrix + product entries. This is because all leaf nodes beneath it are under + a different key that's closer to them in the hierarchy, which is used in + the final product. If a key is used only for grouping and does not appear + in the final product, it should be prefixed with an underscore (``_``) to + indicate that they are hidden. Likewise, keys that appear in the final + product should not be prefixed with an underscore. Failure to follow this + convention will not affect the proper functioning of the algorithm, but a + warning will be emitted unless the respective ``warn_unused`` and/or + ``warn_used`` parameters are set to ``False``. + + Every leaf node in the input document must have at least one dictionary key + in its path, or else a ``NoKeyError`` will be thrown. For example, a + document consisting only of an array of strings is invalid. + + Parameters + ---------- + matrix : Matrix + JSON document on which to compute the product. + warn_unused : bool + Whether or not to warn if unused keys are not prefixed with underscores + (true by default). + warn_used : bool + Whether or not to warn if used keys are prefixed with underscores (true + by default). + + Returns + ------- + Generator[dict[str, MatrixValue]] + Iterator of matrix product entries. + + Raises + ------ + NoKeyError + If a leaf node in the document does not have any key leading to it. + + Warns + ----- + UnusedKeyWarning + If an unused key is not prefixed with an underscore. + UsedKeyWarning + If a used key is prefixed with an underscore. + """ + + def iterate_next_dimension( + queue: "Iterator[tuple[tuple[str | int, ...], str | None, Matrix]]", + entry: "dict[str, MatrixValue]", + ) -> "Generator[tuple[dict[str, MatrixValue], bool]]": + try: + path, key, matrix = next(queue) + except StopIteration: + yield (entry, False) + else: + used = False + for e, u in iterate_impl(path, key, matrix, queue, entry): + if u: + used = True + yield e, u + + try: + last = path[-1] + except IndexError: + pass + else: + if isinstance(last, str): + match = IDENTIFIER_RE.search(last) + assert match + underscores = match.group("underscores") + rest = match.group("rest") + path_repr = "".join( + f"[{json.dumps(i)}]" for i in path[:-1] + ) + + if warn_used and used and underscores: + warnings.warn( + f"Key {json.dumps(last)} at root{path_repr} " + f"is used in a matrix product entry even though it " + f"begins with {json.dumps(underscores)}. Consider " + f"renaming it to {json.dumps(rest)} to indicate this.", + category=UsedKeyWarning, + ) + elif warn_unused and not used and not underscores: + warnings.warn( + f"Key {json.dumps(last)} at root{path_repr} " + f"is never used in a matrix product entry and is used " + f"only for grouping. Consider renaming it to " + f"{json.dumps(f'_{last}')} to indicate this.", + category=UnusedKeyWarning, + ) + + def iterate_impl( + path: tuple[str | int, ...], + key: str | None, + matrix: "Matrix", + queue: "Iterator[tuple[tuple[str | int, ...], str | None, Matrix]]", + entry: "dict[str, MatrixValue]", + ) -> "Generator[tuple[dict[str, MatrixValue], bool]]": + if isinstance(matrix, dict): + yield from ( + (e, False) + for e, _ in iterate_next_dimension( + chain( + ( + ((*path, k), k, v) + for (k, v) in sorted(matrix.items()) + ), + queue, + ), + entry, + ) + ) + elif isinstance(matrix, list): + queue_list = list(queue) + for i, v in enumerate(matrix): + yield from iterate_next_dimension( + chain([((*path, i), key, v)], queue_list), entry + ) + else: + if key is None: + raise NoKeyError + yield from ( + (e, True) + for e, _ in iterate_next_dimension( + queue, {**entry, key: matrix} + ) + ) + + yield from ( + entry for entry, _ in iterate_impl((), None, matrix, chain(), {}) + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--warn-unused", action=argparse.BooleanOptionalAction, default=True + ) + parser.add_argument( + "--warn-used", action=argparse.BooleanOptionalAction, default=True + ) + parser.add_argument("filename", nargs="?", default="-") + namespace = parser.parse_args() + with ( + sys.stdin + if namespace.filename == "-" + else open(namespace.filename) as f + ): + matrix: "Matrix" = json.load(f) + + json.dump( + list( + iterate_matrix_product( + matrix=matrix, + warn_unused=namespace.warn_unused, + warn_used=namespace.warn_used, + ) + ), + sys.stdout, + indent=2, + sort_keys=True, + ) diff --git a/cpp/cmake/Modules/generate_jit_lto_kernels.cmake b/cpp/cmake/Modules/generate_jit_lto_kernels.cmake new file mode 100644 index 000000000000..95e8c314e3c6 --- /dev/null +++ b/cpp/cmake/Modules/generate_jit_lto_kernels.cmake @@ -0,0 +1,140 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +include_guard(GLOBAL) + +include(${CMAKE_CURRENT_LIST_DIR}/compute_matrix_product.cmake) + +# This function compiles a JIT LTO kernel object target and writes a C header embedding its fatbin. +function(add_jit_lto_kernel kernel_target) + set(options) + set(one_value KERNEL_FILE FATBIN_HEADER_FILE) + set(multi_value LINK_LIBRARIES) + + cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN}) + + add_library(${kernel_target} OBJECT EXCLUDE_FROM_ALL "${_JIT_LTO_KERNEL_FILE}") + # Do not modify these properties, options, and libraries. Usage requirements (including CUDA + # version, etc.) should be propagated to the kernel targets via INTERFACE libraries passed in + # through the LINK_LIBRARIES argument. + target_link_libraries(${kernel_target} PRIVATE ${_JIT_LTO_LINK_LIBRARIES}) + target_compile_options(${kernel_target} PRIVATE -Xfatbin=--compress-all --compress-mode=size) + set_target_properties( + ${kernel_target} + PROPERTIES CUDA_SEPARABLE_COMPILATION ON + CUDA_FATBIN_COMPILATION ON + POSITION_INDEPENDENT_CODE ON + INTERPROCEDURAL_OPTIMIZATION ON + ) + + add_custom_command( + OUTPUT "${_JIT_LTO_FATBIN_HEADER_FILE}" + COMMENT "Generate embedded fatbin header for ${kernel_target}" + COMMAND "${bin_to_c}" --const --name embedded_fatbin --static $ + > "${_JIT_LTO_FATBIN_HEADER_FILE}" + DEPENDS $ + ) +endfunction() + +# This function materializes one matrix entry into configured sources and registers its fatbin. +function(process_jit_lto_matrix_entry source_list_var) + set(options) + set(one_value NAME_FORMAT KERNEL_INPUT_FILE OUTPUT_DIRECTORY FRAGMENT_TAG_FORMAT + MATRIX_JSON_ENTRY + ) + set(multi_value KERNEL_LINK_LIBRARIES FRAGMENT_TAG_HEADER_FILES) + + cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN}) + + populate_matrix_variables("${_JIT_LTO_MATRIX_JSON_ENTRY}") + string(CONFIGURE "${_JIT_LTO_NAME_FORMAT}" kernel_name @ONLY) + string(CONFIGURE "${_JIT_LTO_FRAGMENT_TAG_FORMAT}" fragment_tag @ONLY) + + set(fragment_tag_header_files "") + foreach(header_file IN LISTS _JIT_LTO_FRAGMENT_TAG_HEADER_FILES) + if(NOT header_file MATCHES "^(\".*\"|<.*>)$") + set(header_file "\"${header_file}\"") + endif() + string(APPEND fragment_tag_header_files "#include ${header_file}\n") + endforeach() + + set(kernel_file "${_JIT_LTO_OUTPUT_DIRECTORY}/${kernel_name}_kernel.cu") + set(kernel_target "${kernel_name}_kernel") + set(fatbin_header_file "${_JIT_LTO_OUTPUT_DIRECTORY}/${kernel_name}_fatbin.h") + set(fatbin_file "${_JIT_LTO_OUTPUT_DIRECTORY}/${kernel_name}_fatbin.cpp") + configure_file("${_JIT_LTO_KERNEL_INPUT_FILE}" "${kernel_file}" @ONLY) + configure_file("${CMAKE_CURRENT_FUNCTION_LIST_DIR}/register_fatbin.cpp.in" "${fatbin_file}" @ONLY) + + add_jit_lto_kernel( + ${kernel_target} KERNEL_FILE "${kernel_file}" FATBIN_HEADER_FILE "${fatbin_header_file}" + LINK_LIBRARIES ${_JIT_LTO_KERNEL_LINK_LIBRARIES} + ) + list(APPEND ${source_list_var} "${fatbin_header_file}" "${fatbin_file}") + set(${source_list_var} + "${${source_list_var}}" + PARENT_SCOPE + ) +endfunction() + +# This function expands a JIT LTO matrix JSON description into kernel and fatbin source files. +function(generate_jit_lto_kernels source_list_var) + set(options) + set(one_value NAME_FORMAT MATRIX_JSON_FILE MATRIX_JSON_STRING KERNEL_INPUT_FILE + FRAGMENT_TAG_FORMAT OUTPUT_DIRECTORY + ) + set(multi_value KERNEL_LINK_LIBRARIES FRAGMENT_TAG_HEADER_FILES) + + cmake_parse_arguments(_JIT_LTO "${options}" "${one_value}" "${multi_value}" ${ARGN}) + + find_package(CUDAToolkit REQUIRED) + find_program( + bin_to_c + NAMES bin2c + PATHS ${CUDAToolkit_BIN_DIR} + ) + + if(_JIT_LTO_MATRIX_JSON_FILE) + set_property( + DIRECTORY + PROPERTY CMAKE_CONFIGURE_DEPENDS "${_JIT_LTO_MATRIX_JSON_FILE}" + APPEND + ) + compute_matrix_product(matrix_product MATRIX_JSON_FILE "${_JIT_LTO_MATRIX_JSON_FILE}") + else() + compute_matrix_product(matrix_product MATRIX_JSON_STRING "${_JIT_LTO_MATRIX_JSON_STRING}") + endif() + + string(JSON len LENGTH "${matrix_product}") + math(EXPR last "${len} - 1") + + # cmake-lint: disable=C0103,E1120 + foreach(i RANGE "${last}") + string(JSON matrix_json_entry GET "${matrix_product}" "${i}") + process_jit_lto_matrix_entry( + "${source_list_var}" + NAME_FORMAT + "${_JIT_LTO_NAME_FORMAT}" + KERNEL_INPUT_FILE + "${_JIT_LTO_KERNEL_INPUT_FILE}" + FRAGMENT_TAG_FORMAT + "${_JIT_LTO_FRAGMENT_TAG_FORMAT}" + FRAGMENT_TAG_HEADER_FILES + ${_JIT_LTO_FRAGMENT_TAG_HEADER_FILES} + OUTPUT_DIRECTORY + "${_JIT_LTO_OUTPUT_DIRECTORY}" + MATRIX_JSON_ENTRY + "${matrix_json_entry}" + KERNEL_LINK_LIBRARIES + ${_JIT_LTO_KERNEL_LINK_LIBRARIES} + ) + endforeach() + + set(${source_list_var} + "${${source_list_var}}" + PARENT_SCOPE + ) +endfunction() diff --git a/cpp/cmake/Modules/register_fatbin.cpp.in b/cpp/cmake/Modules/register_fatbin.cpp.in new file mode 100644 index 000000000000..6f11ce4bad77 --- /dev/null +++ b/cpp/cmake/Modules/register_fatbin.cpp.in @@ -0,0 +1,22 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "@fatbin_header_file@" +#include + +@fragment_tag_header_files@ + +namespace { + +using fragment_tag = @fragment_tag@; +using fragment_entry = cudf::detail::jit_lto::StaticFatbinFragmentEntry; + +} // namespace + +template <> +const uint8_t* const fragment_entry::data = embedded_fatbin; + +template <> +const size_t fragment_entry::length = sizeof(embedded_fatbin); diff --git a/cpp/include/cudf/detail/jit_lto/AlgorithmLauncher.hpp b/cpp/include/cudf/detail/jit_lto/AlgorithmLauncher.hpp new file mode 100644 index 000000000000..c09592608c56 --- /dev/null +++ b/cpp/include/cudf/detail/jit_lto/AlgorithmLauncher.hpp @@ -0,0 +1,50 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +#include +#include + +#include +#include +#include +#include + +namespace cudf::detail::jit_lto { + +struct AlgorithmLauncher { + AlgorithmLauncher() : kernel{nullptr}, library{nullptr} {} + + AlgorithmLauncher(cudaKernel_t k, cudaLibrary_t lib); + + ~AlgorithmLauncher(); + + AlgorithmLauncher(const AlgorithmLauncher&) = delete; + AlgorithmLauncher& operator=(const AlgorithmLauncher&) = delete; + + AlgorithmLauncher(AlgorithmLauncher&& other) noexcept; + AlgorithmLauncher& operator=(AlgorithmLauncher&& other) noexcept; + + template + void dispatch(cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, Args&&... args) + { + static_assert(std::is_same_v...)>, + "dispatch() argument types do not match the kernel function signature FuncT"); + + void* kernel_args[] = {const_cast(static_cast(&args))...}; + this->call(stream, grid, block, shared_mem, kernel_args); + } + + cudaKernel_t get_kernel() { return this->kernel; } + + private: + void call(cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, void** args); + cudaKernel_t kernel; + cudaLibrary_t library; +}; + +} // namespace cudf::detail::jit_lto diff --git a/cpp/include/cudf/detail/jit_lto/AlgorithmPlanner.hpp b/cpp/include/cudf/detail/jit_lto/AlgorithmPlanner.hpp new file mode 100644 index 000000000000..e0e71421254c --- /dev/null +++ b/cpp/include/cudf/detail/jit_lto/AlgorithmPlanner.hpp @@ -0,0 +1,57 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace cudf::detail::jit_lto { + +struct LauncherJitCache { + std::shared_mutex mutex; + std::unordered_map> launchers; +}; + +struct AlgorithmPlanner { + AlgorithmPlanner(std::string entrypoint, LauncherJitCache& jit_cache) + : entrypoint(std::move(entrypoint)), jit_cache_(jit_cache) + { + } + + std::shared_ptr get_launcher(); + + std::string entrypoint; + std::vector> fragments; + + template >> + void add_fragment(std::unique_ptr fragment) + { + fragments.push_back(std::unique_ptr(std::move(fragment))); + } + + template + void add_static_fragment() + { + add_fragment(std::make_unique>()); + } + + private: + std::string get_fragments_key() const; + std::shared_ptr build(); + + std::shared_ptr read_cache(std::string const& launch_key) const; + + LauncherJitCache& jit_cache_; +}; + +} // namespace cudf::detail::jit_lto diff --git a/cpp/include/cudf/detail/jit_lto/FragmentEntry.hpp b/cpp/include/cudf/detail/jit_lto/FragmentEntry.hpp new file mode 100644 index 000000000000..aed48c242393 --- /dev/null +++ b/cpp/include/cudf/detail/jit_lto/FragmentEntry.hpp @@ -0,0 +1,50 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +#include + +#include +#include +#include +#include +#include + +namespace cudf::detail::jit_lto { + +struct FragmentEntry { + virtual ~FragmentEntry() = default; + + virtual bool add_to(nvJitLinkHandle& handle) const = 0; + + virtual const char* get_key() const = 0; +}; + +struct FatbinFragmentEntry : FragmentEntry { + virtual const uint8_t* get_data() const = 0; + + virtual size_t get_length() const = 0; + + bool add_to(nvJitLinkHandle& handle) const override final; +}; + +template +struct StaticFatbinFragmentEntry final : FatbinFragmentEntry { + const uint8_t* get_data() const override { return StaticFatbinFragmentEntry::data; } + + size_t get_length() const override { return StaticFatbinFragmentEntry::length; } + + const char* get_key() const override + { + return typeid(StaticFatbinFragmentEntry).name(); + } + + static const uint8_t* const data; + static const size_t length; +}; + +} // namespace cudf::detail::jit_lto diff --git a/cpp/include/cudf/detail/jit_lto/nvjitlink_checker.hpp b/cpp/include/cudf/detail/jit_lto/nvjitlink_checker.hpp new file mode 100644 index 000000000000..70755a9adf96 --- /dev/null +++ b/cpp/include/cudf/detail/jit_lto/nvjitlink_checker.hpp @@ -0,0 +1,13 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +namespace cudf::detail::jit_lto { + +void check_nvjitlink_result(nvJitLinkHandle handle, nvJitLinkResult result); + +} // namespace cudf::detail::jit_lto diff --git a/cpp/include/cudf/hashing/detail/murmurhash3_x86_32_jit_tags.hpp b/cpp/include/cudf/hashing/detail/murmurhash3_x86_32_jit_tags.hpp new file mode 100644 index 000000000000..7def96d96bde --- /dev/null +++ b/cpp/include/cudf/hashing/detail/murmurhash3_x86_32_jit_tags.hpp @@ -0,0 +1,56 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +namespace cudf::hashing::detail::jit_lto { + +struct tag_murmur_entry {}; + +struct tag_i8 {}; +struct tag_i16 {}; +struct tag_i32 {}; +struct tag_i64 {}; +struct tag_u8 {}; +struct tag_u16 {}; +struct tag_u32 {}; +struct tag_u64 {}; +struct tag_f32 {}; +struct tag_f64 {}; +struct tag_b8 {}; +struct tag_ts_day {}; +struct tag_ts_s {}; +struct tag_ts_ms {}; +struct tag_ts_us {}; +struct tag_ts_ns {}; +struct tag_du_day {}; +struct tag_du_s {}; +struct tag_du_ms {}; +struct tag_du_us {}; +struct tag_du_ns {}; +struct tag_dict {}; +struct tag_str {}; +struct tag_list {}; +struct tag_dec32 {}; +struct tag_dec64 {}; +struct tag_dec128 {}; +struct tag_struct {}; + +template +struct fragment_tag_murmur_hasher {}; + +/// Strong `murmur_jit_hasher` with a no-op body; used for storage types not present in the +/// table so nvJitLink still sees exactly one definition per `T` (weak overrides are not supported). +template +struct fragment_tag_murmur_hasher_noop {}; + +struct fragment_tag_murmur_entry {}; + +/// Owns ``murmur_jit_hash_dispatcher`` with the full type-id switch (matrix ``suffix`` = ``all``). +struct fragment_tag_murmur_dispatch_all {}; + +/// Owns ``murmur_jit_hash_dispatcher`` with an INT32-only branch (matrix ``suffix`` = ``int32``). +struct fragment_tag_murmur_dispatch_int32 {}; + +} // namespace cudf::hashing::detail::jit_lto diff --git a/cpp/src/hash/jit_lto_kernels/murmurhash3_x86_32_jit_dispatch_impl.cuh b/cpp/src/hash/jit_lto_kernels/murmurhash3_x86_32_jit_dispatch_impl.cuh new file mode 100644 index 000000000000..385dcdf541cf --- /dev/null +++ b/cpp/src/hash/jit_lto_kernels/murmurhash3_x86_32_jit_dispatch_impl.cuh @@ -0,0 +1,88 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include +#include + +namespace cudf::hashing::detail { + +/// Full type-id switch used when the dispatch fatbin is built with ``dispatch_int32_only == 0``. +__device__ inline hash_value_type murmur_jit_hash_dispatcher_all_types(column_device_view col, + uint32_t seed, + bool const nullable, + size_type row_index) +{ + switch (col.type().id()) { + case type_id::INT8: + return murmur_jit_hasher>(col, seed, nullable, row_index); + case type_id::INT16: + return murmur_jit_hasher>(col, seed, nullable, row_index); + case type_id::INT32: + return murmur_jit_hasher>(col, seed, nullable, row_index); + case type_id::INT64: + return murmur_jit_hasher>(col, seed, nullable, row_index); + case type_id::UINT8: + return murmur_jit_hasher>(col, seed, nullable, row_index); + case type_id::UINT16: + return murmur_jit_hasher>(col, seed, nullable, row_index); + case type_id::UINT32: + return murmur_jit_hasher>(col, seed, nullable, row_index); + case type_id::UINT64: + return murmur_jit_hasher>(col, seed, nullable, row_index); + case type_id::FLOAT32: + return murmur_jit_hasher>(col, seed, nullable, row_index); + case type_id::FLOAT64: + return murmur_jit_hasher>(col, seed, nullable, row_index); + case type_id::BOOL8: + return murmur_jit_hasher>(col, seed, nullable, row_index); + case type_id::TIMESTAMP_DAYS: + return murmur_jit_hasher>(col, seed, nullable, row_index); + case type_id::TIMESTAMP_SECONDS: + return murmur_jit_hasher>( + col, seed, nullable, row_index); + case type_id::TIMESTAMP_MILLISECONDS: + return murmur_jit_hasher>( + col, seed, nullable, row_index); + case type_id::TIMESTAMP_MICROSECONDS: + return murmur_jit_hasher>( + col, seed, nullable, row_index); + case type_id::TIMESTAMP_NANOSECONDS: + return murmur_jit_hasher>( + col, seed, nullable, row_index); + case type_id::DURATION_DAYS: + return murmur_jit_hasher>(col, seed, nullable, row_index); + case type_id::DURATION_SECONDS: + return murmur_jit_hasher>( + col, seed, nullable, row_index); + case type_id::DURATION_MILLISECONDS: + return murmur_jit_hasher>( + col, seed, nullable, row_index); + case type_id::DURATION_MICROSECONDS: + return murmur_jit_hasher>( + col, seed, nullable, row_index); + case type_id::DURATION_NANOSECONDS: + return murmur_jit_hasher>( + col, seed, nullable, row_index); + case type_id::DICTIONARY32: + return murmur_jit_hasher>(col, seed, nullable, row_index); + case type_id::STRING: + return murmur_jit_hasher>(col, seed, nullable, row_index); + case type_id::LIST: + return murmur_jit_hasher>(col, seed, nullable, row_index); + case type_id::DECIMAL32: + return murmur_jit_hasher>(col, seed, nullable, row_index); + case type_id::DECIMAL64: + return murmur_jit_hasher>(col, seed, nullable, row_index); + case type_id::DECIMAL128: + return murmur_jit_hasher>(col, seed, nullable, row_index); + case type_id::STRUCT: + return murmur_jit_hasher>(col, seed, nullable, row_index); + default: CUDF_UNREACHABLE("Invalid type_id."); + } +} + +} // namespace cudf::hashing::detail diff --git a/cpp/src/hash/jit_lto_kernels/murmurhash3_x86_32_jit_hasher_decl.cuh b/cpp/src/hash/jit_lto_kernels/murmurhash3_x86_32_jit_hasher_decl.cuh new file mode 100644 index 000000000000..f05b427b089c --- /dev/null +++ b/cpp/src/hash/jit_lto_kernels/murmurhash3_x86_32_jit_hasher_decl.cuh @@ -0,0 +1,34 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include +#include + +namespace cudf::hashing::detail { + +/** + * @brief Forward declarations for the Murmur JIT link graph. + * + * Per-type explicit specializations must appear in the TU before any use of that specialization. + * Hasher / noop fragment TUs include this header, then define `template <> ... + * murmur_jit_hasher`. The ``murmurhash_jit_dispatch`` fatbin defines + * ``murmur_jit_hash_dispatcher``; the entry kernel and hasher fragments only declare it here. + * + * `murmur_jit_hash_dispatcher` is declared here so `murmurhash3_x86_32_lto.cuh` can call it from + * nested/dictionary paths without including the dispatcher body before `murmur_jit_hasher` + * specializations in hasher fragment TUs. + */ +extern __device__ hash_value_type murmur_jit_hash_dispatcher(column_device_view col, + uint32_t seed, + bool nullable, + size_type row_index); + +template +extern __device__ hash_value_type +murmur_jit_hasher(column_device_view col, uint32_t seed, bool nullable, size_type row_index); + +} // namespace cudf::hashing::detail diff --git a/cpp/src/hash/jit_lto_kernels/murmurhash3_x86_32_lto.cuh b/cpp/src/hash/jit_lto_kernels/murmurhash3_x86_32_lto.cuh new file mode 100644 index 000000000000..cc3184f62a58 --- /dev/null +++ b/cpp/src/hash/jit_lto_kernels/murmurhash3_x86_32_lto.cuh @@ -0,0 +1,156 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include "murmurhash3_x86_32_jit_hasher_decl.cuh" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace cudf::hashing::detail { + +using result_type = hash_value_type; + +class element_hasher { + public: + /** + * @brief Constructs an element_hasher object. + * + * @param nulls Indicates whether to check for nulls + * @param seed The seed to use for the hash function + * @param null_hash The hash value to use for nulls + */ + __device__ element_hasher( + bool nulls, + result_type seed = DEFAULT_HASH_SEED, + result_type null_hash = cuda::std::numeric_limits::max()) noexcept + : _check_nulls(nulls), _seed(seed), _null_hash(null_hash) + { + } + + __device__ result_type seed() const noexcept { return _seed; } + + template + __device__ result_type operator()(column_device_view const& col, + size_type row_index) const noexcept + requires(column_device_view::has_element_accessor()) + { + if (_check_nulls && col.is_null(row_index)) { return _null_hash; } + return MurmurHash3_x86_32{_seed}(col.element(row_index)); + } + + template + __device__ result_type operator()(column_device_view const& col, + size_type row_index) const noexcept + requires(not column_device_view::has_element_accessor()) + { + CUDF_UNREACHABLE("Unsupported type in hash."); + } + + bool _check_nulls; + result_type _seed; + result_type _null_hash; +}; + +class element_hasher_adaptor { + static constexpr result_type NULL_HASH = cuda::std::numeric_limits::max(); + static constexpr result_type NON_NULL_HASH = 0; + + public: + __device__ element_hasher_adaptor(bool check_nulls, result_type seed) noexcept + : _element_hasher(check_nulls, seed), _check_nulls(check_nulls) + { + } + + template + __device__ result_type operator()(column_device_view const& col, + size_type row_index) const noexcept + requires(not cudf::is_nested() and not cudf::is_dictionary()) + { + return _element_hasher.template operator()(col, row_index); + } + + template + __device__ result_type operator()(column_device_view const& col, + size_type row_index) const noexcept + requires(cudf::is_dictionary()) + { + if (_check_nulls && col.is_null(row_index)) { return NULL_HASH; } + + auto const keys = col.child(dictionary_column_view::keys_column_index); + return murmur_jit_hash_dispatcher(keys, + _element_hasher.seed(), + _check_nulls, + static_cast(col.element(row_index))); + } + + template + __device__ result_type operator()(column_device_view const& col, + size_type row_index) const noexcept + requires(cudf::is_nested()) + { + auto hash = result_type{0}; + column_device_view curr_col = col.slice(row_index, 1); + while (curr_col.type().id() == type_id::STRUCT || curr_col.type().id() == type_id::LIST) { + if (_check_nulls) { + auto validity_it = cudf::detail::make_validity_iterator(curr_col); + hash = cudf::detail::accumulate( + validity_it, validity_it + curr_col.size(), hash, [](auto h, auto is_valid) { + return cudf::hashing::detail::hash_combine(h, is_valid ? NON_NULL_HASH : NULL_HASH); + }); + } + if (curr_col.type().id() == type_id::STRUCT) { + if (curr_col.num_child_columns() == 0) { return hash; } + curr_col = cudf::detail::structs_column_device_view(curr_col).get_sliced_child(0); + } else if (curr_col.type().id() == type_id::LIST) { + auto list_col = cudf::detail::lists_column_device_view(curr_col); + auto list_sizes = cudf::make_list_size_iterator(list_col); + hash = cudf::detail::accumulate( + list_sizes, list_sizes + list_col.size(), hash, [](auto h, auto size) { + return cudf::hashing::detail::hash_combine(h, MurmurHash3_x86_32{}(size)); + }); + curr_col = list_col.get_sliced_child(); + } + } + for (int i = 0; i < curr_col.size(); ++i) { + hash = cudf::hashing::detail::hash_combine( + hash, murmur_jit_hash_dispatcher(curr_col, _element_hasher.seed(), _check_nulls, i)); + } + return hash; + } + + element_hasher const _element_hasher; + bool const _check_nulls; +}; + +template +__device__ hash_value_type +hasher_impl(bool check_nulls, cudf::column_device_view col, uint32_t seed, size_type row_index) +{ + auto const hasher = element_hasher_adaptor{check_nulls, seed}; + return hasher.template operator()(col, row_index); +} + +template +__device__ hash_value_type +hasher(cudf::column_device_view col, uint32_t seed, bool const nullable, size_type row_index) +{ + return hasher_impl(nullable, col, seed, row_index); +} + +} // namespace cudf::hashing::detail diff --git a/cpp/src/hash/jit_lto_kernels/murmurhash_dispatch_fragment.cu.in b/cpp/src/hash/jit_lto_kernels/murmurhash_dispatch_fragment.cu.in new file mode 100644 index 000000000000..f640a807ccd1 --- /dev/null +++ b/cpp/src/hash/jit_lto_kernels/murmurhash_dispatch_fragment.cu.in @@ -0,0 +1,33 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +// Owns the device definition of ``murmur_jit_hash_dispatcher``. Matrix row injects +// ``dispatch_int32_only`` (0 = full switch, 1 = INT32-only) for ``if constexpr`` so the unused +// branch is not codegen'd in this fatbin. + +#include + +#include "murmurhash3_x86_32_jit_hasher_decl.cuh" +#include "murmurhash3_x86_32_jit_dispatch_impl.cuh" + +namespace cudf::hashing::detail { + +static constexpr bool dispatch_int32_only_v = static_cast(@dispatch_int32_only@); + +__device__ hash_value_type murmur_jit_hash_dispatcher(column_device_view col, + uint32_t seed, + bool const nullable, + size_type row_index) +{ + if constexpr (dispatch_int32_only_v) { + if (col.type().id() == type_id::INT32) { + return murmur_jit_hasher>(col, seed, nullable, row_index); + } + CUDF_UNREACHABLE("Invalid type_id."); + } else { + return murmur_jit_hash_dispatcher_all_types(col, seed, nullable, row_index); + } +} + +} // namespace cudf::hashing::detail diff --git a/cpp/src/hash/jit_lto_kernels/murmurhash_dispatch_matrix.json b/cpp/src/hash/jit_lto_kernels/murmurhash_dispatch_matrix.json new file mode 100644 index 000000000000..756f858393c1 --- /dev/null +++ b/cpp/src/hash/jit_lto_kernels/murmurhash_dispatch_matrix.json @@ -0,0 +1,12 @@ +{ + "_row": [ + { + "suffix": "all", + "dispatch_int32_only": 0 + }, + { + "suffix": "int32", + "dispatch_int32_only": 1 + } + ] +} diff --git a/cpp/src/hash/jit_lto_kernels/murmurhash_entry_kernel.cu.in b/cpp/src/hash/jit_lto_kernels/murmurhash_entry_kernel.cu.in new file mode 100644 index 000000000000..e66c042cacc7 --- /dev/null +++ b/cpp/src/hash/jit_lto_kernels/murmurhash_entry_kernel.cu.in @@ -0,0 +1,35 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include "murmurhash3_x86_32_jit_hasher_decl.cuh" + +extern "C" __global__ void cudf_murmurhash3_x86_32_jit_link_kernel(cudf::mutable_column_device_view output, + uint32_t seed, + cudf::table_device_view input, + bool nullable) +{ + auto const idx = cudf::detail::grid_1d::global_thread_id(); + if (idx >= input.num_rows()) { return; } + + auto const num_cols = input.num_columns(); + if (num_cols == 0) { + output.element(idx) = seed; + return; + } + + cudf::hash_value_type hash_value = + cudf::hashing::detail::murmur_jit_hash_dispatcher(input.column(0), seed, nullable, idx); + for (cudf::size_type i = 1; i < num_cols; ++i) { + hash_value = cudf::hashing::detail::hash_combine( + hash_value, + cudf::hashing::detail::murmur_jit_hash_dispatcher(input.column(i), seed, nullable, idx)); + } + output.element(idx) = hash_value; +} diff --git a/cpp/src/hash/jit_lto_kernels/murmurhash_entry_matrix.json b/cpp/src/hash/jit_lto_kernels/murmurhash_entry_matrix.json new file mode 100644 index 000000000000..c108798fd88a --- /dev/null +++ b/cpp/src/hash/jit_lto_kernels/murmurhash_entry_matrix.json @@ -0,0 +1,7 @@ +{ + "_singleton": [ + { + "k": "entry" + } + ] +} diff --git a/cpp/src/hash/jit_lto_kernels/murmurhash_hasher_fragment.cu.in b/cpp/src/hash/jit_lto_kernels/murmurhash_hasher_fragment.cu.in new file mode 100644 index 000000000000..5956f8010e08 --- /dev/null +++ b/cpp/src/hash/jit_lto_kernels/murmurhash_hasher_fragment.cu.in @@ -0,0 +1,26 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +// Strong explicit specialization of `murmur_jit_hasher` (real hasher); linked when logical +// type T appears in the input table. Other storage types use `murmurhash_jit_hasher_noop_*` +// (strong no-op) so nvJitLink sees exactly one definition per T. + +#include + +// `murmurhash3_x86_32_lto.cuh` pulls in `murmurhash3_x86_32_jit_hasher_decl.cuh` first (dispatcher +// decl + `murmur_jit_hasher` template decl). Dispatcher body lives only in the entry fatbin. +#include "murmurhash3_x86_32_lto.cuh" + +namespace cudf::hashing::detail { + +template <> +__device__ hash_value_type murmur_jit_hasher<@storage_cpp@>(column_device_view col, + uint32_t seed, + bool nullable, + size_type row_index) +{ + return hasher<@storage_cpp@>(col, seed, nullable, row_index); +} + +} // namespace cudf::hashing::detail diff --git a/cpp/src/hash/jit_lto_kernels/murmurhash_hasher_matrix.json b/cpp/src/hash/jit_lto_kernels/murmurhash_hasher_matrix.json new file mode 100644 index 000000000000..5216fcc72c94 --- /dev/null +++ b/cpp/src/hash/jit_lto_kernels/murmurhash_hasher_matrix.json @@ -0,0 +1,32 @@ +{ + "_row": [ + {"storage_cpp": "int8_t", "abbrev": "i8"}, + {"storage_cpp": "int16_t", "abbrev": "i16"}, + {"storage_cpp": "int32_t", "abbrev": "i32"}, + {"storage_cpp": "int64_t", "abbrev": "i64"}, + {"storage_cpp": "uint8_t", "abbrev": "u8"}, + {"storage_cpp": "uint16_t", "abbrev": "u16"}, + {"storage_cpp": "uint32_t", "abbrev": "u32"}, + {"storage_cpp": "uint64_t", "abbrev": "u64"}, + {"storage_cpp": "float", "abbrev": "f32"}, + {"storage_cpp": "double", "abbrev": "f64"}, + {"storage_cpp": "bool", "abbrev": "b8"}, + {"storage_cpp": "cudf::timestamp_D", "abbrev": "ts_day"}, + {"storage_cpp": "cudf::timestamp_s", "abbrev": "ts_s"}, + {"storage_cpp": "cudf::timestamp_ms", "abbrev": "ts_ms"}, + {"storage_cpp": "cudf::timestamp_us", "abbrev": "ts_us"}, + {"storage_cpp": "cudf::timestamp_ns", "abbrev": "ts_ns"}, + {"storage_cpp": "cudf::duration_D", "abbrev": "du_day"}, + {"storage_cpp": "cudf::duration_s", "abbrev": "du_s"}, + {"storage_cpp": "cudf::duration_ms", "abbrev": "du_ms"}, + {"storage_cpp": "cudf::duration_us", "abbrev": "du_us"}, + {"storage_cpp": "cudf::duration_ns", "abbrev": "du_ns"}, + {"storage_cpp": "cudf::dictionary32", "abbrev": "dict"}, + {"storage_cpp": "cudf::string_view", "abbrev": "str"}, + {"storage_cpp": "cudf::list_view", "abbrev": "list"}, + {"storage_cpp": "numeric::decimal32", "abbrev": "dec32"}, + {"storage_cpp": "numeric::decimal64", "abbrev": "dec64"}, + {"storage_cpp": "numeric::decimal128", "abbrev": "dec128"}, + {"storage_cpp": "cudf::struct_view", "abbrev": "struct"} + ] +} diff --git a/cpp/src/hash/jit_lto_kernels/murmurhash_hasher_noop_fragment.cu.in b/cpp/src/hash/jit_lto_kernels/murmurhash_hasher_noop_fragment.cu.in new file mode 100644 index 000000000000..f6284bd55aa8 --- /dev/null +++ b/cpp/src/hash/jit_lto_kernels/murmurhash_hasher_noop_fragment.cu.in @@ -0,0 +1,28 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +// Strong explicit specialization with a no-op body for `T` not used by any column in the table. +// nvJitLink does not coalesce weak + strong device symbols, so unused types use this fatbin +// instead of `fragment_tag_murmur_hasher` (real hasher). + +#include + +#include "murmurhash3_x86_32_jit_hasher_decl.cuh" + +namespace cudf::hashing::detail { + +template <> +__device__ hash_value_type murmur_jit_hasher<@storage_cpp@>(column_device_view col, + uint32_t seed, + bool nullable, + size_type row_index) +{ + (void)col; + (void)seed; + (void)nullable; + (void)row_index; + return hash_value_type{0}; +} + +} // namespace cudf::hashing::detail diff --git a/cpp/src/hash/murmurhash3_x86_32.cu b/cpp/src/hash/murmurhash3_x86_32.cu index bc0f6eb99e09..b493204a1cdd 100644 --- a/cpp/src/hash/murmurhash3_x86_32.cu +++ b/cpp/src/hash/murmurhash3_x86_32.cu @@ -2,16 +2,14 @@ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ -#include +#include "murmurhash3_x86_32_jit_launch.hpp" + #include -#include #include -#include #include -#include -#include +#include namespace cudf { namespace hashing { @@ -22,26 +20,7 @@ std::unique_ptr murmurhash3_x86_32(table_view const& input, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { - auto output = make_numeric_column(data_type(type_to_id()), - input.num_rows(), - mask_state::UNALLOCATED, - stream, - mr); - - // Return early if there's nothing to hash - if (input.num_columns() == 0 || input.num_rows() == 0) { return output; } - - bool const nullable = has_nulls(input); - auto const row_hasher = cudf::detail::row::hash::row_hasher(input, stream); - auto output_view = output->mutable_view(); - - // Compute the hash value for each row - thrust::tabulate(rmm::exec_policy_nosync(stream), - output_view.begin(), - output_view.end(), - row_hasher.device_hasher(nullable, seed)); - - return output; + return murmurhash3_x86_32_jit(input, seed, stream, mr); } } // namespace detail diff --git a/cpp/src/hash/murmurhash3_x86_32_jit_launch.hpp b/cpp/src/hash/murmurhash3_x86_32_jit_launch.hpp new file mode 100644 index 000000000000..cc7170a87a0b --- /dev/null +++ b/cpp/src/hash/murmurhash3_x86_32_jit_launch.hpp @@ -0,0 +1,340 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace cudf::hashing::detail { + +namespace { + +inline cudf::detail::jit_lto::LauncherJitCache& murmur_jit_launcher_cache() +{ + static cudf::detail::jit_lto::LauncherJitCache cache; + return cache; +} + +// Every `type_id` handled by `murmur_jit_hash_dispatcher` must have exactly one device definition +// of `murmur_jit_hasher` in the nvJitLink input for that link. Types present in the table use +// the real hasher fatbin; unused types use a strong no-op fatbin when linking the full dispatcher. +// Flat INT32-only tables use the INT32-only dispatcher and link only the INT32 strong hasher (no +// no-op stubs). +static constexpr std::array murmur_jit_hasher_type_ids{{ + type_id::INT8, + type_id::INT16, + type_id::INT32, + type_id::INT64, + type_id::UINT8, + type_id::UINT16, + type_id::UINT32, + type_id::UINT64, + type_id::FLOAT32, + type_id::FLOAT64, + type_id::BOOL8, + type_id::TIMESTAMP_DAYS, + type_id::TIMESTAMP_SECONDS, + type_id::TIMESTAMP_MILLISECONDS, + type_id::TIMESTAMP_MICROSECONDS, + type_id::TIMESTAMP_NANOSECONDS, + type_id::DURATION_DAYS, + type_id::DURATION_SECONDS, + type_id::DURATION_MILLISECONDS, + type_id::DURATION_MICROSECONDS, + type_id::DURATION_NANOSECONDS, + type_id::DICTIONARY32, + type_id::STRING, + type_id::LIST, + type_id::DECIMAL32, + type_id::DECIMAL64, + type_id::DECIMAL128, + type_id::STRUCT, +}}; + +/** + * @brief Collect logical `type_id`s for fragment selection, including descendants of nested + * columns. + * + * Nested hashing dispatches primitives via `murmur_jit_hash_dispatcher`; strong `murmur_jit_hasher` + * fragments must be linked for every storage type that can appear under struct/list/dictionary + * columns, not only top-level columns. + */ +inline void collect_nested_logical_types(std::unordered_set& ids, column_view const& col) +{ + ids.insert(col.type().id()); + switch (col.type().id()) { + case type_id::STRUCT: + for (size_type i = 0; i < col.num_children(); ++i) { + collect_nested_logical_types(ids, col.child(i)); + } + break; + case type_id::LIST: { + lists_column_view const lcv(col); + collect_nested_logical_types(ids, lcv.child()); + break; + } + case type_id::DICTIONARY32: { + dictionary_column_view const dcv(col); + collect_nested_logical_types(ids, dcv.keys()); + break; + } + default: break; + } +} + +inline void collect_table_logical_types(std::unordered_set& ids, table_view const& table) +{ + for (size_type c = 0; c < table.num_columns(); ++c) { + collect_nested_logical_types(ids, table.column(c)); + } +} + +inline void add_strong_hasher_fragment(cudf::detail::jit_lto::AlgorithmPlanner& planner, type_id id) +{ + using namespace cudf::hashing::detail::jit_lto; + switch (id) { + case type_id::INT8: planner.add_static_fragment>(); break; + case type_id::INT16: planner.add_static_fragment>(); break; + case type_id::INT32: planner.add_static_fragment>(); break; + case type_id::INT64: planner.add_static_fragment>(); break; + case type_id::UINT8: planner.add_static_fragment>(); break; + case type_id::UINT16: planner.add_static_fragment>(); break; + case type_id::UINT32: planner.add_static_fragment>(); break; + case type_id::UINT64: planner.add_static_fragment>(); break; + case type_id::FLOAT32: + planner.add_static_fragment>(); + break; + case type_id::FLOAT64: + planner.add_static_fragment>(); + break; + case type_id::BOOL8: planner.add_static_fragment>(); break; + case type_id::TIMESTAMP_DAYS: + planner.add_static_fragment>(); + break; + case type_id::TIMESTAMP_SECONDS: + planner.add_static_fragment>(); + break; + case type_id::TIMESTAMP_MILLISECONDS: + planner.add_static_fragment>(); + break; + case type_id::TIMESTAMP_MICROSECONDS: + planner.add_static_fragment>(); + break; + case type_id::TIMESTAMP_NANOSECONDS: + planner.add_static_fragment>(); + break; + case type_id::DURATION_DAYS: + planner.add_static_fragment>(); + break; + case type_id::DURATION_SECONDS: + planner.add_static_fragment>(); + break; + case type_id::DURATION_MILLISECONDS: + planner.add_static_fragment>(); + break; + case type_id::DURATION_MICROSECONDS: + planner.add_static_fragment>(); + break; + case type_id::DURATION_NANOSECONDS: + planner.add_static_fragment>(); + break; + case type_id::DICTIONARY32: + planner.add_static_fragment>(); + break; + case type_id::STRING: planner.add_static_fragment>(); break; + case type_id::LIST: planner.add_static_fragment>(); break; + case type_id::DECIMAL32: + planner.add_static_fragment>(); + break; + case type_id::DECIMAL64: + planner.add_static_fragment>(); + break; + case type_id::DECIMAL128: + planner.add_static_fragment>(); + break; + case type_id::STRUCT: + planner.add_static_fragment>(); + break; + default: break; + } +} + +inline void add_noop_hasher_fragment(cudf::detail::jit_lto::AlgorithmPlanner& planner, type_id id) +{ + using namespace cudf::hashing::detail::jit_lto; + switch (id) { + case type_id::INT8: + planner.add_static_fragment>(); + break; + case type_id::INT16: + planner.add_static_fragment>(); + break; + case type_id::INT32: + planner.add_static_fragment>(); + break; + case type_id::INT64: + planner.add_static_fragment>(); + break; + case type_id::UINT8: + planner.add_static_fragment>(); + break; + case type_id::UINT16: + planner.add_static_fragment>(); + break; + case type_id::UINT32: + planner.add_static_fragment>(); + break; + case type_id::UINT64: + planner.add_static_fragment>(); + break; + case type_id::FLOAT32: + planner.add_static_fragment>(); + break; + case type_id::FLOAT64: + planner.add_static_fragment>(); + break; + case type_id::BOOL8: + planner.add_static_fragment>(); + break; + case type_id::TIMESTAMP_DAYS: + planner.add_static_fragment>(); + break; + case type_id::TIMESTAMP_SECONDS: + planner.add_static_fragment>(); + break; + case type_id::TIMESTAMP_MILLISECONDS: + planner.add_static_fragment>(); + break; + case type_id::TIMESTAMP_MICROSECONDS: + planner.add_static_fragment>(); + break; + case type_id::TIMESTAMP_NANOSECONDS: + planner.add_static_fragment>(); + break; + case type_id::DURATION_DAYS: + planner.add_static_fragment>(); + break; + case type_id::DURATION_SECONDS: + planner.add_static_fragment>(); + break; + case type_id::DURATION_MILLISECONDS: + planner.add_static_fragment>(); + break; + case type_id::DURATION_MICROSECONDS: + planner.add_static_fragment>(); + break; + case type_id::DURATION_NANOSECONDS: + planner.add_static_fragment>(); + break; + case type_id::DICTIONARY32: + planner.add_static_fragment>(); + break; + case type_id::STRING: + planner.add_static_fragment>(); + break; + case type_id::LIST: + planner.add_static_fragment>(); + break; + case type_id::DECIMAL32: + planner.add_static_fragment>(); + break; + case type_id::DECIMAL64: + planner.add_static_fragment>(); + break; + case type_id::DECIMAL128: + planner.add_static_fragment>(); + break; + case type_id::STRUCT: + planner.add_static_fragment>(); + break; + default: break; + } +} + +} // namespace + +/// True for a flat ``INT32``-only table: ``collect_table_logical_types`` yields exactly +/// ``{INT32}``. Then we link ``fragment_tag_murmur_dispatch_int32`` and **omit** noop hasher +/// fatbins. Nested or mixed types keep the full dispatcher and noop stubs for unreferenced +/// ``murmur_jit_hasher`` symbols. +inline bool use_murmur_jit_int32_only_dispatch(std::unordered_set const& logical_types) +{ + return logical_types.size() == 1u && logical_types.count(type_id::INT32) == 1u; +} + +inline std::unique_ptr murmurhash3_x86_32_jit(table_view const& input, + uint32_t seed, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + auto output = make_numeric_column(data_type(type_to_id()), + input.num_rows(), + mask_state::UNALLOCATED, + stream, + mr); + + if (input.num_columns() == 0 || input.num_rows() == 0) { return output; } + + bool const nullable = has_nulls(input); + auto const preprocessed = cudf::detail::row::hash::preprocessed_table::create(input, stream); + table_device_view const input_dv{*preprocessed}; + + auto output_view = output->mutable_view(); + auto d_output = mutable_column_device_view::create(output_view, stream); + + cudf::detail::jit_lto::AlgorithmPlanner planner{"cudf_murmurhash3_x86_32_jit_link_kernel", + murmur_jit_launcher_cache()}; + + std::unordered_set logical_types; + collect_table_logical_types(logical_types, input); + + using namespace cudf::hashing::detail::jit_lto; + planner.add_static_fragment(); + bool const int32_only_jit_link = use_murmur_jit_int32_only_dispatch(logical_types); + if (int32_only_jit_link) { + planner.add_static_fragment(); + } else { + planner.add_static_fragment(); + } + for (type_id const id : murmur_jit_hasher_type_ids) { + if (logical_types.count(id) != 0u) { + add_strong_hasher_fragment(planner, id); + } else if (not int32_only_jit_link) { + add_noop_hasher_fragment(planner, id); + } + } + + auto launcher = planner.get_launcher(); + + cudf::detail::grid_1d const grid{input.num_rows(), 256}; + launcher + ->dispatch( + stream.value(), + dim3(grid.num_blocks), + dim3(grid.num_threads_per_block), + 0, + *d_output, + seed, + input_dv, + nullable); + + return output; +} + +} // namespace cudf::hashing::detail diff --git a/cpp/src/jit_lto/AlgorithmLauncher.cpp b/cpp/src/jit_lto/AlgorithmLauncher.cpp new file mode 100644 index 000000000000..def84fb3626d --- /dev/null +++ b/cpp/src/jit_lto/AlgorithmLauncher.cpp @@ -0,0 +1,51 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +namespace cudf::detail::jit_lto { + +AlgorithmLauncher::AlgorithmLauncher(cudaKernel_t k, cudaLibrary_t lib) : kernel{k}, library{lib} {} + +AlgorithmLauncher::~AlgorithmLauncher() +{ + if (library != nullptr) { (void)cudaLibraryUnload(library); } +} + +AlgorithmLauncher::AlgorithmLauncher(AlgorithmLauncher&& other) noexcept + : kernel{other.kernel}, library{other.library} +{ + other.kernel = nullptr; + other.library = nullptr; +} + +AlgorithmLauncher& AlgorithmLauncher::operator=(AlgorithmLauncher&& other) noexcept +{ + if (this != &other) { + if (library != nullptr) { cudaLibraryUnload(library); } + kernel = other.kernel; + library = other.library; + other.kernel = nullptr; + other.library = nullptr; + } + return *this; +} + +void AlgorithmLauncher::call( + cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, void** kernel_args) +{ + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.stream = stream; + config.dynamicSmemBytes = shared_mem; + config.numAttrs = 0; + config.attrs = NULL; + + CUDF_CUDA_TRY(cudaLaunchKernelExC(&config, kernel, kernel_args)); +} + +} // namespace cudf::detail::jit_lto diff --git a/cpp/src/jit_lto/AlgorithmPlanner.cpp b/cpp/src/jit_lto/AlgorithmPlanner.cpp new file mode 100644 index 000000000000..357bdab6cb19 --- /dev/null +++ b/cpp/src/jit_lto/AlgorithmPlanner.cpp @@ -0,0 +1,171 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cudf::detail::jit_lto { + +namespace { + +void emit_jit_lto_build_timing(double const build_ms) +{ + CUDF_LOG_INFO("jit_lto: AlgorithmPlanner::build %.6f ms", build_ms); + if (std::getenv("CUDF_JIT_LTO_LINK_TIMING") != nullptr) { + if (auto const* meta = std::getenv("CUDF_JIT_LTO_LINK_TIMING_META"); + meta != nullptr && meta[0] != '\0') { + std::fprintf(stderr, "CUDF_JIT_LTO_LINK_TIMING build_ms=%.6f %s\n", build_ms, meta); + } else { + std::fprintf(stderr, "CUDF_JIT_LTO_LINK_TIMING build_ms=%.6f\n", build_ms); + } + std::fflush(stderr); + } +} + +std::vector make_nvjitlink_option_strings(std::string arch_sm) +{ + std::vector opts; + opts.reserve(8); + opts.emplace_back("-lto"); + opts.push_back(std::move(arch_sm)); + + // Whitespace-separated extra nvJitLink flags (e.g. "-O3", "-split-compile=0"). When non-empty, + // this replaces compile-time CUDF_JIT_LTO_NVJITLINK_OPTSET extras so one build can sweep flags. + // Note: "-O0" with "-lto" has triggered cudaErrorLaunchFailure on some toolchains; axis + // benchmarks skip that pair (see benchmark_murmur_jit_lto_axes.sh). JIT fragments are LTO-IR; + // -lto is required. + if (auto const* env = std::getenv("CUDF_JIT_LTO_NVJITLINK_OPTIONS"); + env != nullptr && env[0] != '\0') { + std::istringstream iss(env); + std::string tok; + while (iss >> tok) { + opts.push_back(std::move(tok)); + } + return opts; + } + +#if defined(CUDF_JIT_LTO_NVJITLINK_PROFILE_O0) + opts.emplace_back("-O0"); +#elif defined(CUDF_JIT_LTO_NVJITLINK_PROFILE_O3) + opts.emplace_back("-O3"); +#elif defined(CUDF_JIT_LTO_NVJITLINK_PROFILE_SPLIT_COMPILE) + opts.emplace_back("-split-compile=0"); +#elif defined(CUDF_JIT_LTO_NVJITLINK_PROFILE_SPLIT_COMPILE_EXTENDED) + opts.emplace_back("-split-compile-extended=0"); +#endif + return opts; +} + +} // namespace + +std::string AlgorithmPlanner::get_fragments_key() const +{ + std::string key = ""; + for (const auto& fragment : this->fragments) { + key += fragment->get_key(); + } + return key; +} + +std::shared_ptr AlgorithmPlanner::read_cache(std::string const& launch_key) const +{ + auto& launchers = jit_cache_.launchers; + std::shared_lock read_lock(jit_cache_.mutex); + if (auto it = launchers.find(launch_key); it != launchers.end()) { return it->second; } + return nullptr; +} + +std::shared_ptr AlgorithmPlanner::get_launcher() +{ + auto& launchers = jit_cache_.launchers; + auto launch_key = this->get_fragments_key(); + + if (auto hit = read_cache(launch_key)) { return hit; } + + std::unique_lock write_lock(jit_cache_.mutex); + if (auto it = launchers.find(launch_key); it != launchers.end()) { return it->second; } + + auto launcher = this->build(); + launchers[launch_key] = launcher; + return launcher; +} + +std::shared_ptr AlgorithmPlanner::build() +{ + using clock = std::chrono::steady_clock; + using duration_ms = std::chrono::duration; + + auto const t_build_start = clock::now(); + + int device = 0; + int major = 0; + int minor = 0; + CUDF_CUDA_TRY(cudaGetDevice(&device)); + CUDF_CUDA_TRY(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device)); + CUDF_CUDA_TRY(cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device)); + + std::string archs = "-arch=sm_" + std::to_string((major * 10 + minor)); + + auto const opt_strings = make_nvjitlink_option_strings(std::move(archs)); + std::vector opt_ptrs; + opt_ptrs.reserve(opt_strings.size()); + for (auto const& s : opt_strings) { + opt_ptrs.push_back(s.c_str()); + } + + nvJitLinkHandle handle; + auto result = nvJitLinkCreate(&handle, static_cast(opt_ptrs.size()), opt_ptrs.data()); + check_nvjitlink_result(handle, result); + + for (const auto& frag : this->fragments) { + frag->add_to(handle); + } + + result = nvJitLinkComplete(handle); + check_nvjitlink_result(handle, result); + + size_t cubin_size; + result = nvJitLinkGetLinkedCubinSize(handle, &cubin_size); + check_nvjitlink_result(handle, result); + + std::unique_ptr cubin{new char[cubin_size]}; + result = nvJitLinkGetLinkedCubin(handle, cubin.get()); + check_nvjitlink_result(handle, result); + + result = nvJitLinkDestroy(&handle); + CUDF_EXPECTS(result == NVJITLINK_SUCCESS, "nvJitLinkDestroy failed"); + + cudaLibrary_t library; + CUDF_CUDA_TRY( + cudaLibraryLoadData(&library, cubin.get(), nullptr, nullptr, 0, nullptr, nullptr, 0)); + + cudaKernel_t kernel; + CUDF_CUDA_TRY(cudaLibraryGetKernel(&kernel, library, this->entrypoint.c_str())); + + double const build_ms = duration_ms(clock::now() - t_build_start).count(); + emit_jit_lto_build_timing(build_ms); + + return std::make_shared(kernel, library); +} + +} // namespace cudf::detail::jit_lto diff --git a/cpp/src/jit_lto/FragmentEntry.cpp b/cpp/src/jit_lto/FragmentEntry.cpp new file mode 100644 index 000000000000..7c66c33dfd80 --- /dev/null +++ b/cpp/src/jit_lto/FragmentEntry.cpp @@ -0,0 +1,18 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +namespace cudf::detail::jit_lto { + +bool FatbinFragmentEntry::add_to(nvJitLinkHandle& handle) const +{ + auto result = nvJitLinkAddData(handle, NVJITLINK_INPUT_ANY, get_data(), get_length(), get_key()); + + check_nvjitlink_result(handle, result); + return true; +} + +} // namespace cudf::detail::jit_lto diff --git a/cpp/src/jit_lto/nvjitlink_checker.cpp b/cpp/src/jit_lto/nvjitlink_checker.cpp new file mode 100644 index 000000000000..25bdc27cd5b9 --- /dev/null +++ b/cpp/src/jit_lto/nvjitlink_checker.cpp @@ -0,0 +1,31 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include + +#include +#include + +namespace cudf::detail::jit_lto { + +void check_nvjitlink_result(nvJitLinkHandle handle, nvJitLinkResult result) +{ + if (result != NVJITLINK_SUCCESS) { + std::string error_msg = "nvJITLink failed with error " + std::to_string(result); + size_t log_size = 0; + result = nvJitLinkGetErrorLogSize(handle, &log_size); + if (result == NVJITLINK_SUCCESS && log_size > 0) { + std::unique_ptr log{new char[log_size]}; + result = nvJitLinkGetErrorLog(handle, log.get()); + if (result == NVJITLINK_SUCCESS) { error_msg += "\n" + std::string(log.get()); } + } + CUDF_FAIL(error_msg); + } +} + +} // namespace cudf::detail::jit_lto diff --git a/python/pylibcudf/benchmark_murmur_jit_lto_axes.sh b/python/pylibcudf/benchmark_murmur_jit_lto_axes.sh new file mode 100755 index 000000000000..03d4e9289e60 --- /dev/null +++ b/python/pylibcudf/benchmark_murmur_jit_lto_axes.sh @@ -0,0 +1,123 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# +# Murmur JIT+LTO link benchmark: sweep nvJitLink *extra* flags via CUDF_JIT_LTO_NVJITLINK_OPTIONS. +# libcudf always passes -lto; embedded fragments are LTO-IR, so omitting -lto is not supported. +# +# Skipped preset: +# - O0 (nvJitLink ``-lto`` + ``-O0`` → ``cudaErrorLaunchFailure`` on observed stacks) +# +# Each profile clears ~/.nv/ComputeCache once, then runs two Python processes (cold/warm disk); +# each run does one in-process table sweep. All timings append to one CSV. +# +# Usage: +# conda activate cudf_2606 # or set CONDA_ENV +# ./python/pylibcudf/benchmark_murmur_jit_lto_axes.sh [PROFILE ...] +# +# No arguments: lto, split-compile, split-compile-extended. (-O3 is nvJitLink default with -lto; +# pass ``O3`` explicitly if needed.) ``O0`` is not run (see above). +# +# Environment: +# CONDA_ENV conda env name (default: cudf_2606) +# REPO_ROOT cuDF git root (default: inferred from this script) +# OUTDIR output directory (default: $REPO_ROOT/jit_lto_bench_out) +# CSV_OUT raw per-link CSV (default: $OUTDIR/murmur_jit_lto_bench.csv); removed at start +# SUMMARY_CSV pivot summary for analysis (default: $OUTDIR/murmur_jit_lto_bench_summary.csv) + +set -euo pipefail + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +CONDA_ENV="${CONDA_ENV:-cudf_2606}" +OUTDIR="${OUTDIR:-${REPO_ROOT}/jit_lto_bench_out}" +CSV_OUT="${CSV_OUT:-${OUTDIR}/murmur_jit_lto_bench.csv}" +SUMMARY_CSV="${SUMMARY_CSV:-${OUTDIR}/murmur_jit_lto_bench_summary.csv}" +BENCH_PY="${REPO_ROOT}/python/pylibcudf/benchmark_murmur_jit_lto_link.py" + +mkdir -p "${OUTDIR}" +rm -f "${CSV_OUT}" + +declare -A NVJITLINK_PROFILE_OPTS=( + [lto]="" + [O0]="-O0" + [O3]="-O3" + [split-compile]="-split-compile=0" + [split-compile-extended]="-split-compile-extended=0" +) + +DEFAULT_PROFILES=(lto split-compile split-compile-extended) + +filter_profiles() { + local -a in=( "$@" ) + RUN_PROFILES=() + local p + for p in "${in[@]}"; do + if [[ "${p}" == O0 ]]; then + echo "Skipping profile O0: nvJitLink -lto -O0 fails at kernel launch on observed stacks." >&2 + continue + fi + RUN_PROFILES+=( "${p}" ) + done +} + +if [[ ${#} -gt 0 ]]; then + filter_profiles "$@" +else + RUN_PROFILES=( "${DEFAULT_PROFILES[@]}" ) +fi + +if [[ ${#RUN_PROFILES[@]} -eq 0 ]]; then + echo "No profiles left to run after filtering." >&2 + exit 1 +fi + +for p in "${RUN_PROFILES[@]}"; do + if [[ ! -v NVJITLINK_PROFILE_OPTS[$p] ]]; then + echo "Unknown profile '${p}'. Valid: ${!NVJITLINK_PROFILE_OPTS[*]}" >&2 + exit 1 + fi +done + +unset CUDF_JIT_LTO_DISABLE_LTO + +# Conda's cuda-nvcc activate hook uses NVCC_PREPEND_FLAGS without a default; ``set -u`` makes that +# an error. Allow unset variables only while initializing conda. +set +u +# shellcheck source=/dev/null +source "$(conda info --base)/etc/profile.d/conda.sh" +conda activate "${CONDA_ENV}" +set -u + +echo "REPO_ROOT=${REPO_ROOT}" +echo "CONDA_DEFAULT_ENV=${CONDA_DEFAULT_ENV}" +echo "OUTDIR=${OUTDIR}" +echo "CSV_OUT=${CSV_OUT}" +echo "SUMMARY_CSV=${SUMMARY_CSV}" +echo "PROFILES=${RUN_PROFILES[*]}" +echo "" + +for tag in "${RUN_PROFILES[@]}"; do + extra="${NVJITLINK_PROFILE_OPTS[$tag]}" + if [[ -n "${extra}" ]]; then + export CUDF_JIT_LTO_NVJITLINK_OPTIONS="${extra}" + else + unset CUDF_JIT_LTO_NVJITLINK_OPTIONS + fi + + echo "======== Profile ${tag} (CUDF_JIT_LTO_NVJITLINK_OPTIONS=${extra:-}, always -lto) ========" + echo "======== Clear CUDA compute cache ========" + rm -rf "${HOME}/.nv/ComputeCache" + + for inv in 1 2; do + echo "-------- Python: profile=${tag} script_invocation=${inv} (append ${CSV_OUT}) --------" + python "${BENCH_PY}" \ + --nvjitlink-optset "${tag}" \ + --script-invocation "${inv}" \ + --output-csv "${CSV_OUT}" + done + echo "" +done + +python "${BENCH_PY}" --summarize-bench-csv "${CSV_OUT}" --summary-output "${SUMMARY_CSV}" + +echo "Done. Raw: ${CSV_OUT} Summary: ${SUMMARY_CSV}" diff --git a/python/pylibcudf/benchmark_murmur_jit_lto_link.py b/python/pylibcudf/benchmark_murmur_jit_lto_link.py new file mode 100644 index 000000000000..5528d7e22386 --- /dev/null +++ b/python/pylibcudf/benchmark_murmur_jit_lto_link.py @@ -0,0 +1,538 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +""" +Benchmark ``AlgorithmPlanner::build()`` (MurmurHash3 x86_32 JIT+LTO / nvJitLink). + +**Axes (how this script fits in)** + +- **CUDA / nvJitLink disk cache** (~``~/.nv/ComputeCache``): clear before the *first* process + invocation; the *second* process invocation (separate ``python`` run) sees a warm disk cache. +- **In-process sweep**: one full A→E table pass per Python run (``--process-passes`` defaults to + **1**; increase only if you explicitly want repeated sweeps in the same process). +- **nvJitLink link flags**: **``CUDF_JIT_LTO_NVJITLINK_OPTIONS``** — whitespace-separated extras after + ``-lto`` and ``-arch=sm_XX``. Murmur JIT fragments are LTO-IR; **``-lto`` is always required** (the + axis driver does not attempt no-LTO links). The axis script does **not** sweep **``-O0``** with + ``-lto`` (observed ``cudaErrorLaunchFailure``). No redundant default **``-O3``** row (O3 is the + default optimization level with LTO). + +Set ``CUDF_JIT_LTO_LINK_TIMING=1`` (set automatically here) so ``build()`` emits one stderr line per +link. Optional ``CUDF_JIT_LTO_LINK_TIMING_META`` is set per pass for parsing (``script_inv``, +``process_pass``). + +**Suggested sweep** (one libcudf build, then):: + + ./python/pylibcudf/benchmark_murmur_jit_lto_axes.sh + +Use ``--output-csv PATH`` to append one row per ``build_ms`` event (header written if the file is +new or empty). The axis shell driver writes a single CSV for the full sweep, then a **summary CSV** +(``--summarize-bench-csv``) with cold/warm columns and percent deltas vs ``lto`` baseline per +scenario and totals. + +Table order: **A** ``float32`` (full dispatcher first), **B** ``int32``, then struct, list, int64. +""" + +from __future__ import annotations + +import argparse +import csv +import io +import json +import os +import re +import sys +import threading +from pathlib import Path +from typing import Any + +# (letter, description) aligned with table order in ``_benchmark_tables``. +_BENCHMARK_TYPE_LABELS: list[tuple[str, str]] = [ + ("A", "float32 column (first build: full dispatcher)"), + ("B", "int32 column"), + ("C", "struct column (nested)"), + ("D", "list column (nested)"), + ("E", "int64 column"), +] + + +def _drop_script_dir_from_sys_path() -> None: + """If this file lives under ``python/pylibcudf/``, running it prepends that directory to + ``sys.path`` so ``import pylibcudf`` loads the **source** tree (no built Cython extensions) + instead of the conda/site-packages install. Drop that entry so the installed package wins. + """ + script_dir = os.path.dirname(os.path.abspath(__file__)) + if sys.path and os.path.realpath(sys.path[0]) == os.path.realpath( + script_dir + ): + del sys.path[0] + + +_drop_script_dir_from_sys_path() + +# Optional trailing meta: space-separated key=value tokens (from CUDF_JIT_LTO_LINK_TIMING_META). +_TIMING_LINE = re.compile( + r"^CUDF_JIT_LTO_LINK_TIMING build_ms=(?P[\d.]+)(?:\s+(?P.+))?\s*$" +) + + +def _parse_timing_meta(meta: str | None) -> dict[str, str]: + if not meta: + return {} + out: dict[str, str] = {} + for part in meta.split(): + if "=" in part: + k, _, v = part.partition("=") + out[k] = v + return out + + +def _parse_build_timing_lines(captured_stderr: str) -> list[dict[str, Any]]: + events: list[dict[str, Any]] = [] + for line in captured_stderr.splitlines(): + m = _TIMING_LINE.match(line.strip()) + if m: + ev: dict[str, Any] = {"build_ms": float(m.group("build"))} + meta = _parse_timing_meta(m.group("meta")) + if meta: + ev["meta"] = meta + events.append(ev) + return events + + +def _attach_scenario_labels( + events: list[dict[str, Any]], +) -> list[dict[str, Any]]: + """First ``build`` line in a pass is scenario A, second is B, ...""" + out = [] + for i, ev in enumerate(events): + letter, desc = _BENCHMARK_TYPE_LABELS[i] if i < len(_BENCHMARK_TYPE_LABELS) else ( + "?", + "unknown", + ) + row = {**ev, "type_letter": letter, "type_desc": desc} + out.append(row) + return out + + +def _capture_cpp_stderr(work): + """Redirect FD 2 so libcudf ``fprintf(stderr, ...)`` is captured.""" + read_fd, write_fd = os.pipe() + saved = os.dup(2) + os.dup2(write_fd, 2) + os.close(write_fd) + buf = io.StringIO() + + def reader(): + with os.fdopen(read_fd, "r") as pipe: + buf.write(pipe.read()) + + th = threading.Thread(target=reader) + th.start() + try: + work() + finally: + os.dup2(saved, 2) + os.close(saved) + th.join() + return buf.getvalue() + + +def _benchmark_tables(): + """Return ``pylibcudf.Table`` for types A–E (float32 before int32 for full-dispatch first link).""" + import pyarrow as pa + + import pylibcudf as plc + + pa_float_first = pa.table( + {"x": pa.array([1.0, 2.0, 3.0, 4.0, 5.0], type=pa.float32())} + ) + pa_int32 = pa.table({"x": pa.array([1, 2, 3, 4, 5], type=pa.int32())}) + pa_c = pa.table( + { + "s": pa.array( + [{"a": 1, "b": 2}, {"a": 3, "b": 4}, {"a": 5, "b": 6}], + type=pa.struct([("a", pa.int32()), ("b", pa.int32())]), + ) + } + ) + pa_d = pa.table( + { + "lst": pa.array( + [[1, 2], [3, 4, 5], []], + type=pa.list_(pa.int32()), + ) + } + ) + pa_e = pa.table({"x": pa.array([1, 2, 3, 4, 5], type=pa.int64())}) + + return [ + plc.Table.from_arrow(pa_float_first), + plc.Table.from_arrow(pa_int32), + plc.Table.from_arrow(pa_c), + plc.Table.from_arrow(pa_d), + plc.Table.from_arrow(pa_e), + ] + + +def _cache_axis_interpretation( + *, script_invocation: int, process_pass: int +) -> dict[str, str]: + """Human-readable labels for the benchmark grid (not exhaustive of every subsystem).""" + disk = "warm" if script_invocation >= 2 else "cold" + mem = "warm" if process_pass >= 2 else "cold" + return { + "cuda_compute_cache_disk": disk, + "libcudf_jit_lto_in_process": mem, + } + + +def run_benchmark( + *, + script_invocation: int = 1, + process_passes: int = 1, + nvjitlink_optset: str = "lto", +) -> dict[str, Any]: + """Run ``process_passes`` full table sweeps; capture timing lines per pass.""" + os.environ["CUDF_JIT_LTO_LINK_TIMING"] = "1" + + import pylibcudf as plc + + seed = plc.hashing.LIBCUDF_DEFAULT_HASH_SEED + tables = _benchmark_tables() + + pass_results: list[dict[str, Any]] = [] + + for process_pass in range(1, process_passes + 1): + os.environ["CUDF_JIT_LTO_LINK_TIMING_META"] = ( + f"script_inv={script_invocation} process_pass={process_pass} " + f"nvjitlink_optset={nvjitlink_optset}" + ) + + def work(): + for tbl in tables: + plc.hashing.murmurhash3_x86_32(tbl, seed) + + captured = _capture_cpp_stderr(work) + raw_events = _parse_build_timing_lines(captured) + labeled = _attach_scenario_labels(raw_events) + pass_results.append( + { + "process_pass": process_pass, + "axes_interpretation": _cache_axis_interpretation( + script_invocation=script_invocation, + process_pass=process_pass, + ), + "link_events": labeled, + "expected_scenarios": len(_BENCHMARK_TYPE_LABELS), + } + ) + + return { + "nvjitlink_optset": nvjitlink_optset, + "script_invocation": script_invocation, + "process_passes": process_passes, + "passes": pass_results, + } + + +_CSV_COLUMNS = ( + "nvjitlink_optset", + "script_invocation", + "process_pass", + "cuda_compute_cache_disk", + "libcudf_jit_lto_in_process", + "scenario", + "type_desc", + "build_ms", +) + + +_SUMMARY_CSV_COLUMNS = ( + "nvjitlink_optset", + "scenario", + "type_desc", + "cold_ms", + "warm_ms", + "warm_pct_of_cold", + "cold_pct_vs_lto_baseline", + "warm_pct_vs_lto_baseline", +) + + +def summarize_bench_csv(raw_path: Path, summary_path: Path) -> None: + """Read raw benchmark CSV (per-link rows); write a pivot-style summary for analysis. + + Expects ``script_invocation`` 1 = cold disk, 2 = warm disk (as produced by the axis shell). + Percent columns: ``warm_pct_of_cold`` = 100×warm/cold; ``*_pct_vs_lto_baseline`` = 100×(x/lto−1) + for the same scenario and invocation class. Baseline rows (``lto``) use 0 for those deltas. + Appends a **TOTAL** row per ``nvjitlink_optset`` (sum of A–E present in the raw file). + """ + with raw_path.open(newline="", encoding="utf-8") as f: + rows = list(csv.DictReader(f)) + if not rows: + raise SystemExit(f"empty or missing benchmark CSV: {raw_path}") + + # ms[optset][scenario][inv] = float + ms: dict[str, dict[str, dict[int, float]]] = {} + type_desc: dict[tuple[str, str], str] = {} + opt_order: list[str] = [] + + for r in rows: + opt = r["nvjitlink_optset"] + scen = r["scenario"] + inv = int(r["script_invocation"]) + val = float(r["build_ms"]) + if opt not in ms: + ms[opt] = {} + opt_order.append(opt) + ms[opt].setdefault(scen, {})[inv] = val + type_desc[(opt, scen)] = r.get("type_desc", "") + + known_scen = ["A", "B", "C", "D", "E"] + all_scen: set[str] = set() + for d in ms.values(): + all_scen |= set(d.keys()) + scen_order = [s for s in known_scen if s in all_scen] + sorted( + all_scen - set(known_scen) + ) + + def lto_ms(scen: str, inv: int) -> float | None: + if "lto" not in ms or scen not in ms["lto"]: + return None + return ms["lto"][scen].get(inv) + + def pct_vs_baseline( + opt: str, scen: str, inv: int, value: float | None + ) -> str: + if value is None: + return "" + if opt == "lto": + return "0" + base = lto_ms(scen, inv) + if base is None or base <= 0: + return "" + return f"{100.0 * (value / base - 1.0):.2f}" + + out_rows: list[dict[str, str]] = [] + + for opt in opt_order: + cold_sum = 0.0 + warm_sum = 0.0 + for scen in scen_order: + if scen not in ms[opt]: + continue + d = ms[opt][scen] + cold = d.get(1) + warm = d.get(2) + if cold is not None: + cold_sum += cold + if warm is not None: + warm_sum += warm + + warm_pct_of_cold = "" + if cold is not None and warm is not None and cold > 0: + warm_pct_of_cold = f"{100.0 * warm / cold:.2f}" + + out_rows.append( + { + "nvjitlink_optset": opt, + "scenario": scen, + "type_desc": type_desc.get((opt, scen), ""), + "cold_ms": f"{cold:.6f}" if cold is not None else "", + "warm_ms": f"{warm:.6f}" if warm is not None else "", + "warm_pct_of_cold": warm_pct_of_cold, + "cold_pct_vs_lto_baseline": pct_vs_baseline(opt, scen, 1, cold), + "warm_pct_vs_lto_baseline": pct_vs_baseline(opt, scen, 2, warm), + } + ) + + warm_pct_tot = "" + if cold_sum > 0 and warm_sum > 0: + warm_pct_tot = f"{100.0 * warm_sum / cold_sum:.2f}" + lto_cold_tot = sum( + ms["lto"][s][1] + for s in scen_order + if "lto" in ms and s in ms["lto"] and 1 in ms["lto"][s] + ) + lto_warm_tot = sum( + ms["lto"][s][2] + for s in scen_order + if "lto" in ms and s in ms["lto"] and 2 in ms["lto"][s] + ) + cold_pct_tot = "" + warm_pct_lto_tot = "" + if opt != "lto" and lto_cold_tot > 0 and cold_sum > 0: + cold_pct_tot = f"{100.0 * (cold_sum / lto_cold_tot - 1.0):.2f}" + if opt == "lto": + cold_pct_tot = "0" + warm_pct_lto_tot = "0" + elif lto_warm_tot > 0 and warm_sum > 0: + warm_pct_lto_tot = f"{100.0 * (warm_sum / lto_warm_tot - 1.0):.2f}" + + out_rows.append( + { + "nvjitlink_optset": opt, + "scenario": "TOTAL", + "type_desc": "", + "cold_ms": f"{cold_sum:.6f}" if cold_sum > 0 else "", + "warm_ms": f"{warm_sum:.6f}" if warm_sum > 0 else "", + "warm_pct_of_cold": warm_pct_tot, + "cold_pct_vs_lto_baseline": cold_pct_tot, + "warm_pct_vs_lto_baseline": warm_pct_lto_tot, + } + ) + + summary_path.parent.mkdir(parents=True, exist_ok=True) + with summary_path.open("w", newline="", encoding="utf-8") as f: + w = csv.DictWriter(f, fieldnames=_SUMMARY_CSV_COLUMNS) + w.writeheader() + w.writerows(out_rows) + + +def append_report_to_csv(path: Path, report: dict[str, Any]) -> None: + """Append one CSV row per link timing event; write header if *path* is missing or empty.""" + path.parent.mkdir(parents=True, exist_ok=True) + new_file = (not path.exists()) or path.stat().st_size == 0 + with path.open("a", newline="", encoding="utf-8") as f: + w = csv.DictWriter(f, fieldnames=_CSV_COLUMNS) + if new_file: + w.writeheader() + for p in report["passes"]: + ax = p["axes_interpretation"] + for ev in p["link_events"]: + w.writerow( + { + "nvjitlink_optset": report["nvjitlink_optset"], + "script_invocation": report["script_invocation"], + "process_pass": p["process_pass"], + "cuda_compute_cache_disk": ax["cuda_compute_cache_disk"], + "libcudf_jit_lto_in_process": ax["libcudf_jit_lto_in_process"], + "scenario": ev["type_letter"], + "type_desc": ev["type_desc"], + "build_ms": ev["build_ms"], + } + ) + + +def _print_human_report(report: dict[str, Any]) -> None: + opt = report["nvjitlink_optset"] + inv = report["script_invocation"] + print( + "AlgorithmPlanner::build() (Murmur JIT+LTO), CUDF_JIT_LTO_LINK_TIMING=1\n" + f"nvjitlink_optset={opt} script_invocation={inv}\n" + ) + for p in report["passes"]: + pn = p["process_pass"] + ax = p["axes_interpretation"] + evs = p["link_events"] + nexp = p["expected_scenarios"] + print( + f"--- process_pass={pn} " + f"(cuda_compute_cache_disk~{ax['cuda_compute_cache_disk']}, " + f"libcudf_jit_lto_in_process~{ax['libcudf_jit_lto_in_process']}) ---" + ) + if len(evs) != nexp: + print( + f" Expected {nexp} timing lines, got {len(evs)} " + "(fragment collapse, wrong build, or cache hit)." + ) + if not evs: + print(" (no nvJitLink build lines — expected for pass 2 if in-memory cache hit)\n") + continue + for row in evs: + letter = row["type_letter"] + desc = row["type_desc"] + ms = row["build_ms"] + print(f" Type {letter} ({desc}): {ms:.6f} ms") + print("") + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument( + "--summarize-bench-csv", + type=Path, + metavar="RAW.csv", + help="Read raw per-link benchmark CSV and exit after writing summary (see --summary-output)", + ) + p.add_argument( + "--summary-output", + type=Path, + metavar="PATH", + help="Summary CSV path (default: RAW stem + _summary.csv next to RAW)", + ) + p.add_argument( + "--script-invocation", + type=int, + default=1, + metavar="N", + help="1 = first python process after clearing ~/.nv/ComputeCache; 2 = second process (warm disk cache). Default: 1", + ) + p.add_argument( + "--process-passes", + type=int, + default=1, + metavar="N", + help="In-process repetitions of the full table sweep (default: 1).", + ) + p.add_argument( + "--nvjitlink-optset", + default="lto", + metavar="NAME", + help="Label for this run (e.g. matches CUDF_JIT_LTO_NVJITLINK_OPTIONS preset in the shell driver). Default: lto", + ) + p.add_argument( + "--output-json", + type=Path, + metavar="PATH", + help="Write full report as JSON to PATH", + ) + p.add_argument( + "--output-csv", + type=Path, + metavar="PATH", + help="Append link timing rows to PATH (creates file and header if missing or empty)", + ) + p.add_argument( + "-q", + "--quiet", + action="store_true", + help="Suppress human-readable report (use with --output-csv or --output-json)", + ) + args = p.parse_args(argv) + + if args.summarize_bench_csv is not None: + raw = args.summarize_bench_csv + if not raw.is_file(): + p.error(f"not a file: {raw}") + out = args.summary_output + if out is None: + out = raw.with_name(f"{raw.stem}_summary.csv") + summarize_bench_csv(raw, out) + print(f"Wrote summary CSV: {out}") + return 0 + + if args.process_passes < 1: + p.error("--process-passes must be >= 1") + + report = run_benchmark( + script_invocation=args.script_invocation, + process_passes=args.process_passes, + nvjitlink_optset=args.nvjitlink_optset, + ) + + if args.output_json: + args.output_json.write_text( + json.dumps(report, indent=2) + "\n", encoding="utf-8" + ) + + if args.output_csv: + append_report_to_csv(args.output_csv, report) + + if not args.quiet: + _print_human_report(report) + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())