Skip to content

Commit b75bba2

Browse files
committed
Do not reuse AST nodes in generated derivatives.
clad assembles each derivative by hand and must keep it a proper tree -- no node the child (Stmt::children()) of two parents. Across the differentiation rules it instead fed one expression into two places: a primal operand reused as the derived operand (pointer/shift derivatives, product and quotient rules), an assignment lvalue reused in its own adjoint, recorded values shared between the forward and reverse passes, aggregate/InitListExpr subtrees, cached references (adjoints, loop counters, tape and tracker refs, the error index/return expressions), a tape push emitted both in a parallel loop body and its OpenMP residual sweep, and a CUDA cudaMalloc's pointer/size reused in the adjoint's cudaMemset. Such a node has two parents, benign until a later pass rewrites it in place and the edit leaks into unrelated code (and it breaks CodeGen for shared aggregate initializers). This is a property of the hand-built body, not of Clang ASTs in general, which do share nodes across type and template-argument edges, InitListExpr's dual forms, the reduction combiners Sema synthesizes for OpenMP clauses, the OpaqueValueExpr common subexpression of a PseudoObjectExpr (e.g. a threadIdx.x property access), and a call's default-argument expression across its call sites. Clone the reused expression so each occurrence owns its node. The clone lives at the narrowest correct owner: at the reuse site for a one-off local; inside the clad helper that builds the second use (memset, tracker store, central-difference calls); and on read in the accessor of an inherently reused cached reference. Add CloneNode for a reference-preserving (non-remapping) subtree copy, and buildClonedLambda so a copied lambda gets a fresh closure instead of sharing the original's operator() body. StmtClone deep-clones a PseudoObjectExpr -- rebuilding its bound OpaqueValueExprs and remapping the syntactic and semantic forms to them -- so a copied property access shares nothing with its source. Generic builders (BuildOp) stay non-cloning: they receive whole subtrees and would deep-copy them repeatedly. Enforce the invariant in DerivativeBuilder: findSharedNode returns the first node reachable through two parent edges -- skipping the legitimate Clang shares above (type and template-argument locations, OpenMP clauses, OpaqueValueExprs, default-argument expressions, InitListExpr's dual form) -- which the assert rejects on debug builds, catching a rule that stops cloning at generation time rather than as a later codegen crash, and which warns on release builds, where the assert is compiled out, so a regression is surfaced rather than silently shipped. The check runs on clang-17+, where buildClonedLambda can synthesize a fresh closure; below that lambda derivatives legitimately share and it is disabled.
1 parent 395a057 commit b75bba2

18 files changed

Lines changed: 1104 additions & 285 deletions

