EnzymeAD / EnzymeAD/Enzyme-JAX
convert_concat / elementwise_reshape_like / concat_insert_dim_elementwise never terminate on convert(concat(reshape))
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
## Summary
`convert_concat`, `elementwise_reshape_like` and `concat_insert_dim_elementwise` form a
rewrite cycle on `convert(concat(reshape(a), reshape(b)))`. The greedy driver never reaches
a fixed point, and because one of the three steps is a batching rewrite, the module *grows*
on every trip (a fresh `enzymexla_unbatched_ConcatInsertDimToBatch_*` wrapper function each
round), so this is an unbounded loop, not just a slow one.
## Minimal reproducer
```mlir
// repro.mlir
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg0: !transform.any_op) {
%0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %0 {
transform.apply_patterns.enzyme_hlo.concat_insert_dim_elementwise
transform.apply_patterns.enzyme_hlo.elementwise_reshape_like
transform.apply_patterns.enzyme_hlo.convert_concat
} : !transform.any_op
transform.yield
}
func.func @cycle(%arg0: tensor<16x64x2xf32>, %arg1: tensor<16x64x2xf32>) -> tensor<2x16x64x2xbf16> {
%0 = stablehlo.reshape %arg0 : (tensor<16x64x2xf32>) -> tensor<1x16x64x2xf32>
%1 = stablehlo.reshape %arg1 : (tensor<16x64x2xf32>) -> tensor<1x16x64x2xf32>
%2 = stablehlo.concatenate %0, %1, dim = 0 : (tensor<1x16x64x2xf32>, tensor<1x16x64x2xf32>) -> tensor<2x16x64x2xf32>
%3 = stablehlo.convert %2 : (tensor<2x16x64x2xf32>) -> tensor<2x16x64x2xbf16>
return %3 : tensor<2x16x64x2xbf16>
}
}
```
```
enzymexlamlir-opt --transform-interpreter --enzyme-hlo-remove-transform repro.mlir
```
does not terminate. Measured through Reactant.jl (same patterns, same C++, `Reactant_jll`
v0.0.405): no output after **10 minutes** on this six-line function, with RSS climbing
steadily while it spins (1.33 GB -> 1.70 GB over three minutes). All three patterns are
needed; drop any one and it converges immediately.
## The cycle
```
start: convert(concat(reshape(a), reshape(b)))
convert_concat -> concat(convert(reshape(a)), convert(reshape(b)))
elementwise_reshape_like -> concat(reshape(convert(a)), reshape(convert(b)))
concat_insert_dim_elementwise -> convert(concat(reshape(a), reshape(b))) == start
```
The last step is `ConcatInsertDimToBatchBase::matchAndRewriteImpl` (`AutoBatching.cpp`),
which outlines the batched elementwise op into a new wrapper function every time, so each
round leaves more IR behind. `reorder_elementwise_and_shape_op` plays the same role as
`elementwise_reshape_like` for the transpose-propagation group, so the transpose side has
the same loop.
## How it shows up in practice
Any graph with fp32 master weights and bf16 compute: a rotary embedding's
`concat(reshape(...), reshape(...))` feeding a bf16 cast is enough. In Reactant.jl this hung
`@compile` of a small GPT (rotary, QK-norm, GQA, sliding-window attention, ReLU^2) for over
45 minutes with no output as soon as the model had two or more transformer layers, on both
CPU and CUDA. One layer compiled in 91 s; fp32 versions of the same graphs were fine at any
depth. It is not Enzyme-specific -- a hand-written backward pass with no `enzyme.autodiff`
op in the module hangs identically.
Reactant's `:all` pipeline is where this bites, because its post-optimization cleanup stage
re-runs the pattern set with reshape/transpose propagation reversed (`down`, `up`, `down`),
which is the configuration that puts all three patterns in one greedy run.
`optimize=:only_enzyme` and `:before_enzyme` avoid it.
## Fix
PR to follow: `ConvertConcat` declines when every concat operand is a reshape-like op
inserting the concat dimension -- exactly the shape `ConcatInsertDimToBatch` produces, where
pushing the convert inside gains nothing because it is batched straight back. That is one of
three possible cut points; happy to move it if you prefer breaking the loop elsewhere.
Reactant.jl gets a matching change so its default pipeline stops feeding all three patterns
into one run: EnzymeAD/Reactant.jl#3225.
## Environment
Reactant.jl main (v0.2.283 source) with `Reactant_jll` v0.0.405, Julia 1.12.7,
Linux x86_64. Reproduces on the CPU backend and on CUDA (H100, driver 580.105).
## Possibly related
- EnzymeAD/Reactant.jl#2652 -- `GreedyPatternRewriteDriver` non-termination in
`ReduceMulToDotGeneral` (different pattern, same failure mode)
- EnzymeAD/Reactant.jl#2805 -- `ConvertMulConvert` correctness error; that pattern rewrites
the same convert/elementwise chains a layered bf16 graph is made of
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with the repro.mlir example and run enzymexlamlir-opt with the transform interpreter to observe the non-terminating rewrite cycle. Read ConcatInsertDimToBatchBase::matchAndRewriteImpl in AutoBatching.cpp and the ConvertConcat pattern; done means the pattern set reaches a fixed point without growing wrapper functions for the reproducer, including the analogous transpose-propagation case.
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
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 58/100