diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 48a4a69cb575..b715d2203a1a 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -413,6 +413,9 @@ include(cmake/thirdparty/get_nanoarrow.cmake) # find thread_pool include(cmake/thirdparty/get_thread_pool.cmake) +# find xxhash +include(cmake/thirdparty/get_xxhash.cmake) + # find zstd include(cmake/thirdparty/get_zstd.cmake) @@ -422,6 +425,9 @@ add_subdirectory(librtcx) # JIT Embedding helper functions include(librtcx/embed.cmake) +# Pre-compiled fragment management helper functions +include(cmake/Modules/AddFragment.cmake) + # Workaround until https://github.com/rapidsai/rapids-cmake/issues/176 is resolved if(NOT BUILD_SHARED_LIBS) include("${rapids-cmake-dir}/export/find_package_file.cmake") @@ -444,39 +450,39 @@ if(NOT BUILD_SHARED_LIBS) ) endif() -add_embed(cudf_cuda_embed) +rtcx_add_embed(cudf_cuda_embed) -embed_includes( +rtcx_embed_includes( cudf_cuda_embed SOURCE_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/librtcx/libcxx DEST_DIRECTORY librtcx/libcxx INCLUDE_DIRECTORIES librtcx/libcxx ) -embed_includes( +rtcx_embed_includes( cudf_cuda_embed SOURCE_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/include/cudf DEST_DIRECTORY cudf/cpp/include/cudf INCLUDE_DIRECTORIES cudf/cpp/include ) -embed_includes( +rtcx_embed_includes( cudf_cuda_embed SOURCE_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/src/jit DEST_DIRECTORY cudf/cpp/src/jit INCLUDE_DIRECTORIES cudf/cpp/src ) -embed_includes( +rtcx_embed_includes( cudf_cuda_embed SOURCE_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/src/binaryop/jit DEST_DIRECTORY cudf/cpp/src/binaryop/jit INCLUDE_DIRECTORIES cudf/cpp/src ) -embed_includes( +rtcx_embed_includes( cudf_cuda_embed SOURCE_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/src/join/jit DEST_DIRECTORY cudf/cpp/src/join/jit INCLUDE_DIRECTORIES cudf/cpp/src ) -embed_includes( +rtcx_embed_includes( cudf_cuda_embed SOURCE_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/src/rolling DEST_DIRECTORY cudf/cpp/src/rolling INCLUDE_DIRECTORIES cudf/cpp/src ) -embed_includes( +rtcx_embed_includes( cudf_cuda_embed SOURCE_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/src/transform/jit DEST_DIRECTORY cudf/cpp/src/transform/jit INCLUDE_DIRECTORIES cudf/cpp/src ) @@ -486,13 +492,91 @@ get_target_property(LIBCUDACXX_RAW_INCLUDE_DIRS CCCL::libcudacxx INTERFACE_INCLU foreach(INC_DIR IN LISTS LIBCUDACXX_RAW_INCLUDE_DIRS) cmake_path(GET INC_DIR FILENAME INC_DIR_NAME) - embed_includes( + rtcx_embed_includes( cudf_cuda_embed SOURCE_DIRECTORY ${INC_DIR} DEST_DIRECTORY CCCL/libcudacxx/${INC_DIR_NAME} INCLUDE_DIRECTORIES CCCL/libcudacxx/${INC_DIR_NAME} ) endforeach() -embed(cudf_cuda_embed COMPRESSION zstd OUTPUT_DIRECTORY "${CUDF_GENERATED_INCLUDE_DIR}/rtcx_embed") +rtcx_embed( + cudf_cuda_embed COMPRESSION zstd OUTPUT_DIRECTORY "${CUDF_GENERATED_INCLUDE_DIR}/rtcx_embed" +) + +rtcx_add_embed(cudf_fragments) + +list(APPEND CUDF_PRECOMPILE_PHYSICAL_TYPES uint8_t uint16_t uint32_t uint64_t numeric::decimal32 + numeric::decimal64 numeric::decimal128 +) + +foreach(TYPE IN ITEMS ${CUDF_PRECOMPILE_PHYSICAL_TYPES}) + set(FRAGMENT_NAME transform_kernel) + get_property( + FILE_INDEX + TARGET cudf_fragments__embed_props + PROPERTY EMBED_FILE_INDEX + ) + set(VARIANT_NAME transform_kernel_${FILE_INDEX}) + set(INSTANCE + "cudf::jit::transform_kernel>, cudf::jit::type_list>>" + ) + add_fragment( + cudf_fragments + FRAGMENT + ${VARIANT_NAME} + SOURCE + src/transform/jit/kernel.cu + KERNEL_INSTANCE + ${INSTANCE} + UDF_TYPE + "int(${TYPE} *, ${TYPE})" + DEFINITIONS + CUDF_LTO_MODE + ARRAY_IDS + ${FRAGMENT_NAME}_FILE_INDEX + ${FRAGMENT_NAME}_INSTANCE + ARRAY_VALUES + ${FILE_INDEX} + "${INSTANCE}" + ) +endforeach() + +foreach(TYPE IN ITEMS ${CUDF_PRECOMPILE_PHYSICAL_TYPES}) + foreach(RHS_IS_SCALAR IN ITEMS "false" "true") + set(FRAGMENT_NAME transform_kernel) + get_property( + FILE_INDEX + TARGET cudf_fragments__embed_props + PROPERTY EMBED_FILE_INDEX + ) + set(VARIANT_NAME transform_kernel_${FILE_INDEX}) + set(INSTANCE + "cudf::jit::transform_kernel, cudf::jit::column_accessor<1ULL, cudf::column_device_view_core, ${TYPE}, ${RHS_IS_SCALAR}, 0>>, cudf::jit::type_list>>" + ) + add_fragment( + cudf_fragments + FRAGMENT + ${VARIANT_NAME} + SOURCE + src/transform/jit/kernel.cu + KERNEL_INSTANCE + ${INSTANCE} + UDF_TYPE + "int(${TYPE} *, ${TYPE}, ${TYPE})" + DEFINITIONS + CUDF_LTO_MODE + ARRAY_IDS + ${FRAGMENT_NAME}_FILE_INDEX + ${FRAGMENT_NAME}_INSTANCE + ARRAY_VALUES + ${FILE_INDEX} + "${INSTANCE}" + ) + endforeach() +endforeach() + +rtcx_embed( + cudf_fragments COMPRESSION none OUTPUT_DIRECTORY "${CUDF_GENERATED_INCLUDE_DIR}/rtcx_embed" +) # ################################################################################################## # * library targets ------------------------------------------------------------------------------- @@ -1058,9 +1142,10 @@ add_library( src/utilities/type_checks.cpp src/utilities/type_dispatcher.cpp ${cudf_cuda_embed_SOURCE_DIR}/cudf_cuda_embed.s + ${cudf_fragments_SOURCE_DIR}/cudf_fragments.s ) -add_dependencies(cudf cudf_cuda_embed) +add_dependencies(cudf cudf_cuda_embed cudf_fragments) set_property( SOURCE src/io/parquet/writer_impl.cu @@ -1130,7 +1215,9 @@ target_include_directories( "$" "$" "$" + "$" "$" + "$" INTERFACE "$" ) @@ -1172,8 +1259,15 @@ target_compile_definitions(cudf PRIVATE THRUST_FORCE_32_BIT_OFFSET_TYPE=1 CCCL_A target_link_libraries( cudf PUBLIC CCCL::CCCL $ - PRIVATE $ $ ZLIB::ZLIB - ${CUDF_nvcomp_TARGET} kvikio::kvikio ${CUDF_nanoarrow_TARGET} zstd rtcx::rtcx + PRIVATE $ + $ + ZLIB::ZLIB + ${CUDF_nvcomp_TARGET} + kvikio::kvikio + ${CUDF_nanoarrow_TARGET} + zstd + $ + rtcx::rtcx ) # When rmm is a static library being absorbed via whole-archive, strip nvtx3 from its public diff --git a/cpp/benchmarks/CMakeLists.txt b/cpp/benchmarks/CMakeLists.txt index e17fdab03f88..68d7aec9ab17 100644 --- a/cpp/benchmarks/CMakeLists.txt +++ b/cpp/benchmarks/CMakeLists.txt @@ -389,11 +389,24 @@ ConfigureNVBench(CSV_WRITER_NVBENCH io/csv/csv_writer.cpp) # * ast benchmark --------------------------------------------------------------------------------- ConfigureNVBench(AST_NVBENCH ast/polynomials.cpp ast/transform.cpp) +# ################################################################################################## +# * LTO Fragments ---------------------------------------------------------------------------- +rtcx_add_embed(cudf_benchmark_fragments) +add_fragment(cudf_benchmark_fragments FRAGMENT add_f32 SOURCE binaryop/fragments/add_f32.cu) +add_fragment(cudf_benchmark_fragments FRAGMENT mul_f32 SOURCE binaryop/fragments/mul_f32.cu) +rtcx_embed( + cudf_benchmark_fragments COMPRESSION none OUTPUT_DIRECTORY + "${CUDF_GENERATED_INCLUDE_DIR}/rtcx_embed" +) + # ################################################################################################## # * binaryop benchmark ---------------------------------------------------------------------------- ConfigureNVBench( BINARYOP_NVBENCH binaryop/binaryop.cpp binaryop/compiled_binaryop.cpp binaryop/polynomials.cpp + ${cudf_benchmark_fragments_SOURCE_DIR}/cudf_benchmark_fragments.s ) +target_include_directories(BINARYOP_NVBENCH PRIVATE ${cudf_benchmark_fragments_SOURCE_DIR}) +add_dependencies(BINARYOP_NVBENCH cudf_benchmark_fragments) # ################################################################################################## # * transform benchmark diff --git a/cpp/benchmarks/binaryop/compiled_binaryop.cpp b/cpp/benchmarks/binaryop/compiled_binaryop.cpp index 9f6ea706e44f..bc359f825730 100644 --- a/cpp/benchmarks/binaryop/compiled_binaryop.cpp +++ b/cpp/benchmarks/binaryop/compiled_binaryop.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -7,7 +7,9 @@ #include #include +#include +#include #include template @@ -72,6 +74,7 @@ BINARYOP_BENCHMARK_DEFINE(timestamp_s, duration_s, ADD, time BINARYOP_BENCHMARK_DEFINE(duration_s, duration_D, SUB, duration_ms); BINARYOP_BENCHMARK_DEFINE(int64_t, int64_t, SUB, int64_t); BINARYOP_BENCHMARK_DEFINE(float, float, MUL, int64_t); +BINARYOP_BENCHMARK_DEFINE(float, float, MUL, float); BINARYOP_BENCHMARK_DEFINE(duration_s, int64_t, MUL, duration_s); BINARYOP_BENCHMARK_DEFINE(int64_t, int64_t, DIV, int64_t); BINARYOP_BENCHMARK_DEFINE(duration_ms, int32_t, DIV, duration_ms); @@ -101,3 +104,115 @@ BINARYOP_BENCHMARK_DEFINE(duration_ms, duration_ns, NULL_EQUALS, bool BINARYOP_BENCHMARK_DEFINE(duration_ms, duration_ns, NULL_NOT_EQUALS, bool); BINARYOP_BENCHMARK_DEFINE(decimal32, decimal32, NULL_MAX, decimal32); BINARYOP_BENCHMARK_DEFINE(timestamp_D, timestamp_s, NULL_MIN, timestamp_s); +// clang-format on + +template +void BM_jit_binaryop(nvbench::state& state, cudf::binary_operator binop) +{ + constexpr auto const jit_mul_cuda = R"***( +__device__ void transform(float* out, float a, float b) { + *out = a * b; +} +)***"; + + constexpr auto const jit_add_cuda = R"***( +__device__ void transform(float* out, float a, float b) { + *out = a + b; +} +)***"; + + auto const num_rows = static_cast(state.get_int64("num_rows")); + auto const use_lto = state.get_string("use_lto") == "true"; + static_assert(std::is_same_v && std::is_same_v); + static_assert(std::is_same_v); + + auto const source_table = create_random_table( + {cudf::type_to_id(), cudf::type_to_id()}, row_count{num_rows}); + + auto lhs = cudf::column_view(source_table->get_column(0)); + auto rhs = cudf::column_view(source_table->get_column(1)); + + size_t fragment_id = 0; + char const* cuda = nullptr; + + switch (binop) { + case cudf::binary_operator::ADD: { + fragment_id = cudf_benchmark_fragments::add_f32; + cuda = jit_add_cuda; + } break; + case cudf::binary_operator::MUL: { + fragment_id = cudf_benchmark_fragments::mul_f32; + cuda = jit_mul_cuda; + } break; + default: throw std::runtime_error("Unsupported binary operator for JIT benchmark"); + } + + // Call once for hot cache. + cudf::transform_input inputs[] = {lhs, rhs}; + cudf::transform_output outputs[] = { + {cudf::data_type{cudf::type_to_id()}, cudf::output_nullability::ALL_VALID}}; + + auto const range = cudf_benchmark_fragments::file_ranges[fragment_id]; + std::span udf{cudf_benchmark_fragments::files.subspan(range[0], range[1])}; + + auto result = use_lto ? cudf::transform_lto(udf, + cudf::lto_binary_type::FATBIN, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + {}, + std::nullopt) + : cudf::multi_transform(cuda, + cudf::udf_source_type::CUDA, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + {}, + std::nullopt); + + // use number of bytes read and written to global memory + state.add_global_memory_reads(num_rows); + state.add_global_memory_reads(num_rows); + state.add_global_memory_writes(num_rows); + + state.exec(nvbench::exec_tag::sync, [&](nvbench::launch&) { + [[maybe_unused]] auto result = use_lto ? cudf::transform_lto(udf, + cudf::lto_binary_type::FATBIN, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + {}, + std::nullopt) + : cudf::multi_transform(cuda, + cudf::udf_source_type::CUDA, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + {}, + std::nullopt); + }); +} + +#define BM_JIT_BINARYOP_BENCHMARK_DEFINE(name, lhs, rhs, bop, tout) \ + static void name(::nvbench::state& st) \ + { \ + ::BM_jit_binaryop(st, ::cudf::binary_operator::bop); \ + } \ + NVBENCH_BENCH(name) \ + .set_name("jit_binary_op_" BM_STRINGIFY(name)) \ + .add_int64_axis("num_rows", {10'000, 100'000, 1'000'000, 10'000'000, 100'000'000}) \ + .add_string_axis("use_lto", {"true", "false"}) + +#define build_name_jit(a, b, c, d) a##_##b##_##c##_##d##_jit + +#define JIT_BINARYOP_BENCHMARK_DEFINE(lhs, rhs, bop, tout) \ + BM_JIT_BINARYOP_BENCHMARK_DEFINE(build_name_jit(bop, lhs, rhs, tout), lhs, rhs, bop, tout) + +// clang-format off +JIT_BINARYOP_BENCHMARK_DEFINE(float, float, ADD, float); +JIT_BINARYOP_BENCHMARK_DEFINE(float, float, MUL, float); +// clang-format on diff --git a/cpp/benchmarks/binaryop/fragments/add_f32.cu b/cpp/benchmarks/binaryop/fragments/add_f32.cu new file mode 100644 index 000000000000..b2bee38e7df2 --- /dev/null +++ b/cpp/benchmarks/binaryop/fragments/add_f32.cu @@ -0,0 +1,10 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +extern "C" __device__ int transform(float* out, float a, float b) +{ + *out = a + b; + return 0; +} diff --git a/cpp/benchmarks/binaryop/fragments/mul_f32.cu b/cpp/benchmarks/binaryop/fragments/mul_f32.cu new file mode 100644 index 000000000000..a6f4edc7c303 --- /dev/null +++ b/cpp/benchmarks/binaryop/fragments/mul_f32.cu @@ -0,0 +1,10 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +extern "C" __device__ int transform(float* out, float a, float b) +{ + *out = a * b; + return 0; +} diff --git a/cpp/cmake/Modules/AddFragment.cmake b/cpp/cmake/Modules/AddFragment.cmake new file mode 100644 index 000000000000..187d9c36d594 --- /dev/null +++ b/cpp/cmake/Modules/AddFragment.cmake @@ -0,0 +1,87 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +include_guard(GLOBAL) + +# This macro is used to create object libraries for JIT compilation fragments, and embed them as +# fatbins in the final library. It compiles the specified source file with the appropriate flags to +# generate a fatbin containing the specified kernel instance, and then embeds that fatbin in the +# final library with metadata that allows it to be looked up at runtime. +macro(add_fragment) + set(TARGET ${ARGV0}) + set(ONE_VALUE_ARGS FRAGMENT SOURCE KERNEL_ONLY KERNEL_INSTANCE UDF_TYPE) + set(MULTI_VALUE_ARGS DEFINITIONS ARRAY_IDS ARRAY_VALUES) + cmake_parse_arguments(ARG "${OPTIONS}" "${ONE_VALUE_ARGS}" "${MULTI_VALUE_ARGS}" ${ARGN}) + + if(NOT ARG_FRAGMENT) + message(FATAL_ERROR "add_fragment requires FRAGMENT argument") + endif() + + if(NOT ARG_SOURCE) + message(FATAL_ERROR "add_fragment requires SOURCE argument") + endif() + + set(OBJECT_ID ${TARGET}_${ARG_FRAGMENT}) + add_library(${OBJECT_ID} OBJECT ${ARG_SOURCE}) + target_compile_options(${OBJECT_ID} PRIVATE --compress-mode=size) + + if(DEFINED ARG_KERNEL_ONLY AND ARG_KERNEL_ONLY) + # ensure that the FATBIN symbols only contain the specified kernel + target_compile_options(${OBJECT_ID} PRIVATE -Xnvlink=--kernels-used=cudf_kernel_entry) + endif() + + set(INSTANTIATION_DIR "${CUDF_GENERATED_INCLUDE_DIR}/${TARGET}/instantiations/${ARG_FRAGMENT}") + target_include_directories(${OBJECT_ID} PRIVATE ${INSTANTIATION_DIR}) + + if(ARG_KERNEL_INSTANCE) + file( + GENERATE + OUTPUT "${INSTANTIATION_DIR}/cudf/detail/kernel_instance.cuh" + CONTENT "#pragma once\n#define CUDF_KERNEL_INSTANCE ${ARG_KERNEL_INSTANCE}" + ) + endif() + + if(ARG_UDF_TYPE) + file( + GENERATE + OUTPUT "${INSTANTIATION_DIR}/cudf/detail/operation_udf.cuh" + CONTENT "#pragma once\n#define CUDF_UDF_TYPE ${ARG_UDF_TYPE}" + ) + endif() + + target_compile_definitions(${OBJECT_ID} PRIVATE CUDF_DISABLE_EXPORTS ${ARG_DEFINITIONS}) + set_target_properties( + ${OBJECT_ID} + PROPERTIES CUDA_SEPARABLE_COMPILATION ON + CUDA_FATBIN_COMPILATION ON + POSITION_INDEPENDENT_CODE ON + INTERPROCEDURAL_OPTIMIZATION ON + CXX_STANDARD 20 + CXX_STANDARD_REQUIRED ON + CXX_EXTENSIONS ON + CXX_VISIBILITY_PRESET hidden + CUDA_STANDARD 20 + CUDA_STANDARD_REQUIRED ON + CUDA_VISIBILITY_PRESET hidden + ) + target_link_libraries( + ${OBJECT_ID} + PUBLIC CCCL::CCCL rapids_logger::rapids_logger rmm::rmm $ + PRIVATE $ $ + ZLIB::ZLIB nvcomp::nvcomp kvikio::kvikio nanoarrow::nanoarrow zstd + ) + target_include_directories( + ${OBJECT_ID} PRIVATE "$" + "$" + ) + target_compile_options(${OBJECT_ID} PRIVATE "$<$:${CUDF_CUDA_FLAGS}>") + + rtcx_embed_blob( + ${TARGET} FILE $ DEST fragments/${ARG_FRAGMENT}.fatbin ID + ${ARG_FRAGMENT} ARRAY_IDS ${ARG_ARRAY_IDS} ARRAY_VALUES ${ARG_ARRAY_VALUES} + ) +endmacro() diff --git a/cpp/cmake/thirdparty/get_xxhash.cmake b/cpp/cmake/thirdparty/get_xxhash.cmake new file mode 100644 index 000000000000..046d95b23f2a --- /dev/null +++ b/cpp/cmake/thirdparty/get_xxhash.cmake @@ -0,0 +1,48 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +# Use CPM to find or clone xxHash +function(find_and_configure_xxhash) + + include(${rapids-cmake-dir}/find/generate_module.cmake) + rapids_find_generate_module(xxhash HEADER_NAMES xxhash.h xxh3.h) + + set(CPM_DOWNLOAD_xxhash ON) + rapids_cpm_find( + xxhash 0.8.3 + GLOBAL_TARGETS xxhash + CPM_ARGS + GIT_REPOSITORY https://github.com/Cyan4973/xxHash.git + GIT_TAG e626a72bc2321cd320e953a0ccf1584cad60f363 + GIT_SHALLOW TRUE + DOWNLOAD_ONLY TRUE + EXCLUDE_FROM_ALL ${CUDF_EXCLUDE_DEPS_FROM_ALL} + ) + + if(xxhash_ADDED AND NOT TARGET xxhash) + add_library(xxhash INTERFACE) + target_include_directories(xxhash INTERFACE "${xxhash_SOURCE_DIR}") + endif() + + if(DEFINED xxhash_SOURCE_DIR) + set(XXHASH_INCLUDE_DIR + "${xxhash_SOURCE_DIR}" + PARENT_SCOPE + ) + set(xxhash_SOURCE_DIR + "${xxhash_SOURCE_DIR}" + PARENT_SCOPE + ) + endif() + + if(DEFINED xxhash_SOURCE_DIR) + include("${rapids-cmake-dir}/export/find_package_root.cmake") + rapids_export_find_package_root(BUILD xxhash "${xxhash_SOURCE_DIR}" EXPORT_SET cudf-exports) + endif() +endfunction() + +find_and_configure_xxhash() diff --git a/cpp/include/cudf/transform.hpp b/cpp/include/cudf/transform.hpp index 574db028a435..b5519246a6ad 100644 --- a/cpp/include/cudf/transform.hpp +++ b/cpp/include/cudf/transform.hpp @@ -187,6 +187,14 @@ std::unique_ptr multi_transform( rmm::cuda_stream_view stream = cudf::get_default_stream(), rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); +/** + * @brief The type of LTO Binary + */ +enum class lto_binary_type : uint8_t { + LTO_IR, //< LTO-IR binary + FATBIN //< FATBIN binary +}; + /** * @brief Creates a new table by applying a transform function against every * element of the input columns. @@ -208,12 +216,12 @@ std::unique_ptr
multi_transform( * The size of the resulting column is the `row_size` if provided, otherwise it is inferred from * the input and pre-allocated output columns. * - * @param udf The PTX/CUDA string of the transform function to apply - * @param source_type The source type of the UDF (CUDA or PTX) + * @param udf The LTO-IR fragment containing the transform function to apply. The UDF must + * be named `transform` and follow the CUDF UDF ABI + * @param binary_type The type of the LTO binary provided in `udf` * @param is_null_aware Signifies the UDF will receive row inputs as optional values - * @param is_fallible Signifies the UDF may produce errors during execution * @param user_data User-defined device data to pass to the UDF - * @param inputs Immutable views of the inputs to transform (columns and scalar columns) + * @param inputs Immutable views of the inputs to transform * @param outputs Specification of the output columns to be created * @param string_offsets For string output columns, the offsets can be pre-allocated and passed in * to prevent overhead of compacting string views into run-end strings column. @@ -225,11 +233,10 @@ std::unique_ptr
multi_transform( * function to every element of the input according to the output specifications * */ -std::unique_ptr
multi_transform( - std::string const& udf, - udf_source_type source_type, +std::unique_ptr
transform_lto( + std::span udf, + lto_binary_type binary_type, null_aware is_null_aware, - fallible is_fallible, std::optional user_data, std::span inputs, std::span outputs, diff --git a/cpp/include/cudf/types.hpp b/cpp/include/cudf/types.hpp index 188424299a64..414f4c554473 100644 --- a/cpp/include/cudf/types.hpp +++ b/cpp/include/cudf/types.hpp @@ -261,14 +261,6 @@ enum class output_nullability : uint8_t { ///< null values are produced }; -/** - * @brief Indicates whether a function can produce an error during execution. - */ -enum class fallible : bool { - NO = 0, ///< The function is not fallible - YES = 1 ///< The function is fallible -}; - /** * @brief Indicates whether a function nullifies its output on error. * diff --git a/cpp/librtcx/CMakeLists.txt b/cpp/librtcx/CMakeLists.txt index 47662051f7f1..a40bb36d2914 100644 --- a/cpp/librtcx/CMakeLists.txt +++ b/cpp/librtcx/CMakeLists.txt @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= @@ -60,7 +60,7 @@ if(NOT TARGET nvtx3::nvtx3-cpp) rapids_cpm_nvtx3() endif() -add_library(rtcx STATIC rtcx.cpp) +add_library(rtcx STATIC hash.cpp rtcx.cpp) add_library(rtcx::rtcx ALIAS rtcx) set_target_properties( diff --git a/cpp/librtcx/README.md b/cpp/librtcx/README.md index 06ec01b01f82..b8f9898c0d13 100644 --- a/cpp/librtcx/README.md +++ b/cpp/librtcx/README.md @@ -12,7 +12,8 @@ RTCX (runtime-compiler extended) is a wrapper around NVRTC and nvJitLink designe ## Build-time Requirements - CMake >= 4.0 -- LibZSTD - for binary compression +- libzstd +- xxHash - CUDA >= 12.2 # Dependencies diff --git a/cpp/librtcx/embed.cmake b/cpp/librtcx/embed.cmake index b4510b6d1c64..dd1e83190835 100644 --- a/cpp/librtcx/embed.cmake +++ b/cpp/librtcx/embed.cmake @@ -1,19 +1,23 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= if(NOT TARGET zstd) - message(FATAL_ERROR "embed(): zstd target is required for LIBRTCX embedding.") + message(FATAL_ERROR "zstd target is required for LIBRTCX embedding.") +endif() + +if(NOT TARGET xxhash) + message(FATAL_ERROR "xxhash target is required for LIBRTCX embedding.") endif() # This function initializes a target for JIT embedding. It must be called before any calls to # embed_includes() or embed_blob() for the target. It creates a dedicated INTERFACE library target # that is used to track registered files and dependencies via target properties. The TARGET argument # specifies the name of the target being initialized. -function(add_embed TARGET) +function(rtcx_add_embed TARGET) set(OPTIONS "") set(ONE_VALUE_ARGS) set(MULTI_VALUE_ARGS) @@ -24,10 +28,11 @@ function(add_embed TARGET) endif() add_library(${TARGET}__embed_props INTERFACE) + set_property(TARGET ${TARGET}__embed_props PROPERTY EMBED_FILE_INDEX 0) endfunction() # This function registers a directory of include files to be embedded for JIT compilation. -function(embed_includes TARGET) +function(rtcx_embed_includes TARGET) set(OPTIONS "") set(ONE_VALUE_ARGS SOURCE_DIRECTORY # Source directory where files will be copied from DEST_DIRECTORY # Destination directory where files will be copied to @@ -106,17 +111,26 @@ function(embed_includes TARGET) PROPERTY EMBED_INCLUDE_DIRECTORIES ${ARG_INCLUDE_DIRECTORIES} ) + get_property( + SOURCE_FILE_IDS + TARGET ${TARGET}__embed_props + PROPERTY EMBED_SOURCE_FILE_IDS + ) + list(LENGTH SOURCE_FILE_IDS IDX) + + set_property(TARGET ${TARGET}__embed_props PROPERTY EMBED_FILE_INDEX ${IDX}) + endfunction() # This function registers a single file to be embedded for JIT compilation. -function(embed_blob TARGET) +function(rtcx_embed_blob TARGET) set(OPTIONS) set(ONE_VALUE_ARGS ID FILE DEST) set(MULTI_VALUE_ARGS ARRAY_IDS ARRAY_VALUES) cmake_parse_arguments(ARG "${OPTIONS}" "${ONE_VALUE_ARGS}" "${MULTI_VALUE_ARGS}" ${ARGN}) if(NOT TARGET ${TARGET}__embed_props) - message(FATAL_ERROR "embed target '${TARGET}' has not been initialized with add_embed()") + message(FATAL_ERROR "embed target '${TARGET}' has not been initialized with rtcx_add_embed()") endif() if(NOT ARG_ID @@ -179,6 +193,15 @@ function(embed_blob TARGET) PROPERTY EMBED_SOURCE_FILE_DESTS ${ARG_DEST} ) + get_property( + SOURCE_FILE_IDS + TARGET ${TARGET}__embed_props + PROPERTY EMBED_SOURCE_FILE_IDS + ) + list(LENGTH SOURCE_FILE_IDS IDX) + + set_property(TARGET ${TARGET}__embed_props PROPERTY EMBED_FILE_INDEX ${IDX}) + endfunction() #[==[ @@ -186,14 +209,14 @@ endfunction() # for JIT compilation. #]==] # cmake-lint: disable=R0915 -function(embed TARGET) +function(rtcx_embed TARGET) set(OPTIONS "") set(ONE_VALUE_ARGS "COMPRESSION" "OUTPUT_DIRECTORY") set(MULTI_VALUE_ARGS "") cmake_parse_arguments(ARG "${OPTIONS}" "${ONE_VALUE_ARGS}" "${MULTI_VALUE_ARGS}" ${ARGN}) if(NOT TARGET ${TARGET}__embed_props) - message(FATAL_ERROR "embed target '${TARGET}' has not been initialized with add_embed()") + message(FATAL_ERROR "embed target '${TARGET}' has not been initialized with rtcx_add_embed()") endif() if(NOT DEFINED ARG_COMPRESSION) @@ -271,8 +294,10 @@ function(embed TARGET) ) set(RUNNER "${TARGET}__jit_embed_run") - add_executable(${RUNNER} EXCLUDE_FROM_ALL "${EMBED_SCRIPT}") - target_link_libraries(${RUNNER} PRIVATE ${CMAKE_DL_LIBS} zstd) + add_executable( + ${RUNNER} EXCLUDE_FROM_ALL "${EMBED_SCRIPT}" ${CMAKE_CURRENT_FUNCTION_LIST_DIR}/hash.cpp + ) + target_link_libraries(${RUNNER} PRIVATE ${CMAKE_DL_LIBS} xxhash zstd) target_include_directories( ${RUNNER} PRIVATE ${CMAKE_CURRENT_FUNCTION_LIST_DIR} ${ZSTD_INCLUDE_DIR} ) diff --git a/cpp/librtcx/embed.hpp b/cpp/librtcx/embed.hpp index a71b7e1bb240..5df1aa4db40b 100644 --- a/cpp/librtcx/embed.hpp +++ b/cpp/librtcx/embed.hpp @@ -4,11 +4,14 @@ */ #pragma once -#include "sha256.hpp" +#include "hash.hpp" #include #include #include + +#define XXH_INLINE_ALL +#include #include #include @@ -99,17 +102,20 @@ std::vector compress_bytes(std::span bytes, std::string_ return compressed; } -rtcx::sha256 compute_embed_hash(std::span uncompressed_files_bytes, - std::span merged_dests_bytes, - std::span merged_include_dirs_bytes, - std::string_view compression) +rtcx::hash128 compute_embed_hash(std::span uncompressed_files_bytes, + std::span merged_dests_bytes, + std::span merged_include_dirs_bytes, + std::string_view compression) { - rtcx::sha256_context ctx; - ctx.update(uncompressed_files_bytes); - ctx.update(merged_dests_bytes); - ctx.update(merged_include_dirs_bytes); - ctx.update(std::span{reinterpret_cast(compression.data()), compression.size()}); - return ctx.finalize(); + XXH3_state_t state; + XXH3_INITSTATE(&state); + XXH3_128bits_reset(&state); + XXH3_128bits_update(&state, uncompressed_files_bytes.data(), uncompressed_files_bytes.size()); + XXH3_128bits_update(&state, merged_dests_bytes.data(), merged_dests_bytes.size()); + XXH3_128bits_update(&state, merged_include_dirs_bytes.data(), merged_include_dirs_bytes.size()); + XXH3_128bits_update(&state, compression.data(), compression.size()); + auto hash = XXH3_128bits_digest(&state); + return rtcx::hash128{hash.high64, hash.low64}; } template diff --git a/cpp/librtcx/hash.cpp b/cpp/librtcx/hash.cpp new file mode 100644 index 000000000000..fd32d55b2269 --- /dev/null +++ b/cpp/librtcx/hash.cpp @@ -0,0 +1,70 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "hash.hpp" + +#include +#include +#include +#include +#include + +namespace rtcx { + +char const* hash128_hex_string::data() const { return data_; } + +char const* hash128_hex_string::c_str() const { return data_; } + +hash128_hex_string hash128_hex_string::make(std::span input) +{ + constexpr char const HEX_CHARS[] = "0123456789abcdef"; // NOLINT(modernize-avoid-c-arrays) + hash128_hex_string hex; + for (std::size_t i = 0; i < NUM_HEX_BYTES; ++i) { + hex.data_[i * 2] = HEX_CHARS[(input[i] >> 4) & 0x0F]; + hex.data_[i * 2 + 1] = HEX_CHARS[input[i] & 0x0F]; + } + hex.data_[NUM_HEX_DIGITS] = '\0'; + return hex; +} + +hash128_hex_string hash128_hex_string::make(__uint128_t hash) +{ + auto array = std::bit_cast>(hash); + return make(array); +} + +std::uint8_t hash128::operator[](std::size_t index) const +{ + return reinterpret_cast(&value)[15 - index]; +} + +std::size_t hash128::size() const { return 16; } + +std::uint8_t const* hash128::data() const { return reinterpret_cast(&value); } + +hash128_hex_string hash128::to_hex_string() const { return hash128_hex_string::make(value); } + +hash128 hash128::parse(std::string_view hex) +{ + if (hex.size() != hash128_hex_string::NUM_HEX_DIGITS) { + throw std::invalid_argument( + std::format("Invalid hash128 hex string length, expected {} got {} (hash: `{}`)", + hash128_hex_string::NUM_HEX_DIGITS, + hex.size(), + hex)); + } + std::array data{}; + for (std::size_t i = 0; i < hash128_hex_string::NUM_HEX_BYTES; ++i) { + auto hex_byte = hex.substr(i * 2, 2); + auto [ptr, ec] = std::from_chars(hex_byte.begin(), hex_byte.end(), data[i], 16); + if (ec != std::errc()) { + throw std::invalid_argument( + std::format("Invalid hex character {} in HEX string: `{}`", hex_byte, hex)); + } + } + return hash128{std::bit_cast<__uint128_t>(data)}; +} + +} // namespace rtcx diff --git a/cpp/librtcx/hash.hpp b/cpp/librtcx/hash.hpp new file mode 100644 index 000000000000..d1aa12c85c61 --- /dev/null +++ b/cpp/librtcx/hash.hpp @@ -0,0 +1,62 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include + +namespace rtcx { + +struct [[nodiscard]] hash128_hex_string { + static constexpr std::size_t NUM_HEX_DIGITS = 32; + static constexpr std::size_t NUM_HEX_BYTES = NUM_HEX_DIGITS / 2; + + char data_[NUM_HEX_DIGITS + 1]; // NOLINT(modernize-avoid-c-arrays) + + [[nodiscard]] constexpr std::string_view view() const + { + return std::string_view{data_, NUM_HEX_DIGITS}; + } + + [[nodiscard]] constexpr operator std::string_view() const { return view(); } + + [[nodiscard]] char const* data() const; + + [[nodiscard]] char const* c_str() const; + + [[nodiscard]] static constexpr std::size_t size() { return NUM_HEX_DIGITS; } + + static hash128_hex_string make(std::span input); + + static hash128_hex_string make(__uint128_t hash); +}; + +struct hash128 { + __uint128_t value; + + constexpr hash128(__uint128_t v = 0) : value(v) {} + + constexpr hash128(std::uint64_t high, std::uint64_t low) + : value((static_cast<__uint128_t>(high) << 64) | low) + { + } + + [[nodiscard]] constexpr bool operator==(hash128 const&) const = default; + + [[nodiscard]] std::uint8_t operator[](std::size_t index) const; + + [[nodiscard]] std::size_t size() const; + + [[nodiscard]] std::uint8_t const* data() const; + + hash128_hex_string to_hex_string() const; + + static hash128 parse(std::string_view hex); +}; + +} // namespace rtcx diff --git a/cpp/librtcx/rtcx.cpp b/cpp/librtcx/rtcx.cpp index 199f2b563eab..1c3144869a8b 100644 --- a/cpp/librtcx/rtcx.cpp +++ b/cpp/librtcx/rtcx.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -183,6 +183,7 @@ void log_error(std::string_view msg) DO_IT(LibraryUnload) #define FOR_EACH_NVRTC_FUNC(DO_IT) \ + DO_IT(Version) \ DO_IT(GetErrorString) \ DO_IT(CreateProgram) \ DO_IT(DestroyProgram) \ @@ -199,6 +200,7 @@ void log_error(std::string_view msg) DO_IT(GetLoweredName) #define FOR_EACH_NVJITLINK_FUNC(DO_IT) \ + DO_IT(Version) \ DO_IT(Create) \ DO_IT(Destroy) \ DO_IT(AddData) \ @@ -658,6 +660,24 @@ void log_nvJitLink_result(link_params const& params, } // namespace +std::int32_t nvrtc_version() +{ + RTCX_FUNC_RANGE(); + + std::int32_t major, minor; + RTCX_CHECK_NVRTC(nvrtc->Version(&major, &minor)); + return major * 1000 + minor * 10; +} + +std::int32_t nvjitlink_version() +{ + RTCX_FUNC_RANGE(); + + std::uint32_t major, minor; + RTCX_CHECK_NVJITLINK(nvjitlink->Version(&major, &minor)); + return static_cast(major * 1000 + minor * 10); +} + byte_buffer compile(compile_params const& params) { RTCX_FUNC_RANGE(); @@ -979,12 +999,14 @@ std::optional blob_t::from_file(char const* path) namespace { -/// @brief retrieves a blob from disk based on the given sha256 hash and object type (e.g. "blob", +/// @brief retrieves a blob from disk based on the given hash and object type (e.g. "blob", /// "cuLibrary"). Returns nullopt if the file doesn't exist on disk, and throws if any other error /// occurs. -std::optional get_disk_blob(std::string const& cache_dir, object_type type, sha256 const& sha) +std::optional get_disk_blob(std::string const& cache_dir, + object_type type, + hash128 const& hash) { - auto hex = sha.to_hex_string(); + auto hex = hash.to_hex_string(); auto path = std::format("{}/{}.{}.bin", cache_dir, hex.view(), object_tag(type)); auto blob = blob_t::from_file(path.c_str()); @@ -992,11 +1014,11 @@ std::optional get_disk_blob(std::string const& cache_dir, object_type type return std::make_shared(std::move(*blob)); } -std::optional get_disk_library(std::string const& cache_dir, sha256 const& sha) +std::optional get_disk_library(std::string const& cache_dir, hash128 const& hash) { RTCX_FUNC_RANGE(); - auto hex = sha.to_hex_string(); + auto hex = hash.to_hex_string(); auto path = std::format("{}/{}.{}.bin", cache_dir, hex.view(), object_tag(object_type::LIBRARY)); // WAR: avoid a driver API call when the cache file is not present @@ -1033,7 +1055,7 @@ std::vector get_disk_entries(std::string const& cache_dir) void cache_blob_to_disk(std::string const& cache_dir, std::string const& tmp_dir, object_type type, - sha256 const& sha, + hash128 const& hash, std::span binary) { RTCX_FUNC_RANGE(); @@ -1061,7 +1083,7 @@ void cache_blob_to_disk(std::string const& cache_dir, } } - auto hex = sha.to_hex_string(); + auto hex = hash.to_hex_string(); auto final_path = std::format("{}/{}.{}.bin", cache_dir, hex.view(), object_tag(type)); std::filesystem::create_directories(std::filesystem::path{final_path}.parent_path()); @@ -1084,7 +1106,7 @@ void cache_blob_to_disk(std::string const& cache_dir, } // namespace -std::shared_future cache_t::get_or_add_blob(sha256 const& sha, blob_compile_func compile) +std::shared_future cache_t::get_or_add_blob(hash128 const& hash, blob_compile_func compile) { RTCX_FUNC_RANGE(); @@ -1094,7 +1116,7 @@ std::shared_future cache_t::get_or_add_blob(sha256 const& sha, blob_compil std::unique_lock lock{lock_}; // check memory cache - if (auto it = enabled_ ? blobs_cache_.entries_.find(sha) : blobs_cache_.entries_.end(); + if (auto it = enabled_ ? blobs_cache_.entries_.find(hash) : blobs_cache_.entries_.end(); it != blobs_cache_.entries_.end()) { counter_.blob_mem_hits.incr(); @@ -1108,7 +1130,7 @@ std::shared_future cache_t::get_or_add_blob(sha256 const& sha, blob_compil // check disk cache std::optional disk_blob = std::nullopt; - if (enabled_) { disk_blob = get_disk_blob(cache_dir_, object_type::BLOB, sha); } + if (enabled_) { disk_blob = get_disk_blob(cache_dir_, object_type::BLOB, hash); } std::promise promise; auto fut = promise.get_future().share(); @@ -1121,30 +1143,30 @@ std::shared_future cache_t::get_or_add_blob(sha256 const& sha, blob_compil promise.set_value(std::move(*disk_blob)); // insert into cache - blobs_cache_.insert(sha, std::move(cache_fut), current_tick); + blobs_cache_.insert(hash, std::move(cache_fut), current_tick); return ret_fut; } else { counter_.blob_disk_misses.incr(); - blobs_cache_.insert(sha, std::move(cache_fut), current_tick); + blobs_cache_.insert(hash, std::move(cache_fut), current_tick); // we can release the lock while calling the maker function since it may be expensive and we - // have already reserved a spot in the cache for this sha + // have already reserved a spot in the cache for this hash lock.unlock(); auto result = compile(); promise.set_value(result); - cache_blob_to_disk(cache_dir_, tmp_dir_, object_type::BLOB, sha, result->view()); + cache_blob_to_disk(cache_dir_, tmp_dir_, object_type::BLOB, hash, result->view()); return ret_fut; } } } -std::shared_future cache_t::get_or_add_library(sha256 const& sha, +std::shared_future cache_t::get_or_add_library(hash128 const& hash, library_compile_func compile) { RTCX_FUNC_RANGE(); @@ -1155,7 +1177,7 @@ std::shared_future cache_t::get_or_add_library(sha256 const& sha, std::unique_lock lock{lock_}; // check memory cache - if (auto it = enabled_ ? libraries_cache_.entries_.find(sha) : libraries_cache_.entries_.end(); + if (auto it = enabled_ ? libraries_cache_.entries_.find(hash) : libraries_cache_.entries_.end(); it != libraries_cache_.entries_.end()) { counter_.library_mem_hits.incr(); @@ -1169,7 +1191,7 @@ std::shared_future cache_t::get_or_add_library(sha256 const& sha, // check disk cache std::optional disk_library = std::nullopt; - if (enabled_) { disk_library = get_disk_library(cache_dir_, sha); } + if (enabled_) { disk_library = get_disk_library(cache_dir_, hash); } std::promise promise; auto fut = promise.get_future().share(); @@ -1179,10 +1201,10 @@ std::shared_future cache_t::get_or_add_library(sha256 const& sha, if (disk_library.has_value()) { counter_.library_disk_hits.incr(); - libraries_cache_.insert(sha, std::move(cache_fut), current_tick); + libraries_cache_.insert(hash, std::move(cache_fut), current_tick); // we can release the lock while calling the maker function since it may be expensive and we - // have already reserved a spot in the cache for this sha + // have already reserved a spot in the cache for this hash lock.unlock(); promise.set_value(std::move(*disk_library)); @@ -1192,17 +1214,17 @@ std::shared_future cache_t::get_or_add_library(sha256 const& sha, } else { counter_.library_disk_misses.incr(); - libraries_cache_.insert(sha, std::move(cache_fut), current_tick); + libraries_cache_.insert(hash, std::move(cache_fut), current_tick); // we can release the lock while calling the maker function since it may be expensive and we - // have already reserved a spot in the cache for this sha + // have already reserved a spot in the cache for this hash lock.unlock(); auto [library, blob] = compile(); promise.set_value(library); // store result to disk - cache_blob_to_disk(cache_dir_, tmp_dir_, object_type::LIBRARY, sha, blob->view()); + cache_blob_to_disk(cache_dir_, tmp_dir_, object_type::LIBRARY, hash, blob->view()); return ret_fut; } @@ -1290,8 +1312,8 @@ void cache_t::preload_from_disk() for (auto const& path : entries) { try { auto file_name = std::filesystem::path{path}.filename().string(); - auto sha_str = file_name.substr(0, file_name.find('.')); - auto sha = sha256::parse(sha_str); + auto hash_str = file_name.substr(0, file_name.find('.')); + auto hash = hash128::parse(hash_str); if (path.ends_with(".blob.bin")) { auto data = blob_t::from_file(path.c_str()); @@ -1300,14 +1322,14 @@ void cache_t::preload_from_disk() std::promise promise; auto fut = promise.get_future().share(); promise.set_value(std::move(blob)); - blobs_cache_.insert(sha, std::move(fut), tick_); + blobs_cache_.insert(hash, std::move(fut), tick_); } else if (path.ends_with(".cuLibrary.bin")) { - auto lib = get_disk_library(cache_dir_, sha); + auto lib = get_disk_library(cache_dir_, hash); if (!lib.has_value()) { continue; } std::promise promise; auto fut = promise.get_future().share(); promise.set_value(std::move(*lib)); - libraries_cache_.insert(sha, std::move(fut), tick_); + libraries_cache_.insert(hash, std::move(fut), tick_); } } catch (std::exception const& e) { // ignore any errors during preload diff --git a/cpp/librtcx/rtcx.hpp b/cpp/librtcx/rtcx.hpp index bc3ff8fd5569..3a9929c20301 100644 --- a/cpp/librtcx/rtcx.hpp +++ b/cpp/librtcx/rtcx.hpp @@ -1,11 +1,11 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once -#include "sha256.hpp" +#include #include #include @@ -104,26 +104,11 @@ func(void*, R (*)(void*, Args...)) -> func; template func(R (*)(Args...)) -> func; -struct [[nodiscard]] sha256_hasher { - constexpr std::uint64_t operator()(sha256 const& obj) const +struct [[nodiscard]] hash128_hasher { + constexpr std::uint64_t operator()(hash128 const& obj) const { - struct u64x4 { - alignas(alignof(sha256)) std::uint64_t // NOLINT(modernize-avoid-c-arrays) - v[sizeof(sha256) / sizeof(std::uint64_t)]; - }; - - auto value = std::bit_cast(obj); - auto h0 = value.v[0]; - auto h1 = value.v[1]; - auto h2 = value.v[2]; - auto h3 = value.v[3]; - - auto mix = [](std::uint64_t seed, std::uint64_t v) { - seed ^= v + 0x9e3779b97f4a7c15ULL + (seed << 6) + (seed >> 2); - return seed; - }; - - return mix(mix(mix(h0, h1), h2), h3); + // use only the lower 64 bits of the hash for the hash table + return static_cast(obj.value); } }; @@ -477,7 +462,7 @@ struct alignas(CACHELINE_ALIGNMENT) lru_memory_cache { void hit(std::uint64_t tick) { last_touched_tick = tick; } }; - std::unordered_map entries_ = {}; + std::unordered_map entries_ = {}; std::size_t limit_; explicit lru_memory_cache(std::size_t limit) : limit_{limit} @@ -492,7 +477,7 @@ struct alignas(CACHELINE_ALIGNMENT) lru_memory_cache { auto num_to_purge = (entries_.size() + 1) / 2; - std::vector> rankings; + std::vector> rankings; rankings.reserve(entries_.size()); for (auto& [key, entry] : entries_) { @@ -510,7 +495,7 @@ struct alignas(CACHELINE_ALIGNMENT) lru_memory_cache { } } - void insert(sha256 const& sha, T&& value, std::uint64_t tick) + void insert(hash128 const& sha, T&& value, std::uint64_t tick) { if (limit_ == 0) { return; } @@ -648,23 +633,22 @@ struct cache_t { // NOLINT [[nodiscard]] std::string const& get_tmp_dir(); /** - * @brief Query the cache for a compiled blob by its SHA-256 hash, or insert it if not present - * @param sha SHA-256 hash of the blob to query or insert + * @brief Query the cache for a compiled blob by its hash, or insert it if not present + * @param hash hash of the blob to query or insert * @param compile Function to compile the blob if it's not found in the cache * @return A shared future that will hold the compiled blob once it's available */ - [[nodiscard]] std::shared_future get_or_add_blob(sha256 const& sha, + [[nodiscard]] std::shared_future get_or_add_blob(hash128 const& hash, blob_compile_func compile); /** - * @brief Query the cache for a compiled library by its SHA-256 hash and binary type, or insert + * @brief Query the cache for a compiled library by its hash and binary type, or insert * it if not present - * @param sha SHA-256 hash of the library to query or insert - * @param type Binary type of the library (e.g., CUBIN, PTX) + * @param hash hash of the library to query or insert * @param compile Function to compile the library if it's not found in the cache * @return A shared future that will hold the compiled library once it's available */ - [[nodiscard]] std::shared_future get_or_add_library(sha256 const& sha, + [[nodiscard]] std::shared_future get_or_add_library(hash128 const& hash, library_compile_func compile); /** @@ -741,6 +725,20 @@ struct cache_t { // NOLINT [[nodiscard]] bool is_enabled(); }; +/** + * @brief Get the version of the NVRTC library + * @return An integer representing the NVRTC version. With the encoding major * 1000 + minor * 10 + + * patch + */ +[[nodiscard]] std::int32_t nvrtc_version(); + +/** + * @brief Get the version of the NVJITLINK library + * @return An integer representing the NVJITLINK version. With the encoding major * 1000 + minor * + * 10 + patch + */ +[[nodiscard]] std::int32_t nvjitlink_version(); + /** * @brief Compile source code into a binary blob * @@ -819,7 +817,7 @@ std::string reflect(T value) = delete; template <> inline std::string reflect(bool value) { - return std::format("(bool){}", value); + return std::format("{}", value); } /** @@ -830,7 +828,7 @@ inline std::string reflect(bool value) template <> inline std::string reflect(std::uint8_t value) { - return std::format("(unsigned char){}U", value); + return std::format("{}U", value); } /** @@ -841,7 +839,7 @@ inline std::string reflect(std::uint8_t value) template <> inline std::string reflect(std::uint16_t value) { - return std::format("(unsigned short){}U", value); + return std::format("{}U", value); } /** @@ -852,7 +850,7 @@ inline std::string reflect(std::uint16_t value) template <> inline std::string reflect(std::uint32_t value) { - return std::format("(unsigned int){}U", value); + return std::format("{}U", value); } /** @@ -863,7 +861,7 @@ inline std::string reflect(std::uint32_t value) template <> inline std::string reflect(std::uint64_t value) { - return std::format("(unsigned long long int){}ULL", value); + return std::format("{}ULL", value); } /** @@ -874,7 +872,7 @@ inline std::string reflect(std::uint64_t value) template <> inline std::string reflect(std::int8_t value) { - return std::format("(signed char){}", value); + return std::format("{}", value); } /** @@ -885,7 +883,7 @@ inline std::string reflect(std::int8_t value) template <> inline std::string reflect(std::int16_t value) { - return std::format("(signed short){}", value); + return std::format("{}", value); } /** @@ -896,7 +894,7 @@ inline std::string reflect(std::int16_t value) template <> inline std::string reflect(std::int32_t value) { - return std::format("(signed int){}", value); + return std::format("{}", value); } /** @@ -907,7 +905,7 @@ inline std::string reflect(std::int32_t value) template <> inline std::string reflect(std::int64_t value) { - return std::format("(signed long long int){}LL", value); + return std::format("{}LL", value); } /** @@ -918,7 +916,7 @@ inline std::string reflect(std::int64_t value) template <> inline std::string reflect(float value) { - return std::format("(float){}F", value); + return std::format("{}F", value); } /** @@ -929,20 +927,7 @@ inline std::string reflect(float value) template <> inline std::string reflect(double value) { - return std::format("(double){}", value); -} - -/** - * @brief Reflect a value of any type into its CUDA string representation, given the type name as a - * string - * @param type The name of the type to be reflected (e.g., "int", "float", "MyStruct", etc.) - * @param value The string representation of the value to be reflected, which will be used in the - * resulting CUDA code - * @return A string containing the CUDA representation of the value with the specified type - */ -inline std::string reflect_cast(std::string_view type, std::string_view value) -{ - return std::format("(({})({}))", type, value); + return std::format("{}", value); } /** @@ -959,7 +944,7 @@ template requires(std::is_enum_v) std::string reflect_enum(std::string_view type, T value) { - return reflect_cast(type, reflect(static_cast>(value))); + return std::format("{}{}{}{}", type, "{", static_cast>(value), "}"); } /** diff --git a/cpp/librtcx/sha256.hpp b/cpp/librtcx/sha256.hpp deleted file mode 100644 index 1149477f4b38..000000000000 --- a/cpp/librtcx/sha256.hpp +++ /dev/null @@ -1,288 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include -#include -#include -#include -#include -#include - -namespace rtcx { - -namespace functions { - -inline std::uint32_t ror(std::uint32_t x, std::uint32_t n) { return (x >> n) | (x << (32 - n)); } - -inline std::uint32_t ch(std::uint32_t x, std::uint32_t y, std::uint32_t z) -{ - return z ^ (x & (y ^ z)); -} - -inline std::uint32_t maj(std::uint32_t x, std::uint32_t y, std::uint32_t z) -{ - return ((x | y) & z) | (x & y); -} - -inline std::uint32_t sigma0(std::uint32_t x) { return ror(x, 2) ^ ror(x, 13) ^ ror(x, 22); } - -inline std::uint32_t sigma1(std::uint32_t x) { return ror(x, 6) ^ ror(x, 11) ^ ror(x, 25); } - -inline std::uint32_t gamma0(std::uint32_t x) { return ror(x, 7) ^ ror(x, 18) ^ (x >> 3); } - -inline std::uint32_t gamma1(std::uint32_t x) { return ror(x, 17) ^ ror(x, 19) ^ (x >> 10); } - -inline void put_be32(void* ptr, std::uint32_t value) -{ - auto* p = (std::uint8_t*)ptr; - p[0] = (value >> 24) & 0xff; - p[1] = (value >> 16) & 0xff; - p[2] = (value >> 8) & 0xff; - p[3] = (value >> 0) & 0xff; -} - -inline std::uint32_t get_be32(void const* ptr) -{ - auto const* p = (std::uint8_t const*)ptr; - return (std::uint32_t)p[0] << 24 | (std::uint32_t)p[1] << 16 | (std::uint32_t)p[2] << 8 | - (std::uint32_t)p[3] << 0; -} - -} // namespace functions - -struct [[nodiscard]] sha256_hex_string { - char data_[65]; // NOLINT(modernize-avoid-c-arrays) - - [[nodiscard]] constexpr std::string_view view() const { return std::string_view{data_, 64}; } - - [[nodiscard]] constexpr operator std::string_view() const { return view(); } - - [[nodiscard]] char const* data() const { return data_; } - - [[nodiscard]] char const* c_str() const { return data_; } - - [[nodiscard]] static constexpr std::size_t size() { return 64; } - - static sha256_hex_string make(std::span input) - { - constexpr char const HEX_CHARS[] = "0123456789abcdef"; // NOLINT(modernize-avoid-c-arrays) - sha256_hex_string hex; - for (std::size_t i = 0; i < 32; ++i) { - hex.data_[i * 2] = HEX_CHARS[(input[i] >> 4) & 0x0F]; - hex.data_[i * 2 + 1] = HEX_CHARS[input[i] & 0x0F]; - } - hex.data_[64] = '\0'; - return hex; - } -}; - -struct [[nodiscard]] sha256 { - alignas(16) std::uint8_t data_[32]; // NOLINT(modernize-avoid-c-arrays) - - [[nodiscard]] std::uint8_t operator[](std::size_t index) const { return data_[index]; } - - [[nodiscard]] std::size_t size() const { return 32; } - - [[nodiscard]] std::uint8_t const* data() const { return data_; } - - [[nodiscard]] constexpr bool operator==(sha256 const&) const = default; - - sha256_hex_string to_hex_string() const { return sha256_hex_string::make(data_); } - - static sha256 parse(std::string_view hex) - { - if (hex.size() != 64) { - throw std::invalid_argument(std::format( - "Invalid SHA256 hex string length, expected 64 got {} (sha: `{}`)", hex.size(), hex)); - } - sha256 hash; - for (std::size_t i = 0; i < 32; ++i) { - auto hex_byte = hex.substr(i * 2, 2); - auto [ptr, ec] = std::from_chars(hex_byte.begin(), hex_byte.end(), hash.data_[i], 16); - if (ec != std::errc()) { - throw std::invalid_argument( - std::format("Invalid hex character in SHA256 string: `{}`", hex_byte)); - } - } - return hash; - } -}; - -struct sha256_context { - private: - static constexpr std::size_t BLOCK_SIZE = 64; - std::uint32_t state_[8] = // NOLINT(modernize-avoid-c-arrays) - {0x6a09'e667ul, - 0xbb67'ae85ul, - 0x3c6e'f372ul, - 0xa54f'f53aul, - 0x510e'527ful, - 0x9b05'688cul, - 0x1f83'd9abul, - 0x5be0'cd19ul}; - std::uint64_t size_ = 0; - std::uint8_t buf_[BLOCK_SIZE] = {}; // NOLINT(modernize-avoid-c-arrays) - - public: - sha256_context() = default; - sha256_context(sha256_context const&) = delete; - sha256_context& operator=(sha256_context const&) = delete; - sha256_context(sha256_context&&) = delete; - sha256_context& operator=(sha256_context&&) = delete; - ~sha256_context() = default; - - private: - void transform(std::uint8_t const* buf) - { - std::uint32_t S[8], W[64], t0, t1; // NOLINT(modernize-avoid-c-arrays) - int i; - - /* copy state into S */ - for (i = 0; i < 8; i++) - S[i] = state_[i]; - - /* copy the state into 512-bits into W[0..15] */ - for (i = 0; i < 16; i++, buf += sizeof(std::uint32_t)) - W[i] = functions::get_be32(buf); - - /* fill W[16..63] */ - for (i = 16; i < 64; i++) - W[i] = functions::gamma1(W[i - 2]) + W[i - 7] + functions::gamma0(W[i - 15]) + W[i - 16]; - -#define RND(a, b, c, d, e, f, g, h, i, ki) \ - t0 = h + functions::sigma1(e) + functions::ch(e, f, g) + ki + W[i]; \ - t1 = functions::sigma0(a) + functions::maj(a, b, c); \ - d += t0; \ - h = t0 + t1; - - RND(S[0], S[1], S[2], S[3], S[4], S[5], S[6], S[7], 0, 0x428a'2f98); - RND(S[7], S[0], S[1], S[2], S[3], S[4], S[5], S[6], 1, 0x7137'4491); - RND(S[6], S[7], S[0], S[1], S[2], S[3], S[4], S[5], 2, 0xb5c0'fbcf); - RND(S[5], S[6], S[7], S[0], S[1], S[2], S[3], S[4], 3, 0xe9b5'dba5); - RND(S[4], S[5], S[6], S[7], S[0], S[1], S[2], S[3], 4, 0x3956'c25b); - RND(S[3], S[4], S[5], S[6], S[7], S[0], S[1], S[2], 5, 0x59f1'11f1); - RND(S[2], S[3], S[4], S[5], S[6], S[7], S[0], S[1], 6, 0x923f'82a4); - RND(S[1], S[2], S[3], S[4], S[5], S[6], S[7], S[0], 7, 0xab1c'5ed5); - RND(S[0], S[1], S[2], S[3], S[4], S[5], S[6], S[7], 8, 0xd807'aa98); - RND(S[7], S[0], S[1], S[2], S[3], S[4], S[5], S[6], 9, 0x1283'5b01); - RND(S[6], S[7], S[0], S[1], S[2], S[3], S[4], S[5], 10, 0x2431'85be); - RND(S[5], S[6], S[7], S[0], S[1], S[2], S[3], S[4], 11, 0x550c'7dc3); - RND(S[4], S[5], S[6], S[7], S[0], S[1], S[2], S[3], 12, 0x72be'5d74); - RND(S[3], S[4], S[5], S[6], S[7], S[0], S[1], S[2], 13, 0x80de'b1fe); - RND(S[2], S[3], S[4], S[5], S[6], S[7], S[0], S[1], 14, 0x9bdc'06a7); - RND(S[1], S[2], S[3], S[4], S[5], S[6], S[7], S[0], 15, 0xc19b'f174); - RND(S[0], S[1], S[2], S[3], S[4], S[5], S[6], S[7], 16, 0xe49b'69c1); - RND(S[7], S[0], S[1], S[2], S[3], S[4], S[5], S[6], 17, 0xefbe'4786); - RND(S[6], S[7], S[0], S[1], S[2], S[3], S[4], S[5], 18, 0x0fc1'9dc6); - RND(S[5], S[6], S[7], S[0], S[1], S[2], S[3], S[4], 19, 0x240c'a1cc); - RND(S[4], S[5], S[6], S[7], S[0], S[1], S[2], S[3], 20, 0x2de9'2c6f); - RND(S[3], S[4], S[5], S[6], S[7], S[0], S[1], S[2], 21, 0x4a74'84aa); - RND(S[2], S[3], S[4], S[5], S[6], S[7], S[0], S[1], 22, 0x5cb0'a9dc); - RND(S[1], S[2], S[3], S[4], S[5], S[6], S[7], S[0], 23, 0x76f9'88da); - RND(S[0], S[1], S[2], S[3], S[4], S[5], S[6], S[7], 24, 0x983e'5152); - RND(S[7], S[0], S[1], S[2], S[3], S[4], S[5], S[6], 25, 0xa831'c66d); - RND(S[6], S[7], S[0], S[1], S[2], S[3], S[4], S[5], 26, 0xb003'27c8); - RND(S[5], S[6], S[7], S[0], S[1], S[2], S[3], S[4], 27, 0xbf59'7fc7); - RND(S[4], S[5], S[6], S[7], S[0], S[1], S[2], S[3], 28, 0xc6e0'0bf3); - RND(S[3], S[4], S[5], S[6], S[7], S[0], S[1], S[2], 29, 0xd5a7'9147); - RND(S[2], S[3], S[4], S[5], S[6], S[7], S[0], S[1], 30, 0x06ca'6351); - RND(S[1], S[2], S[3], S[4], S[5], S[6], S[7], S[0], 31, 0x1429'2967); - RND(S[0], S[1], S[2], S[3], S[4], S[5], S[6], S[7], 32, 0x27b7'0a85); - RND(S[7], S[0], S[1], S[2], S[3], S[4], S[5], S[6], 33, 0x2e1b'2138); - RND(S[6], S[7], S[0], S[1], S[2], S[3], S[4], S[5], 34, 0x4d2c'6dfc); - RND(S[5], S[6], S[7], S[0], S[1], S[2], S[3], S[4], 35, 0x5338'0d13); - RND(S[4], S[5], S[6], S[7], S[0], S[1], S[2], S[3], 36, 0x650a'7354); - RND(S[3], S[4], S[5], S[6], S[7], S[0], S[1], S[2], 37, 0x766a'0abb); - RND(S[2], S[3], S[4], S[5], S[6], S[7], S[0], S[1], 38, 0x81c2'c92e); - RND(S[1], S[2], S[3], S[4], S[5], S[6], S[7], S[0], 39, 0x9272'2c85); - RND(S[0], S[1], S[2], S[3], S[4], S[5], S[6], S[7], 40, 0xa2bf'e8a1); - RND(S[7], S[0], S[1], S[2], S[3], S[4], S[5], S[6], 41, 0xa81a'664b); - RND(S[6], S[7], S[0], S[1], S[2], S[3], S[4], S[5], 42, 0xc24b'8b70); - RND(S[5], S[6], S[7], S[0], S[1], S[2], S[3], S[4], 43, 0xc76c'51a3); - RND(S[4], S[5], S[6], S[7], S[0], S[1], S[2], S[3], 44, 0xd192'e819); - RND(S[3], S[4], S[5], S[6], S[7], S[0], S[1], S[2], 45, 0xd699'0624); - RND(S[2], S[3], S[4], S[5], S[6], S[7], S[0], S[1], 46, 0xf40e'3585); - RND(S[1], S[2], S[3], S[4], S[5], S[6], S[7], S[0], 47, 0x106a'a070); - RND(S[0], S[1], S[2], S[3], S[4], S[5], S[6], S[7], 48, 0x19a4'c116); - RND(S[7], S[0], S[1], S[2], S[3], S[4], S[5], S[6], 49, 0x1e37'6c08); - RND(S[6], S[7], S[0], S[1], S[2], S[3], S[4], S[5], 50, 0x2748'774c); - RND(S[5], S[6], S[7], S[0], S[1], S[2], S[3], S[4], 51, 0x34b0'bcb5); - RND(S[4], S[5], S[6], S[7], S[0], S[1], S[2], S[3], 52, 0x391c'0cb3); - RND(S[3], S[4], S[5], S[6], S[7], S[0], S[1], S[2], 53, 0x4ed8'aa4a); - RND(S[2], S[3], S[4], S[5], S[6], S[7], S[0], S[1], 54, 0x5b9c'ca4f); - RND(S[1], S[2], S[3], S[4], S[5], S[6], S[7], S[0], 55, 0x682e'6ff3); - RND(S[0], S[1], S[2], S[3], S[4], S[5], S[6], S[7], 56, 0x748f'82ee); - RND(S[7], S[0], S[1], S[2], S[3], S[4], S[5], S[6], 57, 0x78a5'636f); - RND(S[6], S[7], S[0], S[1], S[2], S[3], S[4], S[5], 58, 0x84c8'7814); - RND(S[5], S[6], S[7], S[0], S[1], S[2], S[3], S[4], 59, 0x8cc7'0208); - RND(S[4], S[5], S[6], S[7], S[0], S[1], S[2], S[3], 60, 0x90be'fffa); - RND(S[3], S[4], S[5], S[6], S[7], S[0], S[1], S[2], 61, 0xa450'6ceb); - RND(S[2], S[3], S[4], S[5], S[6], S[7], S[0], S[1], 62, 0xbef9'a3f7); - RND(S[1], S[2], S[3], S[4], S[5], S[6], S[7], S[0], 63, 0xc671'78f2); - -#undef RND - - for (i = 0; i < 8; i++) { - state_[i] += S[i]; - } - } - - public: - void update(std::span span) - { - std::uint32_t len_buf = size_ & 63; - auto data = span.data(); - auto len = span.size(); - - size_ += len; - - /* Read the data into buf and process blocks as they get full */ - if (len_buf) { - std::uint32_t left = 64 - len_buf; - if (len < left) left = len; - std::memcpy(len_buf + buf_, data, left); - len_buf = (len_buf + left) & 63; - len -= left; - data = (data + left); - if (len_buf) return; - transform(buf_); - } - - while (len >= 64) { - transform(data); - data = data + 64; - len -= 64; - } - - if (len) std::memcpy(buf_, data, len); - } - - sha256 finalize() - { - static std::uint8_t const pad[64] = {0x80}; // NOLINT(modernize-avoid-c-arrays) - std::uint32_t padlen[2]; // NOLINT(modernize-avoid-c-arrays) - int i; - - /* Pad with a binary 1 (ie 0x80), then zeroes, then length */ - padlen[0] = __builtin_bswap32((std::uint32_t)(size_ >> 29)); - padlen[1] = __builtin_bswap32((std::uint32_t)(size_ << 3)); - - i = size_ & 63; - update(std::span{pad, (std::size_t)(1 + (63 & (55 - i)))}); - update(std::span{reinterpret_cast(padlen), sizeof(padlen)}); - - sha256 out; - std::uint8_t* digest = out.data_; - - /* copy output */ - for (i = 0; i < 8; i++, digest += sizeof(std::uint32_t)) { - functions::put_be32(digest, state_[i]); - } - return out; - } -}; - -} // namespace rtcx diff --git a/cpp/src/jit/cache.cpp b/cpp/src/jit/cache.cpp index 10a2e5558737..8a33e807be2c 100644 --- a/cpp/src/jit/cache.cpp +++ b/cpp/src/jit/cache.cpp @@ -1,6 +1,5 @@ - /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -14,6 +13,9 @@ #include #include +#define XXH_INLINE_ALL +#include + #include #include #include @@ -23,20 +25,38 @@ namespace CUDF_EXPORT cudf { namespace { -rtcx::sha256 hash(std::span input) +void hash(XXH3_state_t* ctx, std::span input) { - rtcx::sha256_context ctx; - ctx.update(std::span{reinterpret_cast(input.data()), input.size()}); - return ctx.finalize(); + XXH3_128bits_update(ctx, input.data(), input.size()); } -rtcx::sha256 hash(std::span inputs) +void hash(XXH3_state_t* ctx, std::span inputs) { - rtcx::sha256_context ctx; for (auto const* input : inputs) { - ctx.update(std::span{reinterpret_cast(input), std::strlen(input)}); + XXH3_128bits_update(ctx, input, std::strlen(input)); + XXH3_128bits_update(ctx, "\0", 1); // null terminator + } +} + +void hash(XXH3_state_t* ctx, std::span file_fragments) +{ + for (auto const& fragment : file_fragments) { + XXH3_128bits_update(ctx, fragment.path, std::strlen(fragment.path)); + XXH3_128bits_update(ctx, "\0", 1); // null terminator + } +} + +void hash(XXH3_state_t* ctx, std::span memory_fragments) +{ + for (auto const& fragment : memory_fragments) { + if (fragment.name != nullptr) { + XXH3_128bits_update(ctx, fragment.name, std::strlen(fragment.name)); + XXH3_128bits_update(ctx, "\0", 1); // null terminator + } else { + XXH3_128bits_update(ctx, fragment.data.data(), fragment.data.size()); + XXH3_128bits_update(ctx, "\0", 1); // null terminator + } } - return ctx.finalize(); } void install_file_set( @@ -149,7 +169,7 @@ void jit_bundle_t::ensure_installed() const std::string jit_bundle_t::get_hash() const { - auto str = rtcx::sha256_hex_string::make(cudf_cuda_embed::hash); + auto str = rtcx::hash128_hex_string::make(cudf_cuda_embed::hash); return std::string{str.view()}; } @@ -172,34 +192,15 @@ std::vector jit_bundle_t::get_include_directories() const namespace { -constexpr int MIN_CUDA_VERSION_PCH = 12800; // minimum CUDA version for the "--pch" NVRTC flag -constexpr int MIN_CUDA_VERSION_MINIMAL = - 12800; // minimum CUDA version for the "--minimal" NVRTC flag - -int32_t get_driver_version() +constexpr int32_t make_cuda_version(int32_t major, int32_t minor, int32_t patch) { - int32_t driver_version; - CUDF_CUDA_TRY(cudaDriverGetVersion(&driver_version)); - return driver_version; + return major * 1000 + minor * 10 + patch; } -int32_t get_runtime_version() -{ - int32_t runtime_version; - CUDF_CUDA_TRY(cudaRuntimeGetVersion(&runtime_version)); - return runtime_version; -} - -int32_t get_current_device_compute_capability() -{ - int32_t device; - CUDF_CUDA_TRY(cudaGetDevice(&device)); - - cudaDeviceProp props; - CUDF_CUDA_TRY(cudaGetDeviceProperties(&props, device)); - - return props.major * 10 + props.minor; -} +constexpr int32_t MIN_NVRTC_VERSION_PCH = + make_cuda_version(12, 8, 0); // minimum CUDA version for the "--pch" NVRTC flag +constexpr int32_t MIN_NVRTC_VERSION_MINIMAL = + make_cuda_version(12, 8, 0); // minimum CUDA version for the "--minimal" NVRTC flag std::tuple compile_library( char const* name, @@ -210,15 +211,16 @@ std::tuple compile_library( { CUDF_FUNC_RANGE(); - auto& ctx = cudf::get_context(); - auto& cfg = ctx.config(); - auto& bundle = ctx.jit_bundle(); - auto sm = get_current_device_compute_capability(); - auto runtime = get_runtime_version(); + auto& ctx = cudf::get_context(); + auto& cfg = ctx.config(); + auto& bundle = ctx.jit_bundle(); + auto& device_properties = ctx.get_device_properties(); + auto sm = device_properties.compute_capability; + auto nvrtc_version = ctx.nvrtc_version().value(); auto include_dirs = bundle.get_include_directories(); - auto use_pch = runtime >= MIN_CUDA_VERSION_PCH; - auto use_minimal = runtime >= MIN_CUDA_VERSION_MINIMAL; + auto use_pch = nvrtc_version >= MIN_NVRTC_VERSION_PCH; + auto use_minimal = nvrtc_version >= MIN_NVRTC_VERSION_MINIMAL; std::vector options; @@ -250,6 +252,9 @@ std::tuple compile_library( options.emplace_back("--pch-verbose=false"); options.emplace_back("--pch-messages=false"); } + } else { + // suppress warning about nv_hdrstop directive on later CUDA versions + options.emplace_back("--diag-suppress=161"); } if (cfg.disable_cuda_cache) { options.emplace_back("--no-cache"); } @@ -280,6 +285,90 @@ std::tuple compile_library( return std::make_tuple(library, std::make_shared(std::move(blob))); } +rtcx::blob compile_fragment(char const* name, + char const* cuda_code, + std::span extra_header_include_names, + std::span extra_headers, + std::span name_expressions) +{ + CUDF_FUNC_RANGE(); + + auto& ctx = cudf::get_context(); + auto& cfg = ctx.config(); + auto& bundle = ctx.jit_bundle(); + auto& device_properties = ctx.get_device_properties(); + auto sm = device_properties.compute_capability; + + auto include_dirs = bundle.get_include_directories(); + auto pch_dir = ctx.get_jit_pch_dir(); + + auto nvrtc_version = ctx.nvrtc_version().value(); + + auto use_pch = nvrtc_version >= MIN_NVRTC_VERSION_PCH; + auto use_minimal = nvrtc_version >= MIN_NVRTC_VERSION_MINIMAL; + + std::vector options; + + for (auto const& include_dir : include_dirs) { + options.emplace_back(std::format("-I{}", include_dir)); + } + + options.emplace_back(std::format("--gpu-architecture=sm_{}", sm)); + + options.emplace_back("--diag-suppress=47"); + options.emplace_back("--device-int128"); + + if (sm >= 100) { options.emplace_back("--device-float128"); } + + options.emplace_back("-std=c++20"); + options.emplace_back("--device-as-default-execution-space"); + options.emplace_back("--generate-line-info"); + options.emplace_back("--dopt=on"); + options.emplace_back("--dlink-time-opt"); + + if (use_minimal) { options.emplace_back("--minimal"); } + + if (use_pch) { + options.emplace_back("--pch"); + options.emplace_back(std::format("--pch-dir={}", pch_dir)); + + if (cfg.jit_verbose) { + options.emplace_back("--pch-verbose=true"); + options.emplace_back("--pch-messages=true"); + } else { + options.emplace_back("--pch-verbose=false"); + options.emplace_back("--pch-messages=false"); + } + } else { + // suppress warning about nv_hdrstop directive on later CUDA versions + options.emplace_back("--diag-suppress=161"); + } + + if (cfg.disable_cuda_cache) { options.emplace_back("--no-cache"); } + + if (cfg.dump_jit_trace) { options.emplace_back("--time=-"); } + + if (cfg.dump_jit_time_profile) { + options.emplace_back(std::format("--fdevice-time-trace=cudf_kernel_{}_trace", name)); + } + + std::vector options_cstr; + for (auto const& option : options) { + options_cstr.emplace_back(option.c_str()); + } + + auto params = rtcx::compile_params{.name = name, + .source = cuda_code, + .header_include_names = extra_header_include_names, + .headers = extra_headers, + .options = options_cstr, + .name_expressions = name_expressions, + .target_type = rtcx::binary_type::LTO_IR}; + + auto cubin = rtcx::compile(params); + return std::make_shared(rtcx::blob_t::from_buffer(std::move(cubin))); +} + } // namespace kernel get_kernel(std::string const& name, @@ -290,18 +379,18 @@ kernel get_kernel(std::string const& name, { CUDF_FUNC_RANGE(); - auto& cache = cudf::get_context().rtcx_cache(); - auto& bundle = cudf::get_context().jit_bundle(); + auto& ctx = cudf::get_context(); + auto& cache = ctx.rtcx_cache(); + auto& bundle = ctx.jit_bundle(); + auto& device_properties = ctx.get_device_properties(); + auto sm = device_properties.compute_capability; + auto runtime = device_properties.runtime_version; + auto driver = device_properties.driver_version; + auto bundle_hash = bundle.get_hash(); - auto runtime = get_runtime_version(); - auto driver = get_driver_version(); - auto sm = get_current_device_compute_capability(); - auto header_include_names_hash = hash(header_include_names).to_hex_string(); - auto headers_hash = hash(headers).to_hex_string(); - auto bundle_hash = bundle.get_hash(); - auto source_file = std::format("{}/{}", bundle.get_directory(), source_file_id); + auto source_file = std::format("{}/{}", bundle.get_directory(), source_file_id); - auto cache_key = std::format(R"***(cuLibrary + auto spec = std::format(R"***(cuLibrary name={} binary_type=CUBIN cuda_runtime={} @@ -309,21 +398,27 @@ cuda_driver={} arch={} bundle={} source_file={} -header_include_names={} -headers={} kernel_instance={} )***", - name, - runtime, - driver, - sm, - bundle_hash, - source_file, - header_include_names_hash.view(), - headers_hash.view(), - kernel_instance); - - auto cache_key_sha256 = hash(cache_key); + name, + runtime, + driver, + sm, + bundle_hash, + source_file, + kernel_instance); + + XXH3_state_t state; + XXH3_INITSTATE(&state); + XXH3_128bits_reset(&state); + hash(&state, spec); + hash(&state, "header_include_names: "); + hash(&state, header_include_names); + hash(&state, "headers: "); + hash(&state, headers); + + auto digest = XXH3_128bits_digest(&state); + auto key = rtcx::hash128{digest.high64, digest.low64}; auto compile = [&] { auto bundle_dir = cudf::get_context().jit_bundle().get_directory(); @@ -331,8 +426,160 @@ kernel_instance={} return compile_library(name.c_str(), source.c_str(), header_include_names, headers, {}); }; - auto fut = - cache.get_or_add_library(cache_key_sha256, rtcx::library_compile_func::from_functor(compile)); + auto fut = cache.get_or_add_library(key, rtcx::library_compile_func::from_functor(compile)); + + auto lib = fut.get(); + return kernel{lib, lib->get_kernel("cudf_kernel_entry")}; +} + +rtcx::blob get_kernel_fragment(std::string const& name, + std::string const& source_file_id, + std::span header_include_names, + std::span headers, + std::string const& kernel_instance) +{ + CUDF_FUNC_RANGE(); + + auto& ctx = cudf::get_context(); + auto& cache = ctx.rtcx_cache(); + auto& bundle = ctx.jit_bundle(); + auto& device_properties = ctx.get_device_properties(); + auto runtime = device_properties.runtime_version; + auto driver = device_properties.driver_version; + auto sm = device_properties.compute_capability; + auto bundle_hash = bundle.get_hash(); + + auto source_file = std::format("{}/{}", bundle.get_directory(), source_file_id); + + auto spec = std::format(R"***(objectFile +name={} +binary_type=LTO_IR +cuda_runtime={} +cuda_driver={} +arch={} +bundle={} +source_file={} +kernel_instance={} +)***", + name, + runtime, + driver, + sm, + bundle_hash, + source_file, + kernel_instance); + + XXH3_state_t state; + XXH3_INITSTATE(&state); + XXH3_128bits_reset(&state); + hash(&state, spec); + hash(&state, "header_include_names: "); + hash(&state, header_include_names); + hash(&state, "headers: "); + hash(&state, headers); + + auto digest = XXH3_128bits_digest(&state); + auto key = rtcx::hash128{digest.high64, digest.low64}; + + auto compile = [&] { + auto bundle_dir = cudf::get_context().jit_bundle().get_directory(); + auto source = read_file_string(source_file.c_str()); + return compile_fragment(name.c_str(), source.c_str(), header_include_names, headers, {}); + }; + + auto fut = cache.get_or_add_blob(key, rtcx::blob_compile_func::from_functor(compile)); + + return fut.get(); +} + +std::tuple link_library_uncached( + char const* name, + std::span file_fragments, + std::span memory_fragments) +{ + CUDF_FUNC_RANGE(); + + auto& ctx = cudf::get_context(); + auto& device_properties = ctx.get_device_properties(); + auto sm = device_properties.compute_capability; + auto& cfg = ctx.config(); + + std::vector options; + + options.emplace_back("-lto"); + options.emplace_back(std::format("-arch=sm_{}", sm)); + + if (cfg.disable_cuda_cache) { options.emplace_back("--no-cache"); } + + if (cfg.jit_verbose) { options.emplace_back("-verbose"); } + + if (cfg.dump_jit_trace) { options.emplace_back("-time"); } + + std::vector options_cstr; + + for (auto const& option : options) { + options_cstr.emplace_back(option.c_str()); + } + + auto params = rtcx::link_params{.name = name, + .output_type = rtcx::binary_type::CUBIN, + .file_fragments = file_fragments, + .memory_fragments = memory_fragments, + .link_options = options_cstr}; + + auto cubin = rtcx::link_library(params); + auto library = rtcx::load_library(cubin); + auto blob = rtcx::blob_t::from_buffer(std::move(cubin)); + + return std::make_tuple(library, std::make_shared(std::move(blob))); +} + +kernel get_lto_linked_kernel(std::string const& name, + std::span file_fragments, + std::span memory_fragments) +{ + CUDF_FUNC_RANGE(); + + auto& ctx = cudf::get_context(); + auto& cache = ctx.rtcx_cache(); + auto& bundle = ctx.jit_bundle(); + auto& device_properties = ctx.get_device_properties(); + auto runtime = device_properties.runtime_version; + auto driver = device_properties.driver_version; + auto sm = device_properties.compute_capability; + auto bundle_hash = bundle.get_hash(); + + auto spec = std::format(R"***(cuLibrary +name={} +binary_type=CUBIN +cuda_runtime={} +cuda_driver={} +arch={} +bundle={} +)***", + name, + runtime, + driver, + sm, + bundle_hash); + + XXH3_state_t state; + XXH3_INITSTATE(&state); + XXH3_128bits_reset(&state); + hash(&state, spec); + hash(&state, "file_fragments: "); + hash(&state, file_fragments); + hash(&state, "memory_fragments: "); + hash(&state, memory_fragments); + + auto digest = XXH3_128bits_digest(&state); + auto key = rtcx::hash128{digest.high64, digest.low64}; + + auto compile = [&] { + return link_library_uncached(name.c_str(), file_fragments, memory_fragments); + }; + + auto fut = cache.get_or_add_library(key, rtcx::library_compile_func::from_functor(compile)); auto lib = fut.get(); return kernel{lib, lib->get_kernel("cudf_kernel_entry")}; diff --git a/cpp/src/jit/cache.hpp b/cpp/src/jit/cache.hpp index e9b68f9e933b..16696f3d0a71 100644 --- a/cpp/src/jit/cache.hpp +++ b/cpp/src/jit/cache.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -89,4 +89,32 @@ kernel get_kernel(std::string const& name, std::span headers, std::string const& kernel_instance); +/** + * @brief Gets a kernel fragment from an embedded CUDA source file + * @param name Debug name for the kernel fragment (used for caching and logging) + * @param source_file_id Identifier for the embedded source file (used to locate the source and for + * caching) + * @param header_include_names Names of any additional embedded header files to include during + * compilation + * @param headers Contents of any additional embedded header files to include during compilation + * @param kernel_instance String identifier for the specific kernel instance being requested (used + * for caching) + */ +rtcx::blob get_kernel_fragment(std::string const& name, + std::string const& source_file_id, + std::span header_include_names, + std::span headers, + std::string const& kernel_instance); + +/** + * @brief Gets a kernel by linking together embedded binary fragments + * @param name Debug name for the kernel (used for caching and logging) + * @param file_fragments Paths of the fragments to link together to form the kernel + * @param memory_fragments Memory fragments to link + * @param extra_options Additional linking options + */ +kernel get_lto_linked_kernel(std::string const& name, + std::span file_fragments, + std::span memory_fragments); + } // namespace CUDF_EXPORT cudf diff --git a/cpp/src/jit/helpers.cpp b/cpp/src/jit/helpers.cpp index a7620284b05c..a016ca34dce5 100644 --- a/cpp/src/jit/helpers.cpp +++ b/cpp/src/jit/helpers.cpp @@ -99,5 +99,21 @@ kernel get_udf_kernel(std::string const& source_file, return get_kernel(source_file, source_file, include_names, include_headers, kernel_name); } +rtcx::blob get_udf_kernel_fragment(std::string const& source_file, + std::string const& kernel_name, + std::string const& udf_type) +{ + auto kernel_instance_source = std::format(R"***(#define CUDF_KERNEL_INSTANCE {} + #define CUDF_LTO_MODE)***", + kernel_name); + auto kernel_udf_source = std::format(R"***(#define CUDF_UDF_TYPE {})***", udf_type); + char const* include_names[] = // NOLINT(modernize-avoid-c-arrays) + {"cudf/detail/kernel_instance.cuh", "cudf/detail/operation_udf.cuh"}; + char const* include_headers[] = // NOLINT(modernize-avoid-c-arrays) + {kernel_instance_source.c_str(), kernel_udf_source.c_str()}; + + return get_kernel_fragment(source_file, source_file, include_names, include_headers, kernel_name); +} + } // namespace jit } // namespace cudf diff --git a/cpp/src/jit/helpers.hpp b/cpp/src/jit/helpers.hpp index 9d06d7883c33..9def0042525e 100644 --- a/cpp/src/jit/helpers.hpp +++ b/cpp/src/jit/helpers.hpp @@ -76,5 +76,9 @@ kernel get_udf_kernel(std::string const& source_file, std::string const& kernel_name, std::string const& cuda_source); +rtcx::blob get_udf_kernel_fragment(std::string const& source_file, + std::string const& kernel_name, + std::string const& udf_type); + } // namespace jit } // namespace cudf diff --git a/cpp/src/jit/row_ir.cpp b/cpp/src/jit/row_ir.cpp index cb801b30e857..a41bde6bdc0b 100644 --- a/cpp/src/jit/row_ir.cpp +++ b/cpp/src/jit/row_ir.cpp @@ -551,20 +551,6 @@ bool node::is_null_aware() const return std::any_of(args_.begin(), args_.end(), [](auto& a) { return a->is_null_aware(); }); } -bool node::is_fallible() const -{ - if (op_ == opcode::GET_INPUT) { return false; } - - if (get_op_info(op_, error_policy_).is_fallible) { return true; } - - CUDF_EXPECTS(!args_.empty(), - "Unexpectedly found an operator node with no arguments. All operator nodes should " - "have at least one argument.", - std::runtime_error); - - return std::any_of(args_.begin(), args_.end(), [](auto& a) { return a->is_fallible(); }); -} - bool node::is_always_valid() const { if (op_ == opcode::GET_INPUT) { return false; } @@ -789,7 +775,7 @@ bool is_nullable(scalar_input const& in) { return in.scalar_column->view().nulla bool is_nullable(column_input const& in) { return in.column.nullable(); } -std::tuple ast_converter::generate_code( +std::tuple ast_converter::generate_code( target target_id, ast::expression const& expr, std::string_view function_name) { // add 1 auto-deduced output variable @@ -805,9 +791,6 @@ std::tuple ast_converter: bool is_null_aware = std::any_of( output_irs_.cbegin(), output_irs_.cend(), [](auto& ir) { return ir->is_null_aware(); }); - bool is_fallible = std::any_of( - output_irs_.cbegin(), output_irs_.cend(), [](auto& ir) { return ir->is_fallible(); }); - bool output_is_always_valid = std::all_of( output_irs_.cbegin(), output_irs_.cend(), [](auto& ir) { return ir->is_always_valid(); }); @@ -870,10 +853,7 @@ std::tuple ast_converter: ir->emit_code(instance_, target, sink); } sink.emit("return cudf::errc::SUCCESS;\n}"); - return {sink.get_code(), - is_null_aware ? null_aware::YES : null_aware::NO, - is_fallible ? fallible::YES : fallible::NO, - null_policy}; + return {sink.get_code(), is_null_aware ? null_aware::YES : null_aware::NO, null_policy}; } std::variant get_column_view(scalar_input const& in) @@ -901,7 +881,7 @@ transform_args ast_converter::compute_column(target target_id, // TODO(lamarrr): consider deduplicating ast expression's input column references. See // TransformTest/1.DeeplyNestedArithmeticLogicalExpression for reference - auto [code, is_null_aware, is_fallible, output_nullability] = + auto [code, is_null_aware, output_nullability] = converter.generate_code(target_id, expr, function_name); std::vector> inputs; std::vector> scalar_columns; @@ -937,7 +917,6 @@ transform_args ast_converter::compute_column(target target_id, .udf = std::move(code), .source_type = cudf::udf_source_type::CUDA, .is_null_aware = is_null_aware, - .is_fallible = is_fallible, .user_data = std::nullopt, .inputs = inputs, .outputs{output}, diff --git a/cpp/src/jit/row_ir.hpp b/cpp/src/jit/row_ir.hpp index 3fd10037a437..d54ac08b42cd 100644 --- a/cpp/src/jit/row_ir.hpp +++ b/cpp/src/jit/row_ir.hpp @@ -91,7 +91,6 @@ struct [[nodiscard]] transform_args { std::string udf = {}; udf_source_type source_type = cudf::udf_source_type::CUDA; null_aware is_null_aware = null_aware::NO; - fallible is_fallible = fallible::NO; std::optional user_data = std::nullopt; std::vector inputs = {}; std::vector outputs = {}; @@ -363,13 +362,6 @@ struct [[nodiscard]] node { */ [[nodiscard]] bool is_null_aware() const; - /** - * @brief Returns `true` if this node can produce an error during execution. - * e.g. `ADD_OVERFLOW` operator can produce an error if the result of the addition overflows the - * range of the data type. - */ - [[nodiscard]] bool is_fallible() const; - /** * @brief Returns `true` if this node always produces a valid output even if its inputs are * nullable, e.g., `IS_NULL` operator produces a valid boolean output regardless of the @@ -447,7 +439,7 @@ struct [[nodiscard]] ast_converter { [[nodiscard]] std::unique_ptr add_ir_node(ast::jit::detail::operation const& expr); - [[nodiscard]] std::tuple generate_code( + [[nodiscard]] std::tuple generate_code( target target, ast::expression const& expr, std::string_view function_name); /** diff --git a/cpp/src/jit/util.cpp b/cpp/src/jit/util.cpp index b5747226ec47..b19967be8270 100644 --- a/cpp/src/jit/util.cpp +++ b/cpp/src/jit/util.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -7,6 +7,8 @@ #include #include #include +#include +#include namespace cudf { namespace jit { @@ -54,5 +56,42 @@ void const* get_data_ptr(scalar const& s) return type_dispatcher(s.type(), get_data_ptr_functor{}, s); } +data_type physical_type_of(data_type type) +{ + using enum type_id; + switch (type.id()) { + case EMPTY: return data_type{EMPTY}; + case BOOL8: + case INT8: + case UINT8: return data_type{UINT8}; + case INT16: + case UINT16: return data_type{UINT16}; + case INT32: + case UINT32: + case FLOAT32: return data_type{UINT32}; + case INT64: + case UINT64: + case FLOAT64: return data_type{UINT64}; + case TIMESTAMP_DAYS: return data_type{type_to_id()}; + case TIMESTAMP_SECONDS: return data_type{type_to_id()}; + case TIMESTAMP_MILLISECONDS: return data_type{type_to_id()}; + case TIMESTAMP_MICROSECONDS: return data_type{type_to_id()}; + case TIMESTAMP_NANOSECONDS: return data_type{type_to_id()}; + case DURATION_DAYS: return data_type{type_to_id()}; + case DURATION_SECONDS: return data_type{type_to_id()}; + case DURATION_MILLISECONDS: return data_type{type_to_id()}; + case DURATION_MICROSECONDS: return data_type{type_to_id()}; + case DURATION_NANOSECONDS: return data_type{type_to_id()}; + case DICTIONARY32: + case STRING: + case LIST: + case DECIMAL32: + case DECIMAL64: + case DECIMAL128: + case STRUCT: + default: return type; + } +} + } // namespace jit } // namespace cudf diff --git a/cpp/src/jit/util.hpp b/cpp/src/jit/util.hpp index 1e189e115cba..3220a2286a48 100644 --- a/cpp/src/jit/util.hpp +++ b/cpp/src/jit/util.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2023, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -22,5 +22,37 @@ void const* get_data_ptr(column_view const& view); */ void const* get_data_ptr(scalar const& s); +/** + * @brief Determine the physical type of a given element type. + * The Physical type is the type that can be used to alias the element type through a + * `cudf::column_device_view::element`. For example, the physical type of `int32`, `uint32`, + * `float`, `duration_D` is `uint32`. This is because `uint32` can be used to alias all 5 types + * through a `cudf::column_device_view::element`. + * + * It also means these types can safely alias each other across an ABI boundary as they have the + * same register storage type (PTX `b32`). + * + * e.g. + * + * ```cpp + * // PTX: .extern .func (.param .b32 func_retval0) my_udf(.param .b32 a); + * + * extern void my_udf_uint32(uint32_t a); + * extern void my_udf_int32(int32_t a); + * extern void my_udf_float(float a); + * extern void my_udf_duration_D(duration_D a); + * ``` + * + * `my_udf_int32` and `my_udf_float`, `my_udf_duration_D` can safely alias `my_udf` because + * they have the same register storage type (PTX `b32`). + * + * + * This means that some CUDA functions/kernels that are template-specialized on physical types can + * be re-used for other types that have the same physical type, thus reducing the compilation cost + * of code specialization. + * + */ +data_type physical_type_of(data_type type); + } // namespace jit } // namespace cudf diff --git a/cpp/src/runtime/context.cpp b/cpp/src/runtime/context.cpp index 4c31652246e7..7c4ffaf192ec 100644 --- a/cpp/src/runtime/context.cpp +++ b/cpp/src/runtime/context.cpp @@ -20,8 +20,42 @@ namespace cudf { +namespace { + +int32_t get_driver_version() +{ + int32_t driver_version; + CUDF_CUDA_TRY(cudaDriverGetVersion(&driver_version)); + return driver_version; +} + +int32_t get_runtime_version() +{ + int32_t runtime_version; + CUDF_CUDA_TRY(cudaRuntimeGetVersion(&runtime_version)); + return runtime_version; +} + +int32_t get_current_device_compute_capability() +{ + int32_t device; + CUDF_CUDA_TRY(cudaGetDevice(&device)); + + cudaDeviceProp props; + CUDF_CUDA_TRY(cudaGetDeviceProperties(&props, device)); + + return props.major * 10 + props.minor; +} + +} // namespace + context::context(context_config cfg, init_flags flags) - : _config{std::move(cfg)}, _jit_cache_init_flag{} + : _config{std::move(cfg)}, + _jit_cache_init_flag{}, + _device_properties{ + get_driver_version(), get_runtime_version(), get_current_device_compute_capability()}, + _nvrtc_version{0}, + _nvjitlink_version{0} { initialize_components(flags); } @@ -39,6 +73,9 @@ void context::ensure_jit_cache_initialized() rtcx::initialize(); + _nvrtc_version = rtcx::nvrtc_version(); + _nvjitlink_version = rtcx::nvjitlink_version(); + auto limits = rtcx::cache_limits{.num_mem_blobs = _config.kernel_cache_limit_process, .num_mem_libraries = _config.kernel_cache_limit_process}; @@ -81,8 +118,19 @@ bool context::dump_codegen() const { return _config.dump_codegen; } bool context::use_jit() const { return _config.use_jit; } +context_config const& context::config() const { return _config; } + std::string const& context::get_jit_pch_dir() const { return _config.jit_pch_dir; } +context::device_properties const& context::get_device_properties() const +{ + return _device_properties; +} + +std::optional context::nvrtc_version() const { return _nvrtc_version; } + +std::optional context::nvjitlink_version() const { return _nvjitlink_version; } + void context::initialize_components(init_flags flags) { if (has_flag(flags, init_flags::INIT_JIT_CACHE)) { ensure_jit_cache_initialized(); } diff --git a/cpp/src/runtime/context.hpp b/cpp/src/runtime/context.hpp index 8477fcc89f7b..98baeb586385 100644 --- a/cpp/src/runtime/context.hpp +++ b/cpp/src/runtime/context.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -10,6 +10,7 @@ #include #include +#include namespace rtcx { struct cache_t; @@ -45,11 +46,20 @@ struct [[nodiscard]] context_config { /// objects/state across translation units. class context { public: + struct device_properties { + int32_t driver_version = 0; + int32_t runtime_version = 0; + int32_t compute_capability = 0; + }; + private: context_config _config; std::once_flag _jit_cache_init_flag; std::unique_ptr _rtcx_cache; std::unique_ptr _jit_bundle; + device_properties _device_properties; + std::optional _nvrtc_version; + std::optional _nvjitlink_version; private: void ensure_nvcomp_loaded(); @@ -72,10 +82,16 @@ class context { [[nodiscard]] bool use_jit() const; - [[nodiscard]] context_config const& config() const { return _config; } + [[nodiscard]] context_config const& config() const; [[nodiscard]] std::string const& get_jit_pch_dir() const; + [[nodiscard]] device_properties const& get_device_properties() const; + + [[nodiscard]] std::optional nvrtc_version() const; + + [[nodiscard]] std::optional nvjitlink_version() const; + /// @brief Initialize additional components based on the provided flags /// @param flags The initialization flags to process void initialize_components(init_flags flags); diff --git a/cpp/src/transform/jit/kernel.cu b/cpp/src/transform/jit/kernel.cu index 5c2d99b1048b..1ec0c3c3752b 100644 --- a/cpp/src/transform/jit/kernel.cu +++ b/cpp/src/transform/jit/kernel.cu @@ -34,21 +34,35 @@ #include // clang-format on +#ifndef CUDF_LTO_MODE +#define CUDF_UDF_TYPE int() +#endif + +// Use LTO-dispatch for transform operators if we're in LTO mode. This allows the operator to be +// defined in a separate translation unit and compiled with LTO, which can result in better +// performance due to more optimization opportunities +#ifdef CUDF_LTO_MODE +#define GENERIC_TRANSFORM_OP(...) ::cudf::jit::lto::transform(__VA_ARGS__) +#endif + namespace cudf { namespace jit { +namespace lto { + +using transform_type = CUDF_UDF_TYPE; + +extern "C" __device__ transform_type transform; + +} // namespace lto /// @brief The generic transform kernel. Supports all types and nullability combinations. -template -CUDF_KERNEL void transform_kernel(size_type row_size, - bitmask_type const* __restrict__ stencil, - void* __restrict__ user_data, - column_device_view_core const* __restrict__ input_cols, - mutable_column_device_view_core const* __restrict__ output_cols, - int32_t* __restrict__ max_error) +template +__device__ void transform_kernel(size_type row_size, + bitmask_type const* __restrict__ stencil, + void* __restrict__ user_data, + column_device_view_core const* __restrict__ input_cols, + mutable_column_device_view_core const* __restrict__ output_cols, + int32_t* __restrict__ max_error) { auto start = detail::grid_1d::global_thread_id(); auto stride = detail::grid_1d::grid_stride(); @@ -58,8 +72,8 @@ CUDF_KERNEL void transform_kernel(size_type row_size, auto operation = [&](Args args) { // TODO: static assert invocable auto func = [&](auto... a) { - if constexpr (!discard_errors) { - return GENERIC_TRANSFORM_OP(a...); + if constexpr (!cuda::std::is_void_v) { + return static_cast(GENERIC_TRANSFORM_OP(a...)); } else { (void)GENERIC_TRANSFORM_OP(a...); return errc::SUCCESS; @@ -91,7 +105,7 @@ CUDF_KERNEL void transform_kernel(size_type row_size, (A::assign(output_cols, row, cuda::std::get(outs)), ...); }); - if constexpr (!discard_errors) { thread_error = cuda::std::max(thread_error, row_error); } + thread_error = cuda::std::max(thread_error, row_error); } else { auto active_mask = __ballot_sync(__activemask(), row < row_size); @@ -114,14 +128,15 @@ CUDF_KERNEL void transform_kernel(size_type row_size, ...); }); - if constexpr (!discard_errors) { thread_error = cuda::std::max(thread_error, row_error); } + thread_error = cuda::std::max(thread_error, row_error); } } - if constexpr (!discard_errors) { - cuda::atomic_ref ref(*max_error); - ref.fetch_max(static_cast(thread_error), cuda::std::memory_order_relaxed); - } + // early exit if no error occurred + if (thread_error == errc::SUCCESS) { return; } + + cuda::atomic_ref ref(*max_error); + ref.fetch_max(static_cast(thread_error), cuda::std::memory_order_relaxed); } } // namespace jit diff --git a/cpp/src/transform/transform.cu b/cpp/src/transform/transform.cu index bf0137e863e6..a34b8389a774 100644 --- a/cpp/src/transform/transform.cu +++ b/cpp/src/transform/transform.cu @@ -24,6 +24,7 @@ #include +#include #include #include #include @@ -31,6 +32,8 @@ #include #include +#include +#include #include #include @@ -163,7 +166,6 @@ using handle = std::variant< namespace jit_transform { kernel instantiate(bool is_null_aware, - bool discard_errors, bool has_user_data, std::string const& ins, std::string const& outs, @@ -182,7 +184,6 @@ kernel instantiate(bool is_null_aware, auto kernel = rtcx::reflect_template("cudf::jit::transform_kernel", rtcx::reflect(is_null_aware), - rtcx::reflect(discard_errors), rtcx::reflect(has_user_data), ins, outs); @@ -208,29 +209,28 @@ void launch(cudf::kernel const& kernel, kernel.launch({cfg.min_grid_size}, {cfg.block_size}, 0, stream, args); } -namespace { - -std::string get_element_type_name(column_view const& view); +std::string get_element_type_name(column_view const& view, bool use_physical_type); struct element_type_name_fn { template - std::string operator()(column_view const& view) const + std::string operator()(column_view const& view, bool use_physical_type) const requires(is_fixed_width() || std::same_as) { - return type_to_name(view.type()); + return type_to_name(use_physical_type ? jit::physical_type_of(view.type()) : view.type()); } template - std::string operator()(column_view const& view) const + std::string operator()(column_view const& view, bool use_physical_type) const requires(std::same_as) { - return std::format("cudf::dictionary_element<{}, {}>", - get_element_type_name(view.child(cudf::dictionary_indices_column_index)), - get_element_type_name(view.child(cudf::dictionary_keys_column_index))); + return std::format( + "cudf::dictionary_element<{}, {}>", + get_element_type_name(view.child(cudf::dictionary_indices_column_index), use_physical_type), + get_element_type_name(view.child(cudf::dictionary_keys_column_index), use_physical_type)); } template - std::string operator()(column_view const& view) const + std::string operator()(column_view const& view, bool use_physical_type) const requires(!is_fixed_width() && !std::same_as && !std::same_as) { @@ -238,57 +238,63 @@ struct element_type_name_fn { } }; -std::string get_element_type_name(column_view const& view) +std::string get_element_type_name(column_view const& view, bool use_physical_type) { - return cudf::type_dispatcher(view.type(), element_type_name_fn{}, view); + return cudf::type_dispatcher(view.type(), element_type_name_fn{}, view, use_physical_type); } -} // namespace - -std::string reflect_input_element(column_view const& c) { return get_element_type_name(c); } +std::string reflect_input_element(column_view const& c, bool use_physical_type) +{ + return get_element_type_name(c, use_physical_type); +} -std::string reflect_input_element(scalar_column_view const& c) +std::string reflect_input_element(scalar_column_view const& c, bool use_physical_type) { - return get_element_type_name(c.as_column_view()); + return get_element_type_name(c.as_column_view(), use_physical_type); } -std::string reflect_output_element(fixed_width_column const& c) +std::string reflect_output_element(fixed_width_column const& c, bool use_physical_type) { - return get_element_type_name(c._col->view()); + return get_element_type_name(c._col->view(), use_physical_type); } -std::string reflect_output_element(string_views_column const&) { return "cudf::string_view"; } +std::string reflect_output_element(string_views_column const&, + [[maybe_unused]] bool use_physical_type) +{ + return "cudf::string_view"; +} -std::string reflect_output_element(mutable_strings_column const&) +std::string reflect_output_element(mutable_strings_column const&, + [[maybe_unused]] bool use_physical_type) { return "cuda::std::span"; } -std::string reflect_input_value_type(column_view const& c) +std::string reflect_input_value_type(column_view const& c, bool use_physical_type) { - return is_dictionary(c.type()) - ? reflect_input_value_type(c.child(cudf::dictionary_keys_column_index)) - : reflect_input_element(c); + return is_dictionary(c.type()) ? reflect_input_value_type( + c.child(cudf::dictionary_keys_column_index), use_physical_type) + : reflect_input_element(c, use_physical_type); } -std::string reflect_input_value_type(scalar_column_view const& c) +std::string reflect_input_value_type(scalar_column_view const& c, bool use_physical_type) { - return reflect_input_value_type(c.as_column_view()); + return reflect_input_value_type(c.as_column_view(), use_physical_type); } -std::string reflect_output_value_type(fixed_width_column const& c) +std::string reflect_output_value_type(fixed_width_column const& c, bool use_physical_type) { - return reflect_output_element(c); + return reflect_output_element(c, use_physical_type); } -std::string reflect_output_value_type(string_views_column const& c) +std::string reflect_output_value_type(string_views_column const& c, bool use_physical_type) { - return reflect_output_element(c); + return reflect_output_element(c, use_physical_type); } -std::string reflect_output_value_type(mutable_strings_column const& c) +std::string reflect_output_value_type(mutable_strings_column const& c, bool use_physical_type) { - return reflect_output_element(c); + return reflect_output_element(c, use_physical_type); } std::string reflect_input_column(column_view const&) { return "cudf::column_device_view_core"; } @@ -313,16 +319,18 @@ std::string reflect_output_column(mutable_strings_column const&) return "cudf::jit::mutable_strings_column_device_view"; } -auto reflect(udf_source_type source_type, +auto reflect(std::variant source_type, std::span inputs, std::span outputs) { std::vector in_types; + bool use_physical_types = std::holds_alternative(source_type); for (size_t i = 0; i < inputs.size(); i++) { - auto& in = inputs[i]; - auto column = std::visit([](auto& c) { return reflect_input_column(c); }, in); - auto element = std::visit([](auto& c) { return reflect_input_element(c); }, in); + auto& in = inputs[i]; + auto column = std::visit([&](auto& c) { return reflect_input_column(c); }, in); + auto element = + std::visit([&](auto& c) { return reflect_input_element(c, use_physical_types); }, in); bool as_scalar = std::holds_alternative(in); auto accessor = rtcx::reflect_template("cudf::jit::column_accessor", rtcx::reflect(i), @@ -336,9 +344,10 @@ auto reflect(udf_source_type source_type, std::vector out_types; for (size_t i = 0; i < outputs.size(); i++) { - auto& out = outputs[i]; - auto column = std::visit([](auto& c) { return reflect_output_column(c); }, out); - auto element = std::visit([](auto& c) { return reflect_output_element(c); }, out); + auto& out = outputs[i]; + auto column = std::visit([&](auto& c) { return reflect_output_column(c); }, out); + auto element = + std::visit([&](auto& c) { return reflect_output_element(c, use_physical_types); }, out); bool as_scalar = false; // never scalar auto accessor = rtcx::reflect_template("cudf::jit::column_accessor", rtcx::reflect(i), @@ -356,20 +365,87 @@ auto reflect(udf_source_type source_type, std::vector ptx_in_types; std::vector ptx_out_types; - if (source_type == udf_source_type::PTX) { + if (std::holds_alternative(source_type) && + std::get(source_type) == udf_source_type::PTX) { for (auto& in : inputs) { - ptx_in_types.push_back(std::visit([](auto& c) { return reflect_input_value_type(c); }, in)); + ptx_in_types.push_back( + std::visit([&](auto& c) { return reflect_input_value_type(c, use_physical_types); }, in)); } for (auto& out : outputs) { ptx_out_types.push_back( - std::visit([](auto& c) { return reflect_output_value_type(c); }, out)); + std::visit([&](auto& c) { return reflect_output_value_type(c, use_physical_types); }, out)); } } return std::make_tuple(ins, outs, ptx_in_types, ptx_out_types); } +std::string reflect_udf_signature(bool is_null_aware, + bool has_user_data, + std::span inputs, + std::span outputs, + bool use_physical_types) +{ + std::vector in_types; + + for (size_t i = 0; i < inputs.size(); i++) { + auto& in = inputs[i]; + auto element = + std::visit([&](auto& c) { return reflect_input_element(c, use_physical_types); }, in); + in_types.push_back(is_null_aware ? std::format("cuda::std::optional<{}>", element) : element); + } + + std::vector out_types; + + for (size_t i = 0; i < outputs.size(); i++) { + auto& out = outputs[i]; + auto element = + std::visit([&](auto& c) { return reflect_output_element(c, use_physical_types); }, out); + out_types.push_back(is_null_aware ? std::format("cuda::std::optional<{}> *", element) + : std::format("{} *", element)); + } + + std::vector params; + if (has_user_data) { params.push_back("void*"); } + params.insert(params.end(), out_types.begin(), out_types.end()); + params.insert(params.end(), in_types.begin(), in_types.end()); + + auto joined = + params.empty() + ? "" + : std::accumulate(std::next(params.begin()), params.end(), params[0], [](auto a, auto b) { + return std::format("{}, {}", a, b); + }); + + return std::format("int({})", joined); +} + +std::tuple instantiate_fragment( + bool is_null_aware, + bool has_user_data, + std::string const& ins, + std::string const& outs, + std::span inputs, + std::span outputs) +{ + CUDF_FUNC_RANGE(); + // substitutes the `CUDF_KERNEL_INSTANCE` macro + auto kernel = rtcx::reflect_template("cudf::jit::transform_kernel", + rtcx::reflect(is_null_aware), + rtcx::reflect(has_user_data), + ins, + outs); + + // substitutes the `CUDF_UDF_TYPE` macro + auto signature = reflect_udf_signature( + is_null_aware, has_user_data, inputs, outputs, /*use_physical_types=*/true); + + return {jit::get_udf_kernel_fragment("cudf/cpp/src/transform/jit/kernel.cu", kernel, signature), + lto_binary_type::LTO_IR, + kernel}; +} + auto to_args(std::span inputs, std::span outputs, rmm::cuda_stream_view stream, @@ -409,7 +485,6 @@ auto to_args(std::span inputs, } void run(bool is_null_aware, - bool discard_errors, bool has_user_data, size_type row_size, bitmask_type const* d_stencil, @@ -424,7 +499,6 @@ void run(bool is_null_aware, { auto [in_types, out_types, ptx_in_types, ptx_out_types] = reflect(source_type, inputs, outputs); auto kernel = instantiate(is_null_aware, - discard_errors, has_user_data, in_types, out_types, @@ -440,6 +514,71 @@ void run(bool is_null_aware, kernel, row_size, d_stencil, user_data, input_cols, output_cols, d_max_error, stream); } +rtcx::binary_type as_rtcx_binary_type(lto_binary_type type) +{ + switch (type) { + case lto_binary_type::LTO_IR: return rtcx::binary_type::LTO_IR; + case lto_binary_type::FATBIN: return rtcx::binary_type::FATBIN; + default: + CUDF_FAIL( + std::format("Unrecognized LTO binary type {} for LTO transform", static_cast(type)), + std::invalid_argument); + } +} + +void run_lto(std::optional, lto_binary_type, char const*>> + precompiled_kernel_fragment, + bool is_null_aware, + bool has_user_data, + size_type row_size, + bitmask_type const* d_stencil, + void* user_data, + std::span inputs, + std::span outputs, + int32_t* d_max_error, + std::span udf_binary, + lto_binary_type source_type, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + auto [in_types, out_types, ptx_in_types, ptx_out_types] = reflect(source_type, inputs, outputs); + + std::span kernel_fragment; + lto_binary_type kernel_fragment_binary_type = lto_binary_type::FATBIN; + rtcx::blob fragment_blob = nullptr; + std::string kernel_fragment_id; + + if (precompiled_kernel_fragment.has_value()) { + std::tie(kernel_fragment, kernel_fragment_binary_type, kernel_fragment_id) = + *precompiled_kernel_fragment; + } else { + std::tie(fragment_blob, kernel_fragment_binary_type, kernel_fragment_id) = + instantiate_fragment(is_null_aware, has_user_data, in_types, out_types, inputs, outputs); + kernel_fragment = fragment_blob->view(); + } + + rtcx::memory_fragment memory_fragments[] = { + { + .data = kernel_fragment, + .type = as_rtcx_binary_type(kernel_fragment_binary_type), + .name = kernel_fragment_id.c_str(), + }, + { + .data = udf_binary, + .type = as_rtcx_binary_type(source_type), + .name = nullptr // nullptr = unnamed fragment: the binary will be used to hash the UDF + }}; + + auto kernel = get_lto_linked_kernel("cudf/cpp/src/transform/jit/kernel.cu", {}, memory_fragments); + + auto [cols, handles] = to_args(inputs, outputs, stream, mr); + auto* input_cols = reinterpret_cast(cols.data()); + auto* output_cols = + reinterpret_cast(input_cols + inputs.size()); + return launch( + kernel, row_size, d_stencil, user_data, input_cols, output_cols, d_max_error, stream); +} + } // namespace jit_transform CUDF_KERNEL void copy_offset_bitmask(bitmask_type* __restrict__ destination, @@ -583,14 +722,15 @@ auto get_null_transformation(null_aware is_null_aware, return output_may_be_nullable; } -void perform_checks(udf_source_type source_type, +void perform_checks(std::variant source_type, null_aware is_null_aware, std::optional in_row_size, std::span inputs, std::span outputs, std::span const> string_offsets) { - if (source_type == udf_source_type::PTX) { + if (auto* udf_source = std::get_if(&source_type); + udf_source != nullptr && *udf_source == udf_source_type::PTX) { static constexpr auto is_input_value_supported = [](auto const& c) { return is_integral(c.type()) || is_floating_point(c.type()); }; @@ -900,7 +1040,6 @@ auto finalize_outputs(null_aware is_null_aware, std::unique_ptr
execute_transform(std::string const& udf, udf_source_type source_type, null_aware is_null_aware, - fallible is_fallible, std::optional in_row_size, std::optional user_data, std::span inputs, @@ -926,7 +1065,6 @@ std::unique_ptr
execute_transform(std::string const& udf, rmm::device_scalar d_max_error(static_cast(errc::SUCCESS), stream, mr); jit_transform::run(is_null_aware == null_aware::YES, - is_fallible == fallible::NO, user_data.has_value(), row_size, stencil_has_nulls ? stencil_arg : nullptr, @@ -957,7 +1095,6 @@ std::unique_ptr
execute_transform(std::string const& udf, std::unique_ptr
multi_transform(std::string const& udf, udf_source_type source_type, null_aware is_null_aware, - fallible is_fallible, std::optional user_data, std::span inputs, std::span outputs, @@ -971,7 +1108,6 @@ std::unique_ptr
multi_transform(std::string const& udf, return execute_transform(udf, source_type, is_null_aware, - is_fallible, row_size, user_data, inputs, @@ -981,30 +1117,6 @@ std::unique_ptr
multi_transform(std::string const& udf, mr); } -std::unique_ptr
multi_transform(std::string const& udf, - udf_source_type source_type, - null_aware is_null_aware, - std::optional user_data, - std::span inputs, - std::span outputs, - std::vector>&& string_offsets, - std::optional row_size, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) -{ - return multi_transform(udf, - source_type, - is_null_aware, - fallible::NO, - user_data, - inputs, - outputs, - std::move(string_offsets), - row_size, - stream, - mr); -} - std::unique_ptr transform_extended(std::span inputs, std::string const& udf, data_type output_type, @@ -1018,17 +1130,8 @@ std::unique_ptr transform_extended(std::span inpu { transform_output outputs[] = {{.type = output_type, .nullability = null_policy}}; - auto table = multi_transform(udf, - source_type, - is_null_aware, - fallible::NO, - user_data, - inputs, - outputs, - {}, - row_size, - stream, - mr); + auto table = multi_transform( + udf, source_type, is_null_aware, user_data, inputs, outputs, {}, row_size, stream, mr); auto cols = table->release(); return std::move(cols[0]); @@ -1081,7 +1184,6 @@ std::unique_ptr compute_column_jit(table_view const& table, auto result = multi_transform(args.udf, args.source_type, args.is_null_aware, - args.is_fallible, args.user_data, args.inputs, args.outputs, @@ -1093,4 +1195,105 @@ std::unique_ptr compute_column_jit(table_view const& table, return std::move(cols[0]); } +// if we have a matching pre-compiled kernel fragment for the given transform configuration, return +// it to use for LTO linking instead of compiling a new one +std::optional, lto_binary_type, char const*>> +dispatch_lto_kernel_fragment(bool is_null_aware, + bool has_user_data, + std::span inputs, + std::span outputs) +{ + auto strip_whitespace = [](std::string_view str) { + std::string result; + result.reserve(str.size()); + for (char c : str) { + if (!std::isspace(static_cast(c))) { result.push_back(c); } + } + return result; + }; + + // TODO: better and less error-prone symbol mangling, but this is sufficient for now. + + // the contract here is that CMake and this dispatch function agree on symbol mangling of the + // reflected kernel name. + auto [in_types, out_types, ptx_in_types, ptx_out_types] = + jit_transform::reflect(lto_binary_type::FATBIN, inputs, outputs); + auto target = strip_whitespace(rtcx::reflect_template("cudf::jit::transform_kernel", + rtcx::reflect(is_null_aware), + rtcx::reflect(has_user_data), + in_types, + out_types)); + + for (size_t i = 0; i < std::size(cudf_fragments::transform_kernel_FILE_INDEX); i++) { + auto file_index = cudf_fragments::transform_kernel_FILE_INDEX[i]; + auto instance = strip_whitespace(cudf_fragments::transform_kernel_INSTANCE[i]); + if (target == instance) { + auto range = cudf_fragments::file_ranges[file_index]; + return std::make_tuple(cudf_fragments::files.subspan(range[0], range[1]), + lto_binary_type::FATBIN, + cudf_fragments::transform_kernel_INSTANCE[i]); + } + } + + return std::nullopt; +} + +std::unique_ptr
transform_lto(std::span udf, + lto_binary_type binary_type, + null_aware is_null_aware, + std::optional user_data, + std::span inputs, + std::span outputs, + std::vector>&& string_offsets, + std::optional in_row_size, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + CUDF_FUNC_RANGE(); + perform_checks(binary_type, is_null_aware, in_row_size, inputs, outputs, string_offsets); + auto row_size = in_row_size.has_value() ? *in_row_size : jit::get_projection_size(inputs); + auto output_may_be_nullable = get_null_transformation(is_null_aware, inputs, outputs); + + auto [output_columns, stencil] = make_outputs(is_null_aware, + row_size, + inputs, + outputs, + output_may_be_nullable, + std::move(string_offsets), + stream, + mr); + auto stencil_arg = stencil.has_value() ? stencil->first : nullptr; + auto stencil_has_nulls = stencil.has_value() ? (stencil->second > 0) : false; + + auto precompiled_kernel_fragment = dispatch_lto_kernel_fragment( + is_null_aware == null_aware::YES, user_data.has_value(), inputs, output_columns); + + rmm::device_scalar d_max_error(static_cast(errc::SUCCESS), stream, mr); + + jit_transform::run_lto(precompiled_kernel_fragment, + is_null_aware == null_aware::YES, + user_data.has_value(), + row_size, + stencil_has_nulls ? stencil_arg : nullptr, + user_data.value_or(nullptr), + inputs, + output_columns, + d_max_error.data(), + udf, + binary_type, + stream, + mr); + + auto error = static_cast(d_max_error.value(stream)); + switch (error) { + case errc::SUCCESS: break; + default: + throw evaluation_error( + error, std::format("Transform UDF evaluation failed with error `{}`", to_string(error))); + } + + auto finalized = finalize_outputs(is_null_aware, row_size, std::move(output_columns), stream, mr); + return std::make_unique
(std::move(finalized)); +} + } // namespace cudf diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 2a94c04e13c0..ceb308590ee6 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -690,6 +690,38 @@ ConfigureTest(ENCODE_TEST encode/encode_tests.cpp) # * ast tests ------------------------------------------------------------------------------------- ConfigureTest(AST_TEST ast/transform_tests.cpp ast/ast_tree_tests.cpp ast/jit_expressions_tests.cpp) +# ################################################################################################## +# * transform LTO tests +# ------------------------------------------------------------------------------------- +rtcx_add_embed(cudf_test_fragments) + +add_fragment( + cudf_test_fragments FRAGMENT bankers_rounding SOURCE transform/fragments/bankers_rounding.cu +) + +add_fragment(cudf_test_fragments FRAGMENT distance SOURCE transform/fragments/distance.cu) + +add_fragment(cudf_test_fragments FRAGMENT invsqrt SOURCE transform/fragments/invsqrt.cu) + +add_fragment(cudf_test_fragments FRAGMENT lehmer_mean SOURCE transform/fragments/lehmer_mean.cu) + +add_fragment( + cudf_test_fragments FRAGMENT sum_of_squares SOURCE transform/fragments/sum_of_squares.cu +) + +add_fragment(cudf_test_fragments FRAGMENT to_upper SOURCE transform/fragments/to_upper.cu) + +rtcx_embed( + cudf_test_fragments COMPRESSION none OUTPUT_DIRECTORY "${CUDF_GENERATED_INCLUDE_DIR}/rtcx_embed" +) + +ConfigureTest( + TRANSFORM_LTO_TEST transform/transform_lto_test.cpp + ${cudf_test_fragments_SOURCE_DIR}/cudf_test_fragments.s +) +add_dependencies(TRANSFORM_LTO_TEST cudf_test_fragments) +target_include_directories(TRANSFORM_LTO_TEST PRIVATE ${cudf_test_fragments_SOURCE_DIR}) + # ################################################################################################## # * lists tests ---------------------------------------------------------------------------------- ConfigureTest( diff --git a/cpp/tests/transform/fragments/bankers_rounding.cu b/cpp/tests/transform/fragments/bankers_rounding.cu new file mode 100644 index 000000000000..85494bf93abc --- /dev/null +++ b/cpp/tests/transform/fragments/bankers_rounding.cu @@ -0,0 +1,48 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +template +__device__ Decimal bankers_round(Decimal x) +{ + using Rep = typename Decimal::rep; + + auto const scale = x.scale(); + auto const value = x.value(); + + if (scale >= numeric::scale_type{0}) { return x; } + + Rep factor = 1; + for (int32_t i = 0; i < -static_cast(scale); ++i) { + factor *= 10; + } + + Rep q = value / factor; + Rep r = value % factor; + + if (r < 0) r = -r; + + auto const half = factor / 2; + + bool round_up = false; + + if (r > half) { + round_up = true; + } else if (r == half) { + // tie: round to even + round_up = (q % 2) != 0; + } + + if (round_up) { q += value >= 0 ? Rep{1} : Rep{-1}; } + + return Decimal{q, numeric::scale_type{0}}; +} + +extern "C" __device__ int transform(numeric::decimal128* out, numeric::decimal128 in) +{ + *out = bankers_round(in); + return 0; +} diff --git a/cpp/tests/transform/fragments/distance.cu b/cpp/tests/transform/fragments/distance.cu new file mode 100644 index 000000000000..fee1bded29d9 --- /dev/null +++ b/cpp/tests/transform/fragments/distance.cu @@ -0,0 +1,16 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#include + +__device__ float distance(float x1, float y1, float x2, float y2) +{ + return cuda::std::sqrt((x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1)); +} + +extern "C" __device__ int transform(float* out, float x1, float y1, float x2, float y2) +{ + *out = distance(x1, y1, x2, y2); + return 0; +} diff --git a/cpp/tests/transform/fragments/invsqrt.cu b/cpp/tests/transform/fragments/invsqrt.cu new file mode 100644 index 000000000000..7c652b0ce27e --- /dev/null +++ b/cpp/tests/transform/fragments/invsqrt.cu @@ -0,0 +1,12 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +__device__ float invsqrt(float a) { return 1.0F / sqrtf(a); } + +extern "C" __device__ int transform(float* out, float a) +{ + *out = invsqrt(a); + return 0; +} diff --git a/cpp/tests/transform/fragments/lehmer_mean.cu b/cpp/tests/transform/fragments/lehmer_mean.cu new file mode 100644 index 000000000000..e1747d69adc5 --- /dev/null +++ b/cpp/tests/transform/fragments/lehmer_mean.cu @@ -0,0 +1,43 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include + +/** + * @brief Calculates the integer Lehmer mean of two integers using checked arithmetic. + * The Lehmer mean is defined as `(a^2 + b^2) / (a + b)`. + * @param a The first integer. + * @param b The second integer. + * @return The Lehmer mean of the two integers, or an error code if an overflow occurs during the + * calculation. + * + */ +__device__ cuda::std::expected lehmer_mean(int32_t a, int32_t b) +{ + auto a2 = cudf::detail::ops::mul_overflow(a, a); + if (!a2) return cuda::std::unexpected(a2.error()); + + auto b2 = cudf::detail::ops::mul_overflow(b, b); + if (!b2) return cuda::std::unexpected(b2.error()); + + auto a_b_sum = cudf::detail::ops::add_overflow(a, b); + if (!a_b_sum) return cuda::std::unexpected(a_b_sum.error()); + + auto a2_b2_sum = cudf::detail::ops::add_overflow(a2.value(), b2.value()); + if (!a2_b2_sum) return cuda::std::unexpected(a2_b2_sum.error()); + + return cudf::detail::ops::div_overflow(a2_b2_sum.value(), a_b_sum.value()); +} + +extern "C" __device__ int transform(int32_t* out, int32_t a, int32_t b) +{ + auto result = lehmer_mean(a, b); + if (!result) return static_cast(result.error()); + *out = result.value(); + return 0; +} diff --git a/cpp/tests/transform/fragments/sum_of_squares.cu b/cpp/tests/transform/fragments/sum_of_squares.cu new file mode 100644 index 000000000000..e3e7a149f715 --- /dev/null +++ b/cpp/tests/transform/fragments/sum_of_squares.cu @@ -0,0 +1,12 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +__device__ float sum_of_squares(float a, float b) { return a * a + b * b; } + +extern "C" __device__ int transform(float* out, float a, float b) +{ + *out = sum_of_squares(a, b); + return 0; +} diff --git a/cpp/tests/transform/fragments/to_upper.cu b/cpp/tests/transform/fragments/to_upper.cu new file mode 100644 index 000000000000..f809acc63133 --- /dev/null +++ b/cpp/tests/transform/fragments/to_upper.cu @@ -0,0 +1,21 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +__device__ uint8_t to_upper(uint8_t input) +{ + if (input > 96 && input < 123) { + return input - 32; + } else { + return input; + } +} + +extern "C" __device__ int transform(uint8_t* output, uint8_t input) +{ + *output = to_upper(input); + return 0; +} diff --git a/cpp/tests/transform/integration/unary_transform_test.cpp b/cpp/tests/transform/integration/unary_transform_test.cpp index 4518348033fc..086a2ba311d4 100644 --- a/cpp/tests/transform/integration/unary_transform_test.cpp +++ b/cpp/tests/transform/integration/unary_transform_test.cpp @@ -1404,7 +1404,6 @@ __device__ cudf::errc expression ( EXPECT_NO_THROW(result = cudf::multi_transform(cuda, cudf::udf_source_type::CUDA, cudf::null_aware::NO, - cudf::fallible::YES, std::nullopt, inputs, outputs, @@ -1422,7 +1421,6 @@ __device__ cudf::errc expression ( EXPECT_THROW(result = cudf::multi_transform(cuda, cudf::udf_source_type::CUDA, cudf::null_aware::NO, - cudf::fallible::YES, std::nullopt, inputs, outputs, diff --git a/cpp/tests/transform/transform_lto_test.cpp b/cpp/tests/transform/transform_lto_test.cpp new file mode 100644 index 000000000000..194778804cb2 --- /dev/null +++ b/cpp/tests/transform/transform_lto_test.cpp @@ -0,0 +1,222 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include + +struct TransformLTOTest : public cudf::test::BaseFixture {}; + +template +using column_wrapper = cudf::test::fixed_width_column_wrapper; + +template +using decimal_wrapper = cudf::test::fixed_point_column_wrapper; + +TEST_F(TransformLTOTest, InvSqrt) +{ + column_wrapper input{{1.0f, 4.0f, 9.0f, 16.0f}}; + + cudf::transform_input inputs[] = {input}; + cudf::transform_output outputs[] = { + {cudf::data_type{cudf::type_id::FLOAT32}, cudf::output_nullability::ALL_VALID}}; + + auto const range = cudf_test_fragments::file_ranges[cudf_test_fragments::invsqrt]; + std::span udf{cudf_test_fragments::files.subspan(range[0], range[1])}; + + auto result = cudf::transform_lto(udf, + cudf::lto_binary_type::FATBIN, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + {}, + std::nullopt, + cudf::test::get_default_stream()); + + column_wrapper expected{{1.0f, 0.5f, 0.33333334f, 0.25f}}; + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), expected); +} + +TEST_F(TransformLTOTest, Distance) +{ + column_wrapper x1{{0, 0, 0, 0}}; + column_wrapper y1{{0.0F, 1.2F, 2.5F, 3.7F}}; + column_wrapper x2{{1.6F, 2.1F, 3.2f, 4.5f}}; + column_wrapper y2{{0, 0, 0, 0}}; + + cudf::transform_input inputs[] = {x1, y1, x2, y2}; + cudf::transform_output outputs[] = { + {cudf::data_type{cudf::type_id::FLOAT32}, cudf::output_nullability::ALL_VALID}}; + + auto const range = cudf_test_fragments::file_ranges[cudf_test_fragments::distance]; + std::span udf{cudf_test_fragments::files.subspan(range[0], range[1])}; + + auto result = cudf::transform_lto(udf, + cudf::lto_binary_type::FATBIN, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + {}, + std::nullopt, + cudf::test::get_default_stream()); + + auto distance = [](float x1, float y1, float x2, float y2) { + return std::sqrt((x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1)); + }; + + column_wrapper expected{{ + distance(0, 0, 1.6F, 0), + distance(0, 1.2F, 2.1F, 0), + distance(0, 2.5F, 3.2F, 0), + distance(0, 3.7F, 4.5F, 0), + }}; + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), expected); +} + +TEST_F(TransformLTOTest, ToUpper) +{ + column_wrapper input{{65, 66, 97, 98, 48, 49, 32, 33, 127, 255}}; + + cudf::transform_input inputs[] = {input}; + cudf::transform_output outputs[] = { + {cudf::data_type{cudf::type_id::UINT8}, cudf::output_nullability::ALL_VALID}}; + + auto const range = cudf_test_fragments::file_ranges[cudf_test_fragments::to_upper]; + std::span udf{cudf_test_fragments::files.subspan(range[0], range[1])}; + + auto result = cudf::transform_lto(udf, + cudf::lto_binary_type::FATBIN, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + {}, + std::nullopt, + cudf::test::get_default_stream()); + + column_wrapper expected{{65, 66, 65, 66, 48, 49, 32, 33, 127, 255}}; + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), expected); +} + +TEST_F(TransformLTOTest, SumOfSquares) +{ + column_wrapper lhs{{1.0f, 4.0f, 9.0f, 16.0f}}; + column_wrapper rhs{{1.0f, 2.0f, 2.0f, 10.0f}}; + + cudf::transform_input inputs[] = {lhs, rhs}; + cudf::transform_output outputs[] = { + {cudf::data_type{cudf::type_id::FLOAT32}, cudf::output_nullability::ALL_VALID}}; + + auto const range = cudf_test_fragments::file_ranges[cudf_test_fragments::sum_of_squares]; + std::span udf{cudf_test_fragments::files.subspan(range[0], range[1])}; + + auto result = cudf::transform_lto(udf, + cudf::lto_binary_type::FATBIN, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + {}, + std::nullopt, + cudf::test::get_default_stream()); + + column_wrapper expected{{2.0f, 20.0f, 85.0f, 356.0f}}; + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), expected); +} + +TEST_F(TransformLTOTest, FallibleIntegerLehmerMean) +{ + // computes integer lehmer mean: `(a^2 + b^2) / (a + b)` for each row using checked arithmetic and + // throws if an overflow occurs + + column_wrapper a{{2, 3, 4, 6, 8}}; + column_wrapper b{{1, 1, 2, 3, 4}}; + column_wrapper a_fail{{2, 3, -2, 6, 8}}; + + cudf::transform_output outputs[] = { + {cudf::data_type{cudf::type_id::INT32}, cudf::output_nullability::ALL_VALID}}; + + auto const range = cudf_test_fragments::file_ranges[cudf_test_fragments::lehmer_mean]; + std::span udf{cudf_test_fragments::files.subspan(range[0], range[1])}; + + cudf::transform_input inputs[] = {a, b}; + + auto result = cudf::transform_lto(udf, + cudf::lto_binary_type::FATBIN, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + {}, + std::nullopt, + cudf::test::get_default_stream()); + auto lehmer_mean = [](int32_t a, int32_t b) { return (a * a + b * b) / (a + b); }; + + column_wrapper expected{{lehmer_mean(2, 1), + lehmer_mean(3, 1), + lehmer_mean(4, 2), + lehmer_mean(6, 3), + lehmer_mean(8, 4)}}; + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), expected); + + cudf::transform_input inputs_fail[] = {a_fail, b}; + + EXPECT_THROW(cudf::transform_lto(udf, + cudf::lto_binary_type::FATBIN, + cudf::null_aware::NO, + std::nullopt, + inputs_fail, + outputs, + {}, + std::nullopt, + cudf::test::get_default_stream()), + cudf::evaluation_error); +} + +TEST_F(TransformLTOTest, BankersRounding) +{ + using T = numeric::decimal128; + + decimal_wrapper input{{12450, 12550, 12650, 12750, -12450, -12550, -12650, -12750}, + numeric::scale_type{-2}}; + + cudf::transform_output output{cudf::data_type{cudf::type_to_id(), numeric::scale_type{0}}, + cudf::output_nullability::ALL_VALID}; + + auto const range = cudf_test_fragments::file_ranges[cudf_test_fragments::bankers_rounding]; + std::span udf{cudf_test_fragments::files.subspan(range[0], range[1])}; + + cudf::transform_input inputs[] = {input}; + cudf::transform_output outputs[] = {output}; + + auto result = cudf::transform_lto(udf, + cudf::lto_binary_type::FATBIN, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + {}, + std::nullopt, + cudf::test::get_default_stream()); + + decimal_wrapper expected{{124, 126, 126, 128, -124, -126, -126, -128}, numeric::scale_type{0}}; + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), expected); +}