EnzymeAD / EnzymeAD/Enzyme

Compilation error differentiating Eigen::Matrix::inverse() with Enzyme forward-mode when matrix size ≥ 4

Open
#2,679 10 comments 0 reactions 0 assignees View on GitHub
Dominant language
LLVM
Stars
1.7k
Forks
188
Avg merge
1d 22h
Merged PRs (30d)
26

Description

When using Enzyme’s forward-mode API (`__enzyme_fwddiff`) to compute gradients through `Eigen::Matrix::inverse()` for fixed-size matrices, compilation fails when `N >= 4` (reproduced with `N = 4`). The same code compiles and runs correctly for `N = 2` and `N = 3`.

Minimal reproducible example

```cpp
#include
#include
#include
#include
#include

int enzyme_dup;
int enzyme_dupnoneed;
int enzyme_out;
int enzyme_const;

template
return_type __enzyme_fwddiff(void*, T...);

template
return_type __enzyme_autodiff(void*, T...);

const size_t input_size = 2;
const size_t matrix_size = 4; // change to 2 or 3 works fine

using Matrix = Eigen::Matrix;

void enzyme_func_matrix(Matrix* mats, double& sol) {
Matrix M0;
M0 = mats[0].inverse();
sol = M0.sum();
}

template
void compute_J_matrix_forward(
std::vector>> inputs,
std::vector>>& J_matrix_forward
) {
J_matrix_forward.assign(input_size,
std::vector>(
matrix_size,
std::vector(matrix_size, 0.0)
)
);

std::vector mats(input_size), dmats(input_size);

for (size_t mat_idx = 0; mat_idx < input_size; mat_idx++) {
for (size_t i = 0; i < matrix_size; i++) {
for (size_t j = 0; j < matrix_size; j++) {
mats[mat_idx](i, j) = inputs[mat_idx][i][j];
}
}
}

auto zero_gradients = [&]() {
for (auto& M : dmats) {
for (int i = 0; i < (int)matrix_size; i++) {
for (int j = 0; j < (int)matrix_size; j++) {
M(i, j) = 0.0;
}
}
}
};

double result, dresult;
for (size_t mat_idx = 0; mat_idx < input_size; mat_idx++) {
for (size_t i = 0; i < matrix_size; i++) {
for (size_t j = 0; j < matrix_size; j++) {
zero_gradients();
dmats[mat_idx](i, j) = 1.0;

__enzyme_fwddiff((void*)F,
enzyme_dup, mats.data(), dmats.data(),
enzyme_dup, &result, &dresult);

J_matrix_forward[mat_idx][i][j] = dresult;
}
}
}
}

void generate_inputs(size_t matrix_size, std::vector>>& inputs) {
inputs.assign(2,
std::vector>(
matrix_size,
std::vector(matrix_size)
)
);

std::random_device rd;
std::mt19937 gen(rd());
std::uniform_real_distribution dist(0.0, 1.0);

for (auto& mat : inputs) {
for (auto& row : mat) {
for (auto& val : row) {
val = dist(gen);
}
}
}
}

int main() {
std::vector>> J_matrix_forward;
std::vector>> inputs;

generate_inputs(matrix_size, inputs);
compute_J_matrix_forward(inputs, J_matrix_forward);

for (size_t input_idx = 0; input_idx < input_size; input_idx++) {
for (size_t i = 0; i < matrix_size; i++) {
for (size_t j = 0; j < matrix_size; j++) {
std::cout << "J_matrix_forward[" << input_idx << "]["
<< i << "][" << j << "] = "
<< J_matrix_forward[input_idx][i][j] << "\n";
}
}
}
return 0;
}
```

Compilation error log

