EnzymeAD / EnzymeAD/Enzyme-JAX
add_reduce_slice_fusion fuses through shared intermediates, making linear accumulator chains quadratic
- 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