Skip to content

Forward-over-reverse: "mixed activity for jl_new_struct" when a runtime-computed type parameter forces generic construction #3246

Description

@ChrisRackauckas-Claude

Forward-over-reverse: "Not yet implemented, mixed activity for jl_new_struct" when a runtime-computed type parameter forces generic construction

Enzyme forward-over-reverse (an HVP: outer Forward over inner gradient(Reverse), both under set_runtime_activity) throws

Enzyme: Not yet implemented, mixed activity for jl_new_struct constants=Bool[1, 0, 0]
  ... @ijl_new_structv ...
  Tuple{Bool, DataType, GPUCompiler.ArgumentCC}[
    (1, Vector{Float64},          GPUCompiler.BITS_REF),   # u0     inactive
    (1, Tuple{Float64, Float64},  GPUCompiler.MUT_REF),    # tspan  active
    (1, Vector{Float64},          GPUCompiler.MUT_REF)]    # p      active

when a struct is built whose type parameter is computed from a runtime value (so the construction is type-unstable and Enzyme lowers the jl_new_struct through its generic runtime, runtime_generic_fwd/runtime_generic_augfwd), and that struct mixes an inactive pointer field, an inactive isbits field, and the active pointer field in one allocation.

Dependency-free MWE (Enzyme only)

import Enzyme
const RA_R = Enzyme.set_runtime_activity(Enzyme.Reverse)
const RA_F = Enzyme.set_runtime_activity(Enzyme.Forward)

# A runtime Bool Enzyme/Julia cannot constant-fold (stands in for SciML's `isinplace(f, 4)`).
const FLAG = Ref(true)
@noinline runtime_iip() = FLAG[]

# ODEProblem-shaped immutable struct, parameterized by `iip`.
struct Prob{iip, U, T, P}
    u0::U      # inactive array
    tspan::T   # inactive isbits tuple
    p::P       # active array
end
Prob{iip}(u0, tspan, p) where {iip} =
    Prob{iip, typeof(u0), typeof(tspan), typeof(p)}(u0, tspan, p)

const U0  = [1.0, 2.0]
const TSP = (0.0, 1.0)

# Type-unstable construction, behind its own call boundary (like ODEProblem's ctor chain).
@noinline build(u0, tspan, p) = Prob{runtime_iip()}(u0, tspan, p)

loss(p) = sum(abs2, build(U0, TSP, p).p)
grad(x) = Enzyme.gradient(RA_R, loss, x)[1]

x0 = [1.0, 2.0, 3.0]; v = [1.0, 0.0, 0.0]

grad(x0)                                                          # [2.0, 4.0, 6.0]  (first-order OK)
Enzyme.autodiff(RA_F, Enzyme.Const(grad), Enzyme.Duplicated(x0, v))  # *** mixed activity jl_new_struct ***

Control (passes)

Make the type parameter a compile-time constant and the same HVP succeeds, returning the correct result (Hessian of sum(p.^2) is 2I, so 2v = [2.0, 0.0, 0.0]):

@noinline build_static(u0, tspan, p) = Prob{true}(u0, tspan, p)
loss_static(p) = sum(abs2, build_static(U0, TSP, p).p)
grad_static(x) = Enzyme.gradient(RA_R, loss_static, x)[1]
Enzyme.autodiff(RA_F, Enzyme.Const(grad_static), Enzyme.Duplicated(x0, v))   # OK -> [2.0, 0.0, 0.0]

What's essential (ablation)

Each of these is necessary; removing any one makes the HVP succeed:

  1. Runtime-computed type parameter (Prob{runtime_iip()}). A static Prob{true} works (the control). This is what pushes the jl_new_struct onto the generic runtime path.
  2. Construction behind a function-call boundary (@noinline build(...)). Constructing Prob{runtime_iip()}(...) inline inside the differentiated closure works; only when it's a separate call does it fail — matching SciML's deep ODEProblem constructor chain.
  3. A mixed-activity allocation of ≥3 fields: an inactive pointer (u0::Vector) + an inactive isbits field (tspan::Tuple) + the active pointer (p::Vector). Two fields (e.g. just {tspan, p}) is not enough.
  4. Forward-over-reverse nesting, both under set_runtime_activity. (Plain first-order reverse is fine and correct.)

Versions

Enzyme 0.13.169 (latest release; main is at the same commit), Julia 1.11.9, x86_64-linux.

Context

Reduced from SciML/SciMLSensitivity.jl#1427 — an HVP (SecondOrder(AutoEnzyme(Forward), AutoEnzyme(Reverse))) of a neural-ODE loss, where the crash is in building the ODEProblem inside the loss (the ODE solver is never reached). There the failing jl_new_struct is the ODEProblem, whose iip type parameter comes from the runtime isinplace(f, 4) and which stores u0 (array), tspan (isbits tuple) and p (the active ComponentVector) — the same mix as above; the real error's active set was likewise {tspan, p} with the other fields inactive.

This is the last Enzyme-side blocker for that forward-over-reverse use case. It is the follow-on after the earlier cascade in that workload cleared: #3135 (jl_field_isdefined_checked, fixed by #3137) and #3213 (jl_eqtable_get, fixed — reduced MWE now passes on 0.13.169).

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions