Skip to content

Commit

Permalink
Add support for verifying Enzyme Gradients with Clad Gradients
Browse files Browse the repository at this point in the history
This commit generates code that will verify the results of Enzyme Gradients with Clad Gradients.
For example, if previously the following code was generated for differentiating with enzyme for a function:
```cpp
void f1_grad_enzyme(double arr[2], clad::array_ref<double> _d_arr) {
    double *d_arr = _d_arr.ptr();
    __enzyme_autodiff_f1(f1, arr, d_arr);
}

```

The above code will be appended with checks to verify the calculated gradients. Thus the newly generated code would be:
```cpp
void f1_grad_enzyme(double arr[2], clad::array_ref<double> _d_arr) {
    double *d_arr = _d_arr.ptr();
    __enzyme_autodiff_f1(f1, arr, d_arr);
    double cladResult1[2];
    f1_grad(arr, cladResult1);
    EssentiallyEqualArrays(cladResult1, _d_arr.ptr(), 2UL);
}
```

`EssentiallyEqualArrays` and `EssentiallyEqual` are functions defined in Differentiator.h

Only functions with primitive type and ConstantArray type parameters can be verified in this manner.

To trigger this verification one must append the following flag to clang while compiling the function to be generated: `-Xclang -plugin-arg-clad -Xclang -fcheck-enzyme-with-clad`
  • Loading branch information
Nirhar committed Aug 26, 2022
1 parent 09e2ed0 commit 2d05fa0
Show file tree
Hide file tree
Showing 11 changed files with 321 additions and 13 deletions.
3 changes: 3 additions & 0 deletions include/clad/Differentiator/DiffPlanner.h
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,9 @@ namespace clad {
// A flag to enable the use of enzyme for backend instead of clad
bool use_enzyme = false;

// A flag to generate code that verifies clad and enzyme
bool checkEnzymeWithClad = false;

/// Recomputes `DiffInputVarsInfo` using the current values of data members.
///
/// Differentiation parameters info is computed by parsing the argument
Expand Down
16 changes: 16 additions & 0 deletions include/clad/Differentiator/Differentiator.h
Original file line number Diff line number Diff line change
Expand Up @@ -450,6 +450,22 @@ namespace clad {
code);
}

void EssentiallyEqual(double a, double b) {
// FIXME: We should select epsilon value in a more robust way.
const double epsilon = 1e-12;
// printf("a=%.40f, b=%.40f\n",a,b);
bool ans = std::fabs(a - b) <=
((std::fabs(a > b) ? std::fabs(b) : std::fabs(a)) * epsilon);

assert(ans && "Clad Gradient is not equal to Enzyme Gradient");
}
void EssentiallyEqualArrays(double* a, double* b, unsigned size) {
// FIXME: We should select epsilon value in a more robust way.
for (int i = 0; i < size; i++) {
EssentiallyEqual(a[i], b[i]);
}
}

