Approach for differentiating through iterative solvers

[dkytezab]

Consider the following example. A while loop with both a data-dependent condition and an upper bound on the total number of iterations. If we attempt to compile its derivative, we get

using Reactant, Enzyme

using Enzyme: autodiff, Reverse, Active, Duplicated, make_zero

function adaptive(x)
    i = 0
    @trace while (i < 8) & (sum(x) > 1.0)
        x = x ./ 2
        i += 1
    end
    return sum(x)
end

∇adaptive(dx, x) = autodiff(Reverse, adaptive, Active, Duplicated(x, dx))

x = Reactant.to_rarray([100.0])
dx = make_zero(x)

@compile ∇adaptive(dx, x)

ERROR: LoadError: CompilationError: MLIR pass pipeline "all" failed

error: WhileOp does not have known iteration count for cache removal

/Users/danielkz/Aeolus2/Oceananigans.jl/while_reactant_ad_mwe.jl:5:0
4 |
5 | function adaptive(x)
6 |     i = 0
Stacktrace:
  [1] traced_while/while_loop
    @ /Users/danielkz/Aeolus2/Reactant.jl/src/ControlFlow.jl:20:0

Presumbably this is intrinsic to XLA having minimal support for dynamic shapes. This is unfortunate because the loop signature of our conjugate gradient solver follows this pattern. If we attempt to differentiate it in the same fashion (working off of #5836), we hit the same problem:

using Reactant, Enzyme, Oceananigans

using LinearAlgebra: dot
using Enzyme: autodiff, Reverse, Active, Duplicated, set_runtime_activity, make_zero
using Oceananigans.Architectures: ReactantState
using Oceananigans.Fields: CenterField
using Oceananigans.Solvers: ConjugateGradientSolver, solve!

grid = RectilinearGrid(ReactantState(); size = 4, x = (0, 1), topology = (Periodic, Flat, Flat))

A!(q, p) = (parent(q) .= 2 .* parent(p); nothing)  # simplest SPD operator: the MWE is about the loop

solver = ConjugateGradientSolver(A!; template_field = CenterField(grid), maxiter = 3)
b, x = CenterField(grid), CenterField(grid)
parent(b) .= 1

loss!(x, solver, b) = (solve!(x, solver, b); dot(x, x))

∇loss!(dx, ds, db, x, solver, b) = autodiff(
    set_runtime_activity(Reverse), loss!, Active,
    Duplicated(x, dx), Duplicated(solver, ds), Duplicated(b, db)
)

@info "forward"
forward = @compile raise = true loss!(x, solver, b)
@show forward(x, solver, b)

@info "reverse (expected to fail)"
dx, ds, db = make_zero(x), make_zero(solver), make_zero(b)
@compile raise = true raise_first = true ∇loss!(dx, ds, db, x, solver, b)

[ Info: forward
forward(x, solver, b) = ConcreteIFRTNumber{Float64}(1.0)
[ Info: reverse (expected to fail)
┌ Error: Compilation failed, MLIR module written to /var/folders/km/c7cykv9x5s955_48nqv3f4cm0000gn/T/reactant_iA0aOL/module_012_BWYU_post_all_pm.mlir
└ @ Reactant.MLIR.IR ~/Aeolus2/Reactant.jl/src/mlir/IR/Pass.jl:146
ERROR: LoadError: CompilationError: MLIR pass pipeline "all" failed

error: WhileOp does not have known iteration count for cache removal

/Users/danielkz/Aeolus2/Oceananigans.jl/ext/OceananigansReactantExt/Solvers.jl:65:0
     64 |
     65 | function Oceananigans.Solvers.solve!(x, solver::ReactantCGSolver, b, args...)
     66 |     q = solver.linear_operator_product
Stacktrace:
  [1] traced_while/while_loo

I see two possible solutions to this:

1. We write an extension such that if we try to take a gradient through solve!, then we exclude the data-dependent condition so that we are simplify differentiating through a for loop. While this is in a sense the simplest approach, this would require either fixing a specific number of checkpoints or allowing the user to specify the number of checkpoints, giving them axis to a private piece of the code.
2. We write a custom rule for this and other solvers. If we are differentiating something that iteratively solves Ax = b through some sequence x1, ..., xn, then we can directly differentiate the solution condition rather than the successive iterations. The drawback of this approach is that the derivative may not be exact in the sense that the solve may not have converged to the exact solution. For example, we converge to some x where Ax - b = z for some non-zero z, but we differentiate with respect to x* where Ax* - b = 0. cc @Pangoraw I think that you were looking into Enzyme custom derivative rules.

[Pangoraw]

is 2. implicit differentiation?

The way i see it in the current setting is something like with max_steps where we do rewriting in a known form:

for _ in 1:8
  if (sum(x) > 1.0)
        x = x ./ 2
  else
        x = x
  end
  i += 1
end
return sum(x)

Another area of investigation is that of online checkpointing where with the fixed budget for caches, there won't be any dynamic shapes in the IR. It would require compilers tweaks on the current binomial checkpointing implementation (which is essentially offline).

[dkytezab]

This would be implicit differentiation. Ah, so we can effectively use all of the remaining iterations to preserve the existing converged state (and then could these get folded away?)

[Pangoraw]

(and then could these get folded away?)

if the other condition (sum(x) > 1.0) is truly data-dependent then it might not be possible to do since it might consume the full 8 iterations. but there would have no effect on the XLA side.