EnzymeAD / EnzymeAD/Enzyme-JAX

Reverse-mode gradient through a `@trace mincut=true` loop got ~30× slower on GPU with Reactant ≥ 0.2.281 (reverse-loop work is batched over the whole tape)

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

Description

_Note: below is a (very) verbose AI-generated summary of a massive performance slowdown we observed in Breeze.jl when moving from Reactant v0.2.278 to v0.2.281. Look only at the reproducer if you want to skip the noise, but the analysis of the increased dynamic slices may still be useful_

Between Reactant **0.2.278** (Reactant_jll 0.0.401) and **0.2.281** (Reactant_jll 0.0.406) the compiled Enzyme reverse-mode gradient of a time-stepping loop in Breeze.jl ([`test/reactant/advection_diffusion_ad.jl`](https://github.com/NumericalEarth/Breeze.jl/blob/7330145a1e07446bf666f4c4ea6a59125893e388/test/reactant/advection_diffusion_ad.jl), `CompressibleDynamics` + WENO5, `N = 64`, `Nₛ = 229` steps, `CompileOptions(raise_first=true, raise=true, sync=true)`) went from **21 s to 695 s of execution** on an RTX 5090 (compile: 286 s → 404 s), with results identical to all printed digits.

Everything on the Julia side is the same (Julia 1.12.7, Enzyme 0.13.199, EnzymeCore 0.8.21, same Oceananigans/Breeze, environments differ *only* in `Reactant` + `Reactant_jll`; the Reactant pass pipeline is identical modulo 32 added patterns — `cse_*`, `select_to_logical`, `while_scatter_accumulator_no_add` — that were shown not to matter: the 0.2.278 pattern list run in the new jll gives the same result). Both causes are in the jll (Enzyme-JAX 568fe05 → aea4857):

1. **(dominant, ≥ 30×)** in the post-AD pattern set, `greedy_while_loop_batch_fission` now hoists the reverse-step arithmetic out of the reverse `stablehlo.while` and batches it over the whole tape (removing either that pattern or `while_op_induction_replacement` from the set prevents it, see bisection below). In the Breeze case this creates **1 567 distinct full-tape (`229×64×64`) tensors** (XLA buffer assignment 0.35 GB → 10.66 GB) and a reverse loop that `dynamic_slice`s all of them each iteration (47 → 1 580 `dynamic_slice`s). The same StableHLO also ran for ~39 min (killed) on the *old* XLA.
2. **(secondary, 3.6×)** the XLA bump in Reactant_jll 0.0.407 runs the *unchanged* 0.2.278 StableHLO in 76 s instead of 21 s (almost identical fusion structure, different fusion decisions).

A Breeze-free reproducer for (1) is below; it shows the rewrite at the StableHLO level on CPU in seconds.

## Minimal reproducer (Reactant + Enzyme only)

```julia
using Reactant, Enzyme, Printf
using GPUArraysCore: @allowscalar
const Reactant_jll = Reactant.Reactant_jll
n = parse(Int, get(ENV, "REPRO_N", "64")) # grid size
Nₛ = parse(Int, get(ENV, "REPRO_STEPS", "229")) # time steps
Reactant.set_default_backend(get(ENV, "REPRO_BACKEND", "cpu"))
println("Reactant $(pkgversion(Reactant)) / Reactant_jll $(pkgversion(Reactant_jll)) / Enzyme $(pkgversion(Enzyme)); n=$n Nₛ=$Nₛ")

# periodic shifts along dims 1 and 2
shift1(x, k) = cat(x[end-k+1:end, :], x[1:end-k, :]; dims=1)
shift2(x, k) = cat(x[:, end-k+1:end], x[:, 1:end-k]; dims=2)
xm(x) = shift1(x, 1); xp(x) = shift1(x, n-1); ym(x) = shift2(x, 1); yp(x) = shift2(x, n-1)

# one explicit step of nonlinear advection–diffusion with an upwind-like select
function step(u, c, Δt, κ)
fx = ifelse.(u .> 0, u .* c, u .* xp(c))
fy = ifelse.(u .> 0, u .* c, u .* yp(c))
adv = (fx .- xm(fx)) .+ (fy .- ym(fy))
lap = xp(c) .+ xm(c) .+ yp(c) .+ ym(c) .- 4 .* c
return c .- Δt .* adv .+ Δt .* κ .* lap
end

function loss(c0, θ, Δt, Nₛ)
U = @allowscalar θ[1]
κ = @allowscalar θ[2]
u = U .* (0.5 .+ 0.5 .* tanh.(c0))
c = c0
@trace track_numbers = false mincut = true checkpointing = false for _ in 1:Nₛ
c = step(u .* (1 .+ 0.1 .* c), c, Δt, κ)
end
return sum(c .^ 2)
end

function grad_loss(c0, dc0, θ, dθ, Δt, Nₛ)
dc0 .= 0; dθ .= 0
_, J = Enzyme.autodiff(Enzyme.set_strong_zero(Enzyme.ReverseWithPrimal), loss, Enzyme.Active,
Enzyme.Duplicated(c0, dc0), Enzyme.Duplicated(θ, dθ), Enzyme.Const(Δt), Enzyme.Const(Nₛ))
return dθ, J
end

x = range(-1, 1; length=n)
c0 = Reactant.to_rarray([exp(-(xi^2 + yj^2) / 0.1) for xi in x, yj in x])
dc0 = Reactant.to_rarray(zeros(n, n))
θ = Reactant.to_rarray([0.5, 0.01]); dθ = Reactant.to_rarray(zeros(2))
Δt = 1e-3

excluded = String.(split(get(ENV, "REPRO_EXCLUDE", ""), ","; keepempty=false))
compile_options = CompileOptions(; sync=true, excluded_passes=excluded)
hlo = String(Reactant.@code_hlo compile_options=compile_options grad_loss(c0, dc0, θ, dθ, Δt, Nₛ))
get(ENV, "REPRO_SAVE", "") != "" && write(ENV["REPRO_SAVE"], hlo)
tape_types = [m.match for m in eachmatch(Regex("tensor<(?:[0-9]+x)*$(Nₛ)x[0-9x]*f64>"), hlo)]
nbytes(t) = 8 * prod(parse.(Int, split(match(r"tensor<([0-9x]+)xf64>", t).captures[1], "x")))
@printf "optimized StableHLO: %d bytes, dynamic_slice=%d, while=%d, tensors with a %d-long dim: %d (%.2f GB summed over all mentions)\n" sizeof(hlo) count("stablehlo.dynamic_slice", hlo) count("stablehlo.while", hlo) Nₛ length(tape_types) sum(nbytes, tape_types; init=0)/1e9

compiled = Reactant.@compile compile_options=compile_options grad_loss(c0, dc0, θ, dθ, Δt, Nₛ)
compiled(c0, dc0, θ, dθ, Δt, Nₛ)
t = @elapsed g, J = compiled(c0, dc0, θ, dθ, Δt, Nₛ)
@printf "run: %.3f s J = %.6g dJ/dθ = %s\n" t Float64(J) Array(g)
```

Results (CPU backend, same script, only `Reactant`/`Reactant_jll` differ; `J` and `dJ/dθ` identical):

| n × Nₛ | Reactant | `dynamic_slice` | values with an Nₛ-long dim (summed size) | run |
|---|---|---:|---:|---:|
| 64 × 229 | 0.2.278 | 1 | 6 (0.05 GB) | 0.022 s |
| 64 × 229 | 0.2.283 | 7 | **86 (0.48 GB)** | 0.027 s |
| 64 × 229 | 0.2.283, `REPRO_EXCLUDE=greedy_while_loop_batch_fission` | 9 | 14 (0.11 GB) | 0.020 s |
| 128 × 500 | 0.2.278 | 1 | 6 (0.39 GB) | 0.83 s |
| 128 × 500 | 0.2.283 | 7 | **86 (4.19 GB)** | 0.93 s |
| 256 × 500 | 0.2.278 | 1 | 6 (1.57 GB) | 1.14 s |
| 256 × 500 | 0.2.283 | 7 | **86 (16.8 GB)** | 1.51 s |

### What the optimized StableHLO looks like

**0.2.278** — the reverse loop reads one slice of the tape per iteration and does the stencil arithmetic on `n×n` slices:

```mlir
%13:3 = stablehlo.while(...) // forward loop: tape = dynamic_update_slice(tape, c, i)
%18:5 = stablehlo.while(...) // reverse loop
%44 = stablehlo.dynamic_slice %13#2, %iterArg_16, %c_3, %c_3, sizes = [1, 64, 64] : (tensor<229x64x64xf64>, ...) -> tensor<1x64x64xf64>
... 64×64 arithmetic ...
```

**0.2.283** — the reverse-step arithmetic that depends only on tape slices (`reverse` of the tape, the periodic-shift `slice`/`concatenate`s, `multiply`, `compare`, `add`, `reshape`, …) is hoisted out of the reverse loop and batched over all `Nₛ` steps, each op materialising a full-tape temporary; the loop then `dynamic_slice`s those:

```mlir
%13:3 = stablehlo.while(...) // forward loop, unchanged
%18 = stablehlo.reverse %13#2, dims = [0] : tensor<229x64x64xf64>
%19 = stablehlo.reshape %18 : (tensor<229x64x64xf64>) -> tensor<229x1x64x64xf64>
%20 = stablehlo.multiply %cst_1, %18 : tensor<229x64x64xf64>
%22 = stablehlo.add %cst, %20 : tensor<229x64x64xf64>
%24 = stablehlo.multiply %23, %22 : tensor<229x64x64xf64>
%25 = stablehlo.compare GT, %24, %cst_3 : (tensor<229x64x64xf64>, tensor<229x64x64xf64>) -> tensor<229x64x64xi1>
%26 = stablehlo.slice %13#2 [0:229, 1:64, 0:64] : (tensor<229x64x64xf64>) -> tensor<229x63x64xf64>
%27 = stablehlo.reverse %26, dims = [0] : tensor<229x63x64xf64>
%30 = stablehlo.concatenate %27, %29, dim = 1 : (...) -> tensor<229x64x64xf64>
... (40+ such ops on 229×64×64 / 229×1×64×64 tensors) ...
%59 = stablehlo.add %58, %21 : tensor<229x1x64x64xf64>
%60:4 = stablehlo.while(...) // reverse loop
%87 = stablehlo.dynamic_slice %13#2, %86, ... -> tensor<1x64x64xf64>
%89 = stablehlo.dynamic_slice %22, %iterArg, ... -> tensor<1x64x64xf64>
%91 = stablehlo.dynamic_slice %24, %iterArg, ... -> tensor<1x64x64xf64>
%93 = stablehlo.dynamic_slice %30, %iterArg, ... -> tensor<1x64x64xf64>
%95 = stablehlo.dynamic_slice %25, %iterArg, ... -> tensor<1x64x64xi1>
%97 = stablehlo.dynamic_slice %38, %iterArg, ... -> tensor<1x64x64xf64>
%99 = stablehlo.dynamic_slice %59, %iterArg, %c_6, %c_6, %c_6, sizes = [1, 1, 64, 64] : (tensor<229x1x64x64xf64>, ...) -> tensor<1x1x64x64xf64>
```

This trades a few `n×n` ops per iteration for `Nₛ·n²`-sized temporaries. In this toy it is 7 tensors; in the real Breeze program it is 1 567 tensors (10.7 GB) and a 33× slowdown.

## Where in the pipeline (bisection on the Breeze module)

Replaying the `all` pipeline pass-by-pass with `Reactant.Compiler.run_pass_pipeline_on_source` on the *same* input module with both jlls:

* op counts are identical through the `enzyme` (AD) pass and the following `inline`/`canonicalize`/`cse`: 28 `dynamic_slice`;
* the first divergence is the post-AD `enzyme-hlo-generate-td` + `transform-interpreter` pattern set: jll 0.0.401 → **31**, jll 0.0.407 → **156** `dynamic_slice`. Removing single patterns from that set in jll 0.0.407: `while_op_induction_replacement` → **31** (exactly the old result), `greedy_while_loop_batch_fission` → 16,`while_licm` → 141, `reshape_licm` → 154; 20 other while/slice/licm/dus patterns tried: no effect. My reading (not verified in the Enzyme-JAX source): the new `WhileOpInductionReplacement` (cf. #2759 "Recognise backward-stepping counters") now rewrites the reverse loop's counter into an induction variable, which lets `greedy_while_loop_batch_fission` recognise the tape reads as batchable;
* the later `post_op_transpose_reshape` pipeline amplifies 31 → 47 (0.0.401) vs 156 → **1 580** (0.0.407), again only via `greedy_while_loop_batch_fission` (without it: 304). Feeding jll 0.0.401's build of that pipeline the 0.0.407 input
*also* gives 1 580, so the amplification is a consequence of the changed post-AD output rather than of a change in
this stage.

`CompileOptions(excluded_passes=["greedy_while_loop_batch_fission"])` restores the old structure in the reproducer, but for the Breeze program the gradient was still running after ~26 min (killed), i.e. it is not a workaround there.

## Cross-compiling to separate the two effects (Breeze program)

`Reactant.Ops.hlo_call` on one version's *optimized* `stablehlo_optimized.mlir` compiled with `CompileOptions(optimization_passes=:none, sync=true)` in the other version's environment, so no EnzymeXLA passes run and only the jll (XLA/PJRT) plus Reactant's Julia-side runtime differ:

| XLA of | StableHLO of | run |
|---|---|---:|
| 0.2.278 (jll 0.0.401) | 0.2.278 | 21.6 s (= native) |
| **0.2.283 (jll 0.0.407)** | 0.2.278 | **76 s** |
| 0.2.278 (jll 0.0.401) | **0.2.283** | **~39 min, killed** |
| 0.2.283 (jll 0.0.407) | 0.2.283 | 696 s (= native) |

On the identical 0.2.278 input the two XLAs produce almost the same program (1 543 vs 1 522 fusions, 1 359 vs 1 312 kernels per reverse iteration, 333 vs 330 MiB), with different fusion choices (`input_concatenate_fusion` 273 → 530 and `input_transpose_fusion` 289 → 183 per iteration, copies 35 → 56), so the 3.6× is per-kernel code quality of the new emitters; I have not profiled it further.

## Other data points

* Breeze program, StableHLO handed to XLA: `dynamic_slice` 47 → 1 580 (reading 1 567 distinct tensors), `reshape` 1 402 → 2 796, new `reverse` (96) and `and` (91) ops; same forward loop (47 iter-args, 116 MB tape), reverse loop 18 vs 17 iter-args (0.5 MB); both loops compare the counter against the constant 229; no un-raised kernels in either version.
* XLA after optimizations: fusions 1 543 → 2 477, buffer assignment 0.35 GB → 10.66 GB (640 × `f64[64,229,64]`, 325 × `f64[229,64,64]`, 200 × `f64[229,1,64,64]`, …), no host offloading, register spills comparable.
* I have not checked whether a forward-only (no `Enzyme.autodiff`) loop is affected; the rewrite acts on the reverse loop reading the AD tape, and the forward loop is unchanged between versions.
* Enzyme-JAX commits in the range whose *titles* concern while-loop / dynamic-slice rewriting: #2748, #2759, #2787, #2882, #2916, #2741 (I have not inspected their diffs).

## Versions

* Julia 1.12.7, Linux, RTX 5090 (sm_120), CUDA 12.9 artifact of Reactant_jll (`mode=opt`)
* Reactant 0.2.278 / Reactant_jll 0.0.401 (Enzyme-JAX 568fe057) — fast
* Reactant 0.2.283 / Reactant_jll 0.0.407 (Enzyme-JAX aea4857f) — slow. Only these two versions were measured here;
the slowdown was first noticed with 0.2.281 (jll 0.0.406).
* Enzyme 0.13.199, EnzymeCore 0.8.21

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with the minimal Reactant + Enzyme reproducer in the issue and compare its optimized StableHLO across the reported versions. Inspect the post-AD transform pipeline, especially greedy_while_loop_batch_fission and while_op_induction_replacement, then use the reported Breeze entry point test/reactant/advection_diffusion_ad.jl to validate the result. Done means reverse-mode execution no longer materializes the excessive full-tape intermediates or incurs the reported slowdown while preserving the gradient.

Written by the indexing model from the issue text.

Assessment

Domain
compilers, performance
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.