Add support for verifying Enzyme Gradients with Clad Gradients #488
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Nirhar
force-pushed
the
enzyme_integration_check
branch
from
e69ceed to
2a36e90
Compare
|
This PR fails when I attempt to differentiate multiple functions with enzyme and verify it with clad. Then clad outputs the wrong gradients in some runs of the compiled code but not in others. double f1(double arr[2]) { return arr[0] * arr[1]; }
double f2(double x, double y, double z){
return x * y * z;
}
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]);
auto f2_grad=clad::gradient<clad::opts::use_enzyme>(f2);
double f2_res[3];
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]);
}If each function is individually differentiated and verified(in different files), this error does not occur. Only when they are jointly differentiated the error occurs. void f1_grad(double arr[2], clad::array_ref<double> _d_arr) {
double _t0;
double _t1;
_t1 = arr[0];
_t0 = arr[1];
double f1_return = _t1 * _t0;
goto _label0;
_label0:
{
double _r0 = 1 * _t0;
_d_arr[0] += _r0;
double _r1 = _t1 * 1;
_d_arr[1] += _r1;
}
}
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);
//Verification part below
double cladResult1[2];
f1_grad(arr, cladResult1);
EssentiallyEqualArrays(cladResult1, _d_arr.ptr(), 2UL);
}
void f2_grad(double x, double y, double z, clad::array_ref<double> _d_x, clad::array_ref<double> _d_y, clad::array_ref<double> _d_z) {
double _t0;
double _t1;
double _t2;
double _t3;
_t2 = x;
_t1 = y;
_t3 = _t2 * _t1;
_t0 = z;
double f2_return = _t3 * _t0;
goto _label0;
_label0:
{
double _r0 = 1 * _t0;
double _r1 = _r0 * _t1;
* _d_x += _r1;
double _r2 = _t2 * _r0;
* _d_y += _r2;
double _r3 = _t3 * 1;
* _d_z += _r3;
}
}
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) {
clad::EnzymeGradient<3> grad = __enzyme_autodiff_f2(f2, x, y, z);
* _d_x = grad.d_arr[0U];
* _d_y = grad.d_arr[1U];
* _d_z = grad.d_arr[2U];
//Verification part below
double cladResult1;
double cladResult2;
double cladResult3;
f2_grad(x, y, z, &cladResult1, &cladResult2, &cladResult3);
EssentiallyEqual(cladResult1, * _d_x);
EssentiallyEqual(cladResult2, * _d_y);
EssentiallyEqual(cladResult3, * _d_z);
} |
|
cc @vgvassilev @parth-07 for review and some help |
Nirhar
force-pushed
the
enzyme_integration_check
branch
from
2a36e90 to
2d05fa0
Compare
|
The above problem has been fixed. It was simply that I was not initializing declared variables in some places |
Codecov Report
@@ Coverage Diff @@
## master #488 +/- ##
==========================================
+ Coverage 92.56% 92.67% +0.10%
==========================================
Files 37 37
Lines 5528 5610 +82
==========================================
+ Hits 5117 5199 +82
Misses 411 411
|
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`
Nirhar
force-pushed
the
enzyme_integration_check
branch
from
2d05fa0 to
a882316
Compare
vgvassilev
reviewed
Sep 14, 2022
| @@ -450,6 +450,22 @@ namespace clad { | |||
| code); | |||
| } | |||
|
|
|||
| void EssentiallyEqual(double a, double b) { | |||
| // FIXME: We should select epsilon value in a more robust way. | |||
There was a problem hiding this comment.
@grimmmyshini, can you take a look?
parth-07
requested changes
Sep 15, 2022
| 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"); |
There was a problem hiding this comment.
Instead of an assert, can we give a non-fatal error here?
| @@ -413,6 +417,16 @@ namespace clad { | |||
| else | |||
| DifferentiateWithEnzyme(); | |||
|
|
|||
| if (use_enzyme && checkEnzymeWithClad) { | |||
| DiffRequest newRequest = const_cast<DiffRequest&>(request); | |||
There was a problem hiding this comment.
Why do you need const_cast here?
| auto paramType = paramsRef[i]->getOriginalType(); | ||
| llvm::SmallVector<Expr*, 2> equalityCheckArguments; | ||
| equalityCheckArguments.push_back(BuildDeclRef(cladResultDecls[i])); | ||
| if (paramType->isFloatingType()) { |
There was a problem hiding this comment.
What if paramType is neither floating type nor an array type?
| } else { | ||
| resultVar = BuildVarDecl(paramType, finalVarName, nullptr, true); | ||
| } | ||
| addToCurrentBlock(BuildDeclStmt(resultVar), direction::forward); |
There was a problem hiding this comment.
We can move this statement outside the if-else chain.
| } | ||
| addToCurrentBlock(BuildDeclStmt(resultVar), direction::forward); | ||
| cladGradArgs.push_back(BuildOp(UO_AddrOf, BuildDeclRef(resultVar))); | ||
| cladResultDecls.push_back(resultVar); |
There was a problem hiding this comment.
We can move this outside if-else chain as well.
| BuildDeclRef(paramsRef[i + numParams]), "ptr", {})); | ||
| ConstantArrayType* t = dyn_cast<ConstantArrayType>( | ||
| const_cast<Type*>(paramType.getTypePtr())); | ||
| int sizeOfArray = (int)(t->getSize().roundToDouble(false)); |
There was a problem hiding this comment.
Why do you need C-style cast here?
| // REQUIRES: Enzyme | ||
|
|
||
| #include "clad/Differentiator/Differentiator.h" | ||
|
|
There was a problem hiding this comment.
Can you please add a test containing nested function calls?
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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:
The above code will be appended with checks to verify the calculated gradients. Thus the newly generated code would be:
EssentiallyEqualArraysandEssentiallyEqualare functions defined in Differentiator.hOnly 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