// Gradient Structure for Reverse Mode Enzyme
template <unsigned N> struct EnzymeGradient { double d_arr[N]; };
}
Expand Down
4 changes: 4 additions & 0 deletions include/clad/Differentiator/ReverseModeVisitor.h
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ namespace clad {
unsigned numParams = 0;
bool isVectorValued = false;
bool use_enzyme = false;
bool checkEnzymeWithClad = false;
// FIXME: Should we make this an object instead of a pointer?
// Downside of making it an object: We will need to include
// 'MultiplexExternalRMVSource.h' file
Expand Down Expand Up @@ -92,6 +93,9 @@ namespace clad {
// Function to Differentiate with Enzyme as Backend
void DifferentiateWithEnzyme();

// Function that inserts code to verify Enzyme Results with Clad Results
void CheckEnzymeResultsWithClad(clang::FunctionDecl* cladFD);

public:
using direction = rmv::direction;
clang::Expr* dfdx() {
Expand Down
6 changes: 6 additions & 0 deletions include/clad/Differentiator/VisitorBase.h
Original file line number Diff line number Diff line change
Expand Up @@ -374,6 +374,12 @@ namespace clad {
/// \returns The declaration of the class with the name ClassName
clang::TemplateDecl*
LookupTemplateDeclInCladNamespace(llvm::StringRef ClassName);
/// Find declaration of clad::function templated type
///
/// \param[in] FunctionName name of the function to be found
/// \returns The declaration of the function with the name FunctionName
clang::FunctionDecl*
LookupFunctionDeclInCladNamespace(llvm::StringRef FunctionName);
/// Instantiate clad::class<TemplateArgs> type
///
/// \param[in] CladClassDecl the decl of the class that is going to be used
Expand Down
109 changes: 109 additions & 0 deletions lib/Differentiator/ReverseModeVisitor.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -268,6 +268,10 @@ namespace clad {
if (request.use_enzyme)
use_enzyme = true;

if (request.checkEnzymeWithClad) {
checkEnzymeWithClad = true;
}

auto derivativeBaseName = request.BaseFunctionName;
std::string gradientName = derivativeBaseName + funcPostfix();
// To be consistent with older tests, nothing is appended to 'f_grad' if
Expand Down Expand Up @@ -413,6 +417,16 @@ namespace clad {
else
DifferentiateWithEnzyme();

if (use_enzyme && checkEnzymeWithClad) {
DiffRequest newRequest = const_cast<DiffRequest&>(request);
newRequest.checkEnzymeWithClad = false;
newRequest.use_enzyme = false;
FunctionDecl* cladFD =
plugin::ProcessDiffRequest(m_CladPlugin, newRequest);

CheckEnzymeResultsWithClad(cladFD);
}

gradientBody = endBlock();
m_Derivative->setBody(gradientBody);
endScope(); // Function body scope
Expand Down Expand Up @@ -707,6 +721,101 @@ namespace clad {
addToCurrentBlock(enzymeCall);
}
}
void ReverseModeVisitor::CheckEnzymeResultsWithClad(FunctionDecl* cladFD) {
// Prepare Arguments for the clad derivative function
llvm::SmallVector<Expr*, 16> cladGradArgs;
llvm::SmallVector<VarDecl*, 16> cladResultDecls;
unsigned numParams = m_Function->getNumParams();
llvm::ArrayRef<ParmVarDecl*> paramsRef = m_Derivative->parameters();

for (int i = 0; i < numParams; i++) {
cladGradArgs.push_back(BuildDeclRef(paramsRef[i]));
}
std::string varNames = "cladResult";
int varNo = 1;
auto size_type = m_Context.getSizeType();
unsigned size_type_bits = m_Context.getIntWidth(size_type);
for (int i = 0; i < numParams; i++) {
std::string finalVarName = varNames + std::to_string(varNo++);
auto paramType = paramsRef[i]->getOriginalType();

// FIX-ME: Non Constant Array/pointer type parameters can't be dealt with
// as of now because we don't know the size of the array This code will
// break if we use array type parameters. This can be fixed if the
// ReverseModeVisitor keeps track of the maximum index of the array seen
// so far.

if (isArrayOrPointerType(paramType)) {
assert(paramType->isConstantArrayType() &&
"Only Constant type arrays are allowed to be parameters of "
"functions whose gradients we want to verify with clad");

// Create InitList to set all elements of the result array to zero
auto init = FloatingLiteral::Create(
m_Context, llvm::APFloat(0.0), true,
dyn_cast<ConstantArrayType>(paramType)->getElementType(), noLoc);
llvm::SmallVector<Expr*, 2> initListElement{init};
auto initList = dyn_cast<InitListExpr>(
m_Sema.BuildInitList(noLoc, initListElement, noLoc).get());
ImplicitValueInitExpr imp(
dyn_cast<ConstantArrayType>(paramType)->getElementType());
initList->setArrayFiller(&imp);

auto resultVar = BuildVarDecl(paramType, finalVarName, initList, true);
addToCurrentBlock(BuildDeclStmt(resultVar), direction::forward);
cladGradArgs.push_back(BuildDeclRef(resultVar));
cladResultDecls.push_back(resultVar);
} else {
VarDecl* resultVar;
if (paramType->isFloatingType()) {
auto init = FloatingLiteral::Create(m_Context, llvm::APFloat(0.0),
true, paramType, noLoc);
resultVar = BuildVarDecl(paramType, finalVarName, init, true);
} else {
resultVar = BuildVarDecl(paramType, finalVarName, nullptr, true);
}
addToCurrentBlock(BuildDeclStmt(resultVar), direction::forward);
cladGradArgs.push_back(BuildOp(UO_AddrOf, BuildDeclRef(resultVar)));
cladResultDecls.push_back(resultVar);
}
}

Expr* cladCall = BuildCallExprToFunction(cladFD, cladGradArgs);
addToCurrentBlock(cladCall);

// Compare the values
FunctionDecl* equalityFD =
LookupFunctionDeclInCladNamespace("EssentiallyEqual");
FunctionDecl* equalityFDForArrays =
LookupFunctionDeclInCladNamespace("EssentiallyEqualArrays");
for (int i = 0; i < numParams; i++) {
auto paramType = paramsRef[i]->getOriginalType();
llvm::SmallVector<Expr*, 2> equalityCheckArguments;
equalityCheckArguments.push_back(BuildDeclRef(cladResultDecls[i]));
if (paramType->isFloatingType()) {
equalityCheckArguments.push_back(
BuildOp(UO_Deref, BuildDeclRef(paramsRef[i + numParams])));
Expr* checkCall =
BuildCallExprToFunction(equalityFD, equalityCheckArguments);
addToCurrentBlock(checkCall);
} else if (paramType->isConstantArrayType()) {
equalityCheckArguments.push_back(BuildCallExprToMemFn(
BuildDeclRef(paramsRef[i + numParams]), "ptr", {}));
ConstantArrayType* t = dyn_cast<ConstantArrayType>(
const_cast<Type*>(paramType.getTypePtr()));
int sizeOfArray = (int)(t->getSize().roundToDouble(false));
llvm::APInt idxValue(size_type_bits, sizeOfArray);
auto idx =
IntegerLiteral::Create(m_Context, idxValue, size_type, noLoc);
equalityCheckArguments.push_back(idx);

Expr* checkCall = BuildCallExprToFunction(equalityFDForArrays,
equalityCheckArguments);
addToCurrentBlock(checkCall);
}
}
}

StmtDiff ReverseModeVisitor::VisitStmt(const Stmt* S) {
diag(
DiagnosticsEngine::Warning,
Expand Down
14 changes: 14 additions & 0 deletions lib/Differentiator/VisitorBase.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -366,6 +366,20 @@ namespace clad {
return cast<TemplateDecl>(TapeR.getFoundDecl());
}

FunctionDecl*
VisitorBase::LookupFunctionDeclInCladNamespace(llvm::StringRef FunctionName) {
NamespaceDecl* CladNS = GetCladNamespace();
CXXScopeSpec CSS;
CSS.Extend(m_Context, CladNS, noLoc, noLoc);
DeclarationName TapeName = &m_Context.Idents.get(FunctionName);
LookupResult TapeR(m_Sema, TapeName, noLoc, Sema::LookupUsingDeclName,
clad_compat::Sema_ForVisibleRedeclaration);
m_Sema.LookupQualifiedName(TapeR, CladNS, CSS);
assert(!TapeR.empty() && isa<FunctionDecl>(TapeR.getFoundDecl()) &&
"cannot find clad::tape");
return cast<FunctionDecl>(TapeR.getFoundDecl());
}

QualType VisitorBase::InstantiateTemplate(TemplateDecl* CladClassDecl,
TemplateArgumentListInfo& TLI) {
// This will instantiate tape<T> type and return it.
Expand Down
144 changes: 144 additions & 0 deletions test/Enzyme/ReverseModeWithCladCheck.C
Original file line number Diff line number Diff line change
@@ -0,0 +1,144 @@
// RUN: %cladclang %s -I%S/../../include -Xclang -plugin-arg-clad -Xclang -fcheck-enzyme-with-clad -lstdc++ -oReverseModeWithCladCheck.out | FileCheck %s
// RUN: ./ReverseModeWithCladCheck.out | FileCheck -check-prefix=CHECK-EXEC %s
// CHECK-NOT: {{.*error|warning|note:.*}}
// REQUIRES: Enzyme

#include "clad/Differentiator/Differentiator.h"

double f1(double arr[2]) { return arr[0] * arr[1]; }

// CHECK: void f1_grad_enzyme(double arr[2], clad::array_ref<double> _d_arr) {
// CHECK-NEXT: double *d_arr = _d_arr.ptr();
// CHECK-NEXT: __enzyme_autodiff_f1(f1, arr, d_arr);
// CHECK-NEXT: double cladResult1[2] = {0.};
// CHECK-NEXT: f1_grad(arr, cladResult1);
// CHECK-NEXT: EssentiallyEqualArrays(cladResult1, _d_arr.ptr(), 2UL);
// CHECK-NEXT:}

double f2(double x, double y, double z){
return x * y * z;
}

// CHECK: void f2_grad_enzyme(double x, double y, double z, clad::array_ref<double> _d_x, clad::array_ref<double> _d_y, clad::array_ref<double> _d_z) {
// CHECK-NEXT: clad::EnzymeGradient<3> grad = __enzyme_autodiff_f2(f2, x, y, z);
// CHECK-NEXT: * _d_x = grad.d_arr[0U];
// CHECK-NEXT: * _d_y = grad.d_arr[1U];
// CHECK-NEXT: * _d_z = grad.d_arr[2U];
// CHECK-NEXT: double cladResult1 = 0.;
// CHECK-NEXT: double cladResult2 = 0.;
// CHECK-NEXT: double cladResult3 = 0.;
// CHECK-NEXT: f2_grad(x, y, z, &cladResult1, &cladResult2, &cladResult3);
// CHECK-NEXT: EssentiallyEqual(cladResult1, * _d_x);
// CHECK-NEXT: EssentiallyEqual(cladResult2, * _d_y);
// CHECK-NEXT: EssentiallyEqual(cladResult3, * _d_z);
// CHECK-NEXT:}

double f3(double arr[3], int n){
double sum=0;
for(int i=0;i<n;i++){
sum+=arr[i]*arr[i];
}
return sum;
}

// CHECK: void f3_grad_enzyme(double arr[3], int n, clad::array_ref<double> _d_arr, clad::array_ref<int> _d_n) {
// CHECK-NEXT: double *d_arr = _d_arr.ptr();
// CHECK-NEXT: __enzyme_autodiff_f3(f3, arr, d_arr, n);
// CHECK-NEXT: double cladResult1[3] = {0.};
// CHECK-NEXT: int cladResult2;
// CHECK-NEXT: f3_grad(arr, n, cladResult1, &cladResult2);
// CHECK-NEXT: EssentiallyEqualArrays(cladResult1, _d_arr.ptr(), 3UL);
// CHECK-NEXT: }

double f4(double arr1[3], int n, double arr2[2], int m){
double sum=0;
for(int i=0;i<n;i++){
sum+=arr1[i]*arr1[i];
}
for(int i=0;i<m;i++){
sum+=arr2[i]*arr2[i];
}
return sum;
}

// CHECK: void f4_grad_enzyme(double arr1[3], int n, double arr2[2], int m, clad::array_ref<double> _d_arr1, clad::array_ref<int> _d_n, clad::array_ref<double> _d_arr2, clad::array_ref<int> _d_m) {
// CHECK-NEXT: double *d_arr1 = _d_arr1.ptr();
// CHECK-NEXT: double *d_arr2 = _d_arr2.ptr();
// CHECK-NEXT: __enzyme_autodiff_f4(f4, arr1, d_arr1, n, arr2, d_arr2, m);
// CHECK-NEXT: double cladResult1[3] = {0.};
// CHECK-NEXT: int cladResult2;
// CHECK-NEXT: double cladResult3[2] = {0.};
// CHECK-NEXT: int cladResult4;
// CHECK-NEXT: f4_grad(arr1, n, arr2, m, cladResult1, &cladResult2, cladResult3, &cladResult4);
// CHECK-NEXT: EssentiallyEqualArrays(cladResult1, _d_arr1.ptr(), 3UL);
// CHECK-NEXT: EssentiallyEqualArrays(cladResult3, _d_arr2.ptr(), 2UL);
// CHECK-NEXT: }

double f5(double arr[3], double x,int n,double y){
double res=0;
for(int i=0;i<n;i++){
res+=(arr[i]*x*y);
}
return res;
}

// CHECK: void f5_grad_enzyme(double arr[3], double x, int n, double y, clad::array_ref<double> _d_arr, clad::array_ref<double> _d_x, clad::array_ref<int> _d_n, clad::array_ref<double> _d_y) {
// CHECK-NEXT: double *d_arr = _d_arr.ptr();
// CHECK-NEXT: clad::EnzymeGradient<2> grad = __enzyme_autodiff_f5(f5, arr, d_arr, x, n, y);
// CHECK-NEXT: * _d_x = grad.d_arr[0U];
// CHECK-NEXT: * _d_y = grad.d_arr[1U];
// CHECK-NEXT: double cladResult1[3] = {0.};
// CHECK-NEXT: double cladResult2 = 0.;
// CHECK-NEXT: int cladResult3;
// CHECK-NEXT: double cladResult4 = 0.;
// CHECK-NEXT: f5_grad(arr, x, n, y, cladResult1, &cladResult2, &cladResult3, &cladResult4);
// CHECK-NEXT: EssentiallyEqualArrays(cladResult1, _d_arr.ptr(), 3UL);
// CHECK-NEXT: EssentiallyEqual(cladResult2, * _d_x);
// CHECK-NEXT: EssentiallyEqual(cladResult4, * _d_y);
// CHECK-NEXT: }


int main() {
auto f1_grad = clad::gradient<clad::opts::use_enzyme>(f1);
double f1_v[2] = {3, 4};
double f1_g[2] = {0};
f1_grad.execute(f1_v, f1_g);
printf("d_x = %.2f, d_y = %.2f\n", f1_g[0], f1_g[1]);
// CHECK-EXEC: d_x = 4.00, d_y = 3.00

auto f2_grad=clad::gradient<clad::opts::use_enzyme>(f2);
double f2_res[3]={0};
double f2_x=3,f2_y=4,f2_z=5;
f2_grad.execute(f2_x,f2_y,f2_z,&f2_res[0],&f2_res[1],&f2_res[2]);
printf("d_x = %.2f, d_y = %.2f, d_z = %.2f\n", f2_res[0], f2_res[1], f2_res[2]);
//CHECK-EXEC: d_x = 20.00, d_y = 15.00, d_z = 12.00

auto f3_grad=clad::gradient<clad::opts::use_enzyme>(f3);
double f3_list[3]={3,4,5};
double f3_res[3]={0};
int f3_dn=0;
f3_grad.execute(f3_list,3,f3_res,&f3_dn);
printf("d_x1 = %.2f, d_x2 = %.2f, d_x3 = %.2f, d_n = %d\n",f3_res[0],f3_res[1],f3_res[2],f3_dn);
//CHECK-EXEC: d_x1 = 6.00, d_x2 = 8.00, d_x3 = 10.00, d_n = 0

auto f4_grad=clad::gradient<clad::opts::use_enzyme>(f4);
double f4_list1[3]={3,4,5};
double f4_list2[2]={1,2};
double f4_res1[3]={0};
double f4_res2[2]={0};
int f4_dn1=0,f4_dn2=0;
f4_grad.execute(f4_list1,3,f4_list2,2,f4_res1,&f4_dn1,f4_res2,&f4_dn2);
printf("d_x1 = %.2f, d_x2 = %.2f, d_x3 = %.2f, d_n1 = %d\n",f4_res1[0],f4_res1[1],f4_res1[2],f4_dn1);
//CHECK-EXEC: d_x1 = 6.00, d_x2 = 8.00, d_x3 = 10.00, d_n1 = 0
printf("d_y1 = %.2f, d_y2 = %.2f, d_n2 = %d\n",f4_res2[0],f4_res2[1],f4_dn2);
//CHECK-EXEC: d_y1 = 2.00, d_y2 = 4.00, d_n2 = 0

auto f5_grad=clad::gradient<clad::opts::use_enzyme>(f5);
double f5_list[3]={3,4,5};
double f5_res[3]={0};
double f5_x=10.0,f5_dx=0,f5_y=5,f5_dy=0;
int f5_dn=0;
f5_grad.execute(f5_list,f5_x,3,f5_y,f5_res,&f5_dx,&f5_dn,&f5_dy);
printf("d_x1 = %.2f, d_x2 = %.2f, d_x3 = %.2f, d_n1 = %d, d_x = %.2f, d_y = %.2f\n",f5_res[0],f5_res[1],f5_res[2],f5_dn, f5_dx, f5_dy);
//CHECK-EXEC: d_x1 = 50.00, d_x2 = 50.00, d_x3 = 50.00, d_n1 = 0, d_x = 60.00, d_y = 120.00
}
6 changes: 6 additions & 0 deletions tools/ClangPlugin.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -182,6 +182,12 @@ namespace clad {
if (m_DO.DumpSourceFnAST) {
FD->dumpColor();
}

// If enabled, update request to also compare enzyme and clad results
if (m_DO.CheckEnzymeWithClad) {
request.checkEnzymeWithClad = true;
}

// if enabled, load the dynamic library input from user to use
// as a custom estimation model.
if (m_DO.CustomEstimationModel) {
Expand Down
6 changes: 5 additions & 1 deletion tools/ClangPlugin.h
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,8 @@ namespace clad {
: DumpSourceFn(false), DumpSourceFnAST(false), DumpDerivedFn(false),
DumpDerivedAST(false), GenerateSourceFile(false),
ValidateClangVersion(false), CustomEstimationModel(false),
PrintNumDiffErrorInfo(false), CustomModelName("") {}
PrintNumDiffErrorInfo(false), CheckEnzymeWithClad(false),
CustomModelName("") {}

bool DumpSourceFn : 1;
bool DumpSourceFnAST : 1;
Expand All @@ -76,6 +77,7 @@ namespace clad {
bool ValidateClangVersion : 1;
bool CustomEstimationModel : 1;
bool PrintNumDiffErrorInfo : 1;
bool CheckEnzymeWithClad : 1;
std::string CustomModelName;
};

Expand Down Expand Up @@ -157,6 +159,8 @@ namespace clad {
m_DO.CustomModelName = args[i];
} else if (args[i] == "-fprint-num-diff-errors") {
m_DO.PrintNumDiffErrorInfo = true;
} else if (args[i] == "-fcheck-enzyme-with-clad") {
m_DO.CheckEnzymeWithClad = true;
} else if (args[i] == "-help") {
// Print some help info.
llvm::errs()
Expand Down
Loading

0 comments on commit 2d05fa0

Please sign in to comment.