[FEA] Multi-String Output & LTO Transform Examples

conda/recipes/libcudf/recipe.yaml

@@ -319,6 +319,8 @@ outputs:
         - cuda-version =${{ cuda_version }}
         - cuda-nvtx-dev
         - cuda-cudart-dev
+        - cuda-nvrtc-dev
+        - libnvjitlink-dev
       run:
         - ${{ pin_subpackage("libcudf", exact=True) }}
         - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }}

cpp/CMakeLists.txt

@@ -198,10 +198,6 @@ set_property(
 )
 message(VERBOSE "CUDF: LIBCUDF_LOGGING_LEVEL = '${LIBCUDF_LOGGING_LEVEL}'.")
 
-if(NOT CUDF_GENERATED_INCLUDE_DIR)
-  set(CUDF_GENERATED_INCLUDE_DIR ${CUDF_BINARY_DIR})
-endif()
-
 # ##################################################################################################
 # * linter configuration ---------------------------------------------------------------------------
 if(CUDF_CLANG_TIDY)
@@ -431,6 +427,24 @@ if(NOT BUILD_SHARED_LIBS)
   endif()
 endif()
 
+set(LIBCUDF_FRAGMENT_LINK_LIBRARIES
+    CCCL::CCCL
+    rapids_logger::rapids_logger
+    rmm::rmm
+    $<BUILD_LOCAL_INTERFACE:BS::thread_pool>
+    $<BUILD_LOCAL_INTERFACE:nvtx3::nvtx3-cpp>
+    $<BUILD_LOCAL_INTERFACE:cuco::cuco>
+    ZLIB::ZLIB
+    nvcomp::nvcomp
+    kvikio::kvikio
+    nanoarrow::nanoarrow
+    zstd
+)
+set(LIBCUDF_FRAGMENT_INCLUDE_DIRECTORIES "$<BUILD_INTERFACE:${CUDF_SOURCE_DIR}/include>"
+                                         "$<BUILD_INTERFACE:${CUDF_SOURCE_DIR}/src>"
+)
+set(LIBCUDF_FRAGMENT_COMPILE_OPTIONS "$<$<COMPILE_LANGUAGE:CUDA>:${CUDF_CUDA_FLAGS}>")
+
 rtcx_add_embed(cudf_cuda_embed)
 
 rtcx_embed_includes(
@@ -479,50 +493,19 @@ foreach(INC_DIR IN LISTS LIBCUDACXX_RAW_INCLUDE_DIRS)
   )
 endforeach()
 
