Skip to content

[FEA] LTO IR Support (4) - Implement LTO Transform Kernels - #22680

Merged
rapids-bot[bot] merged 384 commits into
NVIDIA:mainfrom
lamarrr:lto-ir-4--lto-kernels
Jul 3, 2026
Merged

[FEA] LTO IR Support (4) - Implement LTO Transform Kernels#22680
rapids-bot[bot] merged 384 commits into
NVIDIA:mainfrom
lamarrr:lto-ir-4--lto-kernels

Conversation

@lamarrr

@lamarrr lamarrr commented May 27, 2026

Copy link
Copy Markdown
Contributor

Description

This PR adds LTO-based transforms to libcudf by introducing a new cudf::transform_lto API that accepts LTO-IR or FATBIN UDF binaries and executes them through the existing transform pipeline. The API supports typed transform inputs/outputs, null-awareness, and optional user data.

The change enables libcudf transforms to execute precompiled device UDF fragments instead of relying only on source/PTX-style runtime compilation. This creates a path for lower-overhead, link-time-optimized transform kernels while preserving the existing transform abstraction around input columns, scalar inputs, output specifications, and null policy.

image

The AOT-compiled transform UDF has a similar ABI signature as NUMBA-CUDA UDFs (https://nvidia.github.io/numba-cuda/user/cuda_ffi.html):

extern "C" __device__ int transform(Output * ... outputs, Inputs... inputs);

The integer return is used for signaling errors (non-zero values) and may be discarded or propagated by the implementation.

This pull request also replaces the software SHA256 implementation with a SIMD-accelerated xxHash implementation (XXH3-128, ported from #22920). SHA256 was initially chosen as the obvious choice for a collision-free hash, and it was implemented to keep dependencies minimal.

For a hot JIT cache, the previous cryptographic SHA256 hash took ~88.1% (~4727ns) of the get_kernel function execution time.
With the new non-cryptographic SIMD-accelerated hash, the hash time is reduced by 16.33x to 288ns, ~31.2% of the get_kernel function execution time.

This pull request also:

  • Hoists the cudaGetDeviceProperties values into the cudf context object; each call takes ~1.8ms per-call (see: [BUG] Repeated calls to cudaGetDeviceProperties causes launch overhead in JIT kernels #23074)
  • Refactors librtcx's CMake functions to comply with RAPIDS' naming standards
  • Makes the nvrtc-related flags be dispatched by the nvrtcVersion and not cudaGetRuntimeVersion. They can be different
  • Removes discard_errors specialization of the transform kernel; this helps reduce the number of kernel instantiations that need to be pre-compiled, and also reduces kernel variance

Closes #19578 & #23074

Benchmarks

Throughput benchmarks are provided in #22680 (comment), and compilation-time benchmarks are provided in #22680 (comment)

Checklist

  • I am familiar with the Contributing Guidelines.
  • New or existing tests cover these changes.
  • The documentation is up to date with these changes.

lamarrr and others added 30 commits April 1, 2026 18:59
Co-authored-by: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com>
Co-authored-by: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com>
@lamarrr

lamarrr commented Jul 1, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 9cfd77f

@lamarrr

lamarrr commented Jul 1, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 19e47b9

@lamarrr

lamarrr commented Jul 2, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 543e226

@lamarrr

lamarrr commented Jul 3, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test f730f36

lamarrr added 4 commits July 3, 2026 09:16
… checks

- Removed the `fallible` enum and related checks from the transform functions and kernels.
- Updated the `multi_transform` and `transform_lto` functions to no longer accept a fallibility parameter.
- Adjusted the kernel implementation to handle errors without the fallibility flag.
- Modified the CMake configuration to reflect changes in the transform kernel instantiation.
- Updated tests to remove fallibility assertions and ensure compatibility with the new implementation.
@lamarrr

lamarrr commented Jul 3, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test a2034e2

@lamarrr

lamarrr commented Jul 3, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 5a85802

@lamarrr

lamarrr commented Jul 3, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test f7d2883

@lamarrr

lamarrr commented Jul 3, 2026

Copy link
Copy Markdown
Contributor Author

/merge

@rapids-bot
rapids-bot Bot merged commit 3652e80 into NVIDIA:main Jul 3, 2026
152 checks passed
rapids-bot Bot pushed a commit that referenced this pull request Jul 6, 2026
#23106)

Since #22680, a from-scratch `./build.sh` with the **Unix Makefiles** generator fails with:

```
gmake[2]: *** No rule to make target 'CMakeFiles/cudf_fragments_transform_kernel_20.dir/src/transform/jit/kernel.fatbin', needed by 'rtcx_embed/cudf_fragments.hpp'.  Stop.
```

The `rtcx_embed()` custom command depends on the fragment object libraries only via `$<TARGET_OBJECTS:...>` generator expressions. With the Makefiles generator this produces file-level prerequisites with no build rule and no target-level ordering — `cudf_fragments.dir/all` only depended on `cudf_fragments__jit_embed_run`, not on the 21 `cudf_fragments_transform_kernel_N` object libraries — so a parallel make races ahead of the fatbin compilations and dies. Ninja resolves the same dependency through its global build graph, which is why CI never hit this.

The fix records the object-library target names in a new `EMBED_TARGET_DEP_NAMES` property alongside the existing `$<TARGET_OBJECTS:...>` genexes, and passes those names to the custom command `DEPENDS`. Naming a real target there makes CMake emit a proper target-level dependency, so the Makefiles generator builds all fragment fatbins before running the embed step.

Verified locally with CMake 4.3.4 + Unix Makefiles: after this change `Makefile2` contains target-level deps from `cudf_fragments.dir/all` on every `cudf_fragments_transform_kernel_N.dir/all`, and a clean-state `-j32` build of `cudf_fragments` (and a full `./build.sh`) succeeds where it previously failed.

Authors:
  - GALI PREM SAGAR (https://github.com/galipremsagar)

Approvers:
  - Bradley Dice (https://github.com/bdice)
  - Kyle Edwards (https://github.com/KyleFromNVIDIA)

URL: #23106
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CMake CMake build issue feature request New feature or request libcudf Affects libcudf (C++/CUDA) code. non-breaking Non-breaking change

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[FEA] Evaluate LTO-IR for JIT-Compilation

7 participants