Skip to content

[ML] De-template the RDataLoader engine to cut the JIT cost - #23642

Merged
siliataider merged 10 commits into
root-project:masterfrom
siliataider:rdataloader-opt
Oct 18, 2026
Merged

siliataider merged 10 commits into
root-project:masterfrom
siliataider:rdataloader-opt

Conversation

@siliataider

@siliataider siliataider commented Oct 7, 2026 •

Copy link
Copy Markdown
Contributor

This Pull request:

Changes or fixes:

De-template the RDataLoader by adding a custom RDataFrame action to load the entries as the RDataLoader needs them.

Benchmarking results

7256529 rows x 11 cols, batch 262144, 10 epochs, baseline RSS 0.26 GiB, run on a local machine

  • uproot-based loader, 1 thread
  • RDataLoader eager, 1 thread
RDataLoader eager before #23473 RDataLoader eager after #23473 (master) RDataLoader eager in this PR
setup 0.54 0.43 0.86
ep1 9.02 7.49 4.08
ep2 0.22 0.15 0.14
ep3 0.22 0.04 0.04
ep4 0.22 0.04 0.04
ep5 0.22 0.04 0.04
ep6 0.22 0.04 0.04
ep7 0.22 0.04 0.04
ep8 0.22 0.04 0.04
ep9 0.23 0.04 0.04
ep10 0.22 0.04 0.04
total 11.56 8.38 5.41
peak RSS 2.69G 1.78G 1.65G

Flow

%%{init: {"flowchart": {"nodeSpacing": 20, "rankSpacing": 30}, "themeVariables": {"fontSize": "13px"}}}%%
flowchart LR
    PY["RDataLoader(...)"] --> ENG
    PY -.->|"before (master)"| OLD["template string<br/>→ cling jits per schema"]

    subgraph ML["[ML] compiled engine + loaders"]
        ENG["engine ctor:<br/>column layout once"] --> LOAD["cluster / dataset loader<br/>(cols, vecSizes, vecPadding)"]
    end

    subgraph RDF["[RDF] type-erased custom loader action"]
        CALL["LoadCustomValues"] --> HELP["once per column:<br/>T → AppendAs⟨T⟩<br/>RVec⟨T⟩ → AppendVectorAs⟨T⟩"]
        HELP --> EXEC["each entry: void* per column<br/>→ handler → floats<br/>(vectors truncated / padded)"]
    end

    LOAD --> CALL
    EXEC --> COPY["copy into RFlat2DMatrix<br/>→ batches"]

    classDef old stroke-dasharray:5 5
    class OLD old

Loading

@siliataider siliataider self-assigned this Oct 7, 2026
@siliataider siliataider added the in:ML Everything under ROOT/ML label Oct 7, 2026
@siliataider
siliataider marked this pull request as draft October 7, 2026 10:27
@github-actions

github-actions Bot commented Oct 7, 2026 •

Copy link
Copy Markdown

Test Results

    24 files      24 suites   4d 2h 43m 24s ⏱️
 3 889 tests  3 888 ✅ 0 💤 1 ❌
83 245 runs  83 244 ✅ 0 💤 1 ❌

For more details on these failures, see this check.

Results for commit 05a7e89.

♻️ This comment has been updated with latest results.

@siliataider
siliataider force-pushed the rdataloader-opt branch 7 times, most recently from 480740d to f9f88f6 Compare October 9, 2026 15:08
@siliataider
siliataider marked this pull request as ready for review October 9, 2026 15:24
@siliataider
siliataider requested a review from hageboeck as a code owner October 9, 2026 15:24

@vepadulano vepadulano left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you! Here are some comments from my side.

Comment thread tree/dataframe/inc/ROOT/RDF/CustomLoaderHelpers.hxx
Comment thread tree/dataframe/inc/ROOT/RDF/CustomLoaderHelpers.hxx
Comment thread tree/dataframe/inc/ROOT/RDF/CustomLoaderHelpers.hxx Outdated
Comment thread tree/dataframe/inc/ROOT/RDF/CustomLoaderHelpers.hxx Outdated
Comment thread tree/dataframe/inc/ROOT/RDF/RActionCustomLoader.hxx Outdated
Comment thread tree/dataframe/inc/ROOT/RDF/RActionCustomLoader.hxx Outdated
Comment thread tree/dataframe/inc/ROOT/RDF/RActionCustomLoader.hxx Outdated
Comment thread tree/dataframe/src/RDFCustomLoaderHelpers.cxx Outdated
Comment thread tree/dataframe/src/RDFCustomLoaderHelpers.cxx Outdated
Comment thread tree/dataframe/inc/ROOT/RDF/CustomLoaderHelpers.hxx Outdated
vepadulano and others added 4 commits October 16, 2026 16:16
Add a RAction-derived class, together with an accompanying action helper, that serve the purpose of loading a variable number of columns into a target array of floating-point values. The action and the helper do not require compile-time knowledge of the column types. These classes are meant for internal use only, specifically to be used by the RDataLoader.
Vector columns are truncated or padded to their maximum size.

Assisted-by: Claude Opus 5.5
The padding was truncated to an unsigned integer in eager mode, so negative or fractional values were wrong.

Assisted-by: Claude Opus 5.5
The loaders read the columns with the custom RDF loader action.

Assisted-by: Claude Opus 5.5
Only `RDataLoaderEngine.hxx` and `RFlat2DMatrix.hxx`, which Python uses, stay public, so fewer headers are parsed when the engine is looked up from Python.

Assisted-by: Claude Opus 5.5
Assisted-by: Claude Opus 5.5
`LoadCustomValues` appends to a buffer owned by the caller instead of returning a new one, so the data loader no longer copies the values into its staging buffers and datasets, nor zero-fills them beforehand.

Assisted-by: Claude Opus 5.5

@vepadulano vepadulano left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you, great work! I don't have further comments on the changes. Let's see it validated with a couple of CI runs.

@siliataider
siliataider merged commit 9c9be6d into root-project:master Oct 18, 2026
34 of 36 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants