Enable Automatic Differentiation of OpenMP Programs with Clad

Jiayang Li

Mentors: Vassil Vassilev, Martin Vassilev

Hello everyone, my name is Jiayang Li.

Today, I’ll be presenting my project “Enable Automatic Differentiation of OpenMP Programs with Clad”, which focuses on enabling automatic differentiation for OpenMP programs within the Clad framework.

If you’d like to view the slides online, you can click the link shown on the slide.

Background

• What is Automatic Differentiation (AD)?

  Exact derivatives by transforming programs; no finite-difference error; supports forward and reverse.

• What is OpenMP?

  Pragmas for shared-memory parallelism (parallel, for, reduction, etc.).

• What is Clad?

  Source-transformation AD library for C++ implemented as a plugin for Clang

Let’s start with some background.

First, what is Automatic Differentiation (AD)? Automatic Differentiation is a method of obtaining exact derivatives through program transformation. Unlike symbolic or numerical differentiation, it has no rounding errors or instability, and it supports both forward mode and reverse mode differentiation.

Next, what is OpenMP? OpenMP is a widely used parallel programming interface for C, C++, and Fortran. It enables shared-memory parallelism through #pragma omp directives such as parallel, for, and reduction.

Finally, what is Clad? Clad is a source-transformation automatic differentiation plugin for Clang. It works directly on the Clang Abstract Syntax Tree (AST) to generate derivative code. In other words, it performs “compile-time automatic differentiation,” producing exact source-level transformations for C++ functions.

The goal of this project, as illustrated on the right, is to apply automatic differentiation to functions annotated with OpenMP directives — so that the generated derivative code is also parallelized, achieving acceleration through parallel computation.

Case Study

Start with a representative example: compute the gradient of a simple parallel loop.

double parallel_sum_of_squares(const double* x, int n) {
    double total = 0.;
    #pragma omp parallel for reduction(+: total)
    for (int i = 0; i < n; i++) {
        total += x[i] * x[i];
    }
    return total;
}

Let’s start with a simple OpenMP example. This function parallel_sum_of_squares computes the sum of squares of all elements in an array. It parallelizes the loop using #pragma omp parallel for reduction(+: total).

Case Study

What might the reverse-mode gradient look like conceptually?

void parallel_sum_of_squares_grad(const double *x, int n, double *_d_x, int *_d_n) {
    double _d_total = 0.;
    double total = 0.;
    #pragma omp parallel for reduction(+: total)
    for (int i = 0; i < n; i++) {
        total += x[i] * x[i];
    }
    _d_total += 1;
    #pragma omp parallel private(total) firstprivate(_d_total)
    for (int i = n - 1; i >= 0; i -= 1) {
        double _r_d0 = _d_total;
        _d_x[i] += _r_d0 * x[i];
        _d_x[i] += x[i] * _r_d0;
    }
}

Now, we’d like to obtain its reverse-mode gradient. Conceptually, we want to automatically generate an equivalent reverse-mode function. The slide shows a conceptual example, and you can see that:

• We first perform the same forward loop as in the original function.
• Then, we initialize the derivative of the result variable total (i.e., _d_total) to 1.
• Finally, in the reverse loop, we accumulate the gradients for each element x[i].

Is that correct? Almost — but there are still some subtleties and implementation details to handle, which we’ll look into next. The basic idea, however, is already quite close.

Technical Challenges

• Clang AST specifics for OpenMP (building AST nodes, capture variables, clause handling)
• Variable scoping across OpenMP regions and clauses
• Thread-safe storage for tapes/intermediates
• Deterministic schedule reversal for the reverse pass

There are several core challenges here:

1. Complexity of OpenMP nodes in the Clang AST We need to correctly build and capture variables within parallel regions and handle different OpenMP clauses.
2. Variable scoping Each thread in an OpenMP region has its own private variables, and we must ensure these are properly matched between the forward and reverse passes.
3. Thread-safe storage of intermediates Reverse-mode AD requires saving intermediate values (so-called “tapes”), and these must be thread-safe.
4. Deterministic schedule reversal We must guarantee that during the reverse pass, each thread’s iteration chunk is identical to that in the forward pass.

The diagram above shows the data dependencies between intermediate results. Each thread maintains its own stack. The key is that the same thread must perform the corresponding forward and reverse iterations — in reverse order — to ensure that data is popped correctly from its local stack.

Implementation Overview

• The theory is based on the paper from Tapenade, which provides detailed proofs and demonstrates the implementation of OpenMP automatic differentiation on the Fortran platform using Tapenade.

• Since OpenMP constructs are transformed into AST nodes during the Clang AST phase, we can override the corresponding OpenMP-related AST nodes Visit methods—just as we do with other Visit methods—to build new function ASTs.

