From 578b8fcc99bb17207cf33c61d9cd8ffa6dc53030 Mon Sep 17 00:00:00 2001 From: Parth Date: Sun, 10 Apr 2022 21:45:09 +0530 Subject: [PATCH] add support for operator overload in reverse mode --- .../ReverseModeForwPassVisitor.h | 1 + .../clad/Differentiator/ReverseModeVisitor.h | 2 +- .../ReverseModeForwPassVisitor.cpp | 30 ++++- lib/Differentiator/ReverseModeVisitor.cpp | 63 ++++++---- test/Gradient/MemberFunctions.C | 115 +++++++++++++++++- 5 files changed, 180 insertions(+), 31 deletions(-) diff --git a/include/clad/Differentiator/ReverseModeForwPassVisitor.h b/include/clad/Differentiator/ReverseModeForwPassVisitor.h index 5d60e6cb6..fc2236306 100644 --- a/include/clad/Differentiator/ReverseModeForwPassVisitor.h +++ b/include/clad/Differentiator/ReverseModeForwPassVisitor.h @@ -30,6 +30,7 @@ class ReverseModeForwPassVisitor : public ReverseModeVisitor { StmtDiff VisitCompoundStmt(const clang::CompoundStmt* CS) override; StmtDiff VisitDeclRefExpr(const clang::DeclRefExpr* DRE) override; StmtDiff VisitReturnStmt(const clang::ReturnStmt* RS) override; + StmtDiff VisitUnaryOperator(const clang::UnaryOperator* UnOp) override; }; } // namespace clad diff --git a/include/clad/Differentiator/ReverseModeVisitor.h b/include/clad/Differentiator/ReverseModeVisitor.h index 3b7c4b1cb..64dcb52af 100644 --- a/include/clad/Differentiator/ReverseModeVisitor.h +++ b/include/clad/Differentiator/ReverseModeVisitor.h @@ -336,7 +336,7 @@ namespace clad { StmtDiff VisitParenExpr(const clang::ParenExpr* PE); virtual StmtDiff VisitReturnStmt(const clang::ReturnStmt* RS); StmtDiff VisitStmt(const clang::Stmt* S); - StmtDiff VisitUnaryOperator(const clang::UnaryOperator* UnOp); + virtual StmtDiff VisitUnaryOperator(const clang::UnaryOperator* UnOp); StmtDiff VisitExprWithCleanups(const clang::ExprWithCleanups* EWC); /// Decl is not Stmt, so it cannot be visited directly. StmtDiff VisitWhileStmt(const clang::WhileStmt* WS); diff --git a/lib/Differentiator/ReverseModeForwPassVisitor.cpp b/lib/Differentiator/ReverseModeForwPassVisitor.cpp index c22f4fadf..606664301 100644 --- a/lib/Differentiator/ReverseModeForwPassVisitor.cpp +++ b/lib/Differentiator/ReverseModeForwPassVisitor.cpp @@ -29,7 +29,7 @@ ReverseModeForwPassVisitor::Derive(const FunctionDecl* FD, DiffParams args{}; std::copy(FD->param_begin(), FD->param_end(), std::back_inserter(args)); - auto fnName = m_Function->getNameAsString() + "_forw"; + auto fnName = clad::utils::ComputeEffectiveFnName(m_Function) + "_forw"; auto fnDNI = utils::BuildDeclarationNameInfo(m_Sema, fnName); auto paramTypes = ComputeParamTypes(args); @@ -86,8 +86,6 @@ ReverseModeForwPassVisitor::Derive(const FunctionDecl* FD, QualType ReverseModeForwPassVisitor::GetParameterDerivativeType(QualType yType, QualType xType) { - assert(yType.getNonReferenceType()->isRealType() && - "yType should be a builtin-numerical scalar type!!"); QualType xValueType = utils::GetValueType(xType); // derivative variables should always be of non-const type. xValueType.removeLocalConst(); @@ -240,4 +238,30 @@ ReverseModeForwPassVisitor::VisitReturnStmt(const clang::ReturnStmt* RS) { Stmt* newRS = m_Sema.BuildReturnStmt(noLoc, returnInitList).get(); return {newRS}; } + +StmtDiff +ReverseModeForwPassVisitor::VisitUnaryOperator(const UnaryOperator* UnOp) { + auto opCode = UnOp->getOpcode(); + StmtDiff diff{}; + // If it is a post-increment/decrement operator, its result is a reference + // and we should return it. + Expr* ResultRef = nullptr; + if (opCode == UnaryOperatorKind::UO_Deref) { + if (const auto* MD = dyn_cast(m_Function)) { + if (MD->isInstance()) { + diff = Visit(UnOp->getSubExpr()); + Expr* cloneE = BuildOp(UnaryOperatorKind::UO_Deref, diff.getExpr()); + Expr* derivedE = diff.getExpr_dx(); + return {cloneE, derivedE}; + } + } + } else if (opCode == UO_Plus) + diff = Visit(UnOp->getSubExpr(), dfdx()); + else if (opCode == UO_Minus) { + auto d = BuildOp(UO_Minus, dfdx()); + diff = Visit(UnOp->getSubExpr(), d); + } + Expr* op = BuildOp(opCode, diff.getExpr()); + return StmtDiff(op, ResultRef); +} } // namespace clad diff --git a/lib/Differentiator/ReverseModeVisitor.cpp b/lib/Differentiator/ReverseModeVisitor.cpp index 6c5d09cc2..fa12f99ad 100644 --- a/lib/Differentiator/ReverseModeVisitor.cpp +++ b/lib/Differentiator/ReverseModeVisitor.cpp @@ -369,7 +369,7 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context, if (m_ExternalSource) m_ExternalSource->ActAfterCreatingDerivedFnScope(); - + auto params = BuildParams(args); if (m_ExternalSource) @@ -411,7 +411,7 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context, m_IndependentVars.push_back(arg); } } - + if (m_ExternalSource) m_ExternalSource->ActBeforeCreatingDerivedFnBodyScope(); @@ -743,7 +743,7 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context, StmtDiff SDiff = DifferentiateSingleStmt(S); addToCurrentBlock(SDiff.getStmt(), direction::forward); addToCurrentBlock(SDiff.getStmt_dx(), direction::reverse); - + if (m_ExternalSource) m_ExternalSource->ActAfterProcessingStmtInVisitCompoundStmt(); } @@ -861,7 +861,7 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context, m_ExternalSource->ActBeforeDifferentiatingSingleStmtBranchInVisitIfStmt(); StmtDiff BranchDiff = DifferentiateSingleStmt(Branch, /*dfdS=*/nullptr); addToCurrentBlock(BranchDiff.getStmt(), direction::forward); - + if (m_ExternalSource) m_ExternalSource->ActBeforeFinalisingVisitBranchSingleStmtInIfVisitStmt(); @@ -1372,7 +1372,8 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context, // If the function has no args and is not a member function call then we // assume that it is not related to independent variables and does not // contribute to gradient. - if (!NArgs && !isa(CE)) + if ((NArgs == 0U) && !isa(CE) && + !isa(CE)) return StmtDiff(Clone(CE)); // Stores the call arguments for the function to be derived @@ -1392,7 +1393,7 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context, // derived function. In the case of member functions, `implicit` // this object is always passed by reference. if (!dfdx() && !utils::HasAnyReferenceOrPointerArgument(FD) && - !isa(CE)) { + !isa(CE) && !isa(CE)) { for (const Expr* Arg : CE->arguments()) { StmtDiff ArgDiff = Visit(Arg, dfdx()); CallArgs.push_back(ArgDiff.getExpr()); @@ -1424,9 +1425,14 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context, // FIXME: We should add instructions for handling non-differentiable // arguments. Currently we are implicitly assuming function call only // contains differentiable arguments. - for (std::size_t i = skipFirstArg, e = CE->getNumArgs(); i != e; ++i) { + bool isCXXOperatorCall = isa(CE); + + for (std::size_t i = static_cast(isCXXOperatorCall), + e = CE->getNumArgs(); + i != e; ++i) { const Expr* arg = CE->getArg(i); - const auto* PVD = FD->getParamDecl(i - skipFirstArg); + const auto* PVD = + FD->getParamDecl(i - static_cast(isCXXOperatorCall)); StmtDiff argDiff{}; bool passByRef = utils::IsReferenceOrPointerType(PVD->getType()); // We do not need to create result arg for arguments passed by reference @@ -1714,11 +1720,18 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context, pullback); // Try to find it in builtin derivatives - std::string customPullback = FD->getNameAsString() + "_pullback"; + if (baseDiff.getExpr()) + pullbackCallArgs.insert( + pullbackCallArgs.begin(), + BuildOp(UnaryOperatorKind::UO_AddrOf, baseDiff.getExpr())); + std::string customPullback = + clad::utils::ComputeEffectiveFnName(FD) + "_pullback"; OverloadedDerivedFn = m_Builder.BuildCallToCustomDerivativeOrNumericalDiff( customPullback, pullbackCallArgs, getCurrentScope(), const_cast(FD->getDeclContext())); + if (baseDiff.getExpr()) + pullbackCallArgs.erase(pullbackCallArgs.begin()); } // should be true if we are using numerical differentiation to differentiate @@ -1749,7 +1762,8 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context, // derive the called function. DiffRequest pullbackRequest{}; pullbackRequest.Function = FD; - pullbackRequest.BaseFunctionName = FD->getNameAsString(); + pullbackRequest.BaseFunctionName = + clad::utils::ComputeEffectiveFnName(FD); pullbackRequest.Mode = DiffMode::experimental_pullback; // Silence diag outputs in nested derivation process. pullbackRequest.VerboseDiags = false; @@ -1882,7 +1896,8 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context, DiffRequest calleeFnForwPassReq; calleeFnForwPassReq.Function = FD; calleeFnForwPassReq.Mode = DiffMode::reverse_mode_forward_pass; - calleeFnForwPassReq.BaseFunctionName = FD->getNameAsString(); + calleeFnForwPassReq.BaseFunctionName = + clad::utils::ComputeEffectiveFnName(FD); calleeFnForwPassReq.VerboseDiags = true; FunctionDecl* calleeFnForwPassFD = plugin::ProcessDiffRequest(m_CladPlugin, calleeFnForwPassReq); @@ -1906,13 +1921,17 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context, // (isCladArrayType(derivedBase->getType())) // CallArgs.push_back(derivedBase); // else + // Currently derivedBase `*d_this` can never be CladArrayType CallArgs.push_back( BuildOp(UnaryOperatorKind::UO_AddrOf, derivedBase, noLoc)); } - for (std::size_t i = 0, e = CE->getNumArgs(); i != e; ++i) { + for (std::size_t i = static_cast(isCXXOperatorCall), + e = CE->getNumArgs(); + i != e; ++i) { const Expr* arg = CE->getArg(i); - const ParmVarDecl* PVD = FD->getParamDecl(i); + const ParmVarDecl* PVD = + FD->getParamDecl(i - static_cast(isCXXOperatorCall)); StmtDiff argDiff = Visit(arg); if ((argDiff.getExpr_dx() != nullptr) && PVD->getType()->isReferenceType()) { @@ -1988,8 +2007,7 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context, // Add it to the body statements. addToCurrentBlock(add_assign, direction::reverse); } - } - else { + } else { // FIXME: This is not adding 'address-of' operator support. // This is just making this special case differentiable that is required // for computing hessian: @@ -2382,13 +2400,13 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context, VDDerivedInit = getZeroInit(VD->getType()); // `specialThisDiffCase` is only required for correctly differentiating - // the following code: + // the following code: // ``` // Class _d_this_obj; // Class* _d_this = &_d_this_obj; // ``` // Computation of hessian requires this code to be correctly - // differentiated. + // differentiated. bool specialThisDiffCase = false; if (auto MD = dyn_cast(m_Function)) { if (VDDerivedType->isPointerType() && MD->isInstance()) { @@ -2507,10 +2525,10 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context, return VarDeclDiff(VDClone, VDDerived); } - + // TODO: 'shouldEmit' parameter should be removed after converting // Error estimation framework to callback style. Some more research - // need to be done to + // need to be done to StmtDiff ReverseModeVisitor::DifferentiateSingleStmt(const Stmt* S, Expr* dfdS) { if (m_ExternalSource) @@ -3121,7 +3139,7 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context, bodyDiff = {bodyDiff.getStmt(), CFSS}; } - + void ReverseModeVisitor::AddExternalSource(ExternalRMVSource& source) { if (!m_ExternalSource) m_ExternalSource = new MultiplexExternalRMVSource(); @@ -3177,13 +3195,10 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context, QualType ReverseModeVisitor::GetParameterDerivativeType(QualType yType, QualType xType) { - + if (m_Mode == DiffMode::reverse) assert(yType->isRealType() && "yType should be a non-reference builtin-numerical scalar type!!"); - else if (m_Mode == DiffMode::experimental_pullback) - assert(yType.getNonReferenceType()->isRealType() && - "yType should be a builtin-numerical scalar type!!"); QualType xValueType = utils::GetValueType(xType); // derivative variables should always be of non-const type. xValueType.removeLocalConst(); diff --git a/test/Gradient/MemberFunctions.C b/test/Gradient/MemberFunctions.C index 03c6b07fe..950889531 100644 --- a/test/Gradient/MemberFunctions.C +++ b/test/Gradient/MemberFunctions.C @@ -703,7 +703,19 @@ public: // CHECK-NEXT: } // CHECK-NEXT: } - double& ref_mem_fn(double i) {return x;} + double& ref_mem_fn(double i) { + x = +i; + x = -i; + return x; + } + SimpleFunctions& operator+=(double value) { + x += value; + return *this; + } + SimpleFunctions& operator++() { + x += 1.0; + return *this; + } void mem_fn_grad(double i, double j, clad::array_ref _d_i, clad::array_ref _d_j); void const_mem_fn_grad(double i, double j, clad::array_ref _d_i, clad::array_ref _d_j); @@ -763,13 +775,29 @@ double fn2(SimpleFunctions& sf, double i) { } // CHECK: void ref_mem_fn_pullback(double i, double _d_y, clad::array_ref _d_this, clad::array_ref _d_i) { +// CHECK-NEXT: this->x = +i; +// CHECK-NEXT: this->x = -i; // CHECK-NEXT: goto _label0; // CHECK-NEXT: _label0: // CHECK-NEXT: (* _d_this).x += _d_y; +// CHECK-NEXT: { +// CHECK-NEXT: double _r_d1 = (* _d_this).x; +// CHECK-NEXT: * _d_i += -_r_d1; +// CHECK-NEXT: (* _d_this).x -= _r_d1; +// CHECK-NEXT: } +// CHECK-NEXT: { +// CHECK-NEXT: double _r_d0 = (* _d_this).x; +// CHECK-NEXT: * _d_i += _r_d0; +// CHECK-NEXT: (* _d_this).x -= _r_d0; +// CHECK-NEXT: } // CHECK-NEXT: } + // CHECK: clad::ValueAndAdjoint ref_mem_fn_forw(double i, clad::array_ref _d_this, clad::array_ref _d_i) { +// CHECK-NEXT: this->x = +i; +// CHECK-NEXT: this->x = -i; // CHECK-NEXT: return {this->x, (* _d_this).x}; // CHECK-NEXT: } + // CHECK: void fn2_grad(SimpleFunctions &sf, double i, clad::array_ref _d_sf, clad::array_ref _d_i) { // CHECK-NEXT: double _t0; // CHECK-NEXT: SimpleFunctions _t1; @@ -791,6 +819,78 @@ double fn3(double x, double y, double i, double j) { return sf.mem_fn(i, j); } +double fn5(SimpleFunctions& v, double value) { + v += value; + return v.x; +} + +// CHECK: void operator_plus_equal_pullback(double value, SimpleFunctions _d_y, clad::array_ref _d_this, clad::array_ref _d_value) { +// CHECK-NEXT: this->x += value; +// CHECK-NEXT: goto _label0; +// CHECK-NEXT: _label0: +// CHECK-NEXT: ; +// CHECK-NEXT: { +// CHECK-NEXT: double _r_d0 = (* _d_this).x; +// CHECK-NEXT: (* _d_this).x += _r_d0; +// CHECK-NEXT: * _d_value += _r_d0; +// CHECK-NEXT: (* _d_this).x -= _r_d0; +// CHECK-NEXT: } +// CHECK-NEXT: } + +// CHECK: clad::ValueAndAdjoint operator_plus_equal_forw(double value, clad::array_ref _d_this, clad::array_ref _d_value) { +// CHECK-NEXT: this->x += value; +// CHECK-NEXT: return {*this, (* _d_this)}; +// CHECK-NEXT: } + +// CHECK: void fn5_grad(SimpleFunctions &v, double value, clad::array_ref _d_v, clad::array_ref _d_value) { +// CHECK-NEXT: double _t0; +// CHECK-NEXT: SimpleFunctions _t1; +// CHECK-NEXT: _t0 = value; +// CHECK-NEXT: _t1 = v; +// CHECK-NEXT: clad::ValueAndAdjoint _t2 = _t1.operator_plus_equal_forw(_t0, &(* _d_v), nullptr); +// CHECK-NEXT: goto _label0; +// CHECK-NEXT: _label0: +// CHECK-NEXT: (* _d_v).x += 1; +// CHECK-NEXT: { +// CHECK-NEXT: double _grad0 = 0.; +// CHECK-NEXT: _t1.operator_plus_equal_pullback(_t0, {}, &(* _d_v), &_grad0); +// CHECK-NEXT: double _r0 = _grad0; +// CHECK-NEXT: * _d_value += _r0; +// CHECK-NEXT: } +// CHECK-NEXT: } + +double fn4(SimpleFunctions& v) { + ++v; + return v.x; +} + +// CHECK: void operator_plus_plus_pullback(SimpleFunctions _d_y, clad::array_ref _d_this) { +// CHECK-NEXT: this->x += 1.; +// CHECK-NEXT: goto _label0; +// CHECK-NEXT: _label0: +// CHECK-NEXT: ; +// CHECK-NEXT: { +// CHECK-NEXT: double _r_d0 = (* _d_this).x; +// CHECK-NEXT: (* _d_this).x += _r_d0; +// CHECK-NEXT: (* _d_this).x -= _r_d0; +// CHECK-NEXT: } +// CHECK-NEXT: } + +// CHECK: clad::ValueAndAdjoint operator_plus_plus_forw(clad::array_ref _d_this) { +// CHECK-NEXT: this->x += 1.; +// CHECK-NEXT: return {*this, (* _d_this)}; +// CHECK-NEXT: } + +// CHECK: void fn4_grad(SimpleFunctions &v, clad::array_ref _d_v) { +// CHECK-NEXT: SimpleFunctions _t0; +// CHECK-NEXT: _t0 = v; +// CHECK-NEXT: clad::ValueAndAdjoint _t1 = _t0.operator_plus_plus_forw(&(* _d_v)); +// CHECK-NEXT: goto _label0; +// CHECK-NEXT: _label0: +// CHECK-NEXT: (* _d_v).x += 1; +// CHECK-NEXT: _t0.operator_plus_plus_pullback({}, &(* _d_v)); +// CHECK-NEXT: } + int main() { auto d_mem_fn = clad::gradient(&SimpleFunctions::mem_fn); auto d_const_mem_fn = clad::gradient(&SimpleFunctions::const_mem_fn); @@ -825,12 +925,21 @@ int main() { printf("%.2f ",result[i]); //CHECK-EXEC: 40.00 16.00 } - SimpleFunctions sf(2, 3); + SimpleFunctions sf1(2, 3), sf2(3, 4), sf3(4, 5); SimpleFunctions d_sf; + auto d_fn2 = clad::gradient(fn2); - d_fn2.execute(sf, 2, &d_sf, &result[0]); + d_fn2.execute(sf1, 2, &d_sf, &result[0]); + printf("%.2f", result[0]); //CHECK-EXEC: 39.00 + + auto d_fn5 = clad::gradient(fn5); + d_fn5.execute(sf2, 3, &d_sf, &result[0]); printf("%.2f", result[0]); //CHECK-EXEC: 40.00 + auto d_fn4 = clad::gradient(fn4); + d_fn4.execute(sf3, &d_sf); + printf("%.2f", d_sf.x); //CHECK-EXEC: 2.00 + auto d_const_volatile_lval_ref_mem_fn_i = clad::gradient(&SimpleFunctions::const_volatile_lval_ref_mem_fn, "i"); // CHECK: void const_volatile_lval_ref_mem_fn_grad_0(double i, double j, clad::array_ref _d_this, clad::array_ref _d_i) const volatile & {