Add experimental Project AST JIT for integral add and multiply [fast-ut] [databricks] - #15312
Add experimental Project AST JIT for integral add and multiply [fast-ut] [databricks]#15312thirtiseven wants to merge 27 commits into
Conversation
|
@greptile full review |
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Greptile SummaryThis PR introduces an experimental, opt-in
Confidence Score: 5/5
Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A["GpuProjectExecMeta.convertToGpu()"] --> B{JIT enabled?}
B -- yes --> C["wrapProjectExpressions\n(pre-mark JIT-eligible)"]
B -- no --> D["gpuExprs (unchanged)"]
C --> E{Legacy AST enabled?}
D --> E
E -- yes --> F["GpuProjectAstExpression.wrap()\n(skip if already JIT-wrapped)"]
E -- no --> G["projectList (jit-marked only)"]
F --> G
G --> H["GpuProjectExec.doExecuteColumnar()"]
H --> I["bindGpuProjectReferencesTiered()\n(Project-specific binder)"]
I --> J["buildExprTiers()"]
J --> J1["unwrap() — strip JIT/AST markers"]
J1 --> J2["CSE / GpuEquivalentExpressions"]
J2 --> J3["getExprTiers()"]
J3 --> J4{hasAstOutputs?}
J4 -- yes --> J5["rewrapAstTiers()\n(legacy AST markers)"]
J4 -- no --> J6["tiers unchanged"]
J5 --> J7
J6 --> J7
J7{JIT enabled?} -- yes --> J8["wrapTierExpression()\n(JIT overrides legacy AST for\nfully supported roots)"]
J7 -- no --> J9["final tiers (no JIT)"]
J8 --> J9
J9 --> K["GpuTieredProject.projectAndCloseWithRetrySingleBatch()"]
K --> L{Expression type?}
L -- GpuAstJitExpression --> M["computeColumnJit(table)"]
L -- GpuProjectAstExpression --> N["computeColumn(table)"]
L -- other --> O["columnarEval(batch)"]
Reviews (8): Last reviewed commit: "add nvtx docs" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
|
Java binding is in 26.10, waiting for main branch to switch... |
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
|
NOTE: release/26.08 has been created from main. Please retarget your PR to release/26.08 if it should be included in the release. |
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
…ut] (#15377) Related to #8954. ### Description The motivation for this PR is to prepare for AST JIT work. Legacy AST Project execution is currently all-or-nothing: if any top-level Project output cannot be represented by the legacy AST backend, eligible sibling outputs cannot use AST either. This change selects the legacy AST backend independently for each top-level Project expression. Supported fixed-width outputs use AST while unsupported outputs in the same Project continue through the regular GPU expression path. Top-level null literals also remain on the regular projection path so their outputs can reuse its cached null vector. It uses the existing `spark.rapids.sql.projectAstEnabled` configuration and does not change result semantics or add a user-facing configuration. This is independent of the experimental AST JIT work in #15312. It refactors the existing legacy AST path and does not depend on AST JIT, LTO, or precompiled fragments. This also changes physical-plan rendering: AST projections now appear under GpuProject as `AST(...) AS x` instead of using a dedicated `GpuProjectAstExec` node. This does not change result semantics, APIs, or configuration. Implementation details: - Represent each eligible output with a lightweight `GpuProjectAstExpression`. - Keep top-level null literals on the regular projection path so multiple null outputs can reuse its cached null vector; non-null literals remain AST-eligible. - Share one cuDF input `Table` across all AST outputs evaluated for the same tier and batch. - Preserve TieredProject common-subexpression elimination for duplicate outputs and shared subtrees. - Route fused higher-order-function projections through the same expression evaluator so AST outputs retain the shared input Table. - Close compiled legacy AST expressions at task completion, including retry execution paths. Performance testing: - Benchmark script: [project_ast_per_expression_perf.scala](https://gist.github.com/thirtiseven/f20355b1e8c32444fb11a39fa3717545) - Environment: Spark 3.5.2 with `local[4]`, plugin `26.08.0-SNAPSHOT`, and 2 NVIDIA RTX 5880 Ada Generation GPUs. - Workload: 20 million Parquet rows in 8 partitions, 84 projected outputs, one warmup, and five measured iterations. A global GPU aggregate materializes every projected output. - `LIBCUDF_JIT_ENABLED=0` isolates the legacy AST backend. AST off/on order and case order are reversed on alternating iterations. - Speedup is AST off / AST on; values above 1 mean AST is faster. The Project metric is `opTimeLegacy`, accumulated across tasks, so it is not directly comparable to E2E wall-clock time. | Case | Expressions | Coverage / CSE shape | AST tiers (off/on) | AST outputs (off/on) | AST off/on E2E median (ms) | E2E speedup | AST off/on Project op median (ms) | Project op speedup | |---|---|---|---|---:|---:|---:|---:|---:| | `broad_unique_ast` | 84 unique AST-compatible | 42 AST operator families | `[0]` / `[84]` | 0 / 84 | 670.249 / 589.848 | 1.136x | 684.796 / 356.251 | 1.922x | | `whole_output_duplicates` | 42 AST expressions, each projected twice | whole-output CSE | `[0,0]` / `[42,0]` | 0 / 84 | 601.863 / 526.478 | 1.143x | 355.906 / 155.668 | 2.286x | | `cheap_partial_cse` | 84 AST-compatible | one shared add | `[0,0]` / `[1,84]` | 0 / 84 | 516.149 / 467.530 | 1.104x | 214.336 / 189.046 | 1.134x | | `expensive_partial_cse` | 84 AST-compatible | one shared transcendental subtree | `[0,0]` / `[1,84]` | 0 / 84 | 476.142 / 440.932 | 1.080x | 239.745 / 175.069 | 1.369x | | `mixed_half` | 42 AST-compatible + 42 regular GPU | mixed execution | `[0]` / `[42]` | 0 / 42 | 793.670 / 712.257 | 1.114x | 1167.327 / 868.695 | 1.344x | No median regression was observed in the five main benchmark cases in either E2E or Project op time. #### Top-level literal routing follow-up A targeted follow-up compared top-level literals on the regular and AST paths. It used 20 million rows, 84 Project outputs, two warmups, and 15 alternating measured iterations. The benchmark directly executed and consumed the columnar `GpuProject`, avoiding Catalyst constant folding and unrelated aggregate, shuffle, or ColumnarToRow work. Speedup is regular / AST; values above 1 mean AST is faster. | Case | Expressions | Regular/AST E2E median (ms) | E2E speedup | Regular/AST Project op median (ms) | Project op speedup | |---|---|---:|---:|---:|---:| | `literal_only` | 84 unique non-null long literals | 61.418 / 60.952 | 1.008x | 117.137 / 104.202 | 1.124x | | `mixed_ast_literals` | 42 AST-compatible expressions + 42 unique non-null long literals | 231.422 / 230.816 | 1.003x | 241.066 / 239.654 | 1.006x | | `null_literal_only` | 84 duplicate double null literals | 33.257 / 137.528 | 0.242x | 8.224 / 407.906 | 0.020x | Non-null literals were neutral in E2E time in both the literal-only and mixed workloads, so there is no evidence for excluding all top-level literals from AST. Null literals were different: the regular projection was 4.1x faster E2E and 49.6x faster in Project op time because it can reuse its cached null vector across outputs, whereas per-expression AST evaluates the null outputs independently. Based on these results, this PR keeps only top-level null literals on the regular projection path and leaves non-null literals AST-eligible. ### TPC-H/TPC-DS coverage We also performed physical-plan sweeps over stock TPC-H/NDS-H and TPC-DS workloads. - TPC-H/NDS-H with the standard Decimal schema produced no legacy Project AST expressions across 24 physical plans. - TPC-DS produced only three AST expressions across 102 physical plans. Two operate on a small filtered date dimension, and the remaining expression is in a post-aggregation Project. These workloads therefore spend nearly all of their time in unrelated scans, joins, and aggregations and do not provide a meaningful performance signal for this change. The performance results above use purpose-built AST-only and mixed-expression Projects so that the affected execution path is exercised directly, while still materializing every projected expression. ### Checklists Documentation - [ ] Updated for new or modified user-facing features or behaviors - [x] No user-facing change Testing - [x] Added or modified tests to cover new code paths - [ ] Covered by existing tests (Please provide the names of the existing tests in the PR description.) - [ ] Not required Performance - [x] Tests ran and results are added in the PR description - [ ] Issue filed with a link in the PR description - [ ] Not required --------- Signed-off-by: Haoyang Li <haoyangl@nvidia.com> Co-authored-by: Igor Peshansky <7594381+igorpeshansky@users.noreply.github.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
thirtiseven
changed the title
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
|
@greptile full review |
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
There was a problem hiding this comment.
Pull request overview
This PR introduces an experimental cuDF AST JIT execution path for GpuProjectExec, selectively wrapping fully-supported projection tiers (currently limited to non-ANSI INT/BIGINT add and multiply) with a new GpuAstJitExpression. The feature is gated behind a new internal config (spark.rapids.sql.projectAstJitEnabled) and integrates with existing tiered projection/CSE and retry semantics, including manual build-side Project evaluation in broadcast joins.
Changes:
- Add
GpuAstJitExpressionand AST-JIT eligibility plumbing (supportsAstJit/ operator tagging) to enable tier-level wrapping for supported Project expressions. - Introduce a new internal config to enable Project AST JIT, and route Project-specific binding paths to allow JIT selection without affecting generic tiered binders (e.g., Filter).
- Add/extend Scala + Python integration tests, and update broadcast-join build-side projection to use the Project-specific tiered binding and retry-aware projection execution.
Reviewed changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tests/src/test/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinRetrySuite.scala | New suite validating build-side shared JIT tier retry behavior under injected GPU OOM. |
| tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala | New unit tests covering wrapping rules, CSE exposure, precedence vs legacy AST, and binder isolation. |
| tests/src/test/scala/com/nvidia/spark/rapids/GpuArrayHofFusionSuite.scala | Update HOF fusion test to validate shared input table behavior across legacy AST + JIT AST expressions. |
| sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinExecBase.scala | Switch build-side post-projection to Project-specific tiered binding and retry-aware single-batch projection. |
| sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastHashJoinExecBase.scala | Same as above for chained build-side Projects in BHJ plan extraction. |
| sql-plugin/src/main/scala/org/apache/spark/sql/rapids/arithmetic.scala | Mark GpuAdd/GpuMultiply as AST-JIT-supported operators for non-ANSI INT/BIGINT. |
| sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala | Add internal config spark.rapids.sql.projectAstJitEnabled and isProjectAstJitEnabled accessor. |
| sql-plugin/src/main/scala/com/nvidia/spark/rapids/literals.scala | Mark GpuLiteral as AST-JIT-compatible leaf (enabling literals within supported JIT expressions). |
| sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuProjectAstExpression.scala | Generalize backend-marker handling across tiers and integrate optional JIT wrapping into tier construction. |
| sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuExpressions.scala | Add AST-JIT support predicates (supportsAstJit, containsAstJitOperator) to GpuExpression. |
| sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuBoundAttribute.scala | Add Project-specific tiered binder entrypoints and mark GpuBoundReference AST-JIT-compatible. |
| sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala | New JIT wrapper expression using CompiledExpression.computeColumnJit and task-completion cleanup. |
| sql-plugin/src/main/scala/com/nvidia/spark/rapids/basicPhysicalOperators.scala | Ensure GpuProjectExec uses Project-specific tiered binding; extend AST extraction to include JIT. |
| integration_tests/src/main/python/ast_test.py | Add integration coverage validating correctness and plan selection for JIT-enabled Projects. |
|
Before delving into this in detail, I'll try take this for a test drive. |
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
| conf: SQLConf): Unit = { | ||
| val explain = RapidsConf.EXPLAIN.get(conf) | ||
| if (!explain.equalsIgnoreCase("NONE")) { | ||
| val explanation = GpuAstJitExpression.explainFinalSelections( |
There was a problem hiding this comment.
[Really optional] Unless the user sets "explain=ALL", none of the post-CSE JIT nodes would appear in the log, and thus it would be hard for them to know that the JIT is actually working. It would be useful for the explainer to show counts of the different kinds of nodes (e.g., "20 JIT nodes, 10 legacy AST nodes, 15 regular project") along with the errors/rejected nodes, either unconditionally or with a new "explain=STATS" setting. Definitely out of this PR's scope, so maybe just file a feature request to track?
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
revans2
left a comment
revans2
left a comment
There was a problem hiding this comment.
Could we run this on databricks too? Just to be sure that this works in that environment too.
It would also be nice to have a bridge + JIT integration test, just to be sure that it is all working as expected if we enable all of these together.
| * @param conf SQL configuration | ||
| * @param metrics Metrics to inject into the bound expressions | ||
| */ | ||
| def bindGpuProjectReferencesTiered[A <: Expression]( |
There was a problem hiding this comment.
I personally don't like the name. This just like bindGpuReferencesTiered returns a GpuTieredProject. All of these are binding for a "Project" operation. Adding Project to the name does not distinguish it from the other in any meaningful way by the name alone.
Why do we need to distinguish between these two APIs? If we can get a speedup on a Regular GpuProjectExec why do we not want to do it also for pre-processing on aggregations, expand and filter operations? Not it looks like join already does use this in some cases.
If there are good reasons to keep them separate, can we rename this or modify the original API to take in the JIT/AST enable param? To me that is much cleaner and less confusing.
| private def finalBackend(expression: Expression): String = { | ||
| GpuProjectAstExpressionBase.extractTopLevel(expression) match { | ||
| case Some(_: GpuAstJitExpression) => "Project AST JIT" | ||
| case Some(_: GpuProjectAstExpression) => "legacy Project AST" |
There was a problem hiding this comment.
nit: I don't think this explains it very well, and I am not sure a customer is going to understand. If this is not for a customer to follow, then can we make sure it is documented and drop the Project from it? Something like "AST JIT" and "AST Interpreted" feel better to me.
| tiers | ||
| } | ||
| if (enableProjectAstJit) { | ||
| // Project binding selects JIT after CSE so newly exposed tiers are eligible. |
There was a problem hiding this comment.
I would like a follow-on issue to try to enable AST and JIT AST more generically. AST and JIT AST have different algorithms to search through a Project operation and enable their respective backends, which can
lead to cases where whether an eligible subexpression uses AST/JIT depends on whether CSE happens to materialize it into a separate tier.
Also, neither backend can select an AST/JIT-compatible GPU subtree when it is below a GPU/CPU bridge. For example, with ANSI off and all values typed as longs, my_udf(a + b) * c can run as a GPU add feeding a
CPU bridge, followed by a GPU multiply. However, a + b is not selected for AST or JIT unless it is independently shared and CSE happens to materialize it into a separate tier. Adding a + b as another Project
output can therefore change whether it uses AST/JIT.
If we combine the AST and JIT labeling, especially if we do it as a two pass like operation similar to the GPU/CPU bridge, I think we can do a cost based optimization to reduce data movement and decide if something should be JIT or not, especially if things are intermixed.
| def selfIsAstJitOperator: Boolean = false | ||
|
|
||
| /** Whether this node and its complete expression subtree support AST JIT. */ | ||
| final def supportsAstJit: Boolean = selfSupportsAstJit && children.forall { |
There was a problem hiding this comment.
I get that you are being conservative right now. But I am concerned that this is adding in a lot of code that we are going to have to rip out when we actually do it right. Currently If I have an expression tree like A + B + C that can all but JIT, then we do the JIT. But if I have (A + B + C) / D why would we not want to do the JIT for A + B + C still? This is why I think a two pass optimization is a much better path for this. First pass would be to go through each expression and see if (it by itself) could be JIT or AST or neither. The second pass would be to do cost reduction estimation. For the Bridge it is all about data movement. Here it would be about data materialization cost and possibly reduce JIT vs execution costs. I know that will take a lot of experimentation to understand these costs, but having the framework in place is much better than just do it if we can with all or nothing.
There was a problem hiding this comment.
Totally agreed. If we could partially enable the AST JIT inside the expression, that would be better. Also, since the multi-output and CSE NVIDIA/cudf#23621 have been merged, we might need to adjust some design decisions here. I'm converting this to a draft now to test more solutions...
thirtiseven
changed the title
igorpeshansky
left a comment
igorpeshansky
left a comment
There was a problem hiding this comment.
Also agree with @revans2's comments…
| } | ||
| if (completed) { | ||
| throw new IllegalStateException( | ||
| s"Task completed while registering the $backendName cleanup callback") |
There was a problem hiding this comment.
I wonder if the backend name is helpful here, or if you could make this message generic (e.g., "Task completed while registering compiled expression cleanup callback") and rely on the stack trace for context… This would avoid an abstract member and both overrides.
If you still want to vary the message by subclass, can you reuse nodeName instead?
| withResource(GpuProjectAstExpression.tableFromBatch(batch)) { table => | ||
| computeColumn(table) | ||
| /** Extracts a legacy Project AST wrapper after unwrapping any top-level aliases. */ | ||
| private[rapids] def extractTopLevel(expression: Expression): Option[GpuProjectAstExpression] = |
There was a problem hiding this comment.
This is currently sitting between wrap and rewrap. Want to move it to the top of the object (and in GpuAstJitExpression, for symmetry)?
| } | ||
| compiledExpression | ||
| } | ||
|
|
There was a problem hiding this comment.
Nit: stray blank line (also at the start of this object)…
|
@revans2 @igorpeshansky Another thing to note is that the cuDF team would like some feedback on the row IR integration, which enables precompiled fragments and LTO linking and can significantly reduce cold run time. Do you think the cold run performance results of this PR are a good reason to implement row IR, or what else can we do to evaluate it? |
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
If this is PR is to explore/understand what is possible, then it would be interesting to do to see what the results are. If you think this PR is going to provide real performance improvements/enhancements, then we need to do something with JIT as effectively all of the operations we do are "cold" so the cold JIT times are very concerning, and we would need/want some heuristics to understand clearly when JIT wins and when it does not. (or we have to spend a lot of time figuring out how to start pre-compiling things very early so hopefully it is ready by the time we need it. |
Contributes to #10640
Description
This PR adds an experimental integration for using cuDF AST JIT in
GpuProjectExec. It builds on the per-expression legacy AST projection infrastructure merged in #15377 and groups compatible Project expressions into multi-output AST JIT execution waves.Operator coverage remains intentionally limited to non-ANSI
IntegerTypeandLongTypeaddition and multiplication. The internalspark.rapids.sql.projectAstJitEnabledconfiguration remains disabled by default.Planning and execution
AST JIT planning is a Project-specific, two-stage pass:
This allows maximal compatible children below an unsupported root to run in an earlier JIT wave, and it allows a later JIT root to consume a regular-GPU wave. Regular-GPU CSE and tiering remain inside regular waves. Same-wave JIT roots stay in one physical tier, so shared JIT-compatible subexpressions remain visible to cuDF rather than being materialized by plugin CSE.
Each JIT root is compiled explicitly with
AstExpression.compileJit. When a wave has multiple JIT roots,GpuProjectExecpasses the resulting expressions toCompiledExpression.computeTableJitin one cuDF call, enabling cuDF Row IR CSE across outputs. A single JIT root, or disabling the internalspark.rapids.sql.projectAstJitMultiOutputEnabledA/B switch, evaluates the same JIT-compiled representation through the unifiedCompiledExpression.computeColumnAPI. The multi-output switch defaults to true, but it has no effect unless the primary experimental AST JIT configuration is enabled.The current policy selects JIT whenever a compatible wave is available; this PR does not add a runtime profitability cost model. Non-deterministic expressions are not moved across backend boundaries.
AST JIT still takes precedence over legacy Project AST for JIT-compatible expressions. Eligible remaining expressions can use legacy AST when both experimental backends are enabled. The binding entry point remains Project-specific so generic tiered binders used by other operators do not enable Project AST JIT implicitly.
Retry and resource lifecycle
Compiled expressions participate in the existing
Retryablecheckpoint/restore lifecycle and are closed at task completion. The lazy JIT compilation checkpoint runs inside the projection retry boundary, including when Project split retry is disabled, so a cold-compilationGpuRetryOOMis retried before evaluation. Multi-output result tables and batches use the existing ARM ownership conventions, including failure cleanup, duplicate expression wrappers, output ordering, and null propagation.The manually evaluated build-side Project paths in BHJ and BNLJ use the Project-specific binder. BNLJ also evaluates its bound tiered Project through the same spillable, OOM-retry-aware helper as BHJ, including when AST JIT is disabled.
cuDF backend status
rapidsai/cudf#23615 added scalar column-view support for AST expressions and JIT execution. rapidsai/cudf#23621 added native multi-output AST JIT execution and Row IR CSE. rapidsai/cudf#23828 adds the backend-specific Java/JNI compilation and multi-output APIs consumed by this PR; its 26.10 snapshot dependency still needs to be published.
The single-output and multi-output Row IR paths do not currently stitch a shipped library of precompiled Row IR operator fragments. NVRTC PCH is process-local, so it cannot remove first-use compilation in a fresh process. The pending rapidsai/cudf#23648 is also outside the local cuDF revision used for these measurements and may affect follow-up performance policy.
Testing
The Scala tests cover:
JIT -> regular GPU -> JITdependency waves;The Python integration coverage compares CPU/GPU results for
INTandBIGINT, shared multi-output subtrees, partial unsupported roots, and mixed JIT/regular/JIT waves.Local validation with the private 26.10 JNI artifact:
mvn clean install -Dbuildver=352 -DskipTestsGpuProjectAstJitSuite: 21/21 passedProjectExprSuite: 12/12 passed, including the multi-output retry testProjectSplitRetrySuite: 19/19 passedGpuBroadcastNestedLoopJoinRetrySuite: 1/1 passedPerformance
Warm throughput
The following is a same-process warm-throughput comparison, not a fresh-process cold-run measurement:
local[4];The modes are regular GPU Project, per-expression AST JIT, and multi-output AST JIT.
Regular/MultiandPer-expression/Multivalues above 1 mean multi-output JIT is faster.Project opis the cumulativeopTimeLegacySQL metric across parallel tasks; it is supporting evidence and is not directly comparable to wall-clock time.single_chainmulti_independentmulti_shared_addmixed_wavessingle_chainhas only one JIT root, so it does not activate multi-output execution and serves as a noise control. Grouping 60 independent roots is approximately neutral versus per-expression JIT at the end-to-end level in this run.When all 60 outputs share
a + b, keeping that subtree visible to cuDF for cross-output CSE and consolidating execution improves E2E by 19.7% over per-expression JIT and 22.5% over regular GPU Project. The cumulative Project-op per-expression/multi ratio is 3.018x. For the three-wave mixed workload, multi-output improves E2E by 10.3% over per-expression JIT and 14.2% over regular GPU Project.Because all modes share one Spark process after warmup, process-local compiler and kernel caches may be shared. These results support the steady-state wave/grouping policy but do not establish a cold-start improvement or a general cost policy for operators beyond the add/multiply scope.
Cold start and amortization
A separate benchmark measures first-use latency in fresh Spark processes after an excluded non-JIT GPU bootstrap. The 20-million-row and 200-million-row runs use the same expression shapes and 60-output workloads. Normal Spark file-split planning was retained, with four and 18 realized scan partitions respectively. Each value is the median of three fresh processes, and cross-mode checksums passed.
Cold Regular/JITandPCH Regular/JITare speedups: values above 1 mean JIT is faster.single_chainmulti_shared_addmixed_wavessingle_chainmulti_shared_addmixed_wavesAt 20 million rows, legacy AST first-use latency was 342.261 ms, 865.344 ms, and 913.246 ms for the three cases, close to regular Project. Empty-cache AST-JIT is 7.1–8.3x slower than regular Project at 20 million rows and remains 2.4–2.7x slower after increasing the input by 10x. PCH removes 25–80% of the cold-to-hot gap depending on expression shape, but does not make first-use JIT faster than regular Project in these workloads.
The estimated cache-miss overhead is
JIT cold - exact disk-cache first. The estimates at 20 million rows are 1,968.051 ms, 6,080.500 ms, and 5,403.616 ms; at 200 million rows they are 2,058.447 ms, 6,035.906 ms, and 5,344.617 ms. Every estimate changes by less than 5% after increasing the row count by 10x, indicating that source-JIT cache-miss overhead is primarily a fixed per-expression-shape cost.Cache-state controls and raw measurements
single_chainmulti_shared_addmixed_wavessingle_chainmulti_shared_addmixed_wavesEvery empty-cache sample uses a new
LIBCUDF_KERNEL_CACHE_PATHwithLIBCUDF_JIT_DISABLE_CUDA_CACHE=1.PCH warmfirst compiles a distinct one-row probe while leaving the target kernel key cold; NVRTC logs confirmed PCH reuse by every target kernel.Exact disk cachestarts a fresh process with the target kernel already present, whileExact in-process hotimmediately repeats the target action in the PCH process. Spark startup, the non-JIT bootstrap, PCH probe, and disk-cache population are excluded from measured target latency.Historical Row IR/LTO POC. An earlier proof of concept precompiled
INT32/INT64ADDandMULfragments and linked them with a thin runtime topology. For a 100K-row depth-4 expression, isolated fresh-process/fresh-cache first-call latency fell from approximately 2.33 seconds with a semantically equivalent source-JIT control to 0.12 seconds with LTO, saving approximately 2.21 seconds (95%, 20.2x). Hot kernel throughput was unchanged: the 100M-row, 30-iteration trimmed means were 4.096 ms for LTO and 4.098 ms for source JIT. This is POC evidence rather than a measurement of the current cudf-spark path; the source control used an outerIDENTITYnode to force source-JIT fallback. The prototype is available onproject-ast-jit-ultimate-perf.Checklists
Documentation
Testing
Performance