EnzymeAD / EnzymeAD/Enzyme

Compilation hangs (mk ii) : Matrix edition

Aperta
#142 13 commenti 0 reazioni 0 assegnatari Vedi su GitHub
Lingua principale
LLVM
Stelle
1.7k
Fork
188
Merge medio
2g 4h
PR unite (30g)
22

Descrizione

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.

Guida per i contributori

Apri la guida per i contributori

Valutazione

Questa issue non è ancora stata valutata.

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.