Sink FPErrorEstimationModel in ErrorEstimationHandler

``FPErrorEstimationModel`` used to be the base class for user-provided error estimation models. However, since we removed custom models in favor of custom error estimation functions, there is no point in having it as a separate entity from ``ErrorEstimationHandler``.

Author: PetroZarytskyi

include/clad/Differentiator/DerivativeBuilder.h

@@ -13,6 +13,7 @@
 #include "clad/Differentiator/DerivedFnCollector.h"
 #include "clad/Differentiator/DiffPlanner.h"
 
+#include "clang/AST/Decl.h"
 #include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/AST/StmtVisitor.h"
 #include "clang/Basic/Diagnostic.h"
@@ -23,6 +24,7 @@
 #include <memory>
 #include <stack>
 #include <unordered_map>
+#include <utility>
 
 namespace clang {
   class ASTContext;
@@ -49,24 +51,23 @@ namespace clad {
 } // namespace clad
 
 namespace clad {
-  class ErrorEstimationHandler;
-  class FPErrorEstimationModel;
-
-  /// A pair of FunctionDecl and potential enclosing context, e.g. a function
-  /// in nested namespaces.
-  // This is the type returned by cloneFunction. Using OverloadedDeclWithContext
-  // instead would lead to unnecessarily returning a nullptr in the overloaded
-  // FD
-  using DeclWithContext = std::pair<clang::FunctionDecl*, clang::Decl*>;
-  /// Stores derivative and the corresponding overload. If no overload exist
-  /// then `second` data member should be `nullptr`.
-  struct DerivativeAndOverload {
-    clang::Decl* derivative = nullptr;
-    clang::FunctionDecl* overload = nullptr;
-    DerivativeAndOverload(clang::Decl* p_derivative = nullptr,
-                          clang::FunctionDecl* p_overload = nullptr)
-        : derivative(p_derivative), overload(p_overload) {}
-  };
+class ErrorEstimationHandler;
+
+/// A pair of FunctionDecl and potential enclosing context, e.g. a function
+/// in nested namespaces.
+// This is the type returned by cloneFunction. Using OverloadedDeclWithContext
+// instead would lead to unnecessarily returning a nullptr in the overloaded
+// FD
+using DeclWithContext = std::pair<clang::FunctionDecl*, clang::Decl*>;
+/// Stores derivative and the corresponding overload. If no overload exist
+/// then `second` data member should be `nullptr`.
+struct DerivativeAndOverload {
+  clang::Decl* derivative = nullptr;
+  clang::FunctionDecl* overload = nullptr;
+  DerivativeAndOverload(clang::Decl* p_derivative = nullptr,
+                        clang::FunctionDecl* p_overload = nullptr)
+      : derivative(p_derivative), overload(p_overload) {}
+};
 
   static clang::SourceLocation noLoc{};
   class VisitorBase;

include/clad/Differentiator/ErrorEstimator.h

@@ -1,13 +1,13 @@
 #ifndef CLAD_ERROR_ESTIMATOR_H
 #define CLAD_ERROR_ESTIMATOR_H
 
-#include "EstimationModel.h"
 #include "clad/Differentiator/ExternalRMVSource.h"
 #include "clad/Differentiator/ReverseModeVisitorDirectionKinds.h"
 
 #include "clang/AST/OperationKinds.h"
 
 #include <stack>
+#include <string>
 
 namespace clang {
 class Stmt;
@@ -29,37 +29,44 @@ class ErrorEstimationHandler : public ExternalRMVSource {
   // `Stmts` is originally defined in `VisitorBase`.
   using Stmts = llvm::SmallVector<clang::Stmt*, 16>;
   /// Reference to the return error expression.
-  clang::Expr* m_RetErrorExpr;
-  /// An instance of the custom error estimation model to be used.
-  FPErrorEstimationModel* m_EstModel; // We do not own this.
+  clang::Expr* m_RetErrorExpr = nullptr;
+  /// An Expr representing a custom `getErrorVal` function, if any.
+  clang::Expr* m_CustomErrorFunction = nullptr;
+  void LookupCustomErrorFunction();
+  /// The function to build the error expression of a
+  /// specific estimation model. The error expression is returned in the form
+  /// of a clang::Expr.
+  /// \param[in] refExpr The reference of the expression to which the error
+  /// has to be assigned, this is a StmtDiff type hence one can use getExpr()
+  /// to get the unmodified expression and getExpr_dx() to get the absolute
+  /// derivative of the same.
+  /// \param [in] name Name of the variable being analysed.
+  ///
+  /// \returns The error expression of the input value.
+  // Return an expression of the following kind:
+  // std::abs(dfdx * delta_x * Em)
+  clang::Expr* AssignError(StmtDiff refExpr, const std::string& name);
   /// A set of assignments resulting for declaration statments.
   Stmts m_ForwardReplStmts;
   /// A vector to keep track of error statements for delayed emission.
   Stmts m_ReverseErrorStmts;
   /// The index expression for emitting final errors for input param errors.
-  clang::Expr* m_IdxExpr;
+  clang::Expr* m_IdxExpr = nullptr;
   /// A map from var decls to their size variables (e.g. `var_size`).
   std::unordered_map<const clang::VarDecl*, clang::Expr*> m_ArrSizes;
   // FIXME: Solve this in a more general way.
   /// A flag signaling if the current error comes from a function call.
   bool m_ErrorFromFunctionCall = false;
 
   std::stack<bool> m_ShouldEmit;
-  ReverseModeVisitor* m_RMV;
+  ReverseModeVisitor* m_RMV = nullptr;
   llvm::SmallVectorImpl<clang::ParmVarDecl*>* m_Params = nullptr;
 
 public:
   using direction = rmv::direction;
-  ErrorEstimationHandler()
-      : m_RetErrorExpr(nullptr), m_EstModel(nullptr), m_IdxExpr(nullptr) {}
+  ErrorEstimationHandler() = default;
   ~ErrorEstimationHandler() override = default;
 
-  /// Function to set the error estimation model currently in use.
-  ///
-  /// \param[in] estModel The error estimation model, can be either
-  /// an in-built one or one provided by the user.
-  void SetErrorEstimationModel(FPErrorEstimationModel* estModel);
-
   /// Builds a reference to the final error parameter of the function.
   clang::DeclRefExpr* BuildFinalErrorExpr();
 
@@ -122,7 +129,7 @@ class ErrorEstimationHandler : public ExternalRMVSource {
   /// \returns The error in the variable 'var'.
   clang::Expr* GetError(clang::Expr* var, clang::Expr* varDiff,
                         const std::string& varName) {
-    return m_EstModel->AssignError({var, varDiff}, varName);
+    return AssignError({var, varDiff}, varName);
   }
 
   /// This function adds the final error and the other parameter errors to the

include/clad/Differentiator/EstimationModel.h

@@ -59,7 +59,6 @@ llvm_add_library(cladDifferentiator
   DerivedFnInfo.cpp
   DiffPlanner.cpp
   ErrorEstimator.cpp
-  EstimationModel.cpp
   JacobianModeVisitor.cpp
   HessianModeVisitor.cpp
   MultiplexExternalRMVSource.cpp

lib/Differentiator/CMakeLists.txt

@@ -592,13 +592,9 @@ static void registerDerivative(Decl* D, Sema& S, const DiffRequest& R) {
     } else if (request.Mode == DiffMode::reverse ||
                request.Mode == DiffMode::pullback) {
       ErrorEstimationHandler handler;
-      std::unique_ptr<FPErrorEstimationModel> model;
       ReverseModeVisitor V(*this, request);
-      if (request.EnableErrorEstimation) {
-        model = std::make_unique<FPErrorEstimationModel>(*this, request);
-        handler.SetErrorEstimationModel(model.get());
+      if (request.EnableErrorEstimation)
         V.AddExternalSource(handler);
-      }
       result = V.Derive();
     } else if (request.Mode == DiffMode::reverse_mode_forward_pass) {
       ReverseModeForwPassVisitor V(*this, request);

lib/Differentiator/DerivativeBuilder.cpp

@@ -6,12 +6,20 @@
 #include "clad/Differentiator/ReverseModeVisitor.h"
 
 #include "clang/AST/Decl.h"
+#include "clang/AST/DeclarationName.h"
 #include "clang/AST/Expr.h"
 #include "clang/AST/OperationKinds.h"
+#include "clang/AST/Type.h"
 #include "clang/Basic/LLVM.h"
+#include "clang/Sema/DeclSpec.h"
+#include "clang/Sema/Lookup.h"
+#include "clang/Sema/TemplateDeduction.h"
 
-#include <llvm/ADT/STLExtras.h>
-#include <llvm/ADT/SmallVector.h>
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+
+#include <limits>
+#include <string>
 
 using namespace clang;
 
@@ -26,11 +34,6 @@ QualType getUnderlyingArrayType(QualType baseType, ASTContext& C) {
   return baseType;
 }
 
-void ErrorEstimationHandler::SetErrorEstimationModel(
-    FPErrorEstimationModel* estModel) {
-  m_EstModel = estModel;
-}
-
 DeclRefExpr* ErrorEstimationHandler::BuildFinalErrorExpr() {
   return m_RMV->BuildDeclRef(m_Params->back());
 }
@@ -42,8 +45,7 @@ void ErrorEstimationHandler::BuildReturnErrorStmt() {
     auto flitr =
         FloatingLiteral::Create(m_RMV->m_Context, llvm::APFloat(1.0), true,
                                 m_RMV->m_Context.DoubleTy, noLoc);
-    Expr* finExpr =
-        m_EstModel->AssignError(StmtDiff(m_RetErrorExpr, flitr), "return_expr");
+    Expr* finExpr = AssignError(StmtDiff(m_RetErrorExpr, flitr), "return_expr");
     m_RMV->addToCurrentBlock(
         m_RMV->BuildOp(BO_AddAssign, BuildFinalErrorExpr(), finExpr),
         direction::forward);
@@ -102,9 +104,9 @@ void ErrorEstimationHandler::EmitNestedFunctionParamError(
     // if (utils::IsReferenceOrPointerType(fnDecl->getParamDecl(i)->getType()))
     //   continue;
     auto* derefExpr = m_RMV->BuildOp(UO_Deref, ArgResult[i]);
-    Expr* errorExpr = m_EstModel->AssignError(
-        {derivedCallArgs[i], derefExpr},
-        fnDecl->getNameInfo().getAsString() + "_param_" + std::to_string(i));
+    Expr* errorExpr = AssignError({derivedCallArgs[i], derefExpr},
+                                  fnDecl->getNameInfo().getAsString() +
+                                      "_param_" + std::to_string(i));
     Expr* FinalError = BuildFinalErrorExpr();
     Expr* errorStmt = m_RMV->BuildOp(BO_AddAssign, FinalError, errorExpr);
     m_ReverseErrorStmts.push_back(errorStmt);
@@ -265,6 +267,7 @@ void ErrorEstimationHandler::EmitDeclErrorStmts(DeclDiff<VarDecl> VDDiff,
 
 void ErrorEstimationHandler::InitialiseRMV(ReverseModeVisitor& RMV) {
   m_RMV = &RMV;
+  LookupCustomErrorFunction();
 }
 
 void ErrorEstimationHandler::ForgetRMV() { m_RMV = nullptr; }
@@ -333,7 +336,7 @@ void ErrorEstimationHandler::ActAfterProcessingArraySubscriptExpr(
           m_RMV->m_Sema.ImpCastExprToType(idx, size->getType(), CK_IntegralCast)
               .get();
       llvm::SmallVector<clang::Expr*, 2> args{size, idx};
-      Expr* extendedSize = m_EstModel->GetFunctionCall("max", "std", args);
+      Expr* extendedSize = m_RMV->GetFunctionCall("max", "std", args);
       size = m_RMV->Clone(size);
       Stmt* updateSize = m_RMV->BuildOp(BO_Assign, size, extendedSize);
       m_RMV->addToCurrentBlock(updateSize, direction::reverse);
@@ -450,6 +453,77 @@ void ErrorEstimationHandler::ActBeforeDifferentiatingCallExpr(
   pullbackArgs.push_back(BuildFinalErrorExpr());
 }
 
+void ErrorEstimationHandler::LookupCustomErrorFunction() {
+  Sema& S = m_RMV->m_Sema;
+  ASTContext& C = m_RMV->m_Context;
+  NamespaceDecl* cladNS = utils::LookupNSD(S, "clad", /*shouldExist=*/true);
+  IdentifierInfo* II = &C.Idents.get("getErrorVal");
+  DeclarationNameInfo DNInfo(DeclarationName(II), utils::GetValidSLoc(S));
+  LookupResult R(S, DNInfo, Sema::LookupOrdinaryName);
+  S.LookupQualifiedName(R, cladNS);
+  if (R.empty())
+    return;
+
+  FunctionProtoType::ExtProtoInfo EPI;
+  QualType ConstCharPtr = C.getPointerType(C.getConstType(C.CharTy));
+  QualType DoubleTy = C.DoubleTy;
+  llvm::SmallVector<QualType, 3> FnTypes = {DoubleTy, DoubleTy, ConstCharPtr};
+  QualType FnTy = C.getFunctionType(DoubleTy, FnTypes, EPI);
+  TemplateSpecCandidateSet FailedCandidates(utils::GetValidSLoc(S),
+                                            /*ForTakingAddress=*/false);
+  if (utils::MatchOverloadType(S, FnTy, R, FailedCandidates)) {
+    // FIXME: MatchOverloadType returns an overload expr without the `clad::`
+    // namespace specifier. Here, we rebuild manually.
+    CXXScopeSpec SS;
+    SS.Extend(C, cladNS, noLoc, noLoc);
+    m_CustomErrorFunction =
+        S.BuildDeclarationNameExpr(SS, R, /*ADL=*/false).get();
+    return;
+  }
+
+  // We did not match the found candidates. Warn and offer the user hints.
+  auto errId = S.Diags.getCustomDiagID(
+      DiagnosticsEngine::Error,
+      "user-defined derivative error function was provided but not used; "
+      "expected signature %0 does not match");
+  S.Diag(m_RMV->m_DiffReq->getLocation(), errId) << FnTy;
+  FailedCandidates.NoteCandidates(S, utils::GetValidSLoc(S));
+  utils::DiagnoseSignatureMismatch(S, FnTy, R);
+}
+
+Expr* ErrorEstimationHandler::AssignError(StmtDiff refExpr,
+                                          const std::string& varName) {
+  Sema& S = m_RMV->m_Sema;
+  ASTContext& C = m_RMV->m_Context;
+  if (m_CustomErrorFunction) {
+    llvm::SmallVector<clang::Expr*, 3> callParams{
+        refExpr.getExpr_dx(), refExpr.getExpr(),
+        clad::utils::CreateStringLiteral(C, varName)};
+    return S
+        .ActOnCallExpr(m_RMV->getCurrentScope(), m_CustomErrorFunction, noLoc,
+                       callParams, noLoc)
+        .get();
+  }
+  // Get the machine epsilon value.
+  double val = std::numeric_limits<float>::epsilon();
+  // Convert it into a floating point literal clang::Expr.
+  Expr* epsExpr =
+      FloatingLiteral::Create(C, llvm::APFloat(val), true, C.DoubleTy, noLoc);
+  // Here, we first build a multiplication operation for the following:
+  // refExpr * <--floating point literal (i.e. machine dependent constant)-->
+  // Build another multiplication operation with above and the derivative
+  Expr* errExpr =
+      m_RMV->BuildOp(BO_Mul, refExpr.getExpr_dx(),
+                     m_RMV->BuildOp(BO_Mul, refExpr.getExpr(), epsExpr));
+  // Next, build a llvm vector-like container to store the parameters
+  // of the function call.
+  llvm::SmallVector<Expr*, 1> params{errExpr};
+  // Finally, build a call to std::abs
+  Expr* absExpr = m_RMV->GetFunctionCall("abs", "std", params);
+  // Return the built error expression.
+  return absExpr;
+}
+
 void ErrorEstimationHandler::ActBeforeFinalizingVisitDeclStmt(
     llvm::SmallVectorImpl<Decl*>& decls,
     llvm::SmallVectorImpl<Decl*>& declsDiff) {