Skip to content

Commit 9a4914d

Browse files
PetroZarytskyiPetroZarytskyi
committed
Schedule reverse_forw functions statically.
Before this PR, all `reverse_forw` functions were scheduled only dynamically. Now, they are always scheduled statically.
1 parent dec617a commit 9a4914d

10 files changed

Lines changed: 120 additions & 154 deletions

File tree

‎include/clad/Differentiator/DiffMode.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ inline const char* DiffModeToString(DiffMode mode) {
3939
case DiffMode::jacobian:
4040
return "jacobian";
4141
case DiffMode::reverse_mode_forward_pass:
42-
return "reverse_mode_forward_pass";
42+
return "reverse_forw";
4343
case DiffMode::error_estimation:
4444
return "error_estimation";
4545
default:

‎include/clad/Differentiator/STLBuiltins.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -668,7 +668,8 @@ void constructor_pullback(T* p, ::std::unique_ptr<T>* dthis, T* dp) noexcept {};
668668

669669
template <typename T>
670670
clad::ValueAndAdjoint<T&, T&>
671-
operator_star_reverse_forw(::std::unique_ptr<T>* u, ::std::unique_ptr<T>* d_u) {
671+
operator_star_reverse_forw(const ::std::unique_ptr<T>* u,
672+
const ::std::unique_ptr<T>* d_u) {
672673
return {**u, **d_u};
673674
}
674675

‎lib/Differentiator/CladUtils.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1002,7 +1002,8 @@ namespace clad {
10021002
mode == DiffMode::pullback ||
10031003
mode == DiffMode::error_estimation ||
10041004
mode == DiffMode::vector_forward_mode;
1005-
if (mode == DiffMode::reverse_mode_forward_pass) {
1005+
if (mode == DiffMode::reverse_mode_forward_pass &&
1006+
!oRetTy->isVoidType()) {
10061007
TemplateDecl* valAndAdjointTempDecl =
10071008
utils::LookupTemplateDeclInCladNamespace(S, "ValueAndAdjoint");
10081009
dRetTy = utils::InstantiateTemplate(S, valAndAdjointTempDecl,

‎lib/Differentiator/DerivativeBuilder.cpp‎

Lines changed: 10 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -325,22 +325,19 @@ static void registerDerivative(Decl* D, Sema& S, const DiffRequest& R) {
325325
}
326326
}
327327

328+
// FIXME: Figure out how to assert here in cases where we have provided
329+
// both a clad-generated derivative and a user-provided one.
330+
// #ifndef NDEBUG
331+
// LookupResult R1 = utils::LookupQualifiedName(Name, m_Sema,
332+
// originalFnDC); assert((R1.empty() || R1.getFoundDecl() ==
333+
// R.getFoundDecl()) &&
334+
// "We clad built a derivative for entity which"
335+
// "has a custom derivative!");
336+
// #endif // NDEBUG
328337
CXXScopeSpec SS;
329338
LookupResult R = LookupCustomDerivativeOrNumericalDiff(
330339
Name, originalFnDC, SS, forCustomDerv, namespaceShouldExist);
331-
bool hasUserDefinedPropagator = false;
332-
if (!R.empty()) {
333-
hasUserDefinedPropagator = true;
334-
#ifndef NDEBUG
335-
// FIXME: Figure out how to assert here in cases where we have provided
336-
// both a clad-generated derivative and a user-provided one.
337-
// LookupResult R1 = utils::LookupQualifiedName(Name, m_Sema,
338-
// originalFnDC); assert((R1.empty() || R1.getFoundDecl() ==
339-
// R.getFoundDecl()) &&
340-
// "We clad built a derivative for entity which"
341-
// "has a custom derivative!");
342-
#endif // NDEBUG
343-
} else {
340+
if (R.empty()) {
344341
// Try to find if clad already built a derivative.
345342
R = utils::LookupQualifiedName(Name, m_Sema, originalFnDC);
346343
if (originalFnDC && !originalFnDC->isRecord())
@@ -360,9 +357,6 @@ static void registerDerivative(Decl* D, Sema& S, const DiffRequest& R) {
360357
// Loc = m_DiffReq->getLocation();
361358
UnqualifiedId Member;
362359
Member.setIdentifier(&m_Context.Idents.get(Name), Loc);
363-
if (auto* UO = dyn_cast<UnaryOperator>(Base))
364-
if (UO->getOpcode() == UO_AddrOf)
365-
Base = UO->getSubExpr();
366360
bool isArrow = Base->getType()->isPointerType();
367361
// FIXME: update SS here?
368362
auto* ME =
@@ -381,12 +375,6 @@ static void registerDerivative(Decl* D, Sema& S, const DiffRequest& R) {
381375
CUDAExecConfig)
382376
.get();
383377
}
384-
// If we have user-defined propagators we pass the first argument by
385-
// pointer.
386-
if (hasUserDefinedPropagator &&
387-
!CallArgs[0]->getType()->isPointerType())
388-
CallArgs[0] =
389-
m_Sema.BuildUnaryOp(S, noLoc, UO_AddrOf, CallArgs[0]).get();
390378
}
391379
Expr* UnresolvedLookup =
392380
m_Sema.BuildDeclarationNameExpr(SS, R, /*ADL*/ false).get();

‎lib/Differentiator/DiffPlanner.cpp‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1213,6 +1213,21 @@ DeclRefExpr* getArgFunction(CallExpr* call, Sema& SemaRef) {
12131213
if (request.Function->getDefinition())
12141214
request.Function = request.Function->getDefinition();
12151215

1216+
QualType returnType = FD->getReturnType();
1217+
bool needsForwPass = utils::isNonConstReferenceType(returnType) ||
1218+
returnType->isPointerType();
1219+
if (request.Mode == DiffMode::pullback ||
1220+
request.Mode == DiffMode::reverse) {
1221+
DiffRequest forwPassRequest = request;
1222+
forwPassRequest.DVI.clear();
1223+
forwPassRequest.Mode = DiffMode::reverse_mode_forward_pass;
1224+
forwPassRequest.EnableTBRAnalysis = false;
1225+
forwPassRequest.EnableVariedAnalysis = false;
1226+
forwPassRequest.EnableUsefulAnalysis = false;
1227+
if (LookupCustomDerivativeDecl(forwPassRequest) || needsForwPass)
1228+
m_DiffRequestGraph.addNode(forwPassRequest, /*isSource=*/true);
1229+
}
1230+
12161231
if (!LookupCustomDerivativeDecl(request)) {
12171232
clang::CFG::BuildOptions Options;
12181233
std::unique_ptr<AnalysisDeclContext> AnalysisDC =

‎lib/Differentiator/ReverseModeVisitor.cpp‎

Lines changed: 43 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -1767,7 +1767,6 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
17671767
// Stores differentiation result of implicit `this` object, if any.
17681768
StmtDiff baseDiff;
17691769
Expr* baseExpr = nullptr;
1770-
Stmt* baseDiffPush = nullptr;
17711770
size_t idx = 0;
17721771

17731772
/// Add base derivative expression in the derived call output args list if
@@ -1797,26 +1796,11 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
17971796
if (baseTy->isPointerType())
17981797
baseTy = baseTy->getPointeeType();
17991798
CXXRecordDecl* baseRD = baseTy->getAsCXXRecordDecl();
1800-
bool shouldStore = utils::isCopyable(baseRD);
1801-
if (shouldStore) {
1802-
if (!isPassedByRef || MD->isConst()) {
1803-
// FIXME: Custom _reverse_forw functions take the base by pointer.
1804-
// This means we cannot pass temporary values and because of that,
1805-
// we need to store. For now, only emit the store stmt if a custom
1806-
// reverse_forw is found.
1807-
beginBlock(direction::forward);
1808-
Expr* baseDiffStore =
1809-
StoreAndRef(baseDiff.getExpr(), direction::forward);
1810-
baseDiffPush =
1811-
utils::unwrapIfSingleStmt(endBlock(direction::forward));
1812-
baseExpr = baseDiffStore;
1813-
} else {
1814-
Expr* baseDiffStore =
1815-
GlobalStoreAndRef(baseDiff.getExpr(), "_t", /*force=*/true);
1816-
Expr* assign =
1817-
BuildOp(BO_Assign, baseDiff.getExpr(), baseDiffStore);
1818-
PreCallStmts.push_back(assign);
1819-
}
1799+
if (isPassedByRef && !MD->isConst() && utils::isCopyable(baseRD)) {
1800+
Expr* baseDiffStore =
1801+
GlobalStoreAndRef(baseDiff.getExpr(), "_t", /*force=*/true);
1802+
Expr* assign = BuildOp(BO_Assign, baseDiff.getExpr(), baseDiffStore);
1803+
PreCallStmts.push_back(assign);
18201804
}
18211805
Expr* baseDerivative = baseDiff.getExpr_dx();
18221806
if (!baseDerivative->getType()->isPointerType())
@@ -2017,61 +2001,52 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
20172001
return StmtDiff(Clone(CE));
20182002

20192003
Expr* call = nullptr;
2020-
// Stores the dx of the call arguments for the function to be derived
2021-
for (std::size_t i = 0, e = CE->getNumArgs() - isMethodOperatorCall; i != e;
2022-
++i) {
2023-
const Expr* arg = CE->getArg(i + isMethodOperatorCall);
2024-
if (!utils::IsReferenceOrPointerArg(arg) || arg->isXValue())
2025-
CallArgDx[i] = getZeroInit(arg->getType());
2026-
}
2027-
if (baseDiff.getExpr_dx() &&
2028-
!baseDiff.getExpr_dx()->getType()->isPointerType())
2029-
CallArgDx.insert(CallArgDx.begin(), BuildOp(UnaryOperatorKind::UO_AddrOf,
2030-
baseDiff.getExpr_dx(), Loc));
2031-
2032-
if (Expr* customForwardPassCE =
2033-
BuildCallToCustomForwPassFn(CE, CallArgs, CallArgDx, baseExpr)) {
2034-
addToCurrentBlock(baseDiffPush, direction::forward);
2035-
if (!needsForwPass)
2036-
return StmtDiff{customForwardPassCE};
2037-
Expr* callRes = nullptr;
2038-
if (isInsideLoop)
2039-
callRes = GlobalStoreAndRef(customForwardPassCE, /*prefix=*/"_t",
2040-
/*force=*/true);
2041-
else
2042-
callRes = StoreAndRef(customForwardPassCE);
2043-
auto* resValue =
2044-
utils::BuildMemberExpr(m_Sema, getCurrentScope(), callRes, "value");
2045-
auto* resAdjoint =
2046-
utils::BuildMemberExpr(m_Sema, getCurrentScope(), callRes, "adjoint");
2047-
return StmtDiff(resValue, resAdjoint);
2048-
}
2049-
if (needsForwPass) {
2050-
DiffRequest calleeFnForwPassReq;
2051-
calleeFnForwPassReq.Function = FD;
2052-
calleeFnForwPassReq.Mode = DiffMode::reverse_mode_forward_pass;
2053-
calleeFnForwPassReq.BaseFunctionName =
2054-
clad::utils::ComputeEffectiveFnName(FD);
2055-
calleeFnForwPassReq.VerboseDiags = true;
2056-
2057-
FunctionDecl* calleeFnForwPassFD =
2058-
m_Builder.HandleNestedDiffRequest(calleeFnForwPassReq);
2059-
2060-
assert(calleeFnForwPassFD &&
2061-
"Clad failed to generate callee function forward pass function");
2062-
2004+
// Lookup a reverse_forw function and build if necessary.
2005+
DiffRequest calleeFnForwPassReq;
2006+
calleeFnForwPassReq.Function = FD;
2007+
calleeFnForwPassReq.Mode = DiffMode::reverse_mode_forward_pass;
2008+
calleeFnForwPassReq.BaseFunctionName =
2009+
clad::utils::ComputeEffectiveFnName(FD);
2010+
calleeFnForwPassReq.VerboseDiags = true;
2011+
2012+
FunctionDecl* calleeFnForwPassFD = FindDerivedFunction(calleeFnForwPassReq);
2013+
if (calleeFnForwPassFD) {
2014+
for (std::size_t i = 0, e = CE->getNumArgs() - isMethodOperatorCall;
2015+
i != e; ++i) {
2016+
const Expr* arg = CE->getArg(i + isMethodOperatorCall);
2017+
if (!utils::IsReferenceOrPointerArg(arg) || arg->isXValue())
2018+
CallArgDx[i] = getZeroInit(arg->getType());
2019+
}
2020+
if (baseDiff.getExpr_dx() &&
2021+
!baseDiff.getExpr_dx()->getType()->isPointerType())
2022+
CallArgDx.insert(
2023+
CallArgDx.begin(),
2024+
BuildOp(UnaryOperatorKind::UO_AddrOf, baseDiff.getExpr_dx(), Loc));
20632025
CallArgs.insert(CallArgs.end(), CallArgDx.begin(), CallArgDx.end());
2064-
if (Expr* baseE = baseDiff.getExpr()) {
2065-
call = BuildCallExprToMemFn(baseE, calleeFnForwPassFD->getName(),
2066-
CallArgs, Loc);
2026+
const auto* forwPassMD = dyn_cast<CXXMethodDecl>(calleeFnForwPassFD);
2027+
Expr* baseE = baseDiff.getExpr();
2028+
if (forwPassMD && forwPassMD->isInstance()) {
2029+
call = BuildCallExprToMemFn(
2030+
baseDiff.getExpr(), calleeFnForwPassFD->getName(), CallArgs, Loc);
20672031
} else {
2032+
if (baseE) {
2033+
baseE = BuildOp(UO_AddrOf, baseE);
2034+
CallArgs.insert(CallArgs.begin(), baseE);
2035+
}
20682036
call = m_Sema
20692037
.ActOnCallExpr(getCurrentScope(),
20702038
BuildDeclRef(calleeFnForwPassFD), Loc,
20712039
CallArgs, Loc, CUDAExecConfig)
20722040
.get();
20732041
}
2074-
auto* callRes = StoreAndRef(call);
2042+
if (call->getType()->isVoidType())
2043+
return StmtDiff(call);
2044+
Expr* callRes = nullptr;
2045+
if (isInsideLoop)
2046+
callRes = GlobalStoreAndRef(call, /*prefix=*/"_t",
2047+
/*force=*/true);
2048+
else
2049+
callRes = StoreAndRef(call);
20752050
auto* resValue =
20762051
utils::BuildMemberExpr(m_Sema, getCurrentScope(), callRes, "value");
20772052
auto* resAdjoint =

‎test/Gradient/FunctionCalls.C‎

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -343,6 +343,14 @@ double fn7(double i, double j) {
343343
// CHECK-NEXT: }
344344
// CHECK-NEXT: }
345345

346+
// CHECK: clad::ValueAndAdjoint<double &, double &> identity_reverse_forw(double &i, double &_d_i) {
347+
// CHECK-NEXT: MyStruct::myFunction();
348+
// CHECK-NEXT: double _d__d_i = 0.;
349+
// CHECK-NEXT: double _d_i0 = i;
350+
// CHECK-NEXT: _d_i0 += 1;
351+
// CHECK-NEXT: return {i, _d_i};
352+
// CHECK-NEXT: }
353+
346354
// CHECK: void identity_pullback(double &i, double _d_y, double *_d_i) {
347355
// CHECK-NEXT: MyStruct::myFunction();
348356
// CHECK-NEXT: double _d__d_i = 0.;
@@ -361,8 +369,6 @@ double fn7(double i, double j) {
361369
// CHECK-NEXT: *_d_i += _d_y;
362370
// CHECK-NEXT: }
363371

364-
// CHECK: clad::ValueAndAdjoint<double &, double &> identity_reverse_forw(double &i, double &_d_i);
365-
366372
// CHECK: void fn7_grad(double i, double j, double *_d_i, double *_d_j) {
367373
// CHECK-NEXT: double _t0 = i;
368374
// CHECK-NEXT: clad::ValueAndAdjoint<double &, double &> _t1 = identity_reverse_forw(i, *_d_i);
@@ -1079,13 +1085,6 @@ double sq_defined_later(double x) {
10791085
double fn25_defined_later(double x) {
10801086
return fn25(x);
10811087
}
1082-
// CHECK: clad::ValueAndAdjoint<double &, double &> identity_reverse_forw(double &i, double &_d_i) {
1083-
// CHECK-NEXT: MyStruct::myFunction();
1084-
// CHECK-NEXT: double _d__d_i = 0.;
1085-
// CHECK-NEXT: double _d_i0 = i;
1086-
// CHECK-NEXT: _d_i0 += 1;
1087-
// CHECK-NEXT: return {i, _d_i};
1088-
// CHECK-NEXT: }
10891088

10901089
// CHECK: void weighted_sum_pullback(double *x, const double *w, double _d_y, double *_d_x) {
10911090
// CHECK-NEXT: {

‎test/Gradient/MemberFunctions.C‎

Lines changed: 16 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -439,6 +439,12 @@ double fn2(SimpleFunctions& sf, double i) {
439439
return sf.ref_mem_fn(i);
440440
}
441441

442+
// CHECK: clad::ValueAndAdjoint<double &, double &> ref_mem_fn_reverse_forw(double i, SimpleFunctions *_d_this, double _d_i) {
443+
// CHECK-NEXT: this->x = +i;
444+
// CHECK-NEXT: this->x = -i;
445+
// CHECK-NEXT: return {this->x, _d_this->x};
446+
// CHECK-NEXT: }
447+
442448
// CHECK: void ref_mem_fn_pullback(double i, double _d_y, SimpleFunctions *_d_this, double *_d_i) {
443449
// CHECK-NEXT: double _t0 = this->x;
444450
// CHECK-NEXT: this->x = +i;
@@ -459,8 +465,6 @@ double fn2(SimpleFunctions& sf, double i) {
459465
// CHECK-NEXT: }
460466
// CHECK-NEXT: }
461467

462-
// CHECK: clad::ValueAndAdjoint<double &, double &> ref_mem_fn_reverse_forw(double i, SimpleFunctions *_d_this, double _d_i);
463-
464468
// CHECK: void fn2_grad(SimpleFunctions &sf, double i, SimpleFunctions *_d_sf, double *_d_i) {
465469
// CHECK-NEXT: SimpleFunctions _t0 = sf;
466470
// CHECK-NEXT: clad::ValueAndAdjoint<double &, double &> _t1 = sf.ref_mem_fn_reverse_forw(i, &(*_d_sf), 0.);
@@ -482,6 +486,11 @@ double fn5(SimpleFunctions& v, double value) {
482486
return v.x;
483487
}
484488

489+
// CHECK: clad::ValueAndAdjoint<SimpleFunctions &, SimpleFunctions &> operator_plus_equal_reverse_forw(double value, SimpleFunctions *_d_this, double _d_value) {
490+
// CHECK-NEXT: this->x += value;
491+
// CHECK-NEXT: return {*this, *_d_this};
492+
// CHECK-NEXT: }
493+
485494
// CHECK: void operator_plus_equal_pullback(double value, SimpleFunctions _d_y, SimpleFunctions *_d_this, double *_d_value) {
486495
// CHECK-NEXT: double _t0 = this->x;
487496
// CHECK-NEXT: this->x += value;
@@ -492,8 +501,6 @@ double fn5(SimpleFunctions& v, double value) {
492501
// CHECK-NEXT: }
493502
// CHECK-NEXT: }
494503

495-
// CHECK: clad::ValueAndAdjoint<SimpleFunctions &, SimpleFunctions &> operator_plus_equal_reverse_forw(double value, SimpleFunctions *_d_this, double _d_value);
496-
497504
// CHECK: void fn5_grad(SimpleFunctions &v, double value, SimpleFunctions *_d_v, double *_d_value) {
498505
// CHECK-NEXT: SimpleFunctions _t0 = v;
499506
// CHECK-NEXT: clad::ValueAndAdjoint<SimpleFunctions &, SimpleFunctions &> _t1 = v.operator_plus_equal_reverse_forw(value, &(*_d_v), 0.);
@@ -511,6 +518,11 @@ double fn4(SimpleFunctions& v) {
511518
return v.x;
512519
}
513520

521+
// CHECK: clad::ValueAndAdjoint<SimpleFunctions &, SimpleFunctions &> operator_plus_plus_reverse_forw(SimpleFunctions *_d_this) {
522+
// CHECK-NEXT: this->x += 1.;
523+
// CHECK-NEXT: return {*this, *_d_this};
524+
// CHECK-NEXT: }
525+
514526
// CHECK: void operator_plus_plus_pullback(SimpleFunctions _d_y, SimpleFunctions *_d_this) {
515527
// CHECK-NEXT: double _t0 = this->x;
516528
// CHECK-NEXT: this->x += 1.;
@@ -520,8 +532,6 @@ double fn4(SimpleFunctions& v) {
520532
// CHECK-NEXT: }
521533
// CHECK-NEXT: }
522534

523-
// CHECK: clad::ValueAndAdjoint<SimpleFunctions &, SimpleFunctions &> operator_plus_plus_reverse_forw(SimpleFunctions *_d_this);
524-
525535
// CHECK: void fn4_grad(SimpleFunctions &v, SimpleFunctions *_d_v) {
526536
// CHECK-NEXT: SimpleFunctions _t0 = v;
527537
// CHECK-NEXT: clad::ValueAndAdjoint<SimpleFunctions &, SimpleFunctions &> _t1 = v.operator_plus_plus_reverse_forw(&(*_d_v));
@@ -935,20 +945,4 @@ int main() {
935945
// CHECK-NEXT: _d_y += _r1;
936946
// CHECK-NEXT: }
937947
// CHECK-NEXT: }
938-
939-
// CHECK: clad::ValueAndAdjoint<double &, double &> ref_mem_fn_reverse_forw(double i, SimpleFunctions *_d_this, double _d_i) {
940-
// CHECK-NEXT: this->x = +i;
941-
// CHECK-NEXT: this->x = -i;
942-
// CHECK-NEXT: return {this->x, _d_this->x};
943-
// CHECK-NEXT: }
944-
945-
// CHECK: clad::ValueAndAdjoint<SimpleFunctions &, SimpleFunctions &> operator_plus_equal_reverse_forw(double value, SimpleFunctions *_d_this, double _d_value) {
946-
// CHECK-NEXT: this->x += value;
947-
// CHECK-NEXT: return {*this, *_d_this};
948-
// CHECK-NEXT: }
949-
950-
// CHECK: clad::ValueAndAdjoint<SimpleFunctions &, SimpleFunctions &> operator_plus_plus_reverse_forw(SimpleFunctions *_d_this) {
951-
// CHECK-NEXT: this->x += 1.;
952-
// CHECK-NEXT: return {*this, *_d_this};
953-
// CHECK-NEXT: }
954948
}

0 commit comments

Comments
 (0)