Skip to content

Commit 5be1fff

Browse files
PetroZarytskyivgvassilev
PetroZarytskyi
authored and
vgvassilev
committed
Pass the final error to error estimation pullbacks directly
Currently, in error estimation, we accumulate errors of a nested function into a separate variable, which we then add to the final error. For example ``` double _t0 = 0.; helper_pullback(..., _t0); _final_error += _t0; ``` However, the only thing we do with `_final_error` and `_t0` is add partial errors to them. It doesn't matter if we first add it to the temporary `_t0` or to `_final_error` directly, i.e. ``` helper_pullback(..., _final_error); ``` Moreover, currently, having a call expression outside a binary operator results in its error `_t0` being ignored. For example ``` double f(double x, double y) { return helper(x, y); } ``` Always produces zero error. This PR removes such temporaries. Fixes #427.
1 parent b769b40 commit 5be1fff

8 files changed

Lines changed: 64 additions & 91 deletions

File tree

‎include/clad/Differentiator/ErrorEstimator.h‎

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -40,9 +40,9 @@ class ErrorEstimationHandler : public ExternalRMVSource {
4040
clang::Expr* m_IdxExpr;
4141
/// A map from var decls to their size variables (e.g. `var_size`).
4242
std::unordered_map<const clang::VarDecl*, clang::Expr*> m_ArrSizes;
43-
/// An expression to match nested function call errors with their
44-
/// assignee (if any exists).
45-
clang::Expr* m_NestedFuncError = nullptr;
43+
// FIXME: Solve this in a more general way.
44+
/// A flag signaling if the current error comes from a function call.
45+
bool m_ErrorFromFunctionCall = false;
4646

4747
std::stack<bool> m_ShouldEmit;
4848
ReverseModeVisitor* m_RMV;
@@ -70,9 +70,7 @@ class ErrorEstimationHandler : public ExternalRMVSource {
7070
/// Function to emit error statements into the derivative body.
7171
///
7272
/// \param[in] errorExpr The error expression (LHS) of the variable.
73-
/// \param[in] addToTheFront A flag to decide whether the error stmts
74-
/// should be added to the beginning of the block or the current position.
75-
void AddErrorStmtToBlock(clang::Expr* errorExpr, bool addToTheFront = true);
73+
void AddErrorStmtToBlock(clang::Expr* errorExpr);
7674

7775
/// Emit the error estimation related statements that were saved to be
7876
/// emitted at later points into specific blocks.
@@ -192,8 +190,7 @@ class ErrorEstimationHandler : public ExternalRMVSource {
192190
void ActBeforeFinalizingDifferentiateSingleStmt(const direction& d) override;
193191
void ActBeforeFinalizingDifferentiateSingleExpr(const direction& d) override;
194192
void ActBeforeDifferentiatingCallExpr(
195-
llvm::SmallVectorImpl<clang::Expr*>& pullbackArgs,
196-
llvm::SmallVectorImpl<clang::Stmt*>& ArgDecls, bool hasAssignee) override;
193+
llvm::SmallVectorImpl<clang::Expr*>& pullbackArgs) override;
197194
void ActBeforeFinalizingVisitDeclStmt(
198195
llvm::SmallVectorImpl<clang::Decl*>& decls,
199196
llvm::SmallVectorImpl<clang::Decl*>& declsDiff) override;

‎include/clad/Differentiator/ExternalRMVSource.h‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -141,8 +141,7 @@ class ExternalRMVSource {
141141
virtual void ActBeforeFinalizingDifferentiateSingleExpr(const direction& d) {}
142142

143143
virtual void ActBeforeDifferentiatingCallExpr(
144-
llvm::SmallVectorImpl<clang::Expr*>& pullbackArgs,
145-
llvm::SmallVectorImpl<clang::Stmt*>& ArgDecls, bool hasAssignee) {}
144+
llvm::SmallVectorImpl<clang::Expr*>& pullbackArgs) {}
146145

147146
virtual void ActBeforeFinalizingVisitDeclStmt(
148147
llvm::SmallVectorImpl<clang::Decl*>& decls,

‎include/clad/Differentiator/MultiplexExternalRMVSource.h‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -61,8 +61,7 @@ class MultiplexExternalRMVSource : public ExternalRMVSource {
6161
void ActBeforeFinalizingDifferentiateSingleStmt(const direction& d) override;
6262
void ActBeforeFinalizingDifferentiateSingleExpr(const direction& d) override;
6363
void ActBeforeDifferentiatingCallExpr(
64-
llvm::SmallVectorImpl<clang::Expr*>& pullbackArgs,
65-
llvm::SmallVectorImpl<clang::Stmt*>& ArgDecls, bool hasAssignee) override;
64+
llvm::SmallVectorImpl<clang::Expr*>& pullbackArgs) override;
6665
void ActBeforeFinalizingVisitDeclStmt(
6766
llvm::SmallVectorImpl<clang::Decl*>& decls,
6867
llvm::SmallVectorImpl<clang::Decl*>& declsDiff) override;

‎lib/Differentiator/ErrorEstimator.cpp‎

Lines changed: 16 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -26,19 +26,6 @@ QualType getUnderlyingArrayType(QualType baseType, ASTContext& C) {
2626
return baseType;
2727
}
2828

29-
Expr* UpdateErrorForFuncCallAssigns(ErrorEstimationHandler* handler,
30-
Expr* savedExpr, Expr* origExpr,
31-
Expr*& callError, const std::string& name) {
32-
Expr* errorExpr = nullptr;
33-
if (!callError)
34-
errorExpr = handler->GetError(savedExpr, origExpr, name);
35-
else {
36-
errorExpr = callError;
37-
callError = nullptr;
38-
}
39-
return errorExpr;
40-
}
41-
4229
void ErrorEstimationHandler::SetErrorEstimationModel(
4330
FPErrorEstimationModel* estModel) {
4431
m_EstModel = estModel;
@@ -63,16 +50,11 @@ void ErrorEstimationHandler::BuildReturnErrorStmt() {
6350
}
6451
}
6552

66-
void ErrorEstimationHandler::AddErrorStmtToBlock(Expr* errorExpr,
67-
bool addToTheFront) {
53+
void ErrorEstimationHandler::AddErrorStmtToBlock(Expr* errorExpr) {
6854
Expr* FinalError = BuildFinalErrorExpr();
6955
Stmt* errorStmt = m_RMV->BuildOp(BO_AddAssign, FinalError, errorExpr);
70-
if (addToTheFront) {
71-
auto& block = m_RMV->getCurrentBlock(direction::reverse);
72-
block.insert(block.begin(), errorStmt);
73-
} else {
74-
m_RMV->addToCurrentBlock(errorStmt, direction::reverse);
75-
}
56+
auto& block = m_RMV->getCurrentBlock(direction::reverse);
57+
block.insert(block.begin(), errorStmt);
7658
}
7759

7860
void ErrorEstimationHandler::EmitErrorEstimationStmts(
@@ -252,10 +234,10 @@ void ErrorEstimationHandler::EmitBinaryOpErrorStmts(Expr* LExpr,
252234
auto decl = GetUnderlyingDeclRefOrNull(LExpr)->getDecl();
253235
if (!ShouldEstimateErrorFor(cast<VarDecl>(decl)))
254236
return;
255-
bool errorFromFunctionCall = (bool)m_NestedFuncError;
256-
Expr* errorExpr = UpdateErrorForFuncCallAssigns(
257-
this, LExpr, oldValue, m_NestedFuncError, decl->getNameAsString());
258-
AddErrorStmtToBlock(errorExpr, /*addToTheFront=*/!errorFromFunctionCall);
237+
if (m_ErrorFromFunctionCall)
238+
return;
239+
Expr* errorExpr = GetError(LExpr, oldValue, decl->getNameAsString());
240+
AddErrorStmtToBlock(errorExpr);
259241
// If there are assign statements to emit in reverse, do that.
260242
EmitErrorEstimationStmts(direction::reverse);
261243
}
@@ -265,17 +247,18 @@ void ErrorEstimationHandler::EmitDeclErrorStmts(DeclDiff<VarDecl> VDDiff,
265247
auto VD = VDDiff.getDecl();
266248
if (!ShouldEstimateErrorFor(VD))
267249
return;
250+
if (m_ErrorFromFunctionCall)
251+
return;
268252
// Build the delta expresion for the variable to be registered.
269253
DeclRefExpr* VDRef = m_RMV->BuildDeclRef(VD);
270254
// FIXME: We should do this for arrays too.
271255
if (!VD->getType()->isArrayType()) {
272256
// If the VarDecl has an init, we should assign it with an error.
273257
if (VD->getInit() && !GetUnderlyingDeclRefOrNull(VD->getInit())) {
274-
bool errorFromFunctionCall = (bool)m_NestedFuncError;
275-
Expr* errorExpr = UpdateErrorForFuncCallAssigns(
276-
this, VDRef, m_RMV->BuildDeclRef(VDDiff.getDecl_dx()),
277-
m_NestedFuncError, VD->getNameAsString());
278-
AddErrorStmtToBlock(errorExpr, /*addToTheFront=*/!errorFromFunctionCall);
258+
Expr* errorExpr =
259+
GetError(VDRef, m_RMV->BuildDeclRef(VDDiff.getDecl_dx()),
260+
VD->getNameAsString());
261+
AddErrorStmtToBlock(errorExpr);
279262
}
280263
}
281264
}
@@ -470,20 +453,9 @@ void ErrorEstimationHandler::ActBeforeFinalizingDifferentiateSingleExpr(
470453
}
471454

472455
void ErrorEstimationHandler::ActBeforeDifferentiatingCallExpr(
473-
llvm::SmallVectorImpl<clang::Expr*>& pullbackArgs,
474-
llvm::SmallVectorImpl<Stmt*>& ArgDecls, bool hasAssignee) {
475-
auto errorRef =
476-
m_RMV->BuildVarDecl(m_RMV->m_Context.DoubleTy, "_t",
477-
m_RMV->getZeroInit(m_RMV->m_Context.DoubleTy));
478-
ArgDecls.push_back(m_RMV->BuildDeclStmt(errorRef));
479-
auto finErr = m_RMV->BuildDeclRef(errorRef);
480-
pullbackArgs.push_back(finErr);
481-
if (hasAssignee) {
482-
if (m_NestedFuncError)
483-
m_NestedFuncError = m_RMV->BuildOp(BO_Add, m_NestedFuncError, finErr);
484-
else
485-
m_NestedFuncError = finErr;
486-
}
456+
llvm::SmallVectorImpl<clang::Expr*>& pullbackArgs) {
457+
m_ErrorFromFunctionCall = true;
458+
pullbackArgs.push_back(BuildFinalErrorExpr());
487459
}
488460

489461
void ErrorEstimationHandler::ActBeforeFinalizingVisitDeclStmt(

‎lib/Differentiator/MultiplexExternalRMVSource.cpp‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -179,11 +179,9 @@ void MultiplexExternalRMVSource::ActBeforeFinalizingDifferentiateSingleExpr(
179179
}
180180

181181
void MultiplexExternalRMVSource::ActBeforeDifferentiatingCallExpr(
182-
llvm::SmallVectorImpl<clang::Expr*>& pullbackArgs,
183-
llvm::SmallVectorImpl<clang::Stmt*>& ArgDecls, bool hasAssignee) {
182+
llvm::SmallVectorImpl<clang::Expr*>& pullbackArgs) {
184183
for (auto source : m_Sources)
185-
source->ActBeforeDifferentiatingCallExpr(pullbackArgs, ArgDecls,
186-
hasAssignee);
184+
source->ActBeforeDifferentiatingCallExpr(pullbackArgs);
187185
}
188186

189187
void MultiplexExternalRMVSource::ActBeforeFinalizingVisitDeclStmt(

‎lib/Differentiator/ReverseModeVisitor.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2062,7 +2062,7 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
20622062
pullbackRequest.DVI.push_back(PVD);
20632063
}
20642064
m_ExternalSource->ActBeforeDifferentiatingCallExpr(
2065-
pullbackCallArgs, PreCallStmts, dfdx());
2065+
pullbackCallArgs);
20662066
pullbackFD =
20672067
plugin::ProcessDiffRequest(m_CladPlugin, pullbackRequest);
20682068
} else

‎test/ErrorEstimation/BasicOps.C‎

Lines changed: 37 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -205,11 +205,9 @@ float func6(float x, float y) {
205205
//CHECK-NEXT: {
206206
//CHECK-NEXT: double _r0 = 0.;
207207
//CHECK-NEXT: double _r1 = 0.;
208-
//CHECK-NEXT: double _t0 = 0.;
209-
//CHECK-NEXT: helper_pullback(x, y, _d_z, &_r0, &_r1, _t0);
208+
//CHECK-NEXT: helper_pullback(x, y, _d_z, &_r0, &_r1, _final_error);
210209
//CHECK-NEXT: *_d_x += _r0;
211210
//CHECK-NEXT: *_d_y += _r1;
212-
//CHECK-NEXT: _final_error += _t0;
213211
//CHECK-NEXT: }
214212
//CHECK-NEXT: _final_error += std::abs(*_d_x * x * {{.+}});
215213
//CHECK-NEXT: _final_error += std::abs(*_d_y * y * {{.+}});
@@ -262,10 +260,8 @@ float func8(float x, float y) {
262260
//CHECK-NEXT: z = _t0;
263261
//CHECK-NEXT: *_d_y += _d_z;
264262
//CHECK-NEXT: x = _t1;
265-
//CHECK-NEXT: double _t2 = 0.;
266-
//CHECK-NEXT: helper2_pullback(x, _d_z, _d_x, _t2);
263+
//CHECK-NEXT: helper2_pullback(x, _d_z, _d_x, _final_error);
267264
//CHECK-NEXT: _d_z = 0.F;
268-
//CHECK-NEXT: _final_error += _t2;
269265
//CHECK-NEXT: _final_error += std::abs(*_d_x * x * {{.+}});
270266
//CHECK-NEXT: }
271267
//CHECK-NEXT: _final_error += std::abs(*_d_x * x * {{.+}});
@@ -279,45 +275,58 @@ float func9(float x, float y) {
279275
}
280276

281277
//CHECK: void func9_grad(float x, float y, float *_d_x, float *_d_y, double &_final_error) {
282-
//CHECK-NEXT: float _t1 = x;
278+
//CHECK-NEXT: float _t0 = x;
283279
//CHECK-NEXT: float _d_z = 0.F;
284280
//CHECK-NEXT: float z = helper(x, y) + helper2(x);
285-
//CHECK-NEXT: float _t3 = z;
286-
//CHECK-NEXT: float _t5 = x;
287-
//CHECK-NEXT: double _t7 = helper2(x);
288-
//CHECK-NEXT: float _t8 = y;
289-
//CHECK-NEXT: double _t4 = helper2(y);
290-
//CHECK-NEXT: z += _t7 * _t4;
281+
//CHECK-NEXT: float _t1 = z;
282+
//CHECK-NEXT: float _t3 = x;
283+
//CHECK-NEXT: double _t4 = helper2(x);
284+
//CHECK-NEXT: float _t5 = y;
285+
//CHECK-NEXT: double _t2 = helper2(y);
286+
//CHECK-NEXT: z += _t4 * _t2;
291287
//CHECK-NEXT: _d_z += 1;
292288
//CHECK-NEXT: {
293-
//CHECK-NEXT: z = _t3;
294-
//CHECK-NEXT: x = _t5;
295-
//CHECK-NEXT: double _t6 = 0.;
296-
//CHECK-NEXT: helper2_pullback(x, _d_z * _t4, _d_x, _t6);
297-
//CHECK-NEXT: y = _t8;
298-
//CHECK-NEXT: double _t9 = 0.;
299-
//CHECK-NEXT: helper2_pullback(y, _t7 * _d_z, _d_y, _t9);
300-
//CHECK-NEXT: _final_error += _t6 + _t9;
289+
//CHECK-NEXT: z = _t1;
290+
//CHECK-NEXT: x = _t3;
291+
//CHECK-NEXT: helper2_pullback(x, _d_z * _t2, _d_x, _final_error);
292+
//CHECK-NEXT: y = _t5;
293+
//CHECK-NEXT: helper2_pullback(y, _t4 * _d_z, _d_y, _final_error);
301294
//CHECK-NEXT: _final_error += std::abs(*_d_y * y * {{.+}});
302295
//CHECK-NEXT: _final_error += std::abs(*_d_x * x * {{.+}});
303296
//CHECK-NEXT: }
304297
//CHECK-NEXT: {
305298
//CHECK-NEXT: double _r0 = 0.;
306299
//CHECK-NEXT: double _r1 = 0.;
307-
//CHECK-NEXT: double _t0 = 0.;
308-
//CHECK-NEXT: helper_pullback(x, y, _d_z, &_r0, &_r1, _t0);
300+
//CHECK-NEXT: helper_pullback(x, y, _d_z, &_r0, &_r1, _final_error);
309301
//CHECK-NEXT: *_d_x += _r0;
310302
//CHECK-NEXT: *_d_y += _r1;
311-
//CHECK-NEXT: x = _t1;
312-
//CHECK-NEXT: double _t2 = 0.;
313-
//CHECK-NEXT: helper2_pullback(x, _d_z, _d_x, _t2);
314-
//CHECK-NEXT: _final_error += _t0 + _t2;
303+
//CHECK-NEXT: x = _t0;
304+
//CHECK-NEXT: helper2_pullback(x, _d_z, _d_x, _final_error);
315305
//CHECK-NEXT: _final_error += std::abs(*_d_x * x * {{.+}});
316306
//CHECK-NEXT: }
317307
//CHECK-NEXT: _final_error += std::abs(*_d_x * x * {{.+}});
318308
//CHECK-NEXT: _final_error += std::abs(*_d_y * y * {{.+}});
319309
//CHECK-NEXT: }
320310

311+
double func10(double x, double y) {
312+
return helper(x, y);
313+
}
314+
315+
// CHECK: void func10_grad(double x, double y, double *_d_x, double *_d_y, double &_final_error) {
316+
// CHECK-NEXT: double _ret_value0 = 0.;
317+
// CHECK-NEXT: _ret_value0 = helper(x, y);
318+
// CHECK-NEXT: {
319+
// CHECK-NEXT: double _r0 = 0.;
320+
// CHECK-NEXT: double _r1 = 0.;
321+
// CHECK-NEXT: helper_pullback(x, y, 1, &_r0, &_r1, _final_error);
322+
// CHECK-NEXT: *_d_x += _r0;
323+
// CHECK-NEXT: *_d_y += _r1;
324+
// CHECK-NEXT: }
325+
// CHECK-NEXT: _final_error += std::abs(*_d_x * x * {{.+}});
326+
// CHECK-NEXT: _final_error += std::abs(*_d_y * y * {{.+}});
327+
// CHECK-NEXT: _final_error += std::abs(1. * _ret_value0 * {{.+}});
328+
// CHECK-NEXT: }
329+
321330
int main() {
322331
clad::estimate_error(func);
323332
clad::estimate_error(func2);
@@ -328,4 +337,5 @@ int main() {
328337
clad::estimate_error(func7);
329338
clad::estimate_error(func8);
330339
clad::estimate_error(func9);
340+
clad::estimate_error(func10);
331341
}

‎test/ErrorEstimation/LoopsAndArrays.C‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -321,10 +321,8 @@ double func6(double x) {
321321
//CHECK-NEXT: i--;
322322
//CHECK-NEXT: sum = clad::pop(_t1);
323323
//CHECK-NEXT: double _r0 = 0.;
324-
//CHECK-NEXT: double _t2 = 0.;
325-
//CHECK-NEXT: fun_pullback(x, _d_sum, &_r0, _t2);
324+
//CHECK-NEXT: fun_pullback(x, _d_sum, &_r0, _final_error);
326325
//CHECK-NEXT: *_d_x += _r0;
327-
//CHECK-NEXT: _final_error += _t2;
328326
//CHECK-NEXT: }
329327
//CHECK-NEXT: _final_error += std::abs(_d_sum * sum * {{.+}});
330328
//CHECK-NEXT: _final_error += std::abs(*_d_x * x * {{.+}});

0 commit comments

Comments
 (0)