Efficient Transforms in cuDF with JIT Compilation

TL;DR

• cuDF uses JIT compilation (NVRTC) to fuse transformations into single kernels, improving throughput and reducing intermediate materialization.
• The 25.08 release adds support for a ternary operator (if-else) and string functions such as find and substring.
• JIT transforms typically reduce the number of kernels and data movements vs precompiled approaches, yielding better cache locality and lower memory bandwidth pressure.
• First-run JIT compilation is about 600 ms per kernel; once cached, loading takes ~3 ms. Pre-population of the kernel cache can shift these costs further in favor of JIT.
• Observed speedups range from 2x to 4x for some string transforms, with scalability advantages that increase with data size. JIT also enables processing larger data before hitting GPU memory limits due to fewer intermediates.

Context and background

RAPIDS cuDF provides a broad set of ETL algorithms for processing data on GPUs. For users coming from pandas, cuDF offers accelerated APIs with zero-code-change through cudf.pandas. For C++ developers and advanced users, cuDF exposes a C++ submodule that operates with non-owning input views and returning owning outputs, which simplifies lifetime and ownership reasoning for GPU data and improves API composability. A key challenge in the C++ model is that some operations materialize many intermediates, causing excessive GPU memory transfers. Kernel fusion—where a single GPU kernel performs multiple computations on the same input data—addresses this by reducing memory traffic and kernel launch overhead. This post explains how JIT compilation brings kernel fusion into the cuDF C++ programming model, delivering higher data processing throughput and more efficient use of GPU memory and compute resources. In typical data processing, expressions are treated as a tree: leaves are columns or scalars, and internal nodes are operators. Scalar expressions often implement a row-wise mapping from inputs to outputs. cuDF supports three evaluation approaches for arithmetic expressions: precompiled, AST (abstract syntax tree), and JIT transform. The precompiled approach calls libcudf APIs operator-by-operator, but materializes intermediates in global memory. The AST approach uses compute_column to evaluate an entire tree with a dedicated kernel and relies on a thread-per-row model; while helpful for kernel fusion, it can be limited in data type and operator support. JIT transform uses NVRTC to compile a custom kernel at runtime that can perform arbitrary transformations. This runtime compilation yields fused kernels tailored to the expression, enabling more efficient resource allocation and better data locality. As of cuDF 25.08, JIT transform adds support for key operators not yet available in AST, including a ternary if-else operator and string functions such as find and substring.

What’s new

The JIT transform feature represents a significant evolution in cuDF’s expression evaluation. In cuDF 25.08, you gain:

• Ternary operator support (if-else) for more complex conditional expressions.
• String function support, including find and substring, expanding the set of operations that can be fused into a single kernel.
• A runtime JIT compilation pathway that creates a fused kernel for each transformation, reducing the need to materialize intermediates and potentially lowering global memory traffic. The main performance advantage of JIT transforms comes from a reduced total kernel count. When a single kernel can perform multiple steps of an expression, intermediates can remain in fast GPU storage (registers and shared memory) rather than being spilled to global memory for subsequent kernels. The result is improved cache locality and more efficient use of GPU registers, which translates into faster runtimes for the same workload. In practice, JIT transforms show performance benefits that scale with data size. Smaller datasets may exhibit overhead from the initial kernel launches and compilation, while larger datasets benefit from the reduced kernel count and better memory bandwidth usage. In published observations, JIT transforms yielded speedups in the range of 2x to 4x for certain string processing cases and 1x to 2x for simpler string transforms, with overall advantages increasing as data size grows. JIT also enables processing roughly 30% larger data before reaching the ~100 GB GPU memory limit on the Grace Hopper Superchip due to fewer materialized intermediates. First-run costs for JIT compilation are notable. If a cached kernel is not found in the kernel cache path (LIBCUDF_KERNEL_CACHE_PATH), compilation takes around 600 ms per kernel. If the kernel is found, loading takes about 3 ms. Once compiled and loaded, subsequent calls to the kernel in the same process incur no additional overhead. If the application defers JIT compilation until runtime, breakeven data sizes for the string_transforms examples are typically in the 1–3 billion rows range when processing batches of ~100 million rows; pre-populating the cache with previously compiled kernels can yield benefits even from the very first million rows. To support experimentation, the rapidsai/cudf repository offers string_transforms examples that illustrate the difference between precompiled and JIT-transformed UDFs. For instance, the extract_email_jit example demonstrates computing an email provider by validating the format of an input string and slicing at the @ and . positions, while the extract_email_precompiled path shows how the same result is achieved with multiple intermediate columns and steps. These examples highlight how JIT can simplify logic and reduce materialization overhead.

