Reuse a runtime-descriptor LTO driver for AST JIT evaluation

cpp/src/jit/cache.cpp

@@ -434,7 +434,8 @@ rtcx::blob get_kernel_fragment(std::string const& name,
                                std::string const& source_file_id,
                                std::span<char const* const> header_include_names,
                                std::span<char const* const> headers,
-                               std::string const& kernel_instance)
+                               std::string const& kernel_instance,
+                               std::span<char const* const> name_expressions)
 {
   CUDF_FUNC_RANGE();
 
@@ -474,14 +475,17 @@ kernel_instance={}
   hash(&state, header_include_names);
   hash(&state, "headers: ");
   hash(&state, headers);
+  hash(&state, "name_expressions: ");
+  hash(&state, name_expressions);
 
   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, {});
+    return compile_fragment(
+      name.c_str(), source.c_str(), header_include_names, headers, name_expressions);
   };
 
   auto fut = cache.get_or_add_blob(key, rtcx::blob_compile_func::from_functor(compile));
@@ -531,6 +535,56 @@ std::tuple<rtcx::library, rtcx::blob> link_library_uncached(
   return std::make_tuple(library, std::make_shared<rtcx::blob_t>(std::move(blob)));
 }
 
+kernel get_lto_linked_kernel(std::string const& name,
+                             std::string const& source_identity,
+                             std::function<std::vector<rtcx::blob>()> const& get_fragments)
+{
+  CUDF_FUNC_RANGE();
+
+  auto& ctx               = cudf::get_context();
+  auto& cache             = ctx.rtcx_cache();
+  auto& device_properties = ctx.get_device_properties();
+  auto spec               = std::format(R"***(cuLibrarySourceLTOv1
+name={}
+binary_type=CUBIN
+cuda_runtime={}
+cuda_driver={}
+arch={}
+lto_arch={}
+bundle={}
+)***",
+                          name,
+                          device_properties.runtime_version,
+                          device_properties.driver_version,
+                          device_properties.compute_capability,
+                          LTO_ARCHITECTURE,
+                          ctx.jit_bundle().get_hash());
+
+  XXH3_state_t state;
+  XXH3_INITSTATE(&state);
+  XXH3_128bits_reset(&state);
+  hash(&state, spec);
+  hash(&state, source_identity);
+  auto digest = XXH3_128bits_digest(&state);
+  auto key    = rtcx::hash128{digest.high64, digest.low64};
+
+  // Resolving fragments before this lookup would compile IR even when a linked library is cached.
+  auto compile = [&] {
+    auto blobs = get_fragments();
+    std::vector<rtcx::memory_fragment> fragments;
+    fragments.reserve(blobs.size());
+    for (auto const& blob : blobs) {
+      fragments.push_back(
+        {.data = blob->view(), .type = rtcx::binary_type::LTO_IR, .name = nullptr});
+    }
+    return link_library_uncached(name.c_str(), {}, 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")};
+}
+
 kernel get_lto_linked_kernel(std::string const& name,
                              std::span<rtcx::file_fragment const> file_fragments,
                              std::span<rtcx::memory_fragment const> memory_fragments)

cpp/src/jit/cache.hpp

@@ -10,6 +10,8 @@
 
 #include <rtcx/rtcx.hpp>
 
+#include <functional>
+
 namespace CUDF_EXPORT cudf {
 
 struct [[nodiscard]] jit_bundle_t {
@@ -99,22 +101,36 @@ kernel get_kernel(std::string const& name,
  * @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)
+ * @param name_expressions Global kernel expressions that must be retained in the fragment
  */
 rtcx::blob get_kernel_fragment(std::string const& name,
                                std::string const& source_file_id,
                                std::span<char const* const> header_include_names,
                                std::span<char const* const> headers,
-                               std::string const& kernel_instance);
+                               std::string const& kernel_instance,
+                               std::span<char const* const> name_expressions = {});
 
 /**
  * @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<rtcx::file_fragment const> file_fragments,
                              std::span<rtcx::memory_fragment const> memory_fragments);
 
+/**
+ * @brief Gets an LTO-linked kernel without resolving its fragments on an executable-cache hit.
+ *
+ * The source identity must cover every generated fragment; the embedded bundle covers static
+ * sources and headers. The supplier's owning blobs remain alive until linking completes.
+ * @param name Debug and cache identity for the linked kernel
+ * @param source_identity Complete identity of generated fragment sources
+ * @param get_fragments Supplier of LTO IR blobs, invoked only on an executable-cache miss
+ */
+kernel get_lto_linked_kernel(std::string const& name,
+                             std::string const& source_identity,
+                             std::function<std::vector<rtcx::blob>()> const& get_fragments);
+
 }  // namespace CUDF_EXPORT cudf

cpp/src/jit/column_accessor.cuh

@@ -65,18 +65,6 @@ struct column_accessor {
     return column(cols).template nullable_element<element_type>(map_index(row));
   }
 
-  static __device__ void set_null_mask_word(auto const* __restrict__ cols,
-                                            size_type word_index,
-                                            bitmask_type word)
-    requires(!as_scalar)
-  {
-    auto* mask = column(cols).null_mask();
-
-    if (mask == nullptr) { return; }
-
-    mask[word_index] = word;
-  }
-
   static __device__ void assign(auto const* __restrict__ cols, size_type row, element_type value)
     requires(!as_scalar)
   {

cpp/src/jit/sync.cuh

@@ -1,10 +1,9 @@
 /*
- * 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 <cudf/column/column_device_view_base.cuh>
 #include <cudf/detail/utilities/cuda.cuh>
 #include <cudf/types.hpp>
 
@@ -28,18 +27,26 @@ __device__ inline bool warp_elect(unsigned int mask)
 
 namespace jit {
 
-template <typename Out>
-__device__ void warp_compact_validity(unsigned int active_mask,
-                                      mutable_column_device_view_core const* outcols,
-                                      size_type row,
-                                      bool is_valid)
+__device__ inline void warp_compact_validity(unsigned int active_mask,
+                                             bitmask_type* null_mask,
+                                             size_type row,
+                                             bool is_valid)
 {
   auto null_word = __ballot_sync(active_mask, is_valid);
   // use warp-elect to make sure we only issue one memory transaction per warp
-  if (warp_elect(active_mask)) {
-    Out::set_null_mask_word(outcols, row / cudf::detail::warp_size, null_word);
+  if (null_mask != nullptr && warp_elect(active_mask)) {
+    null_mask[row / cudf::detail::warp_size] = null_word;
   }
 }
 
+template <typename Out, typename Column>
+__device__ void warp_compact_validity(unsigned int active_mask,
+                                      Column const* outcols,
+                                      size_type row,
+                                      bool is_valid)
+{
+  warp_compact_validity(active_mask, Out::column(outcols).null_mask(), row, is_valid);
+}
+
 }  // namespace jit
 }  // namespace cudf

cpp/src/transform/jit/ast_runtime_descriptor.hpp

@@ -0,0 +1,26 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <cuda/std/cstdint>
+
+namespace cudf::jit::ast_runtime {
+
+// Keep this internal ABI independent of column types and schema width so one row-loop fragment
+// can serve every supported fixed-width AST expression.
+struct input_descriptor {
+  void const* data;
+  cuda::std::uint32_t const* null_mask;
+  cuda::std::int32_t offset;
+  bool scalar;
+};
+
+struct output_descriptor {
+  void* data;
+  cuda::std::uint32_t* null_mask;
+};
+
+}  // namespace cudf::jit::ast_runtime

cpp/src/transform/jit/ast_runtime_driver.cuh

@@ -0,0 +1,42 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <cudf/errc.hpp>
+
+#include <transform/jit/ast_runtime_descriptor.hpp>
+
+namespace cudf::jit::ast_runtime {
+
+extern "C" __device__ cudf::errc transform_row(cuda::std::int32_t row,
+                                               cuda::std::uint32_t active_mask,
+                                               input_descriptor const* inputs,
+                                               output_descriptor const* outputs);
+
+}  // namespace cudf::jit::ast_runtime
+
+extern "C" __global__ void cudf_kernel_entry(
+  cuda::std::int32_t row_size,
+  cuda::std::uint32_t const* stencil,
+  void*,
+  cudf::jit::ast_runtime::input_descriptor const* inputs,
+  cudf::jit::ast_runtime::output_descriptor const* outputs,
+  cuda::std::int32_t* max_error)
+{
+  auto const start  = static_cast<cuda::std::int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  auto const stride = static_cast<cuda::std::int64_t>(gridDim.x) * blockDim.x;
+  // The final warp must reach the active-lane ballot together, even when its last rows are absent.
+  auto const padded_rows = (static_cast<cuda::std::int64_t>(row_size) + 31) / 32 * 32;
+  auto thread_error      = cudf::errc::SUCCESS;
+  for (auto row = start; row < padded_rows; row += stride) {
+    auto const active =
+      row < row_size && (stencil == nullptr || (stencil[row / 32] & (1u << (row % 32))));
+    auto const active_mask = __ballot_sync(0xffffffffu, active);
+    if (!active) { continue; }
+    auto const error = cudf::jit::ast_runtime::transform_row(
+      static_cast<cuda::std::int32_t>(row), active_mask, inputs, outputs);
+    if (static_cast<int>(error) > static_cast<int>(thread_error)) { thread_error = error; }
+  }
+  if (thread_error != cudf::errc::SUCCESS) { atomicMax(max_error, static_cast<int>(thread_error)); }
+}

cpp/src/transform/jit/ast_runtime_udf.cuh

@@ -0,0 +1,66 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <cudf/detail/row_ir/opcode.hpp>
+#include <cudf/errc.hpp>
+#include <cudf/utilities/bit.hpp>
+
+#include <cuda/std/optional>
+
+#include <jit/sync.cuh>
+#include <transform/jit/ast_runtime_descriptor.hpp>
+
+// Keep generated operations outside the reusable precompiled-header prefix.
+// https://docs.nvidia.com/cuda/archive/12.9.0/nvrtc/index.html#determining-the-header-stop-point
+#pragma nv_hdrstop
+
+namespace cudf::jit::ast_runtime {
+
+template <typename T, bool Nullable>
+__device__ auto load(input_descriptor const& column, cuda::std::int32_t row)
+{
+  auto const index = column.offset + (column.scalar ? 0 : row);
+  if constexpr (Nullable) {
+    if (!cudf::bit_value_or(column.null_mask, index, true)) { return cuda::std::optional<T>{}; }
+    return cuda::std::optional<T>{static_cast<T const*>(column.data)[index]};
+  } else {
+    return static_cast<T const*>(column.data)[index];
+  }
+}
+
+template <typename T>
+__device__ void store(output_descriptor const& column,
+                      cuda::std::int32_t row,
+                      cuda::std::uint32_t,
+                      T value)
+{
+  static_cast<T*>(column.data)[row] = value;
+}
+
+template <typename T>
+__device__ void store(output_descriptor const& column,
+                      cuda::std::int32_t row,
+                      cuda::std::uint32_t active_mask,
+                      cuda::std::optional<T> value)
+{
+  // Null positions have no defined value, but maskless outputs still need initialized data.
+  if (value.has_value() || column.null_mask == nullptr) {
+    static_cast<T*>(column.data)[row] = value.value_or(T{});
+  }
+  // All active lanes must participate, including lanes whose value store was skipped.
+  cudf::jit::warp_compact_validity(active_mask, column.null_mask, row, value.has_value());
+}
+
+#include <cudf/detail/operation_udf.cuh>
+
+}  // namespace cudf::jit::ast_runtime
+
+// Registering a global entry makes retention of this otherwise device-only fragment explicit
+// across NVRTC versions before the driver is linked. This anchor is never launched.
+extern "C" __global__ void cudf_ast_runtime_entry()
+{
+  auto const operation = &cudf::jit::ast_runtime::transform_row;
+  asm volatile("" : : "l"(operation));
+}