Skip to content

Store adjoint decls in m_Variables and rebuild references on read. NFC #1896

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vgvassilev merged 1 commit into from
Jul 17, 2026
Merged

Store adjoint decls in m_Variables and rebuild references on read. NFC #1896

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 35 additions & 7 deletions include/clad/Differentiator/VisitorBase.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
#include <array>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <stack>
#include <unordered_map>
Expand Down Expand Up @@ -66,11 +67,13 @@ namespace clad {
const clang::Stmt* m_StmtDxSrc = nullptr;
const clang::Stmt* m_RevSweepSrc = nullptr;
utils::StmtClone* m_Cloner = nullptr;
// Deferred build for the forward value (data[1]): produces the node on
// first read, so a forward value no consumer reads constructs nothing.
// Only the forward slot needs it (the primal DeclRef of a reverse-mode
// leaf), so it is not carried for the adjoint or reverse-sweep slots.
// Deferred build for the forward value (data[1]) and the adjoint (data[0]):
// produces the node on first read, so a representation no consumer reads
// constructs nothing (unlike a Lazy clone, it holds no template node). The
// adjoint slot uses it for a reverse-mode leaf's rebuilt m_Variables ref,
// which a terminal product-rule leaf never reads.
std::function<clang::Stmt*()> m_StmtBuild;
std::function<clang::Stmt*()> m_StmtDxBuild;

// Clone Src into Slot on first read; a no-op when Src is null (eager slot).
clang::Stmt* materialize(clang::Stmt*& Slot, const clang::Stmt*& Src);
Expand Down Expand Up @@ -126,7 +129,8 @@ namespace clad {
return diff.Deferred.Cloner;
return valueForRevSweep.Deferred.Cloner;
}()),
m_StmtBuild(std::move(orig.Build)) {
m_StmtBuild(std::move(orig.Build)),
m_StmtDxBuild(std::move(diff.Build)) {
m_Data[1] = orig.Node;
m_Data[0] = diff.Node;
}
Expand All @@ -139,7 +143,14 @@ namespace clad {
}
return materialize(m_Data[1], m_StmtSrc);
}
clang::Stmt* getStmt_dx() { return materialize(m_Data[0], m_StmtDxSrc); }
clang::Stmt* getStmt_dx() {
// Run the deferred build once, on the first read of the adjoint.
if (!m_Data[0] && m_StmtDxBuild) {
m_Data[0] = m_StmtDxBuild();
m_StmtDxBuild = nullptr;
}
return materialize(m_Data[0], m_StmtDxSrc);
}
clang::Expr* getExpr() {
return llvm::cast_or_null<clang::Expr>(getStmt());
}
Expand All @@ -155,6 +166,7 @@ namespace clad {
void updateStmtDx(clang::Stmt* S) {
m_Data[0] = S;
m_StmtDxSrc = nullptr;
m_StmtDxBuild = nullptr;
}
void updateRevSweep(clang::Stmt* S) {
m_ValueForRevSweep = S;
Expand Down Expand Up @@ -221,9 +233,16 @@ namespace clad {
clang::FunctionDecl* m_Derivative;
/// The differentiation request that is being currently processed.
const DiffRequest& m_DiffReq;
/// A cached adjoint reference, stored as its declaration plus how the
/// reference wraps it, so every read rebuilds a fresh expression instead of
/// caching one node that consumers must clone.
struct AdjointInfo {
clang::VarDecl* Decl = nullptr;
enum WrapKind : std::uint8_t { Plain, Deref, ParenDeref } Wrap = Plain;
};
/// Map used to keep track of variable declarations and match them
/// with their derivatives.
std::unordered_map<const clang::ValueDecl*, clang::Expr*> m_Variables;
std::unordered_map<const clang::ValueDecl*, AdjointInfo> m_Variables;

@github-actions github-actions Bot Jul 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

warning: member variable 'm_Variables' has protected visibility [cppcoreguidelines-non-private-member-variables-in-classes]

    std::unordered_map<const clang::ValueDecl*, AdjointInfo> m_Variables;
                                                             ^

/// Map contains variable declarations replacements. If the original
/// function contains a declaration which name collides with something
/// already created inside derivative's body, the declaration is replaced
Expand Down Expand Up @@ -800,6 +819,15 @@ namespace clad {
clang::Expr* cloneThisExprDerivative() {
return CloneNode(m_ThisExprDerivative);
}
/// Rebuild the adjoint reference described by \p A: a fresh reference to
/// A.Decl, dereferenced/parenthesized per A.Wrap. \p Ref's qualifier is
/// reused when the adjoint decl lives in another context (e.g. a lambda).
clang::Expr* buildAdjoint(const AdjointInfo& A,
const clang::DeclRefExpr* Ref = nullptr);
/// Decompose an already-built adjoint expression (a DeclRefExpr or `*ref`)
/// into the AdjointInfo m_Variables stores. Used where the expression is
/// also needed elsewhere; otherwise construct AdjointInfo directly.
static AdjointInfo adjointInfoFrom(clang::Expr* E);
/// A deferred CloneNode(\p N): the clone is produced only if the StmtDiff
/// representation it is stored in is actually read, so a representation no
/// consumer needs allocates no orphaned clone. Drop-in for CloneNode(N) in
Expand Down
52 changes: 14 additions & 38 deletions lib/Differentiator/BaseForwardModeVisitor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -311,7 +311,7 @@ void BaseForwardModeVisitor::SetupDerivativeParameters(
auto* dPVD = utils::BuildParmVarDecl(m_Sema, m_Derivative, II, diffTy,
PVD->getStorageClass());
params.push_back(dPVD);
m_Variables[PVD] = BuildDeclRef(dPVD);
m_Variables[PVD] = {dPVD};
}
}

Expand Down Expand Up @@ -378,7 +378,7 @@ void BaseForwardModeVisitor::GenerateSeeds(const clang::FunctionDecl* dFD) {
// Memorize the derivative of param, i.e. whenever the param is visited
// in the future, it's derivative dParam is found (unless reassigned with
// something new).
m_Variables[param] = dParam;
m_Variables[param] = {dParamDecl};
}
if (const auto* MD = dyn_cast<CXXMethodDecl>(dFD)) {
// We cannot create derivative of lambda yet because lambdas default
Expand Down Expand Up @@ -454,7 +454,7 @@ void BaseForwardModeVisitor::GenerateSeeds(const clang::FunctionDecl* dFD) {
BuildVarDecl(fieldType.getNonReferenceType(),
"_d_" + fieldDecl->getNameAsString(), dInitializer);
addToCurrentBlock(BuildDeclStmt(derivedFieldDecl));
m_Variables.emplace(fieldDecl, BuildDeclRef(derivedFieldDecl));
m_Variables.emplace(fieldDecl, AdjointInfo{derivedFieldDecl});
}
}
}
Expand Down Expand Up @@ -809,11 +809,8 @@ StmtDiff BaseForwardModeVisitor::VisitMemberExpr(const MemberExpr* ME) {
// variable
auto memberDecl = ME->getMemberDecl();
auto it = m_Variables.find(memberDecl);
if (it != std::end(m_Variables)) {
// m_Variables caches one derivative ref per member; clone so repeated
// member uses do not share the node.
return StmtDiff(clonedME, CloneNode(it->second));
}
if (it != std::end(m_Variables))
return StmtDiff(clonedME, buildAdjoint(it->second));
}
// Is not a real variable. Therefore, derivative is 0.
auto zero =
Expand Down Expand Up @@ -892,10 +889,8 @@ BaseForwardModeVisitor::VisitArraySubscriptExpr(const ArraySubscriptExpr* ASE) {
// If the original field is of constant array type, then,
// the derived variable of `arr[i]` is `_d_arr[i]`.
if (it != m_Variables.end() && decl->getType()->isConstantArrayType()) {
// m_Variables caches one adjoint ref per array; clone the base so
// repeated element accesses do not share the node.
auto* result_at_i =
BuildArraySubscript(CloneNode(it->second), derivedIndices());
BuildArraySubscript(buildAdjoint(it->second), derivedIndices());
return StmtDiff{cloned, result_at_i};
}

Expand Down Expand Up @@ -942,16 +937,12 @@ BaseForwardModeVisitor::VisitArraySubscriptExpr(const ArraySubscriptExpr* ASE) {
// Is not an independent variable, ignored.
return StmtDiff(cloned, zero);

Expr* target = it->second;
// FIXME: fix when adding array inputs
if (!isArrayOrPointerType(target->getType()))
// FIXME: fix when adding array inputs. Forward-mode adjoints are plain refs,
// so the decl's type is the adjoint's -- check it before building anything.
if (!isArrayOrPointerType(it->second.Decl->getType().getNonReferenceType()))
return StmtDiff(cloned, zero);
// llvm::APSInt IVal;
// if (!I->EvaluateAsInt(IVal, m_Context))
// return;
// Create the _result[idx] expression. target is the cached adjoint ref;
// clone it so repeated element accesses do not share the base node.
auto* result_at_is = BuildArraySubscript(CloneNode(target), derivedIndices());
auto* result_at_is =
BuildArraySubscript(buildAdjoint(it->second), derivedIndices());
return StmtDiff(cloned, result_at_is);
}

Expand Down Expand Up @@ -989,23 +980,8 @@ StmtDiff BaseForwardModeVisitor::VisitDeclRefExpr(const DeclRefExpr* DRE) {
// If DRE references a variable, try to find if we know something about
// how it is related to the independent variable.
auto it = m_Variables.find(VD);
if (it != std::end(m_Variables)) {
clang::Expr* dExpr = it->second;
// If a record was found, use the recorded derivative.
if (auto dVarDRE = dyn_cast<DeclRefExpr>(dExpr)) {
auto dVar = cast<VarDecl>(dVarDRE->getDecl());
if (dVar->getDeclContext() != m_Sema.CurContext) {
clad_compat::NestedNameSpecifierTy NNS = DRE->getQualifier();
dExpr = BuildDeclRef(dVar, clad_compat::hasQualifier(NNS)
? NNS
: clad_compat::nullNNS());
}
}
// m_Variables caches one derivative reference per variable; hand out a
// fresh clone so callers that combine it (product rule) never parent the
// cached node twice.
return StmtDiff(clonedDRE, CloneNode(dExpr));
}
if (it != std::end(m_Variables))
return StmtDiff(clonedDRE, buildAdjoint(it->second, DRE));
}
// Is not a variable or is a reference to something unrelated to independent
// variable. Derivative is 0.
Expand Down Expand Up @@ -1639,7 +1615,7 @@ BaseForwardModeVisitor::DifferentiateVarDecl(const VarDecl* VD,
initDx, VD->isDirectInit());

if (VDDerived)
m_Variables.emplace(VDClone, BuildDeclRef(VDDerived));
m_Variables.emplace(VDClone, AdjointInfo{VDDerived});
return DeclDiff<VarDecl>(VDClone, VDDerived);
}

Expand Down
11 changes: 6 additions & 5 deletions lib/Differentiator/ErrorEstimator.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -186,13 +186,14 @@ void ErrorEstimationHandler::EmitFinalErrorStmts(
// Finally emit the error.
// Clone the operands: GetError embeds them in the error expression,
// which must not share nodes with the derivative statements.
auto* errorExpr = GetError(m_RMV->CloneNode(paramClone),
m_RMV->CloneNode(m_RMV->m_Variables[decl]),
params[i]->getNameAsString());
auto* errorExpr =
GetError(m_RMV->CloneNode(paramClone),
m_RMV->buildAdjoint(m_RMV->m_Variables[decl]),
params[i]->getNameAsString());
m_RMV->addToCurrentBlock(
m_RMV->BuildOp(BO_AddAssign, BuildFinalErrorExpr(), errorExpr));
} else {
auto LdiffExpr = m_RMV->m_Variables[decl];
Expr* LdiffExpr = m_RMV->buildAdjoint(m_RMV->m_Variables[decl]);
Expr* size = getSizeExpr(decl);
VarDecl* idxExprDecl = nullptr;
// Save our index expression so it can be used later.
Expand Down Expand Up @@ -344,7 +345,7 @@ void ErrorEstimationHandler::ActAfterProcessingArraySubscriptExpr(
if (const auto* DRE =
dyn_cast<DeclRefExpr>(ASE->getBase()->IgnoreImplicit())) {
const auto* VD = cast<VarDecl>(DRE->getDecl());
Expr* VDdiff = m_RMV->m_Variables[VD];
Expr* VDdiff = m_RMV->buildAdjoint(m_RMV->m_Variables[VD]);
// We only need to track sizes for arrays and pointers.
if (!utils::isArrayOrPointerType(VDdiff->getType()))
return;
Expand Down
34 changes: 20 additions & 14 deletions lib/Differentiator/JacobianModeVisitor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,9 @@
#include "clad/Differentiator/CladUtils.h"
#include "clad/Differentiator/DerivativeBuilder.h"

#include "clang/AST/Decl.h"
#include "clang/AST/OperationKinds.h"

#include "llvm/Support/SaveAndRestore.h"

using namespace clang;
Expand Down Expand Up @@ -89,17 +92,17 @@ DerivativeAndOverload JacobianModeVisitor::Derive() {
continue;
auto derivedPVDName = "_d_vector_" + std::string(PVDII->getName());
IdentifierInfo* derivedPVDII = CreateUniqueIdentifier(derivedPVDName);
Expr* derivedExpr = nullptr;
VarDecl* adjointDecl = nullptr;

@github-actions github-actions Bot Jul 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

warning: no header providing "clang::VarDecl" is directly included [misc-include-cleaner]

    VarDecl* adjointDecl = nullptr;
    ^

AdjointInfo::WrapKind wrap = AdjointInfo::Plain;
if (utils::isArrayOrPointerType(PVD->getType())) {
ParmVarDecl* derivedPVD =
utils::BuildParmVarDecl(m_Sema, m_Derivative, derivedPVDII,
utils::GetParameterDerivativeType(
m_Sema, m_DiffReq.Mode, PVD->getType()),
PVD->getStorageClass());
derivedParams.push_back(derivedPVD);
derivedExpr =
BuildOp(UO_Deref, BuildDeclRef(derivedPVD), PVD->getBeginLoc());
derivedExpr = utils::BuildParenExpr(m_Sema, derivedExpr);
adjointDecl = derivedPVD;
wrap = AdjointInfo::ParenDeref;
Expr* getSize = BuildCallExprToMemFn(BuildDeclRef(derivedPVD),
/*MemberFunctionName=*/"rows", {});
llvm::StringRef PVDName = PVD->getName();
Expand All @@ -117,8 +120,8 @@ DerivativeAndOverload JacobianModeVisitor::Derive() {
m_Sema, m_DiffReq.Mode, PVD->getType()),
PVD->getStorageClass());
derivedParams.push_back(derivedPVD);
derivedExpr =
BuildOp(UO_Deref, BuildDeclRef(derivedPVD), PVD->getBeginLoc());
adjointDecl = derivedPVD;
wrap = AdjointInfo::Deref;
nonArrayIndVarCount += 1;
} else {
VarDecl* derivedPVD =
Expand All @@ -127,10 +130,10 @@ DerivativeAndOverload JacobianModeVisitor::Derive() {
->getPointeeType(),
derivedPVDII);
adjointDecls.push_back(BuildDeclStmt(derivedPVD));
derivedExpr = BuildDeclRef(derivedPVD);
adjointDecl = derivedPVD;
nonArrayIndVarCount += 1;
}
m_Variables[newPVD] = derivedExpr;
m_Variables[newPVD] = {adjointDecl, wrap};
}

params.insert(params.end(), derivedParams.begin(), derivedParams.end());
Expand Down Expand Up @@ -178,7 +181,6 @@ DerivativeAndOverload JacobianModeVisitor::Derive() {
bool is_array =
utils::isArrayOrPointerType(m_DiffReq->getParamDecl(i)->getType());
ParmVarDecl* param = params[i];
Expr* paramDiff = m_Variables[param]->IgnoreParens();
QualType dParamType = clad::utils::GetValueType(param->getType());
// Desugaring the type is necessary to pass it to other templates
dParamType = dParamType.getDesugaredType(m_Context);
Expand All @@ -198,9 +200,11 @@ DerivativeAndOverload JacobianModeVisitor::Derive() {
nonArrayIndVarCountExpr);

if (is_array) {
Expr* base = cast<UnaryOperator>(paramDiff)->getSubExpr();
// The adjoint is `(*_d_p)`; the array whose size we need is the bare
// `_d_p` reference (m_Variables stores its decl).
Expr* base = BuildDeclRef(m_Variables[param].Decl);
// Get size of the array.
Expr* getSize = BuildCallExprToMemFn(Clone(base),
Expr* getSize = BuildCallExprToMemFn(base,
/*MemberFunctionName=*/"rows", {});
// Create an identity matrix for the parameter,
// with number of rows equal to the size of the array,
Expand Down Expand Up @@ -246,12 +250,14 @@ DerivativeAndOverload JacobianModeVisitor::Derive() {
// -> clad::array<double> _d_vector_z = {0, 1};
if (utils::isArrayOrPointerType(param->getType()) ||
param->getType()->isReferenceType()) {
// The store target is `*_d_p`; strip the parens the ParenDeref adjoint
// carries for element access elsewhere.
Expr* paramAssignment =
BuildOp(BO_Assign, Clone(paramDiff), dVectorParam);
BuildOp(BO_Assign, buildAdjoint(m_Variables[param])->IgnoreParens(),

@github-actions github-actions Bot Jul 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

warning: no header providing "clang::BO_Assign" is directly included [misc-include-cleaner]

          BuildOp(BO_Assign, buildAdjoint(m_Variables[param])->IgnoreParens(),
                  ^

dVectorParam);
addToCurrentBlock(paramAssignment);
} else {
auto* paramDecl = cast<VarDecl>(cast<DeclRefExpr>(paramDiff)->getDecl());
SetDeclInit(paramDecl, dVectorParam);
SetDeclInit(m_Variables[param].Decl, dVectorParam);
}
}

Expand Down
8 changes: 3 additions & 5 deletions lib/Differentiator/ReverseModeForwPassVisitor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,7 @@ ReverseModeForwPassVisitor::BuildParams(DiffParams& diffParams) {
if (dPVD->getIdentifier())
m_Sema.PushOnScopeChains(dPVD, getCurrentScope(),
/*AddToContext=*/false);
m_Variables[*it] = BuildDeclRef(dPVD), m_DiffReq->getLocation();
m_Variables[*it] = {dPVD};
}
}
if (m_DiffReq.UseRestoreTracker) {
Expand Down Expand Up @@ -291,10 +291,8 @@ StmtDiff ReverseModeForwPassVisitor::VisitDeclRefExpr(const DeclRefExpr* DRE) {
auto* decl = dyn_cast<VarDecl>(clonedDRE->getDecl());
auto foundAdjoint = m_Variables.find(decl);
Expr* adjoint = nullptr;
// m_Variables caches a single adjoint reference per variable; hand out a
// fresh copy so repeated uses do not share the same node.
if (foundAdjoint != m_Variables.end())
adjoint = CloneNode(foundAdjoint->second);
adjoint = buildAdjoint(foundAdjoint->second, DRE);

return StmtDiff(clonedDRE, adjoint);
}
Expand Down Expand Up @@ -335,7 +333,7 @@ ReverseModeForwPassVisitor::DifferentiateVarDecl(const clang::VarDecl* VD,
auto* VDDerived =
BuildGlobalVarDecl(DerivedType, "_d_" + VD->getNameAsString(),
initDiff.getExpr_dx(), VD->isDirectInit());
m_Variables.emplace(VDCloned, BuildDeclRef(VDDerived));
m_Variables.emplace(VDCloned, AdjointInfo{VDDerived});
// Register the primal clone unconditionally so references rebind by map
// lookup rather than the scope-dependent name-lookup fallback (see the
// matching change in ReverseModeVisitor::DifferentiateVarDecl).
Expand Down
Loading
Loading