EnzymeAD / EnzymeAD/Enzyme-JAX

add_reduce_slice_fusion fuses through shared intermediates, making linear accumulator chains quadratic

Open
#2,830 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
MLIR
Stars
131
Forks
53
Avg merge
1d 10h
Merged PRs (30d)
193

Description

`add_reduce_slice_fusion` fuses through binary-op intermediates that have other users. Each intermediate then gets its own `reduce` over a growing prefix, so a linear accumulator chain becomes quadratic work.

Found via EnzymeAD/Reactant.jl#3166, where a plain unrolled Julia loop got ~200x slower than it should at N=1024.

### Minimal reproducer

```mlir
module {
func.func @main(%x: tensor<4xf32>) -> tensor {
%zero = stablehlo.constant dense<0.000000e+00> : tensor
%s0 = stablehlo.slice %x [0:1] : (tensor<4xf32>) -> tensor<1xf32>
%e0 = stablehlo.reshape %s0 : (tensor<1xf32>) -> tensor
%st1 = stablehlo.add %zero, %e0 : tensor
%q1 = stablehlo.multiply %st1, %st1 : tensor
%s1 = stablehlo.slice %x [1:2] : (tensor<4xf32>) -> tensor<1xf32>
%e1 = stablehlo.reshape %s1 : (tensor<1xf32>) -> tensor
%st2 = stablehlo.add %st1, %e1 : tensor
%q2 = stablehlo.multiply %st2, %st2 : tensor
%s2 = stablehlo.slice %x [2:3] : (tensor<4xf32>) -> tensor<1xf32>
%e2 = stablehlo.reshape %s2 : (tensor<1xf32>) -> tensor
%st3 = stablehlo.add %st2, %e2 : tensor
%q3 = stablehlo.multiply %st3, %st3 : tensor
%s3 = stablehlo.slice %x [3:4] : (tensor<4xf32>) -> tensor<1xf32>
%e3 = stablehlo.reshape %s3 : (tensor<1xf32>) -> tensor
%st4 = stablehlo.add %st3, %e3 : tensor
%q4 = stablehlo.multiply %st4, %st4 : tensor
%l1 = stablehlo.add %q1, %q2 : tensor
%l2 = stablehlo.add %l1, %q3 : tensor
%l3 = stablehlo.add %l2, %q4 : tensor
return %l3 : tensor
}
}
```

This is `state_i = state_{i-1} + x[i]; loss += state_i^2` unrolled — a linear chain in which every `state_i` has **two** users: the next `add`, and its own `multiply`.

Running just this one pattern:

```
enzyme-hlo-generate-td{patterns=add_reduce_slice_fusion},transform-interpreter,enzyme-hlo-remove-transform
```

gives:

```mlir
func.func @main(%arg0: tensor<4xf32>) -> tensor {
%cst = stablehlo.constant dense<0.000000e+00> : tensor
%0 = stablehlo.slice %arg0 [0:1] : (tensor<4xf32>) -> tensor<1xf32>
%1 = stablehlo.reshape %0 : (tensor<1xf32>) -> tensor
%2 = stablehlo.add %cst, %1 : tensor
%3 = stablehlo.multiply %2, %2 : tensor
%4 = stablehlo.slice %arg0 [0:2] : (tensor<4xf32>) -> tensor<2xf32>
%5 = stablehlo.reduce(%4 init: %cst) applies stablehlo.add across dimensions = [0] : (tensor<2xf32>, tensor) -> tensor
%6 = stablehlo.add %5, %cst : tensor
%7 = stablehlo.multiply %6, %6 : tensor
%8 = stablehlo.slice %arg0 [0:3] : (tensor<4xf32>) -> tensor<3xf32>
%9 = stablehlo.reduce(%8 init: %cst) applies stablehlo.add across dimensions = [0] : (tensor<3xf32>, tensor) -> tensor
...
%12 = stablehlo.slice %arg0 [0:4] : (tensor<4xf32>) -> tensor<4xf32>
%13 = stablehlo.reduce(%12 init: %cst) applies stablehlo.add across dimensions = [0] : (tensor<4xf32>, tensor) -> tensor
...
}
```

`N-1` independent reduces over prefixes `[0:2]`, `[0:3]`, `[0:4]`, … The original chain is *not* removed, because each `state_i` still has its `multiply` user — so the prefix sums are recomputed from scratch instead of shared, and total work goes from `O(N)` to `O(N^2)`.

(I ran this through Reactant's `run_pass_pipeline!` with the pipeline string above rather than `enzymexlamlir-opt`, which I don't have built locally — the pattern list is the same.)

### Cause

`ReduceSliceFusionBase::collectSlicesInChain` (`src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp`) walks the binary-op chain and descends into nested binary ops unconditionally:

```cpp
if (auto binaryOp = current.template getDefiningOp()) {
worklist.push_back(binaryOp.getLhs());
worklist.push_back(binaryOp.getRhs());
}
```

There is no check that `current` is used only within the chain being fused. Fusing is only a win when the intermediates are dead afterwards; when an intermediate has an outside user, the rewrite duplicates all the work leading up to it, once per user.

A single-use guard on descent should be enough — the root's result may have arbitrary users, but any deeper intermediate must have exactly one:

```cpp
if (auto binaryOp = current.template getDefiningOp()) {
if (current != startOp.getResult() && !current.hasOneUse()) {
extraValues.push_back(current); // treat as opaque leaf, do not fuse through
continue;
}
worklist.push_back(binaryOp.getLhs());
worklist.push_back(binaryOp.getRhs());
}
```

The same reasoning applies to the other members of the family (`Mul`/`Min`/`Max`/`And`/`Or`/`Xor`), since they share `ReduceSliceFusionBase`.

### Why a slightly different loop is unaffected

From the Reactant issue: changing the recurrence to `state = state * 0.9f + x[i]` does not blow up. That is consistent with the above — a `stablehlo.multiply` sits between consecutive adds, the walk hits a non-`BinaryOpType` and stops, so no long chain is ever collected. It is not that the shared-intermediate case is handled; it is that the chain never forms.

### Impact

Measured through Reactant on CPU, excluding only `add_reduce_slice_fusion` from the pattern list (results identical either way):

| N | default | without the pattern |
|---|---|---|
| 16 | 15 reduces, 5.6 µs | 0 reduces, 3.0 µs |
| 64 | 63 reduces, 24.9 µs | 0, 2.3 µs |
| 256 | 255 reduces, 114.3 µs | 0, 2.5 µs |
| 1024 | 1023 reduces, 708.9 µs | 0, 3.4 µs |

Julia-level reproducer, for completeness:

```julia
using Reactant
function PrimalUnrolled(x)
state = zero(x[begin]); loss = zero(x[begin])
for i in eachindex(x)
state += x[i]
loss += state * state
end
return loss
end
x = Reactant.to_rarray(Float32.(1:1024))
Reactant.@allowscalar @code_hlo PrimalUnrolled(x) # 1023 stablehlo.reduce ops
```

cc @avik-pal @mofeing @jumerckx

Contributor guide

No contributing guide indexed for this repository

Research direction

Start in src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp, at ReduceSliceFusionBase::collectSlicesInChain and the add_reduce_slice_fusion pipeline entry. Run the supplied MLIR reproducer through the listed pipeline and inspect the generated reduces. Done means shared intermediate states are not fused through when they have outside users, while the other ReduceSliceFusionBase operations retain their intended behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
compilers, performance
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
72/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.