Support differentiation of lambdas with active captures

include/clad/Differentiator/BaseForwardModeVisitor.h

@@ -174,6 +174,11 @@ class BaseForwardModeVisitor
       llvm::SmallVectorImpl<clang::Expr*>& clonedArgs,
       llvm::SmallVectorImpl<clang::Expr*>& 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

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<clang::FunctionProtoType>();
+      llvm::SmallVector<clang::QualType, 8> paramTypes(
+          FPT->param_types().begin(), FPT->param_types().end());
+      for (const clang::LambdaCapture& Capture : LE->captures()) {
+        const auto* capVD =
+            llvm::cast<clang::VarDecl>(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<clang::ParmVarDecl*>& params,

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 <cstddef>
 #include <stack>
 #include <unordered_map>
+#include <utility>
 
 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<const clang::VarDecl*, clang::VarDecl*>
         m_DeclReplacements;
+    /// Maps a (lambda, by value captured variable) to its creation snapshot
+    llvm::DenseMap<std::pair<const clang::LambdaExpr*, const clang::VarDecl*>,
+                   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<Stmts> 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);

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 @@
   }
 }
 
+clang::QualType
+BaseForwardModeVisitor::GetLambdaPushforwardType(const LambdaExpr* LE) {
+  FunctionDecl* FD = LE->getCallOperator();
+  llvm::SmallVector<const ValueDecl*, 4> 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<FunctionProtoType>();
+  llvm::SmallVector<QualType, 8> paramTypes(FPT->param_types().begin(),
+                                            FPT->param_types().end());
+  for (const LambdaCapture& Capture : LE->captures()) {
+    const auto* capVD = cast<VarDecl>(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<DeclContext*>(LE->getCallOperator()->getDeclContext());
+  // Parent in the enclosing derivative not the original closure.
+  DeclContext* enclosingDC = m_Sema.CurContext;
+  llvm::SaveAndRestore<DeclContext*> SaveContext(m_Sema.CurContext);
+  llvm::SaveAndRestore<FunctionDecl*> SaveDerivative(m_Derivative);
+  llvm::SaveAndRestore<Scope*> 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<DeclaratorChunk::ParamInfo> ParamInfoLambda;
+  llvm::SmallVector<ParmVarDecl*, 8> 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<VarDecl>(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<FunctionProtoType>()->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<CompoundStmt>(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 @@
   // 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<CXXOperatorCallExpr>(CE)) {
+      const Expr* origCallee = OCE->getArg(0)->IgnoreParenImpCasts();
+      const LambdaExpr* calledLambda = nullptr;
+      const VarDecl* lambdaVD = nullptr;
+      if (const auto* DRE = dyn_cast<DeclRefExpr>(origCallee))
+        if ((lambdaVD = dyn_cast<VarDecl>(DRE->getDecl())))
+          if (const Expr* vinit = lambdaVD->getInit()) {
+            if (const auto* EWC = dyn_cast<ExprWithCleanups>(vinit))
+              vinit = EWC->getSubExpr();
+            calledLambda =
+                dyn_cast_or_null<LambdaExpr>(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<VarDecl>(D))
+            if (vd->getNameAsString() == dgName) {
+              dgVD = vd;
+              break;
+            }
+        if (dgVD) {
+          llvm::SmallVector<Expr*, 8> args;
+          llvm::SmallVector<Expr*, 8> 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<VarDecl>(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<CUDAKernelCallExpr>(CE))
     CUDAExecConfig = Clone(KCE->getConfig());
@@ -1669,6 +1879,68 @@
       // 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<ExprWithCleanups>(vinit))
+            vinit = EWC->getSubExpr();
+          LE = dyn_cast_or_null<LambdaExpr>(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<Decl*, 4> lambdaDecls;
+          for (const LambdaCapture& Capture : LE->captures()) {
+            if (!Capture.capturesVariable() ||
+                Capture.getCaptureKind() != LCK_ByCopy)
+              continue;
+            auto* capVD = cast<VarDecl>(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<VarDecl>(D) && "Mixed decl types in a single decl stmt is "
                                   "not standard c++ syntax");

lib/Differentiator/DiffPlanner.cpp

@@ -1208,6 +1208,11 @@
           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;