include/clad/Differentiator/ErrorEstimator.h

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,11 @@ class ErrorEstimationHandler : public ExternalRMVSource {
5252
Stmts m_ReverseErrorStmts;
5353
/// The index expression for emitting final errors for input param errors.
5454
clang::Expr* m_IdxExpr = nullptr;
55+
/// Clone-on-read accessors for the cached references above: each is reused
56+
/// across several generated subscripts/conditions, so every read must own
57+
/// its own node rather than share the cached one.
58+
clang::Expr* cloneIdxExpr();
59+
clang::Expr* cloneRetErrorExpr();
5560
/// A map from var decls to their size variables (e.g. `var_size`).
5661
std::unordered_map<const clang::VarDecl*, clang::Expr*> m_ArrSizes;
5762
// FIXME: Solve this in a more general way.

include/clad/Differentiator/ReverseModeVisitor.h

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -583,23 +583,27 @@ namespace clad {
583583

584584
/// Returns reference to the last object of the clad tape if clad tape
585585
/// is used as the counter; otherwise returns reference to the counter
586-
/// variable.
587-
clang::Expr* getRef() const { return m_Ref; }
586+
/// variable. The reference is cloned on every read so the forward
587+
/// increment, reverse decrement and loop condition each own their
588+
/// counter DeclRefExpr node instead of sharing it.
589+
[[nodiscard]] clang::Expr* cloneRef() const {
590+
return m_RMV.CloneNode(m_Ref);
591+
}
588592

589593
/// Returns counter post-increment expression (`counter++`).
590594
clang::Expr* getCounterIncrement() {
591-
return m_RMV.BuildOp(clang::UnaryOperatorKind::UO_PostInc, m_Ref);
595+
return m_RMV.BuildOp(clang::UnaryOperatorKind::UO_PostInc, cloneRef());
592596
}
593597

594598
/// Returns counter post-decrement expression (`counter--`)
595599
clang::Expr* getCounterDecrement() {
596-
return m_RMV.BuildOp(clang::UnaryOperatorKind::UO_PostDec, m_Ref);
600+
return m_RMV.BuildOp(clang::UnaryOperatorKind::UO_PostDec, cloneRef());
597601
}
598602

599603
/// Returns `ConditionResult` object for the counter.
600604
clang::Sema::ConditionResult getCounterConditionResult() {
601605
return m_RMV.m_Sema.ActOnCondition(m_RMV.getCurrentScope(), noLoc,
602-
m_Ref,
606+
cloneRef(),
603607
clang::Sema::ConditionKind::Boolean);
604608
}
605609

include/clad/Differentiator/StmtClone.h

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,12 @@ namespace utils {
4141
clang::Sema& m_Sema;
4242
clang::ASTContext& Ctx;
4343
Mapping* m_OriginalToClonedStmts;
44+
// While cloning a PseudoObjectExpr, maps each original OpaqueValueExpr to
45+
// its clone so the syntactic form and the semantic expressions reference
46+
// the same fresh OVE (null outside such a clone). See
47+
// VisitPseudoObjectExpr.
48+
llvm::DenseMap<clang::OpaqueValueExpr*, clang::OpaqueValueExpr*>*
49+
m_OVESubst = nullptr;
4450

4551
clang::Decl* CloneDecl(clang::Decl* Node);
4652
clang::VarDecl* CloneDeclOrNull(clang::VarDecl* Node);
@@ -124,6 +130,8 @@ namespace utils {
124130
DECLARE_CLONE_FN(CXXTemporaryObjectExpr)
125131
DECLARE_CLONE_FN(MaterializeTemporaryExpr)
126132
DECLARE_CLONE_FN(PseudoObjectExpr)
133+
DECLARE_CLONE_FN(OpaqueValueExpr)
134+
DECLARE_CLONE_FN(MSPropertyRefExpr)
127135
DECLARE_CLONE_FN(SubstNonTypeTemplateParmExpr)
128136
DECLARE_CLONE_FN(CXXScalarValueInitExpr)
129137
DECLARE_CLONE_FN(ConstantExpr)

include/clad/Differentiator/VisitorBase.h

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -643,6 +643,13 @@ namespace clad {
643643
bool pushOnScopeChains = false,
644644
bool cloneDefaultArg = true,
645645
clang::SourceLocation Loc = noLoc);
646+
/// Build a primal copy of a lambda expression with a *fresh*
647+
/// closure type, rather than reusing the original closure (as a plain
648+
/// StmtClone does). Reusing the closure makes two clones share the
649+
/// operator() body, violating the one-parent-per-node invariant. Only
650+
/// captureless lambdas get a fresh closure; captured lambdas fall back to
651+
/// a plain clone. Inner lambda-init declarations are rebuilt recursively.
652+
clang::Expr* buildClonedLambda(const clang::LambdaExpr* LE);
646653
/// A function to get the single argument "forward_central_difference"
647654
/// call expression for the given arguments.
648655
///
@@ -687,6 +694,15 @@ namespace clad {
687694
clang::Stmt* Clone(const clang::Stmt* S);
688695
/// A shorthand to simplify cloning of expressions.
689696
clang::Expr* Clone(const clang::Expr* E);
697+
/// Structural copy that does NOT run updateReferencesOf. `Clone` re-points
698+
/// references (name lookup, constant folding, type fix-ups), which is
699+
/// correct when copying original-function code into the derivative but
700+
/// corrupts already-generated derivative expressions (e.g. folds `(x + y)`
701+
/// into a garbage literal). Use this to split a reused generated node into
702+
/// a distinct-but-identical copy.
703+
clang::Expr* CloneNode(const clang::Expr* E);
704+
/// Statement overload of the structural copy above.
705+
clang::Stmt* CloneNode(const clang::Stmt* S);
690706
/// Cloning types is necessary since VariableArrayType
691707
/// store a pointer to their size expression.
692708
clang::QualType CloneType(clang::QualType T);
Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
//--------------------------------------------------------------------*- C++ -//
2+
// clad - the C++ Clang-based Automatic Differentiator
3+
//
4+
// See ASTIntegrity.h for the rationale.
5+
//----------------------------------------------------------------------------//
6+
7+
#include "ASTIntegrity.h"
8+
9+
#include "clang/AST/Expr.h"
10+
#include "clang/AST/ExprCXX.h"
11+
#include "clang/AST/RecursiveASTVisitor.h"
12+
#include "clang/AST/Stmt.h"
13+
#include "clang/AST/StmtOpenMP.h"
14+
15+
#include "llvm/ADT/DenseMap.h"
16+
#include "llvm/ADT/SmallPtrSet.h"
17+
18+
using namespace clang;
19+
20+
namespace clad {
21+
22+
const Stmt* findSharedNode(const Stmt* Root) {
23+
if (!Root)
24+
return nullptr;
25+
// clad's hand-built body must be a tree in its Stmt CHILD-edge structure: no
26+
// node is a child (Stmt::children()) of two parents. Count child edges, NOT
27+
// RAV visits -- RecursiveASTVisitor also follows type and template-argument
28+
// edges, and Clang legitimately shares nodes across those (e.g. the
29+
// substituted `N` in a `Foo<5>` instantiation is one ConstantExpr reachable
30+
// from both the type and the body). Counting visits would flag those valid
31+
// shares; counting child edges flags only genuine multiple-parent reuse,
32+
// which is what clad must not produce.
33+
struct Counter : RecursiveASTVisitor<Counter> {
34+
llvm::DenseMap<const Stmt*, unsigned> ParentEdges;
35+
llvm::SmallPtrSet<const Stmt*, 32> SeenParents;
36+
// Visit implicit code: some genuine reuses have a second parent reachable
37+
// only through implicit nodes.
38+
bool shouldVisitImplicitCode() const { return true; }
39+
// But visit only the semantic form of an InitListExpr: with implicit code
40+
// on, RAV visits both the syntactic and semantic forms, which legitimately
41+
// share their element nodes (Clang's representation, not a reuse).
42+
bool TraverseInitListExpr(InitListExpr* ILE,
43+
DataRecursionQueue* = nullptr) {
44+
InitListExpr* Sem = ILE->isSemanticForm() ? ILE : ILE->getSemanticForm();
45+
if (!Sem)
46+
Sem = ILE;
47+
WalkUpFromInitListExpr(Sem);
48+
for (Stmt* Ch : Sem->children())
49+
if (Ch)
50+
TraverseStmt(Ch);
51+
return true;
52+
}
53+
// Skip OpenMP clauses. Lowering `reduction(+:x)` makes Sema synthesize a
54+
// combiner (`.reduction.lhs = .reduction.lhs + .reduction.rhs`) that
55+
// references one helper DeclRefExpr from two parents -- Clang's own AST,
56+
// not clad reuse. TraverseOMPExecutableDirective walks only the clauses;
57+
// the associated loop body clad generates is still traversed by the
58+
// DEF_TRAVERSE_STMT children walk.
59+
bool TraverseOMPExecutableDirective(OMPExecutableDirective* /*D*/) {
60+
return true;
61+
}
62+
// An OpaqueValueExpr is Clang's mechanism for sharing a common
63+
// subexpression -- e.g. the `threadIdx` base of a `threadIdx.x`
64+
// __declspec(property) access, which Clang models as a PseudoObjectExpr
65+
// that references one OpaqueValueExpr from both its syntactic form and its
66+
// semantic expressions. It is shared by design, not clad reuse, so do not
67+
// descend into its source (which the OVE also shares); VisitStmt likewise
68+
// does not count edges to it.
69+
bool TraverseOpaqueValueExpr(OpaqueValueExpr* /*OVE*/,
70+
DataRecursionQueue* = nullptr) {
71+
return true;
72+
}
73+
// Skip default-argument expressions. Clang stores a call's default argument
74+
// once on the ParmVarDecl and every call site references that one
75+
// expression through a CXXDefaultArgExpr (e.g. the `0` in
76+
// thrust::reduce_by_key's default operators, shared across clad's cloned
77+
// calls). That is Clang's representation, not clad reuse.
78+
bool TraverseCXXDefaultArgExpr(CXXDefaultArgExpr* /*E*/,
79+
DataRecursionQueue* = nullptr) {
80+
return true;
81+
}
82+
bool TraverseCXXDefaultInitExpr(CXXDefaultInitExpr* /*E*/,
83+
DataRecursionQueue* = nullptr) {
84+
return true;
85+
}
86+
// Do not descend into type locations. A generated VarDecl's type can embed
87+
// expressions (a template argument, an array bound, a `std::enable_if<N <
88+
// _Dt>` SFINAE condition) that Clang shares across same-typed decls -- e.g.
89+
// the mersenne engine's `_Dt` word-size constant across the pushforward's
90+
// ValueAndPushforward<> temporaries. Those are Clang's shared AST, not a
91+
// clad reuse; only the statement tree is clad's hand-built output.
92+
// clang-22 added a trailing bool (TraverseQualifier) to these; a defaulted
93+
// parameter matches both the older one-argument call and the new one.
94+
bool TraverseTypeLoc(TypeLoc /*TL*/, bool = false) { return true; }
95+
bool TraverseType(QualType /*T*/, bool = false) { return true; }
96+
bool TraverseTemplateArgumentLoc(const TemplateArgumentLoc& /*ArgLoc*/) {
97+
return true;
98+
}
99+
bool TraverseTemplateArgument(const TemplateArgument& /*Arg*/) {
100+
return true;
101+
}
102+
bool VisitStmt(Stmt* S) {
103+
// RAV can visit one node twice when it is embedded in two type locs
104+
// (e.g. a template-argument ConstantExpr shared by several same-typed
105+
// VarDecls). Count each parent's children once, so a single parent seen
106+
// twice is not mistaken for two parents.
107+
if (!SeenParents.insert(S).second)
108+
return true;
109+
for (const Stmt* Ch : S->children())
110+
if (Ch && !isa<OpaqueValueExpr>(Ch))
111+
++ParentEdges[Ch];
112+
return true;
113+
}
114+
} C;
115+
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
116+
C.TraverseStmt(const_cast<Stmt*>(Root));
117+
for (const auto& Entry : C.ParentEdges)
118+
if (Entry.second >= 2)
119+
return Entry.first;
120+
return nullptr;
121+
}
122+
123+
} // namespace clad

lib/Differentiator/ASTIntegrity.h

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
//--------------------------------------------------------------------*- C++ -//
2+
// clad - the C++ Clang-based Automatic Differentiator
3+
//
4+
// AST-integrity checks for hand-synthesized derivative code.
5+
//
6+
// clad assembles a derivative's statements by hand and must keep them a proper
7+
// tree: no node the child (Stmt::children()) of two parents. A node with two
8+
// parents ("shared") lets a later in-place edit leak across contexts and
9+
// silently corrupt codegen (and a shared aggregate initializer breaks CodeGen
10+
// outright). This is a property of the hand-built body, not of Clang ASTs in
11+
// general -- Clang shares nodes across type and template-argument edges and
12+
// InitListExpr's syntactic/semantic forms -- so the check counts child edges,
13+
// not visits (see findSharedNode).
14+
//----------------------------------------------------------------------------//
15+
16+
#ifndef CLAD_AST_INTEGRITY_H
17+
#define CLAD_AST_INTEGRITY_H
18+
19+
namespace clang {
20+
class Stmt;
21+
} // namespace clang
22+
23+
namespace clad {
24+
25+
/// Return the first Stmt that occurs more than once under \p Root (i.e. is
26+
/// shared with itself through two parent edges), or nullptr if \p Root is a
27+
/// proper tree.
28+
const clang::Stmt* findSharedNode(const clang::Stmt* Root);
29+
30+
} // namespace clad
31+
32+
#endif // CLAD_AST_INTEGRITY_H

0 commit comments

Comments
 (0)