IllegalTypeAnalysisException differentiating SimpleChains stack-memory path

[ChrisRackauckas-Claude]

Description

While isolating SciML/SciMLSensitivity.jl#1534, I reduced the failure to a SciML-free MWE involving only Enzyme differentiating a small SimpleChains.SimpleChain call. The original failure appeared in a docs example using Optimization.AutoEnzyme, but the ODE/Optimization/SciMLSensitivity layers are not needed to reproduce the Enzyme error.

MWE

Project dependencies:

[deps]
Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
SimpleChains = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5"
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

import Enzyme
import SimpleChains
using StaticArrays: @SArray

sc = SimpleChains.SimpleChain(SimpleChains.static(2),
    SimpleChains.TurboDense{true}(identity, SimpleChains.static(2)))

p = Array(SimpleChains.init_params(sc))
dp = zero(p)
u = @SArray Float32[2.0, 0.0]

function loss(sc, u, p)
    y = sc(u, p)
    return sum(abs2, y)
end

println("loss = ", loss(sc, u, p))
Enzyme.autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse),
    loss,
    Enzyme.Active,
    Enzyme.Const(sc),
    Enzyme.Const(u),
    Enzyme.Duplicated(p, dp))
println("gradient norm = ", sqrt(sum(abs2, dp)))

Run command used locally:

timeout 3600 env JULIA_DEPOT_PATH="$PWD/.julia_depot" /home/crackauc/.juliaup/bin/julia +1.12.6 --project=mwe_env simplechains_enzyme_mwe.jl

Observed output

The forward pass succeeds, then Enzyme reverse compilation fails:

loss = 2.4585857
ERROR: LoadError: IllegalTypeAnalysisException: Enzyme compilation failed due to illegal type analysis.
 This usually indicates the use of a Union type, which is not fully supported with Enzyme.API.strictAliasing set to true [the default].
 Ideally, remove the union (which will also make your code faster), or try setting Enzyme.API.strictAliasing!(false) before any autodiff call.
 Failure within method: MethodInstance for loss(::SimpleChains.SimpleChain{Tuple{Static.StaticInt{2}}, Tuple{SimpleChains.TurboDense{true, Static.StaticInt{2}, typeof(identity)}}}, ::StaticArraysCore.SVector{2, Float32}, ::Vector{Float32})
Caused by:
Stacktrace:
 [1] reinterpret
   @ ./essentials.jl:733
 [2] align
   @ ~/.julia/packages/VectorizationBase/F4R96/src/alignment.jl:13
 [3] with_stack_memory
   @ ~/.julia/packages/SimpleChains/ozCs8/src/memory.jl:25
 [4] with_memory
   @ ~/.julia/packages/SimpleChains/ozCs8/src/memory.jl:48
 [5] SimpleChain
   @ ~/.julia/packages/SimpleChains/ozCs8/src/simple_chain.jl:287
 [6] loss

Local versions

Verified locally on Julia 1.12.6 with:

Enzyme v0.13.176
SimpleChains v0.4.8
StaticArrays v1.9.18

VectorizationBase v0.21.74 is pulled transitively through SimpleChains.

Notes

A larger original example used Optimization.AutoEnzyme(mode = Enzyme.set_runtime_activity(Enzyme.Reverse)) in a SciMLSensitivity docs example. This reduced MWE shows the failing core is Enzyme's reverse pass through SimpleChains' stack-memory setup, with the error pointing at VectorizationBase.align(::Ptr).

[wsmoses]

@ChrisRackauckas-Claude / @ChrisRackauckas This stack allocation (https://github.com/PumasAI/SimpleChains.jl/blob/69cae9013804781a21125f9b52fcb09179cba079/src/memory.jl#L21) very intentionally is the wrong type, which this yells about. This is not an enzyme-side error but expected behavior. You can write a custom rule for anything above that call.