-rtcx_embed(
-  cudf_cuda_embed COMPRESSION zstd OUTPUT_DIRECTORY "${CUDF_GENERATED_INCLUDE_DIR}/rtcx_embed"
-)
+rtcx_embed(cudf_cuda_embed COMPRESSION zstd OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/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<false, false, cudf::jit::type_list<cudf::jit::column_accessor<0ULL, cudf::column_device_view_core, ${TYPE}, false, 0>>, cudf::jit::type_list<cudf::jit::column_accessor<0ULL, cudf::mutable_column_device_view_core, ${TYPE}, false, 0>>>"
+# This function precompiles the unary and binary transform kernel fragments for all fixed-width
+# types. This helps amortize the cost of JIT compilation for kernels that have matching signatures
+# and reduces the amount of time spent on source-based runtime JIT compilation for LTO-based JIT.
+function(precompile_fixed_width_kernel_fragments)
+  list(APPEND CUDF_PRECOMPILE_PHYSICAL_TYPES uint8_t uint16_t uint32_t uint64_t numeric::decimal32
+       numeric::decimal64 numeric::decimal128
   )
-  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")
+  # Pre-compile unary-op kernel fragments for fixed-width types.
+  foreach(TYPE IN ITEMS ${CUDF_PRECOMPILE_PHYSICAL_TYPES})
     set(FRAGMENT_NAME transform_kernel)
     get_property(
       FILE_INDEX
@@ -531,7 +514,7 @@ foreach(TYPE IN ITEMS ${CUDF_PRECOMPILE_PHYSICAL_TYPES})
     )
     set(VARIANT_NAME transform_kernel_${FILE_INDEX})
     set(INSTANCE
-        "cudf::jit::transform_kernel<false, false, cudf::jit::type_list<cudf::jit::column_accessor<0ULL, cudf::column_device_view_core, ${TYPE}, false, 0>, cudf::jit::column_accessor<1ULL, cudf::column_device_view_core, ${TYPE}, ${RHS_IS_SCALAR}, 0>>, cudf::jit::type_list<cudf::jit::column_accessor<0ULL, cudf::mutable_column_device_view_core, ${TYPE}, false, 0>>>"
+        "cudf::jit::transform_kernel<false, false, cudf::jit::type_list<cudf::jit::column_accessor<0ULL, cudf::column_device_view_core, ${TYPE}, false, 0>>, cudf::jit::type_list<cudf::jit::column_accessor<0ULL, cudf::mutable_column_device_view_core, ${TYPE}, false, 0>>>"
     )
     add_fragment(
       cudf_fragments
@@ -542,7 +525,7 @@ foreach(TYPE IN ITEMS ${CUDF_PRECOMPILE_PHYSICAL_TYPES})
       KERNEL_INSTANCE
       ${INSTANCE}
       UDF_TYPE
-      "int(${TYPE} *, ${TYPE}, ${TYPE})"
+      "int(${TYPE} *, ${TYPE})"
       DEFINITIONS
       CUDF_LTO_MODE
       ARRAY_IDS
@@ -551,13 +534,137 @@ foreach(TYPE IN ITEMS ${CUDF_PRECOMPILE_PHYSICAL_TYPES})
       ARRAY_VALUES
       ${FILE_INDEX}
       "${INSTANCE}"
+      COMPILE_OPTIONS
+      ${LIBCUDF_FRAGMENT_COMPILE_OPTIONS}
+      INCLUDE_DIRECTORIES
+      ${LIBCUDF_FRAGMENT_INCLUDE_DIRECTORIES}
+      LINK_LIBRARIES
+      ${LIBCUDF_FRAGMENT_LINK_LIBRARIES}
     )
+
   endforeach()
-endforeach()
 
-rtcx_embed(
-  cudf_fragments COMPRESSION none OUTPUT_DIRECTORY "${CUDF_GENERATED_INCLUDE_DIR}/rtcx_embed"
-)
+  # Precompile binary-op kernel fragments for fixed-width types. To minimize binary size impact,
+  # only the RHS-scalar variants are precompiled.
+  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<false, false, cudf::jit::type_list<cudf::jit::column_accessor<0ULL, cudf::column_device_view_core, ${TYPE}, false, 0>, cudf::jit::column_accessor<1ULL, cudf::column_device_view_core, ${TYPE}, ${RHS_IS_SCALAR}, 0>>, cudf::jit::type_list<cudf::jit::column_accessor<0ULL, cudf::mutable_column_device_view_core, ${TYPE}, false, 0>>>"
+      )
+      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}"
+        COMPILE_OPTIONS
+        ${LIBCUDF_FRAGMENT_COMPILE_OPTIONS}
+        INCLUDE_DIRECTORIES
+        ${LIBCUDF_FRAGMENT_INCLUDE_DIRECTORIES}
+        LINK_LIBRARIES
+        ${LIBCUDF_FRAGMENT_LINK_LIBRARIES}
+      )
+    endforeach()
+  endforeach()
+endfunction()
+
+# This function precompiles kernel fragments for string transforms. It precompiles kernel fragments
+# for common string transform operations: multi-extract, multi-split, string-parsing. This helps
+# amortize the cost of JIT compilation for kernels that have matching signatures
+function(precompile_string_kernel_fragments)
+  list(APPEND CUDF_PRECOMPILE_STRING_OUTPUT_PHYSICAL_TYPES uint8_t uint16_t uint32_t uint64_t
+       cuda::std::span<char>
+  )
+  list(
+    APPEND
+    CUDF_PRECOMPILE_STRING_OUTPUT_COLUMN_TYPES
+    cudf::mutable_column_device_view_core
+    cudf::mutable_column_device_view_core
+    cudf::mutable_column_device_view_core
+    cudf::mutable_column_device_view_core
+    cudf::jit::mutable_strings_column_device_view
+  )
+  foreach(OUTPUT_TYPE OUTPUT_COLUMN_TYPE IN ZIP_LISTS CUDF_PRECOMPILE_STRING_OUTPUT_PHYSICAL_TYPES
+                      CUDF_PRECOMPILE_STRING_OUTPUT_COLUMN_TYPES
+  )
+    # Pre-compile for up to 8 output columns
+    foreach(OUTPUT_INDICES IN ITEMS "0" "0;1" "0;1;2" "0;1;2;3" "0;1;2;3;4" "0;1;2;3;4;5"
+                                    "0;1;2;3;4;5;6" "0;1;2;3;4;5;6;7"
+    )
+      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(OUTPUT_ACCESSORS "")
+      set(OUTPUT_POINTERS "")
+      foreach(OUTPUT_INDEX IN LISTS OUTPUT_INDICES)
+        list(
+          APPEND
+          OUTPUT_ACCESSORS
+          "cudf::jit::column_accessor<${OUTPUT_INDEX}ULL, ${OUTPUT_COLUMN_TYPE}, ${OUTPUT_TYPE}, false, 0>"
+        )
+        list(APPEND OUTPUT_POINTERS "${OUTPUT_TYPE} *")
+      endforeach()
+      list(JOIN OUTPUT_ACCESSORS " ," OUTPUT_ACCESSORS_STR)
+      list(JOIN OUTPUT_POINTERS " ," OUTPUT_POINTERS_STR)
+      set(INSTANCE
+          "cudf::jit::transform_kernel<false, false, cudf::jit::type_list<cudf::jit::column_accessor<0ULL, cudf::column_device_view_core, cudf::string_view, false, 0>>, cudf::jit::type_list<${OUTPUT_ACCESSORS_STR}>>"
+      )
+      add_fragment(
+        cudf_fragments
+        FRAGMENT
+        ${VARIANT_NAME}
+        SOURCE
+        src/transform/jit/kernel.cu
+        KERNEL_INSTANCE
+        ${INSTANCE}
+        UDF_TYPE
+        "int(${OUTPUT_POINTERS_STR}, cudf::string_view)"
+        DEFINITIONS
+        CUDF_LTO_MODE
+        ARRAY_IDS
+        ${FRAGMENT_NAME}_FILE_INDEX
+        ${FRAGMENT_NAME}_INSTANCE
+        ARRAY_VALUES
+        ${FILE_INDEX}
+        "${INSTANCE}"
+        COMPILE_OPTIONS
+        ${LIBCUDF_FRAGMENT_COMPILE_OPTIONS}
+        INCLUDE_DIRECTORIES
+        ${LIBCUDF_FRAGMENT_INCLUDE_DIRECTORIES}
+        LINK_LIBRARIES
+        ${LIBCUDF_FRAGMENT_LINK_LIBRARIES}
+      )
+    endforeach()
+  endforeach()
+endfunction()
+
+precompile_fixed_width_kernel_fragments()
+precompile_string_kernel_fragments()
+
+rtcx_embed(cudf_fragments COMPRESSION none OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/embed")
 
 # ##################################################################################################
 # * library targets -------------------------------------------------------------------------------
@@ -1207,7 +1314,7 @@ target_compile_options(
 target_include_directories(
   cudf
   PUBLIC "$<BUILD_INTERFACE:${DLPACK_INCLUDE_DIR}>" "$<BUILD_INTERFACE:${CUDF_SOURCE_DIR}/include>"
-         "$<BUILD_INTERFACE:${CUDF_GENERATED_INCLUDE_DIR}/include>"
+         "$<BUILD_INTERFACE:${CMAKE_CURRENT_BINARY_DIR}/include>"
   PRIVATE "$<BUILD_INTERFACE:${CUDF_SOURCE_DIR}/src>"
           "$<BUILD_INTERFACE:${nanoarrow_SOURCE_DIR}/src>"
           "$<BUILD_INTERFACE:${FlatBuffers_SOURCE_DIR}/include>"

cpp/benchmarks/CMakeLists.txt

@@ -399,11 +399,34 @@ 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)
+add_fragment(
+  cudf_benchmark_fragments
+  FRAGMENT
+  add_f32
+  SOURCE
+  binaryop/fragments/add_f32.cu
+  LINK_LIBRARIES
+  ${LIBCUDF_FRAGMENT_LINK_LIBRARIES}
+  INCLUDE_DIRECTORIES
+  ${LIBCUDF_FRAGMENT_INCLUDE_DIRECTORIES}
+  COMPILE_OPTIONS
+  ${LIBCUDF_FRAGMENT_COMPILE_OPTIONS}
+)
+add_fragment(
+  cudf_benchmark_fragments
+  FRAGMENT
+  mul_f32
+  SOURCE
+  binaryop/fragments/mul_f32.cu
+  LINK_LIBRARIES
+  ${LIBCUDF_FRAGMENT_LINK_LIBRARIES}
+  INCLUDE_DIRECTORIES
+  ${LIBCUDF_FRAGMENT_INCLUDE_DIRECTORIES}
+  COMPILE_OPTIONS
+  ${LIBCUDF_FRAGMENT_COMPILE_OPTIONS}
+)
 rtcx_embed(
-  cudf_benchmark_fragments COMPRESSION none OUTPUT_DIRECTORY
-  "${CUDF_GENERATED_INCLUDE_DIR}/rtcx_embed"
+  cudf_benchmark_fragments COMPRESSION none OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/embed"
 )
 
 # ##################################################################################################

cpp/cmake/Modules/AddFragment.cmake

@@ -13,8 +13,11 @@ include_guard(GLOBAL)
 # final library with metadata that allows it to be looked up at runtime.
 macro(add_fragment)
   set(TARGET ${ARGV0})
+  set(OPTIONS)
   set(ONE_VALUE_ARGS FRAGMENT SOURCE KERNEL_ONLY KERNEL_INSTANCE UDF_TYPE)
-  set(MULTI_VALUE_ARGS DEFINITIONS ARRAY_IDS ARRAY_VALUES)
+  set(MULTI_VALUE_ARGS DEFINITIONS ARRAY_IDS ARRAY_VALUES INCLUDE_DIRECTORIES LINK_LIBRARIES
+                       COMPILE_OPTIONS
+  )
   cmake_parse_arguments(ARG "${OPTIONS}" "${ONE_VALUE_ARGS}" "${MULTI_VALUE_ARGS}" ${ARGN})
 
   if(NOT ARG_FRAGMENT)
@@ -34,7 +37,7 @@ macro(add_fragment)
     target_compile_options(${OBJECT_ID} PRIVATE -Xnvlink=--kernels-used=cudf_kernel_entry)
   endif()
 
-  set(INSTANTIATION_DIR "${CUDF_GENERATED_INCLUDE_DIR}/${TARGET}/instantiations/${ARG_FRAGMENT}")
+  set(INSTANTIATION_DIR "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}/instantiations/${ARG_FRAGMENT}")
   target_include_directories(${OBJECT_ID} PRIVATE ${INSTANTIATION_DIR})
 
   if(ARG_KERNEL_INSTANCE)
@@ -53,7 +56,22 @@ macro(add_fragment)
     )
   endif()
 
-  target_compile_definitions(${OBJECT_ID} PRIVATE CUDF_DISABLE_EXPORTS ${ARG_DEFINITIONS})
+  if(ARG_INCLUDE_DIRECTORIES)
+    target_include_directories(${OBJECT_ID} PRIVATE ${ARG_INCLUDE_DIRECTORIES})
+  endif()
+
+  if(ARG_LINK_LIBRARIES)
+    target_link_libraries(${OBJECT_ID} PRIVATE ${ARG_LINK_LIBRARIES})
+  endif()
+
+  if(ARG_COMPILE_OPTIONS)
+    target_compile_options(${OBJECT_ID} PRIVATE ${ARG_COMPILE_OPTIONS})
+  endif()
+
+  if(ARG_DEFINITIONS)
+    target_compile_definitions(${OBJECT_ID} PRIVATE ${ARG_DEFINITIONS})
+  endif()
+
   set_target_properties(
     ${OBJECT_ID}
     PROPERTIES CUDA_SEPARABLE_COMPILATION ON
@@ -68,17 +86,6 @@ macro(add_fragment)
                CUDA_STANDARD_REQUIRED ON
                CUDA_VISIBILITY_PRESET hidden
   )
-  target_link_libraries(
-    ${OBJECT_ID}
-    PUBLIC CCCL::CCCL rapids_logger::rapids_logger rmm::rmm $<BUILD_LOCAL_INTERFACE:BS::thread_pool>
-    PRIVATE $<BUILD_LOCAL_INTERFACE:nvtx3::nvtx3-cpp> $<BUILD_LOCAL_INTERFACE:cuco::cuco>
-            ZLIB::ZLIB nvcomp::nvcomp kvikio::kvikio nanoarrow::nanoarrow zstd
-  )
-  target_include_directories(
-    ${OBJECT_ID} PRIVATE "$<BUILD_INTERFACE:${CUDF_SOURCE_DIR}/include>"
-                         "$<BUILD_INTERFACE:${CUDF_SOURCE_DIR}/src>"
-  )
-  target_compile_options(${OBJECT_ID} PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:${CUDF_CUDA_FLAGS}>")
 
   rtcx_embed_blob(
     ${TARGET} FILE $<TARGET_OBJECTS:${OBJECT_ID}> DEST fragments/${ARG_FRAGMENT}.fatbin ID