EnzymeAD / EnzymeAD/Enzyme

`__enzyme_batch` over a function containing `__enzyme_fwddiff` returns wrong values

Open
#3,003 0 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

While writing some CFD-based integration tests, I found that batching a function (that itself performs a forward-mode derivative) produces incorrect results. Batching the same function's primal is correct, so the problem is somehow specific to the composition of `__enzyme_batch` over `__enzyme_fwddiff`. This looks like undefined behaviour in the nesting rather than a single arithmetic slip:

* through `opt -passes="enzyme"`, every optimisation level is wrong (see table below)
* through the clang plugin, `-O0` yields `nan` and `-O1`/`-O2`/`-O3` are correct.

Reproduction code:

`f(g, x) = g * x^3`, so `df/dx = 3*g*x^2`. Batched over `x = (1, 3)` with
`g = 2` held common, the primal is `(2, 54)` and the derivative is `(6, 54)`.

```cpp
#include

struct Batch2 {
double a, b;
};

extern Batch2 __enzyme_batch(...);
extern double __enzyme_fwddiff(void *, ...);
extern int enzyme_width;
extern int enzyme_vector;
extern int enzyme_scalar;

double f(double g, double x) { return g * x * x * x; }

// A forward-mode derivative: df/dx = 3*g*x^2.
double dfdx(double g, double x) {
return __enzyme_fwddiff((void *)f, g, 0.0, x, 1.0);
}

int main() {
const double g = 2.0, x0 = 1.0, x1 = 3.0;

// Batching the primal is correct.
Batch2 p = __enzyme_batch((void *)f, enzyme_width, 2, enzyme_scalar, g,
enzyme_vector, x0, x1);

// Batching the derivative is not.
Batch2 dv = __enzyme_batch((void *)dfdx, enzyme_width, 2, enzyme_vector, g, g,
enzyme_vector, x0, x1);
Batch2 ds = __enzyme_batch((void *)dfdx, enzyme_width, 2, enzyme_scalar, g,
enzyme_vector, x0, x1);

printf("primal : (%g, %g) want (%g, %g)\n", p.a, p.b, f(g, x0),
f(g, x1));
printf("df/dx, all vector : (%g, %g) want (%g, %g)\n", dv.a, dv.b,
dfdx(g, x0), dfdx(g, x1));
printf("df/dx, enzyme_scalar: (%g, %g) want (%g, %g)\n", ds.a, ds.b,
dfdx(g, x0), dfdx(g, x1));
return 0;
}
```

Built two ways:

```bash
# 1. opt pipeline, as the Integration tests run it
clang -O2 repro.cpp -S -emit-llvm -o - \
| opt - -load-pass-plugin=./LLVMEnzyme-22.dylib -load=./LLVMEnzyme-22.dylib \
--enzyme-attributor=0 -passes="enzyme" -S \
| lli --jit-kind=mcjit -

# 2. clang plugin
clang++ -fno-exceptions -O2 repro.cpp -o repro \
-fpass-plugin=./ClangEnzyme-22.dylib -Xclang -load -Xclang ./ClangEnzyme-22.dylib
```

Expected:

```
primal : (2, 54) want (2, 54)
df/dx, all vector : (6, 54) want (6, 54)
df/dx, enzyme_scalar: (6, 54) want (6, 54)
```

Actual:

The primal row is `(2, 54)` — correct — in every configuration below. Only the
two derivative rows differ.

**`opt -passes="enzyme"`** — wrong at every -O:

| | `-O0` | `-O1` | `-O2` | `-O3` |
|---|---|---|---|---|
| `df/dx`, all `enzyme_vector` | `(54, 0)` | `(54, 1944)` | `(54, 1944)` | `(54, 1944)` |
| `df/dx`, `enzyme_scalar` | `(54, 0)` | `(54, 54)` | `(54, 54)` | `(54, 54)` |

**clang plugin (`-fpass-plugin`)** — wrong only at `-O0`:

| | `-O0` | `-O1` | `-O2` | `-O3` |
|---|---|---|---|---|
| `df/dx`, all `enzyme_vector` | `(nan, nan)` | `(6, 54)` | `(6, 54)` | `(6, 54)` |
| `df/dx`, `enzyme_scalar` | `(nan, nan)` | `(6, 54)` | `(6, 54)` | `(6, 54)` |

I note that, in every wrong non-`nan` case, lane 0 returns `54`, which is lane 1's answer. The lanes are not being kept apart. Best guess, the configurations that work are those where the inner `__enzyme_fwddiff` call has been inlined into `dfdx` before the Enzyme pass runs. Under `opt`, the pass sees the module once, after clang has finished, and teh inner call is still an opaque external call inside the function being batched

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.