```bash
In file included from ./main.cpp:3:
In file included from ../../../build/debug/include/Eigen/Dense:1:
In file included from ../../../build/debug/include/Eigen/Core:205:
../../../build/debug/include/Eigen/src/Core/arch/SSE/PacketMath.h:423:103: error: Enzyme: ; Function Attrs: mustprogress noinline optnone willreturn uwtable
define linkonce_odr dso_local noundef <2 x double> @preprocess__ZN5Eigen8internal4pxorIDv2_dEET_RKS3_S5_(ptr noundef nonnull align 16 dereferenceable(16) %a, ptr noundef nonnull align 16 dereferenceable(16) %b) #31 !dbg !15509 {
entry:
call void @llvm.dbg.value(metadata ptr %a, metadata !15510, metadata !DIExpression()), !dbg !15511
call void @llvm.dbg.value(metadata ptr %b, metadata !15512, metadata !DIExpression()), !dbg !15511
%0 = load <2 x double>, ptr %a, align 16, !dbg !15513
%1 = load <2 x double>, ptr %b, align 16, !dbg !15514
%2 = bitcast <2 x double> %0 to <2 x i64>, !dbg !15515
%3 = bitcast <2 x double> %1 to <2 x i64>, !dbg !15515
%xor.i = xor <2 x i64> %2, %3, !dbg !15515
%4 = bitcast <2 x i64> %xor.i to <2 x double>, !dbg !15515
ret <2 x double> %4, !dbg !15516
}

constantarg[ptr %a] = 0 type: {[-1]:Pointer, [-1,-1]:Float@double} - vals: {}
constantarg[ptr %b] = 0 type: {[-1]:Pointer} - vals: {}
constantinst[ call void @llvm.dbg.value(metadata ptr %a, metadata !15510, metadata !DIExpression()), !dbg !15511] = 1 val:1 type: {}
constantinst[ call void @llvm.dbg.value(metadata ptr %b, metadata !15512, metadata !DIExpression()), !dbg !15511] = 1 val:1 type: {}
constantinst[ %0 = load <2 x double>, ptr %a, align 16, !dbg !7229] = 0 val:0 type: {[-1]:Float@double}
constantinst[ %1 = load <2 x double>, ptr %b, align 16, !dbg !7230] = 0 val:0 type: {}
constantinst[ %2 = bitcast <2 x double> %0 to <2 x i64>, !dbg !7231] = 0 val:0 type: {[-1]:Float@double}
constantinst[ %3 = bitcast <2 x double> %1 to <2 x i64>, !dbg !7231] = 0 val:0 type: {}
constantinst[ %xor.i = xor <2 x i64> %2, %3, !dbg !7231] = 0 val:0 type: {[-1]:Float@double}
constantinst[ %4 = bitcast <2 x i64> %xor.i to <2 x double>, !dbg !7231] = 0 val:0 type: {[-1]:Float@double}
constantinst[ ret <2 x double> %4, !dbg !7232] = 1 val:1 type: {}
cannot handle unknown binary operator: %xor.i = xor <2 x i64> %2, %3, !dbg !7231

template<> EIGEN_STRONG_INLINE Packet2d pxor(const Packet2d& a, const Packet2d& b) { return _mm_xor_pd(a,b); }
^
../../../build/debug/include/Eigen/src/Core/arch/SSE/PacketMath.h:423:103: error: Enzyme: ; Function Attrs: mustprogress noinline optnone willreturn uwtable
define linkonce_odr dso_local noundef <2 x double> @preprocess__ZN5Eigen8internal4pxorIDv2_dEET_RKS3_S5_(ptr noundef nonnull align 16 dereferenceable(16) %a, ptr noundef nonnull align 16 dereferenceable(16) %b) #31 !dbg !15509 {
entry:
call void @llvm.dbg.value(metadata ptr %a, metadata !15510, metadata !DIExpression()), !dbg !15511
call void @llvm.dbg.value(metadata ptr %b, metadata !15512, metadata !DIExpression()), !dbg !15511
%0 = load <2 x double>, ptr %a, align 16, !dbg !15513
%1 = load <2 x double>, ptr %b, align 16, !dbg !15514
%2 = bitcast <2 x double> %0 to <2 x i64>, !dbg !15515
%3 = bitcast <2 x double> %1 to <2 x i64>, !dbg !15515
%xor.i = xor <2 x i64> %2, %3, !dbg !15515
%4 = bitcast <2 x i64> %xor.i to <2 x double>, !dbg !15515
ret <2 x double> %4, !dbg !15516
}

constantarg[ptr %a] = 0 type: {[-1]:Pointer, [-1,-1]:Float@double} - vals: {}
constantarg[ptr %b] = 0 type: {[-1]:Pointer, [-1,0]:Integer} - vals: {}
constantinst[ call void @llvm.dbg.value(metadata ptr %a, metadata !15510, metadata !DIExpression()), !dbg !15511] = 1 val:1 type: {}
constantinst[ call void @llvm.dbg.value(metadata ptr %b, metadata !15512, metadata !DIExpression()), !dbg !15511] = 1 val:1 type: {}
constantinst[ %0 = load <2 x double>, ptr %a, align 16, !dbg !7229] = 0 val:0 type: {[-1]:Float@double}
constantinst[ %1 = load <2 x double>, ptr %b, align 16, !dbg !7230] = 0 val:0 type: {[0]:Integer}
constantinst[ %2 = bitcast <2 x double> %0 to <2 x i64>, !dbg !7231] = 0 val:0 type: {[-1]:Float@double}
constantinst[ %3 = bitcast <2 x double> %1 to <2 x i64>, !dbg !7231] = 0 val:0 type: {[0]:Integer}
constantinst[ %xor.i = xor <2 x i64> %2, %3, !dbg !7231] = 0 val:0 type: {[-1]:Float@double}
constantinst[ %4 = bitcast <2 x i64> %xor.i to <2 x double>, !dbg !7231] = 0 val:0 type: {[-1]:Float@double}
constantinst[ ret <2 x double> %4, !dbg !7232] = 1 val:1 type: {}
cannot handle unknown binary operator: %xor.i = xor <2 x i64> %2, %3, !dbg !7231

2 errors generated.
```

Contributor guide

Open the contributing guide

Research direction

Start with the minimal reproducer in main.cpp and reproduce the __enzyme_fwddiff compilation failure with matrix_size = 4, comparing it with sizes 2 and 3. Inspect the Eigen/src/Core/arch/SSE/PacketMath.h pxor path and Enzyme's handling of the emitted vector xor instruction. Done means the fixed-size inverse example compiles for N >= 4 without regressing N = 2 or N = 3.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.