Why it matters (impact for developers/enterprises)

• Developer productivity: JIT transforms enable more complex transformations to be fused into single kernels, reducing boilerplate code and avoiding manual optimization of intermediate representations. This simplifies writing efficient UDFs for string processing and other custom transformations.
• Throughput and efficiency: Fused kernels improve data processing throughput by lowering kernel launch overhead and improving cache locality. Fewer intermediates mean reduced GPU memory bandwidth pressure, which is critical for large-scale ETL tasks on GPUs.
• Scalability: The observed speedups grow with data size, making JIT transforms particularly valuable for big data workloads typical in analytics and machine learning preprocessing pipelines.
• Deployment considerations: JIT transforms rely on runtime compilation and a kernel cache. When the cache is cold, there is an upfront cost for compilation, but subsequent runs benefit from fast kernel loading. Teams can pre-populate the cache to maximize early gains.

Technical details or Implementation

• Approaches to expression evaluation:
• Precompiled: Uses libcudf API, supports many types and operators but materializes intermediates in global memory.
• AST: The compute_column API accepts the full expression tree and uses a specialized kernel with GPU thread-per-row parallelism. Useful for kernel fusion but limited by data type and operator coverage.
• JIT transform: Uses NVRTC to compile a custom kernel at runtime for arbitrary transformations. This approach yields fused kernels per expression with optimized resource allocation, avoiding worst-case register usage and reducing global memory traffic.
• Key operators and support: As of cuDF 25.08, JIT transform adds support for a ternary operator (if-else) and string functions such as find and substring, expanding the set of operations that can be fused via JIT.
• String transforms and UDFs: Examples in the rapidsai/cudf repository demonstrate string processing with both precompiled and JIT approaches. The extract_email_jit example uses a raw UDF string to define the transformation, performing steps such as locating the @ and . characters and slicing accordingly. In contrast, the precompiled version materializes positions and additional Boolean columns, increasing memory and compute usage.
• Performance characteristics:
• Fewer kernels generally lead to better cache locality and reduced global memory transfers.
• JIT can deliver faster runtimes due to optimized kernel generation tailored to the expression being evaluated.
• On Grace Hopper hardware with 200 million input rows (about 12.5 GB input), the observed timing and scaling show that JIT benefits become more pronounced as data size grows.
• Kernel cache and warmup: The first execution may incur a warmup time due to NVRTC compilation. The kernel cache path is determined by LIBCUDF_KERNEL_CACHE_PATH. If a kernel is cached, loading is quick (~3 ms). If not cached, compilation may take ~600 ms per kernel. Repeated executions in the same process avoid this overhead.
• Practical adoption notes: If JIT kernels are pre-populated in advance, JIT transforms typically yield benefits even from the first million rows. For testing and deployment, the libcudf example and string_transforms samples provide practical references for comparing precompiled and JIT-based implementations. Users can install prebuilt binaries via the rapidsai-nightly Conda channel and run unit tests and microbenchmarks with the libcudf-tests package to profile expression evaluators.
• References and further reading:
• The NVIDIA blog post describing efficient transforms in cuDF using JIT compilation provides the background and detailed figures referenced here: https://developer.nvidia.com/blog/efficient-transforms-in-cudf-using-jit-compilation

Key takeaways

• JIT compilation enables fused kernels for arbitrary transformations in cuDF, reducing memory traffic and improving throughput.
• The 25.08 release expands JIT capabilities with ternary and string operations, addressing gaps in AST support.
• First-run compilation costs exist, but caching significantly mitigates them; pre-population can yield immediate benefits.
• Real-world performance gains (2x–4x) have been observed for string transforms, with better scaling as data size increases.
• For large ETL workloads, JIT transforms can extend the amount of data processed before hitting GPU memory limits by reducing intermediates.

References

• https://developer.nvidia.com/blog/efficient-transforms-in-cudf-using-jit-compilation