For implementation, we took inspiration from the Tapenade team’s paper (linked in the slides). Tapenade has already implemented OpenMP automatic differentiation for Fortran and provides solid theoretical foundations.

Our approach in Clad is similar. Since OpenMP constructs are represented as specific AST nodes in Clang, we only need to extend the corresponding Visit methods in the AST visitor to produce differentiated versions. In other words, we reuse Clang’s AST mechanisms so that Clad can correctly identify and reconstruct OpenMP statements.

The slide also links to a Compiler Explorer example, where we can see the actual OpenMP AST structure — including a node type called CapturedStmt, which represents captured variables in a parallel region.

Forward Mode (#1491)

• Since the execution order of the forward-mode derivatives is consistent with that of the original code, the scopes of the differentiated variables can be directly inherited from the original variables, and the parallel structure of the entire program can also be reused.

• For the previous example, the actually generated code looks like this:

// auto d_fn_arr = clad::differentiate(parallel_sum_of_squares, "x[1]");
// d_fn_arr.dump();
double parallel_sum_of_squares_darg0_1(const double *x, int n) {
    int _d_n = 0;
    double _d_total = 0.;
    double total = 0.;
    #pragma omp parallel for reduction(+: _d_total,total)
        for (int i = 0; i < n; i++) {
            _d_total += (i == 1) * x[i] + x[i] * (i == 1);
            total += x[i] * x[i];
        }
    return _d_total;
}

I already presented the Forward Mode implementation during the midterm evaluation. Forward mode is relatively straightforward because its execution order matches the original program. Therefore, we can directly reuse the same parallel structure, with identical variable scopes.

In this example, the generated forward-mode code is shown below. You can see that:

• Both total and _d_total appear in the same parallel region.
• We simply add an additional reduction clause reduction(+: _d_total, total) to accumulate the derivatives.

The structure is almost identical to the original function, making it very intuitive.

Reverse Mode — Clang AST

• The previous VisitForStmt method was too invasive for omp for; we implement a dedicated DifferentiateCanonicalLoop for OpenMP loops.

// auto fn_grad = clad::gradient(sum_of_squares);
// fn_grad.dump();
void sum_of_squares_grad(const double *x, int n, double *_d_x, int *_d_n) {
    int _d_i = 0;
    int i = 0;
    double _d_total = 0.;
    double total = 0.;
    unsigned long _t0 = 0;
    for (i = 0; i < n; i++) {
        _t0++;
        total += x[i] * x[i];
    }
    _d_total += 1;
    for (; _t0; _t0--) {
        i--;
        {
            double _r_d0 = _d_total;
            _d_x[i] += _r_d0 * x[i];
            _d_x[i] += x[i] * _r_d0;
        }
    }
}

Reverse mode, however, is much more complex. We can’t just modify the existing VisitForStmt, because its loop counter is embedded within the loop and the generated ForStmt cannot be correctly divided among threads.

Therefore, we introduced a new helper function called DifferentiateCanonicalLoop, which specifically handles OpenMP loops and returns a canonical ForStmt suitable for OpenMP task partitioning and scheduling.

Reverse Mode — Clang AST

• It should be noted that, for Clang’s OpenMP implementation, we need to capture the related variables when building DeclRefExpr (which corresponds to the CaptureStmt’s subnodes in the Clang AST).

  So we also need visit twice to generate two OpenMP regions and their capture lists: one for the forward pass, one for the reverse pass.

CLAD_COMPAT_CLANG19_SemaOpenMP(m_Sema).ActOnOpenMPRegionStart(OMPD_parallel, nullptr);
StmtDiff BodyDiff;
{
    Sema::CompoundScopeRAII CompoundScope(m_Sema);
    if (isOpenMPLoopDirective(D->getDirectiveKind())) {
        const auto* FS = cast<ForStmt>(CS);
        BodyDiff = DifferentiateCanonicalLoop(FS);
    } else {
        BodyDiff = Visit(CS);
    }
}
Stmt* Forward = CLAD_COMPAT_CLANG19_SemaOpenMP(m_Sema)
                    .ActOnOpenMPRegionEnd(BodyDiff.getStmt(), OrigClauses)
                    .get();

CLAD_COMPAT_CLANG19_SemaOpenMP(m_Sema).ActOnOpenMPRegionStart(OMPD_parallel, nullptr);
// Visit twice, but use only the result of the first visit, for capture variables only.
{
    Sema::CompoundScopeRAII CompoundScope(m_Sema);
    Stmts temp;
    m_Globals.swap(temp);
    if (isOpenMPLoopDirective(D->getDirectiveKind())) {
        const auto* FS = cast<ForStmt>(CS);
        DifferentiateCanonicalLoop(FS);
    } else {
        Visit(CS);
    }
    m_Globals.swap(temp);
}
Stmt* Reverse =
    CLAD_COMPAT_CLANG19_SemaOpenMP(m_Sema)
        .ActOnOpenMPRegionEnd(BodyDiff.getStmt_dx(), DiffClauses)
        .get();

In Clang’s OpenMP implementation, each parallel region uses a CapturedStmt to capture external variables. Thus, when constructing the AST for the reverse pass, we must traverse twice: once to generate the forward region and its capture list, and again to generate the reverse region and its capture list.

This ensures that each OpenMP region has the correct variable context.

Reverse Mode — Scope conversion

This diagram shows how scope conversion works. In short, Clad maps variables from each scope in the forward pass to the corresponding scope in the reverse pass. This mapping is especially important in OpenMP, where we distinguish between private, shared, and firstprivate variables.

The example shows a loop with 9 iterations (represented as boxes with iteration numbers), parallelized with two threads and a hypothetical scheduling scheme assigning two iteration chunks per thread. In the multithreaded view, colors represent thread assignment, arrows indicate data flow, and the markers denote variable creation, zero initialization, or accumulation operations. The top half shows the behavior of the original and tangent (forward) modes, while the bottom half shows the adjoint (reverse) behavior.

The paper provides detailed theoretical proofs of correctness, which I won’t cover here.

Reverse Mode — Intermediates and schedule replay

• Intermediates require tapes; the tape itself must become private to each thread; so we make these threadprivate.
• For data consistency, we must ensure that the parallel scheduling of the backward pass strictly reproduces the thread–task allocation of the forward pass. To achieve this, we reverse the schedule while preserving chunk boundaries by a small runtime helper:
  • Forward pass: each thread computes its static chunk from thread ID and executes normally.
  • Reverse pass: the same helper returns identical chunks, then we iterate the chunk in reverse order.

void clad::GetStaticSchedule(int lo, int hi, int stride, int* threadlo, int* threadhi)

Another key issue is schedule replay. Since OpenMP loop scheduling is often static, we can compute each thread’s iteration range using a small runtime helper: clad::GetStaticSchedule.

During the reverse pass, we call the same helper again to obtain the same chunk boundaries and then iterate them in reverse order. This ensures that each thread precisely “replays” its original forward work during the backward pass, maintaining consistency at the thread level.

Reverse Mode (#1641)

The resulting generated code:

// auto fn_grad = clad::gradient(parallel_sum_of_squares);
// fn_grad.dump();
void parallel_sum_of_squares_grad(const double *x, int n, double *_d_x, int *_d_n) {
    double _d_total = 0.;
    double total = 0.;
    #pragma omp parallel reduction(+: total)
        {
            int _t_chunklo0 = 0;
            int _t_chunkhi0 = 0;
            clad::GetStaticSchedule(0, n - 1, 1, &_t_chunklo0, &_t_chunkhi0);
            for (int i = _t_chunklo0; i <= _t_chunkhi0; i += 1) {
                total += x[i] * x[i];
            }
        }
    _d_total += 1;
    #pragma omp parallel private(total) firstprivate(_d_total)
        {
            int _t_chunklo1 = 0;
            int _t_chunkhi1 = 0;
            clad::GetStaticSchedule(0, n - 1, 1, &_t_chunklo1, &_t_chunkhi1);
            for (int i = _t_chunkhi1; i >= _t_chunklo1; i -= 1) {
                {
                    double _r_d0 = _d_total;
                    _d_x[i] += _r_d0 * x[i];
                    _d_x[i] += x[i] * _r_d0;
                }
            }
        }
}

Finally, this slide shows the complete generated reverse-mode code. As you can see:

• The forward pass and backward pass are placed in two separate parallel regions.
• The GetStaticSchedule helper ensures identical chunk boundaries.
• Each thread iterates its chunk in reverse during the backward pass.
• Gradients for each element x[i] are accumulated into _d_x[i].

The structure fully matches the theoretical expectation and works correctly under multithreaded execution.

Future Work

• Support dynamic schedules (record executed chunks at runtime and replay).

• Support more OpenMP clauses and directives (atomic, simd, etc.).

  We can even try to enable OpenMP offloading to target accelerators (e.g., GPUs).

For future improvements:

1. Support dynamic scheduling For dynamically scheduled loops, we’ll need to record the executed chunks at runtime and replay them during the backward pass.
2. Support additional OpenMP clauses and directives Such as atomic, simd, and potentially OpenMP target offloading, enabling automatic differentiation on accelerators like GPUs.

Thank You!

Q & A

That concludes my presentation — thank you for your attention! I’ll be happy to take any questions.