EnzymeAD / EnzymeAD/Enzyme

Compilation hangs (mk ii) : Matrix edition

Open
#142 13 comments 0 reactions 0 assignees View on GitHub
Dominant language
LLVM
Stars
1.7k
Forks
188
Avg merge
2d 4h
Merged PRs (30d)
22

Description

Hello,

I tried to use Eigen3 with enzyme (not sure if it's planned to be supported), and even though it works in the simple cases, it seems to hang the compilation in some more advanced cases. Like matrix inversion, matrix-vector solve, matrix exponentiation.

In https://enzyme.mit.edu/getting_started/CallingConvention/ you describe how to add some custom gradients, but it doesn't seem straight-forward to add them from c++.

For example for the adjoint of the inverse of the matrix
d K^-1/dp = - K^-1 * dK/dp * K^-1 (https://math.stackexchange.com/questions/1471825/derivative-of-the-inverse-of-a-matrix)
Which may be easier and faster to compute, and more numerically stable than the automatically derived one.

Can you please advise ?
Thanks

Here is my test file, where compilation hangs when some of the __enzyme_autodiff lines are present.

`testmatrix.cpp`
```
#include
#include
#include
#include
#include
#include
#include

#include
#include
using Eigen::MatrixXd;
using namespace std;
using namespace Eigen;

int enzyme_dup;
int enzyme_out;
int enzyme_const;

void __enzyme_autodiff(...);

template
double normVector( const Matrix& m )
{
double out = 0.0;
for( int i = 0 ; i < m.rows() ; i++ )
{
out += m(i,0)* m(i,0);
}
return out;
}

template
double normMatrix( const Matrix& m )
{
double out = 0.0;
for( int i = 0 ; i < m.rows() ; i++ )
{
for( int j = 0 ; j < m.cols(); j++)
{
out += m(i,j)* m(i,j);
}
}
return out;
}

double normMatrixXd( const MatrixXd& m )
{
double out = 0.0;
for( int i = 0 ; i < m.rows() ; i++ )
{
for( int j = 0 ; j < m.cols(); j++)
{
out += m(i,j)* m(i,j);
}
}
return out;
}

template
double normInverseMatrix( const Matrix& m )
{
return normMatrix(m.inverse());
}

double normInverseMatrixXd(const MatrixXd& m )
{
MatrixXd inv = m.inverse();
return normMatrixXd(inv);
}

template
double normSolveMatrix( const Matrix& m)
{
Matrix v;
for( int i = 0 ; i < m.cols() ; i++)
{
v(i,0) = i;
}
Matrix sol = m.fullPivLu().solve(v);
return normVector(sol);
}

double normSolveMatrixXd( const MatrixXd& m)
{
MatrixXd v(m.cols(),1);
for( int i = 0 ; i < m.cols() ; i++)
{
v(i,0) = i;
}
return normMatrixXd(m.fullPivLu().solve(v));
}

double normExpMatrixXd( const MatrixXd& m)
{
return normMatrixXd(m.exp());
}

template
double normExpMatrix( const Matrix& m )
{
return normMatrix(m.exp());
}

template< int T>
void testMatrix()
{
Matrix m;
for( int i = 0; i < T ; i++)
{
for( int j = 0 ; j < T ;j++)
{
m(i,j) = (i+j)*(i+j);
}
}

Matrix dm;
for( int i = 0; i < T ; i++)
{
for( int j = 0 ; j < T ;j++)
{
dm(i,j) = 0.0;
}
}
std::cout <<"m : " << std::endl;
std::cout << m << std::endl;
std::cout <<"m.inverse() : " << std::endl;
std::cout << m.inverse() << std::endl;
std::cout << "normSolveMatrix "<< std::endl;
std::cout << normSolveMatrix(m) << std::endl;
std::cout << "normExpMatrix "<< std::endl;
std::cout << normExpMatrix(m) << std::endl;
__enzyme_autodiff(normMatrix, enzyme_dup, &m,&dm); // Works
__enzyme_autodiff(normInverseMatrix, enzyme_dup, &m,&dm);//Hangs compilation
__enzyme_autodiff(normSolveMatrix, enzyme_dup, &m,&dm);//Hangs compilation
__enzyme_autodiff(normExpMatrix, enzyme_dup, &m,&dm);//Hangs compilation

std::cout << dm << std::endl;

}

void testMatrixXd( int T )
{

MatrixXd m(T,T);
for( int i = 0; i < T ; i++)
{
for( int j = 0 ; j < T ;j++)
{
m(i,j) = (i+j)*(i+j);
}
}

MatrixXd dm(T,T);
for( int i = 0; i < T ; i++)
{
for( int j = 0 ; j < T ;j++)
{
dm(i,j) = 0.0;
}
}
std::cout <<"m : " << std::endl;
std::cout << m << std::endl;
std::cout <<"m.inverse() : " << std::endl;
std::cout << m.inverse() << std::endl;
std::cout << "normSolveMatrix "<< std::endl;
std::cout << normSolveMatrixXd(&m) << std::endl;
std::cout << "normExpMatrix "<< std::endl;
std::cout << normExpMatrixXd(&m) << std::endl;
__enzyme_autodiff(normMatrixXd, enzyme_dup, &m,&dm); // Works
__enzyme_autodiff(normInverseMatrixXd, enzyme_dup, &m,&dm); //Hangs compilation
__enzyme_autodiff(normSolveMatrixXd, enzyme_dup, &m,&dm); //Hangs compilation
__enzyme_autodiff(normExpMatrixXd, enzyme_dup, &m,&dm); //Hangs compilation

std::cout << dm << std::endl;

}

int main()
{
testMatrix<3>();
testMatrix<4>();
testMatrix<5>(); //There are no more formulas for matrix inversion in eigen when n = 5

testMatrixXd(3);
testMatrixXd(4);
testMatrixXd(5);

return 0;
}
```
Compilation with provided that you have eigen3 installed by ubuntu (`apt-get install libeigen3-dev`) :
`clang testmatrix.cpp -I/usr/include/eigen3/ -lstdc++ -lm -Xclang -load -Xclang /usr/local/lib/ClangEnzyme-11.so -O2 -o testMatrix -fno-exceptions `

I also tried to add the following flag (`-mllvm -enzyme-max-type-offset=20` ) which helped in the past when compilation was hanging but it didn't have any effect this time.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.