Skip to content

[tmva][sofie] Improve AD-friendlieness of emitted code for Clad#21896

Merged
guitargeek merged 1 commit intoroot-project:masterfrom
guitargeek:sofie_clad_continued
Apr 14, 2026
Merged

[tmva][sofie] Improve AD-friendlieness of emitted code for Clad#21896
guitargeek merged 1 commit intoroot-project:masterfrom
guitargeek:sofie_clad_continued

Conversation

@guitargeek
Copy link
Copy Markdown
Contributor

This commit refactors SOFIE-generated inference code to enable correct and efficient reverse-mode automatic differentiation with Clad.

Key changes:

  • Introduce explicit primitive operations (Copy, Fill, Relu) in SOFIE_common.hxx and provide corresponding custom pullbacks in CladDerivator.h. This replaces previously inlined loops and allows Clad to generate efficient gradient code without relying on tapes or loop-level differentiation.

  • Update Gemm code generation to emit Copy/Fill instead of manually expanding bias initialization loops. This better exposes the intent and improves AD performance and correctness.

  • Replace manual ReLU loops with a dedicated Relu() call, enabling a custom pullback that avoids tape-based condition tracking.

  • Generate an additional "unoptimized" model variant in the SOFIE test suite (OptimizationLevel::kBasic), and use it for AD tests. This disables memory reuse of intermediate tensors. Opaque memory reuse is safe for inference but breaks source-transformation AD.

  • Improve gradient test diagnostics in SOFIE Clad tests by reporting mismatched indices instead of only checking a global max difference.

With these changes, Clad-generated gradients for SOFIE models are both correct and significantly faster, reaching performance comparable to frameworks such as PyTorch and JAX on the CPU for the tested cases (fully-connected neural networks with multiple layers).

@guitargeek guitargeek self-assigned this Apr 12, 2026
@guitargeek guitargeek requested a review from lmoneta as a code owner April 12, 2026 11:57
@github-actions
Copy link
Copy Markdown

github-actions bot commented Apr 12, 2026

Test Results

    22 files      22 suites   3d 7h 19m 23s ⏱️
 3 833 tests  3 831 ✅  1 💤 1 ❌
76 552 runs  76 533 ✅ 18 💤 1 ❌

For more details on these failures, see this check.

Results for commit c558240.

♻️ This comment has been updated with latest results.

Copy link
Copy Markdown
Member

@lmoneta lmoneta left a comment

Choose a reason for hiding this comment

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

LGTM!
Thank you, Jonas, for these improvements, which help CLAD. I have only a question on whether is better to impelment a Copy function or using directly std::copy.

This commit refactors SOFIE-generated inference code to enable correct and
efficient reverse-mode automatic differentiation with Clad.

Key changes:

* Introduce explicit primitive operations (`Copy`, `Fill`, `Relu`) in
  SOFIE_common.hxx and provide corresponding custom pullbacks in
  CladDerivator.h. This replaces previously inlined loops and allows
  Clad to generate efficient gradient code without relying on tapes or
  loop-level differentiation.

* Update Gemm code generation to emit Copy/Fill instead of manually
  expanding bias initialization loops. This better exposes the intent
  and improves AD performance and correctness.

* Replace manual ReLU loops with a dedicated Relu() call, enabling a
  custom pullback that avoids tape-based condition tracking.

* Generate an additional "unoptimized" model variant in the SOFIE test suite
  (`OptimizationLevel::kBasic`), and use it for AD tests. This disables
  memory reuse of intermediate tensors. Opaque memory reuse is safe for
  inference but breaks source-transformation AD.

* Improve gradient test diagnostics in SOFIE Clad tests by reporting
  mismatched indices instead of only checking a global max difference.

With these changes, Clad-generated gradients for SOFIE models are both correct
and significantly faster, reaching performance comparable to frameworks such
as PyTorch and JAX on the CPU for the tested cases (fully-connected
neural networks with multiple layers).
@guitargeek guitargeek force-pushed the sofie_clad_continued branch from c558240 to b7655c7 Compare April 14, 2026 09:30
@guitargeek guitargeek merged commit 47bd6d4 into root-project:master Apr 14, 2026
28 of 29 checks passed
@guitargeek guitargeek deleted the sofie_clad_continued branch April 14, 2026 11:22
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