Compilation hangs (mk ii) : Matrix edition
- 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.