Replace VisitReturnStmt goto/label with named-lambda call - #1861
Replace VisitReturnStmt goto/label with named-lambda call#1861vgvassilev wants to merge 3 commits into
Conversation
VisitorBase::wrapInLambda used to build a lambda and immediately call it (IIFE-style), returning the CallExpr. Construction and invocation were entangled, so the helper could not be used by callers that need to bind the lambda to a VarDecl and invoke it from multiple sites. Split the helper. The new buildLambda performs the lambda construction and returns the LambdaExpr; wrapInLambda is now a thin wrapper that calls buildLambda followed by ActOnCallExpr to preserve the existing IIFE return shape. The single existing caller in BaseForwardModeVisitor goes through the wrapper unchanged. Scaffolding for replacing the goto/label emission in VisitReturnStmt with a [&] lambda bound to a VarDecl and called from each return path (#367).
wrapInLambda produces an IIFE call and buildLambda returns a bare LambdaExpr; neither covers the case where the same lambda body must be invoked from multiple sites. The next step in removing goto/label emission from VisitReturnStmt requires exactly that. Add buildAndBindLambda: routes through buildLambda, binds the result to a fresh VarDecl with a unique identifier, returns the decl so the caller can construct DeclStmt and CallExpr nodes at appropriate sites. Templated and unused in this commit; no instantiation cost. Scaffolding for #367.
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| /// where `_rev` is a [&] lambda wrapping the master reverse sweep. | ||
| /// Replaces the goto/label encoding (vgvassilev/clad#367) so the body | ||
| /// stays well-formed under [stmt.dcl]/3 and Sema::ActOnFinishFunctionBody. | ||
| llvm::SmallPtrSet<clang::Stmt*, 4> m_EarlyReturnMarkers; |
There was a problem hiding this comment.
warning: member variable 'm_EarlyReturnMarkers' has protected visibility [cppcoreguidelines-non-private-member-variables-in-classes]
llvm::SmallPtrSet<clang::Stmt*, 4> m_EarlyReturnMarkers;
^| /// where `_rev` is a [&] lambda wrapping the master reverse sweep. | ||
| /// Replaces the goto/label encoding (vgvassilev/clad#367) so the body | ||
| /// stays well-formed under [stmt.dcl]/3 and Sema::ActOnFinishFunctionBody. | ||
| llvm::SmallPtrSet<clang::Stmt*, 4> m_EarlyReturnMarkers; |
There was a problem hiding this comment.
warning: no header providing "llvm::SmallPtrSet" is directly included [misc-include-cleaner]
include/clad/Differentiator/ReverseModeVisitor.h:34:
- #include <memory>
+ #include <llvm/ADT/SmallPtrSet.h>
+ #include <memory>| /// also has early returns so the seed can be emitted on the natural-tail | ||
| /// path only (outside the lambda) — applying it inside the lambda would | ||
| /// double-seed on every early-return path. | ||
| clang::Stmt* m_NaturalReturnSeed = nullptr; |
There was a problem hiding this comment.
warning: member variable 'm_NaturalReturnSeed' has protected visibility [cppcoreguidelines-non-private-member-variables-in-classes]
.
^| const clang::Expr* E, F&& func) { | ||
| static clang::Expr* | ||
| buildLambda(VisitorBase& V, clang::Sema& S, const clang::Stmt* LocSrc, | ||
| F&& func, |
There was a problem hiding this comment.
warning: forwarding reference parameter 'func' is never forwarded inside the function body [cppcoreguidelines-missing-std-forward]
,
^| clang::LambdaIntroducer Intro; | ||
| Intro.Default = clang::LCD_ByRef; | ||
| if (ExplicitCaptures.empty()) { | ||
| Intro.Default = clang::LCD_ByRef; |
There was a problem hiding this comment.
warning: no header providing "clang::LCD_ByRef" is directly included [misc-include-cleaner]
include/clad/Differentiator/VisitorBase.h:33:
- #include <cstddef>
+ #include <clang/Basic/Lambda.h>
+ #include <cstddef>| if (ExplicitCaptures.empty()) { | ||
| Intro.Default = clang::LCD_ByRef; | ||
| } else { | ||
| Intro.Default = clang::LCD_None; |
There was a problem hiding this comment.
warning: no header providing "clang::LCD_None" is directly included [misc-include-cleaner]
{
^| Intro.Default = clang::LCD_None; | ||
| for (clang::VarDecl* VD : ExplicitCaptures) | ||
| Intro.Captures.emplace_back( | ||
| /*Kind=*/clang::LCK_ByRef, |
There was a problem hiding this comment.
warning: no header providing "clang::LCK_ByRef" is directly included [misc-include-cleaner]
(
^| /// added by addToCurrentBlock from func invocation. | ||
| template <typename F> | ||
| static clang::Expr* wrapInLambda(VisitorBase& V, clang::Sema& S, | ||
| const clang::Stmt* LocSrc, F&& func) { |
There was a problem hiding this comment.
warning: forwarding reference parameter 'func' is never forwarded inside the function body [cppcoreguidelines-missing-std-forward]
,
^| template <typename F> | ||
| clang::VarDecl* | ||
| buildAndBindLambda(const clang::Stmt* LocSrc, llvm::StringRef NameHint, | ||
| F&& func, |
There was a problem hiding this comment.
warning: forwarding reference parameter 'func' is never forwarded inside the function body [cppcoreguidelines-missing-std-forward]
nt,
^| // explicit captures so Sema synthesizes one closure FieldDecl per | ||
| // VarDecl up front and codegen translates each body reference into | ||
| // the corresponding field access ([expr.prim.lambda.capture]/12). | ||
| llvm::SmallSetVector<VarDecl*, 16> Referenced; |
There was a problem hiding this comment.
warning: no header providing "llvm::SmallSetVector" is directly included [misc-include-cleaner]
lib/Differentiator/ReverseModeVisitor.cpp:65:
- #include <memory>
+ #include <llvm/ADT/SetVector.h>
+ #include <memory>The previous reverse-mode encoding emitted a label inside the master reverse and a goto into it from each early-return point. This violates [stmt.dcl]/3 (https://eel.is/c++draft/stmt.dcl#3) whenever a local with a non-trivial initializer or destructor sits between the goto source and its target — the reason DifferentiateWithClad and the forward-mode counterpart have to bypass Sema::ActOnFinishFunctionBody on the derived function (#367). Replace the encoding. At each early return, emit a NullStmt marker and add the adjoint-seed Reverse to the current reverse block (no label wrapping). At function-body finalization, when m_EarlyReturnMarkers is non-empty: materialise the master reverse as a [&]-capture lambda bound to an `auto`-typed VarDecl; walk the forward block to patch each marker with `{ _rev(); return; }`; emit the saved tail-return seed and `_rev();` at the natural-return position. The auto deduction makes the generated source render as `auto _rev0 = [&] {...};` rather than the unspellable closure type. Captures: the body's DeclRefExprs were built outside the lambda scope during the original Visit, so Sema's [&] capture-default scan does not detect them. Collect every outer VarDecl referenced in the body, pass them as explicit captures (Intro.Captures), and inside the lambda scope rebuild each captured DeclRefExpr via Sema::BuildDeclRefExpr so the RefersToEnclosingVariableOrCapture bit is set on the new DRE — codegen relies on that bit to translate references into closure FieldDecl accesses per [expr.prim.lambda.capture]/12. m_EarlyReturnMarkers and m_NaturalReturnSeed are reset at the start of each DifferentiateWithClad call to avoid leaking state across multiple functions handled by the same visitor instance. Status: verified correct under an asan-instrumented clad build (the f3 early-return case in test/Gradient/Loops.C produces gradient {6.00}). A release-build crash in EmitDeclRefLValue does not reproduce under +asan, indicating a layout-sensitive bug elsewhere rather than an issue with the AD logic; tracked as a follow-up. Scaffolding for #367 — Sema::ActOnFinishFunctionBody restoration lands in a follow-up.
7375672 to
70b2687
Compare
| // the corresponding field access ([expr.prim.lambda.capture]/12). | ||
| llvm::SmallSetVector<VarDecl*, 16> Referenced; | ||
| llvm::SmallPtrSet<VarDecl*, 16> LocalToBody; | ||
| class CaptureCollector : public RecursiveASTVisitor<CaptureCollector> { |
There was a problem hiding this comment.
warning: no header providing "clang::RecursiveASTVisitor" is directly included [misc-include-cleaner]
lib/Differentiator/ReverseModeVisitor.cpp:63:
- #include <cstddef>
+ #include <clang/AST/RecursiveASTVisitor.h>
+ #include <cstddef>| llvm::SmallPtrSet<VarDecl*, 16> LocalToBody; | ||
| class CaptureCollector : public RecursiveASTVisitor<CaptureCollector> { | ||
| public: | ||
| llvm::SmallSetVector<VarDecl*, 16>& Ref; |
There was a problem hiding this comment.
warning: member 'Ref' of type 'llvm::SmallSetVector<VarDecl *, 16> &' is a reference [cppcoreguidelines-avoid-const-or-ref-data-members]
llvm::SmallSetVector<VarDecl*, 16>& Ref;
^| class CaptureCollector : public RecursiveASTVisitor<CaptureCollector> { | ||
| public: | ||
| llvm::SmallSetVector<VarDecl*, 16>& Ref; | ||
| llvm::SmallPtrSet<VarDecl*, 16>& Local; |
There was a problem hiding this comment.
warning: member 'Local' of type 'llvm::SmallPtrSet<VarDecl *, 16> &' is a reference [cppcoreguidelines-avoid-const-or-ref-data-members]
llvm::SmallPtrSet<VarDecl*, 16>& Local;
^| class CapRefRebuilder | ||
| : public RecursiveASTVisitor<CapRefRebuilder> { | ||
| public: | ||
| const llvm::SmallPtrSetImpl<VarDecl*>& CapSet; |
There was a problem hiding this comment.
warning: member 'CapSet' of type 'const llvm::SmallPtrSetImpl<VarDecl *> &' is a reference [cppcoreguidelines-avoid-const-or-ref-data-members]
ic:
^| class CapRefRebuilder | ||
| : public RecursiveASTVisitor<CapRefRebuilder> { | ||
| public: | ||
| const llvm::SmallPtrSetImpl<VarDecl*>& CapSet; |
There was a problem hiding this comment.
warning: no header providing "llvm::SmallPtrSetImpl" is directly included [misc-include-cleaner]
lib/Differentiator/ReverseModeVisitor.cpp:65:
- #include <memory>
+ #include <llvm/ADT/SmallPtrSet.h>
+ #include <memory>| : public RecursiveASTVisitor<CapRefRebuilder> { | ||
| public: | ||
| const llvm::SmallPtrSetImpl<VarDecl*>& CapSet; | ||
| llvm::function_ref<DeclRefExpr*(VarDecl*, ExprValueKind)> Rebuild; |
There was a problem hiding this comment.
warning: no header providing "llvm::function_ref" is directly included [misc-include-cleaner]
lib/Differentiator/ReverseModeVisitor.cpp:65:
- #include <memory>
+ #include <llvm/ADT/STLFunctionalExtras.h>
+ #include <memory>| // returns a writable range of Stmt*&. | ||
| class Patcher : public RecursiveASTVisitor<Patcher> { | ||
| public: | ||
| const llvm::SmallPtrSetImpl<Stmt*>& Markers; |
There was a problem hiding this comment.
warning: member 'Markers' of type 'const llvm::SmallPtrSetImpl<Stmt *> &' is a reference [cppcoreguidelines-avoid-const-or-ref-data-members]
public:
^1 participant
The previous reverse-mode encoding emitted a label inside the master reverse and a goto into it from each early-return point. This violates [stmt.dcl]/3 (https://eel.is/c++draft/stmt.dcl#3) whenever a local with a non-trivial initializer or destructor sits between the goto source and its target — the reason DifferentiateWithClad and the forward-mode counterpart have to bypass Sema::ActOnFinishFunctionBody on the derived function (#367).
Replace the encoding. At each early return, emit a NullStmt marker and add the adjoint-seed Reverse to the current reverse block (no label wrapping). At function-body finalization, when m_EarlyReturnMarkers is non-empty: materialise the master reverse as a [&]-capture lambda bound to an
auto-typed VarDecl; walk the forward block to patch each marker with{ _rev(); return; }; emit the saved tail-return seed and_rev();at the natural-return position. The auto deduction makes the generated source render asauto _rev0 = [&] {...};rather than the unspellable closure type.Captures: the body's DeclRefExprs were built outside the lambda scope during the original Visit, so Sema's [&] capture-default scan does not detect them. Collect every outer VarDecl referenced in the body, pass them as explicit captures (Intro.Captures), and inside the lambda scope rebuild each captured DeclRefExpr via Sema::BuildDeclRefExpr so the
RefersToEnclosingVariableOrCapture bit is set on the new DRE — codegen relies on that bit to translate references into closure FieldDecl accesses per [expr.prim.lambda.capture]/12.
m_EarlyReturnMarkers and m_NaturalReturnSeed are reset at the start of each DifferentiateWithClad call to avoid leaking state across multiple functions handled by the same visitor instance.
Status: verified correct under an asan-instrumented clad build (the f3 early-return case in test/Gradient/Loops.C produces gradient {6.00}). A release-build crash in EmitDeclRefLValue does not reproduce under+asan, indicating a layout-sensitive bug elsewhere rather than an issue with the AD logic; tracked as a follow-up.
Scaffolding for #367 — Sema::ActOnFinishFunctionBody restoration lands in a follow-up.