diff --git a/include/clad/Differentiator/BaseForwardModeVisitor.h b/include/clad/Differentiator/BaseForwardModeVisitor.h index 7720e2d3f..5f31ea9e2 100644 --- a/include/clad/Differentiator/BaseForwardModeVisitor.h +++ b/include/clad/Differentiator/BaseForwardModeVisitor.h @@ -174,6 +174,11 @@ class BaseForwardModeVisitor llvm::SmallVectorImpl& clonedArgs, llvm::SmallVectorImpl& derivedArgs); + clang::QualType GetLambdaPushforwardType(const clang::LambdaExpr* LE); + + /// Builds the pushforward lambda for LE + clang::Expr* buildPushforwardLambda(const clang::LambdaExpr* LE); + private: /// Prepares the derivative function parameters. void diff --git a/include/clad/Differentiator/ReverseModeVisitor.h b/include/clad/Differentiator/ReverseModeVisitor.h index a725fa188..18f70ed3a 100644 --- a/include/clad/Differentiator/ReverseModeVisitor.h +++ b/include/clad/Differentiator/ReverseModeVisitor.h @@ -754,11 +754,28 @@ namespace clad { for (const auto* param : FD->parameters()) diffParams.push_back(param); - return utils::GetDerivativeType(m_Sema, FD, DiffMode::pullback, - diffParams, - /*forCustomDerv=*/false, - /*shouldUseRestoreTracker=*/false); + clang::QualType baseTy = + utils::GetDerivativeType(m_Sema, FD, DiffMode::pullback, diffParams, + /*forCustomDerv=*/false, + /*shouldUseRestoreTracker=*/false); + if (LE->capture_size() == 0) + return baseTy; + clang::ASTContext& C = m_Sema.getASTContext(); + const auto* FPT = baseTy->castAs(); + llvm::SmallVector paramTypes( + FPT->param_types().begin(), FPT->param_types().end()); + for (const clang::LambdaCapture& Capture : LE->captures()) { + const auto* capVD = + llvm::cast(Capture.getCapturedVar()); + clang::QualType valTy = utils::getNonConstType( + capVD->getType().getNonReferenceType(), m_Sema); + paramTypes.push_back(valTy); + paramTypes.push_back(C.getPointerType(valTy)); + } + return C.getFunctionType(FPT->getReturnType(), paramTypes, + FPT->getExtProtoInfo()); } + /// Builds the pullback lambda for LE clang::Expr* buildDerivedLambda(const clang::LambdaExpr* LE); /// Builds and returns the sequence of derived function parameters. void BuildParams(llvm::SmallVectorImpl& params, diff --git a/include/clad/Differentiator/VisitorBase.h b/include/clad/Differentiator/VisitorBase.h index d2d694eb5..f4100c48a 100644 --- a/include/clad/Differentiator/VisitorBase.h +++ b/include/clad/Differentiator/VisitorBase.h @@ -26,6 +26,7 @@ #include "clang/Sema/Sema.h" #include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/DenseMap.h" #include "llvm/ADT/StringRef.h" #include "llvm/Support/PrettyStackTrace.h" @@ -34,9 +35,11 @@ #include #include #include +#include namespace clang { class NestedNameSpecifier; +class LambdaExpr; } // namespace clang namespace llvm { @@ -139,6 +142,10 @@ namespace clad { /// See the example inside ForwardModeVisitor::VisitDeclStmt. std::unordered_map m_DeclReplacements; + /// Maps a (lambda, by value captured variable) to its creation snapshot + llvm::DenseMap, + clang::VarDecl*> + m_LambdaCaptureSnapshots; /// A stack of all the blocks where the statements of the gradient function /// are stored (e.g., function body, if statement blocks). std::vector m_Blocks; @@ -210,6 +217,8 @@ namespace clad { return S.ActOnCallExpr(V.getCurrentScope(), lambda, noLoc, {}, noLoc) .get(); } + /// Resolve a captured variable to its clone in the enclosing derivative. + clang::VarDecl* getEnclosingCaptureClone(const clang::VarDecl* capVD); /// For a qualtype QT returns if it's type is Array or Pointer Type static bool isArrayOrPointerType(const clang::QualType QT) { return utils::isArrayOrPointerType(QT); diff --git a/lib/Differentiator/BaseForwardModeVisitor.cpp b/lib/Differentiator/BaseForwardModeVisitor.cpp index c0d120535..b2ea4f6c2 100644 --- a/lib/Differentiator/BaseForwardModeVisitor.cpp +++ b/lib/Differentiator/BaseForwardModeVisitor.cpp @@ -14,6 +14,7 @@ #include "clad/Differentiator/DiffPlanner.h" #include "clad/Differentiator/ErrorEstimator.h" #include "clad/Differentiator/ParseDiffArgsTypes.h" +#include "clad/Differentiator/PushForwardModeVisitor.h" #include "clad/Differentiator/VisitorBase.h" #include "clang/AST/ASTContext.h" @@ -315,6 +316,143 @@ void BaseForwardModeVisitor::SetupDerivativeParameters( } } +clang::QualType +BaseForwardModeVisitor::GetLambdaPushforwardType(const LambdaExpr* LE) { + FunctionDecl* FD = LE->getCallOperator(); + llvm::SmallVector diffParams; + for (const auto* param : FD->parameters()) + diffParams.push_back(param); + QualType baseTy = + utils::GetDerivativeType(m_Sema, FD, DiffMode::pushforward, diffParams, + /*forCustomDerv=*/false, + /*shouldUseRestoreTracker=*/false); + const auto* FPT = baseTy->castAs(); + llvm::SmallVector paramTypes(FPT->param_types().begin(), + FPT->param_types().end()); + for (const LambdaCapture& Capture : LE->captures()) { + const auto* capVD = cast(Capture.getCapturedVar()); + QualType valTy = + utils::getNonConstType(capVD->getType().getNonReferenceType(), m_Sema); + paramTypes.push_back(valTy); + paramTypes.push_back(utils::GetParameterDerivativeType( + m_Sema, DiffMode::pushforward, valTy)); + } + return m_Context.getFunctionType(FPT->getReturnType(), paramTypes, + FPT->getExtProtoInfo()); +} + +#if CLANG_VERSION_MAJOR > 16 +clang::Expr* +BaseForwardModeVisitor::buildPushforwardLambda(const LambdaExpr* LE) { + LambdaIntroducer Intro; + Intro.Default = LCD_None; + Intro.Range.setBegin(noLoc); + Intro.Range.setEnd(noLoc); + AttributeFactory AttrFactory; + const DeclSpec DS(AttrFactory); + Declarator D(DS, clang::ParsedAttributesView::none(), + clang::DeclaratorContext::LambdaExpr); + + QualType dFnType = GetLambdaPushforwardType(LE); + + auto* DC = const_cast(LE->getCallOperator()->getDeclContext()); + // Parent in the enclosing derivative not the original closure. + DeclContext* enclosingDC = m_Sema.CurContext; + llvm::SaveAndRestore SaveContext(m_Sema.CurContext); + llvm::SaveAndRestore SaveDerivative(m_Derivative); + llvm::SaveAndRestore SaveFunctionScope(m_DerivativeFnScope); + beginScope(Scope::LambdaScope | Scope::DeclScope | + Scope::FunctionDeclarationScope | Scope::FunctionPrototypeScope); + m_Sema.PushLambdaScope(); + m_Sema.ActOnLambdaExpressionAfterIntroducer(Intro, getCurrentScope()); + beginScope(Scope::FunctionPrototypeScope | Scope::FunctionDeclarationScope | + Scope::DeclScope); + llvm::SmallVector ParamInfoLambda; + llvm::SmallVector params; + m_Sema.ActOnLambdaClosureQualifiers(Intro, noLoc); + sema::LambdaScopeInfo* LSI = m_Sema.getCurLambda(); + LSI->CallOperator->setType(dFnType); + m_Sema.PushDeclContext(getCurrentScope(), LSI->CallOperator); + LSI->Lambda->setDeclContext(LE->capture_size() ? enclosingDC : DC); + m_Derivative = LSI->CallOperator; + + SetupDerivativeParameters(params); + for (const LambdaCapture& Capture : LE->captures()) { + auto* capVD = cast(Capture.getCapturedVar()); + QualType valTy = + utils::getNonConstType(capVD->getType().getNonReferenceType(), m_Sema); + auto* valPVD = utils::BuildParmVarDecl( + m_Sema, m_Derivative, CreateUniqueIdentifier(capVD->getNameAsString()), + valTy); + m_Sema.PushOnScopeChains(valPVD, getCurrentScope(), /*AddToContext=*/false); + params.push_back(valPVD); + m_DeclReplacements[capVD] = valPVD; + auto* dPVD = utils::BuildParmVarDecl( + m_Sema, m_Derivative, + CreateUniqueIdentifier("_d_" + capVD->getNameAsString()), + utils::GetParameterDerivativeType(m_Sema, DiffMode::pushforward, + valTy)); + m_Sema.PushOnScopeChains(dPVD, getCurrentScope(), /*AddToContext=*/false); + params.push_back(dPVD); + m_Variables[valPVD] = BuildDeclRef(dPVD); + } + + m_Derivative->setBody(MakeCompoundStmt({})); + m_Sema.PopDeclContext(); + for (auto* PVD : params) + ParamInfoLambda.emplace_back(PVD->getIdentifier(), PVD->getLocation(), PVD, + nullptr); + SourceLocation paramListLoc = utils::GetValidSLoc(m_Sema); + QualType retTy = dFnType->castAs()->getReturnType(); + ParsedType trailingReturnType = m_Sema.CreateParsedType( + retTy, m_Context.getTrivialTypeSourceInfo(retTy, paramListLoc)); + D.AddTypeInfo( + DeclaratorChunk::getFunction( + /*hasProto=*/true, /*isAmbiguous=*/false, + /*LParenLoc=*/paramListLoc, /*Params=*/ParamInfoLambda.data(), + /*NumParams=*/ParamInfoLambda.size(), + /*EllipsisLoc=*/SourceLocation(), /*RParenLoc=*/paramListLoc, + /*RefQualifierIsLValueRef=*/true, + /*RefQualifierLoc=*/SourceLocation(), + /*MutableLoc=*/SourceLocation(), /*ESpecType=*/EST_None, + /*ESpecRange=*/SourceRange(), /*Exceptions=*/nullptr, + /*ExceptionRanges=*/nullptr, /*NumExceptions=*/0, + /*NoexceptExpr=*/nullptr, /*ExceptionSpecTokens=*/nullptr, + /*DeclsInPrototype=*/{}, /*LocalRangeBegin=*/noLoc, + /*LocalRangeEnd=*/noLoc, /*Declarator=*/D, + /*TrailingReturnType=*/trailingReturnType, + /*TrailingReturnTypeLoc=*/paramListLoc), + /*EndLoc=*/SourceLocation()); + m_Sema.ActOnLambdaClosureParameters(getCurrentScope(), ParamInfoLambda); + beginScope(Scope::BlockScope | Scope::FnScope | Scope::DeclScope | + Scope::CompoundStmtScope); + m_Sema.ActOnStartOfLambdaDefinition( + Intro, D, + clad_compat::Sema_ActOnStartOfLambdaDefinition_ScopeOrDeclSpec( + getCurrentScope(), DS)); + m_DerivativeFnScope = getCurrentScope(); + + beginBlock(); + StmtDiff bodyDiff = Visit(LE->getCallOperator()->getBody()); + for (auto* S : cast(bodyDiff.getStmt())->body()) + addToCurrentBlock(S); + CompoundStmt* body = endBlock(); + + Expr* lambda = + m_Sema + .ActOnLambdaExpr( + noLoc, + body /*,*/ + CLAD_COMPAT_CLANG17_ActOnLambdaExpr_getCurrentScope_ExtraParam( + *this)) + .get(); + endScope(); + endScope(); + endScope(); + return lambda; +} +#endif // CLANG_VERSION_MAJOR + void BaseForwardModeVisitor::GenerateSeeds(const clang::FunctionDecl* dFD) { // For each function parameter variable, store its derivative value. for (const ParmVarDecl* param : dFD->parameters()) { @@ -1055,6 +1193,78 @@ StmtDiff BaseForwardModeVisitor::VisitCallExpr(const CallExpr* CE) { // Calls to lambda functions are processed differently bool isLambda = isLambdaCallOperator(FD); +#if CLANG_VERSION_MAJOR > 16 + // Bind the captured values and derivatives to lambda's extra parameters. + if (isLambda) { + if (const auto* OCE = dyn_cast(CE)) { + const Expr* origCallee = OCE->getArg(0)->IgnoreParenImpCasts(); + const LambdaExpr* calledLambda = nullptr; + const VarDecl* lambdaVD = nullptr; + if (const auto* DRE = dyn_cast(origCallee)) + if ((lambdaVD = dyn_cast(DRE->getDecl()))) + if (const Expr* vinit = lambdaVD->getInit()) { + if (const auto* EWC = dyn_cast(vinit)) + vinit = EWC->getSubExpr(); + calledLambda = + dyn_cast_or_null(vinit->IgnoreImplicit()); + } + if (calledLambda && calledLambda->capture_size() && lambdaVD) { + VarDecl* dgVD = nullptr; + std::string dgName = "_d_" + lambdaVD->getNameAsString(); + for (Decl* D : m_Derivative->decls()) + if (auto* vd = dyn_cast(D)) + if (vd->getNameAsString() == dgName) { + dgVD = vd; + break; + } + if (dgVD) { + llvm::SmallVector args; + llvm::SmallVector derivs; + for (unsigned i = 1, e = CE->getNumArgs(); i < e; ++i) { + StmtDiff argDiff = Visit(CE->getArg(i)); + args.push_back(argDiff.getExpr()); + derivs.push_back(argDiff.getExpr_dx()); + } + args.append(derivs.begin(), derivs.end()); + for (const LambdaCapture& Capture : calledLambda->captures()) { + if (!Capture.capturesVariable()) + continue; + auto* capVD = cast(Capture.getCapturedVar()); + if (Capture.getCaptureKind() == LCK_ByCopy) { + auto snapIt = + m_LambdaCaptureSnapshots.find({calledLambda, capVD}); + if (snapIt != m_LambdaCaptureSnapshots.end()) { + args.push_back(BuildDeclRef(snapIt->second)); + auto dit = m_Variables.find(snapIt->second); + assert(dit != m_Variables.end() && + "snapshot adjoint must exist"); + args.push_back(Clone(dit->second)); + continue; + } + } + VarDecl* encVD = getEnclosingCaptureClone(capVD); + args.push_back(BuildDeclRef(encVD)); + auto vit = m_Variables.find(encVD); + assert(vit != m_Variables.end() && + "captured variable has no tracked adjoint"); + args.push_back(Clone(vit->second)); + } + Expr* call = m_Sema + .ActOnCallExpr(getCurrentScope(), BuildDeclRef(dgVD), + validLoc, args, validLoc) + .get(); + Expr* vap = StoreAndRef(call, "_t", /*forceDeclCreation=*/true); + Expr* value = + utils::BuildMemberExpr(m_Sema, getCurrentScope(), vap, "value"); + Expr* pushforward = utils::BuildMemberExpr(m_Sema, getCurrentScope(), + vap, "pushforward"); + return StmtDiff(value, pushforward); + } + } + } + } +#endif // CLANG_VERSION_MAJOR + Expr* CUDAExecConfig = nullptr; if (const auto* KCE = dyn_cast(CE)) CUDAExecConfig = Clone(KCE->getConfig()); @@ -1669,6 +1879,68 @@ StmtDiff BaseForwardModeVisitor::VisitDeclStmt(const DeclStmt* DS) { // operator() body -- into the derivative. Rebuild it with a fresh closure // instead; the call's pushforward is regenerated for that closure in // VisitCallExpr, so it still resolves. +#if CLANG_VERSION_MAJOR > 16 + // Synthesize a captureless primal and its pushforward. + if (typeDecl->isLambda() && + !clad::utils::hasNonDifferentiableAttribute(typeDecl)) { + const LambdaExpr* LE = nullptr; + if (const Expr* vinit = VD->getInit()) { + if (const auto* EWC = dyn_cast(vinit)) + vinit = EWC->getSubExpr(); + LE = dyn_cast_or_null(vinit->IgnoreImplicit()); + } + if (LE && LE->capture_size()) { + // Snapshot each by value capture at creation so a later mutation of + // the enclosing variable is not observed. + llvm::SmallVector lambdaDecls; + for (const LambdaCapture& Capture : LE->captures()) { + if (!Capture.capturesVariable() || + Capture.getCaptureKind() != LCK_ByCopy) + continue; + auto* capVD = cast(Capture.getCapturedVar()); + if (m_LambdaCaptureSnapshots.count({LE, capVD})) + continue; + VarDecl* encVD = getEnclosingCaptureClone(capVD); + QualType valTy = utils::getNonConstType( + capVD->getType().getNonReferenceType(), m_Sema); + VarDecl* snap = BuildVarDecl(valTy, "_t", BuildDeclRef(encVD)); + Expr* encDx = nullptr; + auto vit = m_Variables.find(encVD); + if (vit != m_Variables.end()) + encDx = Clone(vit->second); + VarDecl* dsnap = + BuildVarDecl(valTy, "_d_t", encDx ? encDx : getZeroInit(valTy)); + m_Variables[snap] = BuildDeclRef(dsnap); + m_LambdaCaptureSnapshots[{LE, capVD}] = snap; + lambdaDecls.push_back(snap); + lambdaDecls.push_back(dsnap); + } + + DiffRequest lambdaReq = m_DiffReq; + lambdaReq.Function = LE->getCallOperator(); + lambdaReq.Functor = LE->getLambdaClass(); + + Expr* primalE = buildClonedLambda(LE); + QualType AutoTy = m_Context.getAutoDeductType(); + TypeSourceInfo* TSI = m_Context.getTrivialTypeSourceInfo(AutoTy); + VarDecl* g = BuildVarDecl(AutoTy, VD->getNameAsString(), primalE, + VD->isDirectInit(), TSI); + m_DeclReplacements[VD] = g; + + DiffRequest pushforwardReq = lambdaReq; + pushforwardReq.Mode = GetPushForwardMode(); + PushForwardModeVisitor PushforwardVisitor(m_Builder, pushforwardReq); + Expr* pushforwardE = PushforwardVisitor.buildPushforwardLambda(LE); + VarDecl* dg = BuildVarDecl(pushforwardE->getType(), + "_d_" + VD->getNameAsString(), + pushforwardE, VD->isDirectInit()); + + lambdaDecls.push_back(g); + lambdaDecls.push_back(dg); + return StmtDiff(BuildDeclStmt(lambdaDecls), nullptr); + } + } +#endif // CLANG_VERSION_MAJOR for (auto* D : DS->decls()) { assert(isa(D) && "Mixed decl types in a single decl stmt is " "not standard c++ syntax"); diff --git a/lib/Differentiator/DiffPlanner.cpp b/lib/Differentiator/DiffPlanner.cpp index 61d290dc8..58da86adc 100644 --- a/lib/Differentiator/DiffPlanner.cpp +++ b/lib/Differentiator/DiffPlanner.cpp @@ -1208,6 +1208,11 @@ static QualType GetDerivedFunctionType(const CallExpr* CE) { request.EnableVariedAnalysis = false; return true; } + // Captured lambda pushforward is built directly + if (isLambdaCallOperator(MD) && !MD->getParent()->field_empty() && + (m_TopMostReq->Mode == DiffMode::forward || + m_TopMostReq->Mode == DiffMode::vector_forward_mode)) + return true; const CXXRecordDecl* CD = MD->getParent(); if (clad::utils::hasNonDifferentiableAttribute(CD)) nonDiff = true; diff --git a/lib/Differentiator/ReverseModeVisitor.cpp b/lib/Differentiator/ReverseModeVisitor.cpp index 08ef5fa4a..4e4cc8074 100644 --- a/lib/Differentiator/ReverseModeVisitor.cpp +++ b/lib/Differentiator/ReverseModeVisitor.cpp @@ -1818,6 +1818,8 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { auto* DC = const_cast(LE->getCallOperator()->getDeclContext()); + // Parent in the enclosing derivative not the original closure + DeclContext* enclosingDC = m_Sema.CurContext; llvm::SaveAndRestore SaveContext(m_Sema.CurContext); llvm::SaveAndRestore SaveDerivative(m_Derivative); @@ -1841,7 +1843,7 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { m_Sema.PushDeclContext(getCurrentScope(), LSI->CallOperator); - LSI->Lambda->setDeclContext(DC); + LSI->Lambda->setDeclContext(LE->capture_size() ? enclosingDC : DC); m_Derivative = LSI->CallOperator; @@ -1908,6 +1910,21 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { CompoundStmt* DerivedBody = endBlock(); + llvm::SmallVector WL{DerivedBody}; + while (!WL.empty()) { + Stmt* s = WL.pop_back_val(); + if (!s) + continue; + for (auto it = s->child_begin(), e = s->child_end(); it != e; ++it) { + if (auto* dre = dyn_cast_or_null(*it)) + if (dre->refersToEnclosingVariableOrCapture()) + if (auto* vd = dyn_cast(dre->getDecl())) + if (vd->getDeclContext() == m_Derivative) + *it = BuildDeclRef(vd); + WL.push_back(*it); + } + } + if (!m_Globals.empty()) { llvm::SmallVector NewBodyStmts; NewBodyStmts.append(m_Globals.begin(), m_Globals.end()); @@ -1937,6 +1954,29 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { } StmtDiff ReverseModeVisitor::VisitLambdaExpr(const LambdaExpr* LE) { + // Snapshot by value captures at creation; the reverse sweep + // feeds the snapshot adjoint back into the capture + for (const LambdaCapture& Capture : LE->captures()) { + if (!Capture.capturesVariable() || Capture.getCaptureKind() != LCK_ByCopy) + continue; + auto* capVD = cast(Capture.getCapturedVar()); + if (m_LambdaCaptureSnapshots.count({LE, capVD})) + continue; + VarDecl* encVD = getEnclosingCaptureClone(capVD); + QualType ty = utils::getNonConstType( + capVD->getType().getNonReferenceType(), m_Sema); + VarDecl* snap = BuildVarDecl(ty, "_t", BuildDeclRef(encVD)); + addToCurrentBlock(BuildDeclStmt(snap), direction::forward); + VarDecl* dsnap = BuildVarDecl(ty, "_d_t", getZeroInit(ty)); + AddToGlobalBlock(BuildDeclStmt(dsnap)); + m_Variables[snap] = BuildDeclRef(dsnap); + auto it = m_Variables.find(encVD); + if (it != m_Variables.end()) + addToCurrentBlock( + BuildOp(BO_AddAssign, Clone(it->second), BuildDeclRef(dsnap)), + direction::reverse); + m_LambdaCaptureSnapshots[{LE, capVD}] = snap; + } DiffRequest LambdaReq = m_DiffReq; LambdaReq.Function = LE->getCallOperator(); LambdaReq.Functor = LE->getLambdaClass(); @@ -1965,6 +2005,7 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { const auto* MD = dyn_cast(FD); Expr* LambdaCallOpExpr = nullptr; + const LambdaExpr* calledLambda = nullptr; if (MD) { if (isLambdaCallOperator(MD)) { const auto* CallE = CE->getArg(0)->IgnoreParenImpCasts(); @@ -1979,6 +2020,12 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { if (const auto* DRE = llvm::dyn_cast(CallE)) { const auto* VD = dyn_cast(DRE->getDecl()); + if (VD && VD->getInit()) { + const Expr* init = VD->getInit(); + if (const auto* EWC = dyn_cast(init)) + init = EWC->getSubExpr(); + calledLambda = dyn_cast(init->IgnoreImplicit()); + } for (auto* D : m_Derivative->decls()) if (auto* lookupVD = dyn_cast(D)) if (lookupVD->getNameAsString() == @@ -2329,6 +2376,30 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { OverloadedDerivedFn = BuildCallExprToFunction( pullbackFD, pullbackCallArgs, CUDAExecConfig); } else if (MD && isLambdaCallOperator(MD)) { + if (calledLambda) { + for (const LambdaCapture& Capture : calledLambda->captures()) { + if (!Capture.capturesVariable()) + continue; + auto* capVD = cast(Capture.getCapturedVar()); + if (Capture.getCaptureKind() == LCK_ByCopy) { + auto snapIt = + m_LambdaCaptureSnapshots.find({calledLambda, capVD}); + if (snapIt != m_LambdaCaptureSnapshots.end()) { + VarDecl* snap = snapIt->second; + pullbackCallArgs.push_back(BuildDeclRef(snap)); + pullbackCallArgs.push_back( + BuildOp(UO_AddrOf, Clone(m_Variables[snap]))); + continue; + } + } + VarDecl* encVD = getEnclosingCaptureClone(capVD); + pullbackCallArgs.push_back(BuildDeclRef(encVD)); + auto it = m_Variables.find(encVD); + assert(it != m_Variables.end() && + "captured variable has no tracked adjoint"); + pullbackCallArgs.push_back(BuildOp(UO_AddrOf, Clone(it->second))); + } + } OverloadedDerivedFn = m_Sema .ActOnCallExpr( @@ -5141,6 +5212,30 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { params.push_back(dPVD); } + + if (LE && LE->capture_size()) { + for (const LambdaCapture& Capture : LE->captures()) { + auto* capVD = cast(Capture.getCapturedVar()); + QualType valTy = utils::getNonConstType( + capVD->getType().getNonReferenceType(), m_Sema); + auto* valPVD = utils::BuildParmVarDecl( + m_Sema, m_Derivative, + CreateUniqueIdentifier(capVD->getNameAsString()), valTy); + m_Sema.PushOnScopeChains(valPVD, getCurrentScope(), + /*AddToContext=*/false); + params.push_back(valPVD); + m_DeclReplacements[capVD] = valPVD; + auto* adjPVD = utils::BuildParmVarDecl( + m_Sema, m_Derivative, + CreateUniqueIdentifier("_d_" + capVD->getNameAsString()), + m_Context.getPointerType(valTy)); + m_Sema.PushOnScopeChains(adjPVD, getCurrentScope(), + /*AddToContext=*/false); + params.push_back(adjPVD); + m_Variables[valPVD] = + BuildOp(UO_Deref, BuildDeclRef(adjPVD), capVD->getLocation()); + } + } } void ReverseModeVisitor::MarkDeclThreadPrivate(VarDecl* decl) { auto* Init = decl->getInit(); diff --git a/lib/Differentiator/VisitorBase.cpp b/lib/Differentiator/VisitorBase.cpp index 8c5227808..7b2b9f985 100644 --- a/lib/Differentiator/VisitorBase.cpp +++ b/lib/Differentiator/VisitorBase.cpp @@ -21,6 +21,7 @@ #include "clang/AST/Decl.h" #include "clang/AST/DeclCXX.h" #include "clang/AST/Expr.h" +#include "clang/AST/ExprCXX.h" #include "clang/AST/NestedNameSpecifier.h" #include "clang/AST/OperationKinds.h" #include "clang/AST/Stmt.h" @@ -1283,4 +1284,17 @@ namespace clad { } } + clang::VarDecl* + VisitorBase::getEnclosingCaptureClone(const clang::VarDecl* capVD) { + auto it = m_DeclReplacements.find(capVD); + if (it != m_DeclReplacements.end() && it->second) + return it->second; + LookupResult R(m_Sema, capVD->getDeclName(), capVD->getLocation(), + Sema::LookupOrdinaryName); + if (m_Sema.LookupName(R, getCurrentScope()) && R.isSingleResult()) + if (auto* found = dyn_cast(R.getFoundDecl())) + return found; + return const_cast(capVD); + } + } // end namespace clad diff --git a/test/ForwardMode/Lambdas.C b/test/ForwardMode/Lambdas.C index 0e84ec77f..b1b1d3709 100644 --- a/test/ForwardMode/Lambdas.C +++ b/test/ForwardMode/Lambdas.C @@ -1,5 +1,6 @@ // RUN: %cladclang %s -I%S/../../include -oLambdas.out 2>&1 | %filecheck %s // RUN: ./Lambdas.out | %filecheck_exec %s +// UNSUPPORTED: clang-11, clang-12, clang-13, clang-14, clang-15, clang-16 #include "clad/Differentiator/Differentiator.h" @@ -42,6 +43,54 @@ double fn4(double x) { return _f(x); } +double fn5(double x) { + double c = 5.0; + auto _f = [c](double t) { return t * c; }; + return _f(x); +} + +double fn6(double x) { + double c = 2.0; + auto _f = [c](double t) { return t * t * c; }; + return _f(x); +} + +double fn7(double x) { + double d = 4.0; + auto _f = [&d](double t) { return t * d; }; + return _f(x); +} + +double fn8(double x) { + double a = 3.0, b = 7.0; + auto _f = [a, b](double t) { return a * t + b; }; + return _f(x); +} + +double fn9(double x) { + double c = 2.0; + auto _f = [c](double t) { return t * c; }; + c = 100.0; + return _f(x); +} + +double fn10(double x) { + double c = x * x; + auto _f = [c](double) { return c; }; + return _f(x); +} + +double fn11(double x) { + double c = 2.0; + auto _f = [c](double t) { + double s = 0; + for (int k = 0; k < 3; ++k) + s += t * c; + return s; + }; + return _f(x); +} + int main() { auto fn0_dx = clad::differentiate(fn0, 0); printf("Result is = %.2f\n", fn0_dx.execute(7)); // CHECK-EXEC: Result is = 14.00 @@ -62,4 +111,25 @@ int main() { auto fn4_dx = clad::differentiate(fn4, 0); printf("Result is = %.2f\n", fn4_dx.execute(7)); // CHECK-EXEC: Result is = 28.00 printf("Result is = %.2f\n", fn4_dx.execute(-1)); // CHECK-EXEC: Result is = -4.00 + + auto fn5_dx = clad::differentiate(fn5, 0); + printf("Result is = %.2f\n", fn5_dx.execute(3)); // CHECK-EXEC: Result is = 5.00 + + auto fn6_dx = clad::differentiate(fn6, 0); + printf("Result is = %.2f\n", fn6_dx.execute(3)); // CHECK-EXEC: Result is = 12.00 + + auto fn7_dx = clad::differentiate(fn7, 0); + printf("Result is = %.2f\n", fn7_dx.execute(3)); // CHECK-EXEC: Result is = 4.00 + + auto fn8_dx = clad::differentiate(fn8, 0); + printf("Result is = %.2f\n", fn8_dx.execute(3)); // CHECK-EXEC: Result is = 3.00 + + auto fn9_dx = clad::differentiate(fn9, 0); + printf("Result is = %.2f\n", fn9_dx.execute(3)); // CHECK-EXEC: Result is = 2.00 + + auto fn10_dx = clad::differentiate(fn10, 0); + printf("Result is = %.2f\n", fn10_dx.execute(3)); // CHECK-EXEC: Result is = 6.00 + + auto fn11_dx = clad::differentiate(fn11, 0); + printf("Result is = %.2f\n", fn11_dx.execute(2)); // CHECK-EXEC: Result is = 6.00 } diff --git a/test/Gradient/Lambdas.C b/test/Gradient/Lambdas.C index 694094aa9..c37b49aeb 100644 --- a/test/Gradient/Lambdas.C +++ b/test/Gradient/Lambdas.C @@ -76,23 +76,26 @@ double f3(double i, double j) { } // CHECK: void f3_grad(double i, double j, double *_d_i, double *_d_j) { +// CHECK-NEXT: double _d_t = 0.; // CHECK-NEXT: double _d_c = 0.; // CHECK-NEXT: double c = 3.; +// CHECK-NEXT: double _t0 = c; // CHECK-NEXT: auto _f0 = [c](double t) { // CHECK-NEXT: return t * t; // CHECK-NEXT: }; -// CHECK-NEXT: auto _d__f = [](double t, double _d_y, double *_d_t) { +// CHECK-NEXT: auto _d__f = [](double t, double _d_y, double *_d_t0, double c0, double *_d_c0) { // CHECK-NEXT: { -// CHECK-NEXT: *_d_t += _d_y * t; -// CHECK-NEXT: *_d_t += t * _d_y; +// CHECK-NEXT: *_d_t0 += _d_y * t; +// CHECK-NEXT: *_d_t0 += t * _d_y; // CHECK-NEXT: } // CHECK-NEXT: }; // CHECK-NEXT: { // CHECK-NEXT: *_d_i += 1; // CHECK-NEXT: double _r0 = 0.; -// CHECK-NEXT: _d__f(j, 1, &_r0); +// CHECK-NEXT: _d__f(j, 1, &_r0, _t0, &_d_t); // CHECK-NEXT: *_d_j += _r0; // CHECK-NEXT: } +// CHECK-NEXT: _d_c += _d_t; // CHECK-NEXT: } double f4(double i, double j) { @@ -105,22 +108,25 @@ double f4(double i, double j) { } // CHECK: void f4_grad(double i, double j, double *_d_i, double *_d_j) { +// CHECK-NEXT: double _d_t = 0.; // CHECK-NEXT: double _d_c = 0.; // CHECK-NEXT: double c = 3.; // CHECK-NEXT: double _d_d = 0.; // CHECK-NEXT: double d = 4.; +// CHECK-NEXT: double _t0 = c; // CHECK-NEXT: auto _f0 = [c, &d](double t) { // CHECK-NEXT: return t * 2.; // CHECK-NEXT: }; -// CHECK-NEXT: auto _d__f = [](double t, double _d_y, double *_d_t) { -// CHECK-NEXT: *_d_t += _d_y * 2.; +// CHECK-NEXT: auto _d__f = [](double t, double _d_y, double *_d_t0, double c0, double *_d_c0, double d0, double *_d_d0) { +// CHECK-NEXT: *_d_t0 += _d_y * 2.; // CHECK-NEXT: }; // CHECK-NEXT: { // CHECK-NEXT: *_d_i += 1; // CHECK-NEXT: double _r0 = 0.; -// CHECK-NEXT: _d__f(j, 1, &_r0); +// CHECK-NEXT: _d__f(j, 1, &_r0, _t0, &_d_t, d, &_d_d); // CHECK-NEXT: *_d_j += _r0; // CHECK-NEXT: } +// CHECK-NEXT: _d_c += _d_t; // CHECK-NEXT: } const auto _global_f = [](double t) { @@ -173,6 +179,35 @@ double f6(double i) { // CHECK-NEXT: } +double f7(double i, double j) { + double c = 2.0; + auto _f = [c](double t) { return t * c; }; + return i + _f(j); +} + +// CHECK: void f7_grad(double i, double j, double *_d_i, double *_d_j) { +// CHECK-NEXT: double _d_t = 0.; +// CHECK-NEXT: double _d_c = 0.; +// CHECK-NEXT: double c = 2.; +// CHECK-NEXT: double _t0 = c; +// CHECK-NEXT: auto _f0 = [c](double t) { +// CHECK-NEXT: return t * c; +// CHECK-NEXT: }; +// CHECK-NEXT: auto _d__f = [](double t, double _d_y, double *_d_t0, double c0, double *_d_c0) { +// CHECK-NEXT: { +// CHECK-NEXT: *_d_t0 += _d_y * c0; +// CHECK-NEXT: *_d_c0 += t * _d_y; +// CHECK-NEXT: } +// CHECK-NEXT: }; +// CHECK-NEXT: { +// CHECK-NEXT: *_d_i += 1; +// CHECK-NEXT: double _r0 = 0.; +// CHECK-NEXT: _d__f(j, 1, &_r0, _t0, &_d_t); +// CHECK-NEXT: *_d_j += _r0; +// CHECK-NEXT: } +// CHECK-NEXT: _d_c += _d_t; +// CHECK-NEXT: } + int main() { auto df1 = clad::gradient(f1); double di = 0, dj = 0; @@ -203,4 +238,9 @@ int main() { di = 0; df6.execute(3, &di); printf("%.2f\n", di); // CHECK-EXEC: 8.00 + + auto df7 = clad::gradient(f7); + di = 0, dj = 0; + df7.execute(3, 4, &di, &dj); + printf("%.2f %.2f\n", di, dj); // CHECK-EXEC: 1.00 2.00 }