EnzymeAD / EnzymeAD/Enzyme.jl

EnzymeNoTypeError on Julia 1.12: type analysis cannot type the 0xFF pointer-slot poison memset emitted for a by-value read of a large inline struct out of a heap object

Open
#3,433 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
586
Forks
108
Avg merge
1d 5h
Merged PRs (30d)
44

Description

On Julia 1.12 (not 1.10/1.11), a reverse-mode gradient fails with `EnzymeNoTypeError` when a large inline immutable field containing GC pointers is read out of a heap object with `getfield` inside a split-mode thunk. Julia 1.12 lowers that read via the sret + `return_roots` ABI and **poisons the in-struct pointer slots with 0xFF** (`llvm.memset(..., i8 -1, ...)` / `store i64 -1`); Enzyme's type analysis cannot type that memset, so the reconstructed type of the buffer has holes exactly where the poison lands.

Filed from https://github.com/SciML/SciMLSensitivity.jl/issues/1583. Shape is close to the (fixed) #2961 / #3182: also 1.12-only, also a type-analysis hole reached through a by-value struct return under split reverse mode.

## Minimal reproducible example

I have **not** yet been able to remove the SciML packages — see "Status of a SciML-free MWE" at the bottom for how far the reduction got and what is missing. This is the reproducer as it stands. Set `SKIP_A=1` to skip the first gradient; that control passes.

```julia
# ENV: SKIP_A=1 to skip the first gradient (control -- passes).
using ModelingToolkit
using ModelingToolkit: D_nounits as D, t_nounits as t
using OrdinaryDiffEqTsit5
using ADTypes, Enzyme, DifferentiationInterface
using SymbolicIndexingInterface
using SciMLBase, SciMLSensitivity

function loss(x, ps)
(setter, getter, prob0, tspan) = ps
(u0, p) = setter(prob0, x)
if u0 === prob0.u0
u0 = copy(prob0.u0)
end
newprob = remake(prob0; u0, p, tspan)
sol = solve(newprob, Tsit5(); saveat = 0.02, abstol = 1e-6, reltol = 1e-3)
sum(vals -> sum(abs2, vals), getter(sol))
end

adtype = AutoEnzyme(; mode = set_runtime_activity(Reverse))

# --- system A: any other MTK model ---
begin
sts = @variables s1(t)=2.0 s1s2(t)=2.0 s2(t)=2.0
ps_a = @parameters k1=1.0 c1=2.0
eqs_a = [D(s1) ~ -0.25 * c1 * k1 * s1 * s2
D(s1s2) ~ 0.25 * c1 * k1 * s1 * s2
D(s2) ~ -0.25 * c1 * k1 * s1 * s2]
sys_a = mtkcompile(System(eqs_a, t; name = :reactionsystem))
prob_a = ODEProblem(sys_a, [], (0.0, 1.0))
ps_A = (setsym_oop(prob_a, [sys_a.c1]),
getsym(prob_a, [sys_a.s1, sys_a.s1s2, sys_a.s2]), prob_a, (0.0, 1.0))
xA = [3.5]
end

# --- system B: FullSpecialize + trivial initialization ---
begin
@parameters m=1.5 d=9.0
@variables s(t) v(t)
eqs_b = [D(s) ~ v
m * D(v) ~ 1 - d * v]
sys_b = mtkcompile(System(eqs_b, t; name = :model,
initialization_eqs = [s ~ 0, v ~ 0]))
prob_b = ODEProblem{true, SciMLBase.FullSpecialize}(sys_b, [], (0.0, 200.0))
ps_B = (setsym_oop(prob_b, [sys_b.m, sys_b.d]),
getsym(prob_b, [sys_b.v]), prob_b, (0.0, 1.0))
xB = [1.0, 4.3]
end

get(ENV, "SKIP_A", "0") == "1" ||
@show DifferentiationInterface.gradient(loss, adtype, xA, Constant(ps_A))
@show DifferentiationInterface.gradient(loss, adtype, xB, Constant(ps_B))
```

| Julia | grad A | grad B after A | grad B alone |
|---|---|---|---|
| 1.11.9 | OK 459 s | **OK** 211 s | OK |
| 1.12.6 | OK 470 s | **EnzymeNoTypeError** 50 s | OK 778 s |

(The 50 s vs 778 s is not a caching signal: once TA gives up in the `remake` thunk the error is raised at run time before `solve`'s expensive thunk is ever compiled.)

## The failing instruction

With `Enzyme.Compiler.VERBOSE_ERRORS[] = true`:

```
Cannot deduce type of memset
call void @llvm.memset.p0.i64(ptr noundef nonnull align 8 dereferenceable(16) %10,
i8 -1, i64 16, i1 false), !dbg !2441
```

inside `preprocess_julia_maybe_eager_initialize_problem`, at `prob.f` in `SciMLBase.maybe_eager_initialize_problem`:

```llvm
%"prob::ODEProblem.f" = call noalias ptr @malloc(i64 304), !enzyme_fromstack ; alloca -> heap (split mode)
%9 = load i64, ptr addrspace(11) %8 ; prob+0
store i64 %9, ptr %"prob::ODEProblem.f"
%10 = getelementptr inbounds i8, ptr %"prob::ODEProblem.f", i64 8
call void @llvm.memset.p0.i64(ptr %10, i8 -1, i64 16, i1 false) ; <-- poison
%13 = load i64, ptr addrspace(11) %11 ; prob+24
store i64 %13, ptr %12
store i64 -1, ptr %"prob::ODEProblem.f"+32 ; <-- poison
...
```

Enzyme's reconstructed type for that 304-byte buffer has holes exactly where the poison writes land:

```
%"prob::ODEProblem.f" = call noalias ... @malloc(i64 304), !enzyme_fromstack:
{[-1]:Pointer, [-1,0]:Integer, [-1,24]:Integer, [-1,25]:Integer, ... [-1,40]:Integer,
[-1,64]:Pointer, [-1,72]:Integer, ...}
```

— nothing for bytes 8–23, 48–71, 152–175, 248–263, 272–303, which are precisely the memset/`store i64 -1` ranges.

## Where the 0xFF poison comes from (plain Julia 1.12, no Enzyme, no SciML)

It is ordinary Julia 1.12 codegen for "read a large inline immutable field containing GC pointers out of a **heap** object and return it by value": the pointer fields are loaded `atomic unordered` into `return_roots`, and the corresponding slots of the sret body are poisoned with 0xFF.

```julia
using InteractiveUtils
struct Sing end
struct Inner
p1::Vector{Float64}; p2::Vector{Float64}; b1::Bool; b2::Bool
p3::Vector{Float64}; p4::Vector{Float64}; b3::Bool; n::Nothing
end
struct Fn
f::Sing; mass::Bool; a1::Nothing; a2::Nothing
o1::Inner; o2::Inner; o3::Inner; o4::Inner; o5::Inner; o6::Inner
tail::Vector{Float64}
end # sizeof(Fn) == 304, same as the real ODEFunction
struct IProb; f::Fn; u0::Vector{Float64}; end
mutable struct MProb; f::Fn; u0::Vector{Float64}; end
@noinline grab(p) = p.f

# count `memset ... i8 -1` in the optimized IR
nmemset(T) = (io = IOBuffer();
code_llvm(io, grab, Tuple{T}; optimize = true, raw = true, debuginfo = :none);
count(l -> occursin("memset", l) && occursin("i8 -1", l),
split(String(take!(io)), '\n')))

nmemset(IProb) # 0 -- stack container, plain memcpy
nmemset(MProb) # 12 -- heap container, sret + return_roots + 0xFF poison
```

The same thing on the real type confirms the mapping: `@noinline grab(prob) = prob.f` on the FullSpecialize MTK problem emits `memset(sret+8, 0xff, 16)`, `memset(sret+48, 0xff, 24)`, `memset(sret+152, 0xff, 24)`, `memset(sret+248, 0xff, 16)`, `memset(sret+272, 0xff, 32)` plus a `store i64 -1` — byte-for-byte the ranges Enzyme fails to type.

## Why the SciML side needs `FullSpecialize`

Verified layouts (identical on 1.11 and 1.12):

```
ODEProblem (FullSpecialize) sizeof=344 ODEFunction (FullSpecialize) sizeof=304
f off=0 size=304 isptr=false observed off=8 size=40
u0 off=304 size=8 isptr=true sys off=48 size=8 isptr
tspan off=312 size=16 initialization_data off=56 size=248 isptr=false
p off=328 size=16
ODEFunction (AutoSpecialize) sizeof=72
initialization_data off=56 size=8 isptr=true
```

With `FullSpecialize` + trivial initialization the whole `OverrideInitData` (248 B) is stored inline, so `prob.f` is a 304-byte by-value aggregate copy with many GC-pointer slots to poison. With the boxed `Union{Nothing, OverrideInitData}` the field is 72 B and the pattern does not arise.

## Ablations (Julia 1.12.6, each run to completion)

| Variant | Result |
|---|---|
| A then B (above) | **ERROR** 50 s |
| B alone | OK 778 s |
| B first, then A, then B again | OK / OK / OK (3rd is a thunk-cache hit, 0 s) |
| A, then `empty!` every Enzyme.jl global cache, then B | **ERROR** 51 s |
| `Enzyme.API.maxtypeoffset!(65536)` | **ERROR** 51 s |
| `Enzyme.API.looseTypeAnalysis!(true)` | **OK** 492 s |
| first gradient = trivial Enzyme gradient | B OK 723 s |
| first gradient = plain (non-MTK) `solve` | B OK 641 s |
| first gradient = plain (non-MTK) `remake` + `solve` | B OK 628 s |

Notes:

* `looseTypeAnalysis!(true)` fixing it confirms this is purely TA giving up, and is a usable workaround.
* `maxtypeoffset!` not helping matches the #2961 finding that the "register too large" hint in the error text is a red herring here.
* Caches cleared in the clearcache run (sizes after the first gradient, to show they were populated): `Compiler.cache` 29, `autodiff_cache` 57, `ActivityCache` 272, `RRULE_CACHE` 109 808, `INACTIVE_CACHE` 110 928, `EASY_RULE_CACHE` 68, `NOALIAS_CACHE` 68, `Interpreter.SigCache` 2. Clearing all of them does not change the outcome, so the carried-over state is not Enzyme.jl's.
* The first gradient has to be over *another MTK model*; no amount of non-MTK Enzyme compilation reproduces it. That points at Julia's own inference/native-code cache for MethodInstances shared between the two MTK problems.

## Order dependence — hypothesis

Everything above is measured; this last link is not yet verified. Once the shared MTK/SciMLBase MethodInstances have been inferred *and natively compiled* during the first gradient, Julia 1.12 no longer has their inferred source available for inlining, so on the second compile `get_initial_values(prob, prob, prob.f, …)` stays an out-of-line call instead of being inlined and SROA'd. That forces `prob.f` to be materialized as a real 304-byte by-value copy, which is when the 0xFF pointer-slot poisoning is emitted. In the passing run the copy is inlined away and the memset never exists. This is consistent with every observation above (MTK-only first gradient, Enzyme caches irrelevant, B-first fine, 1.12-only).

If that is right, the order dependence is incidental and the real fix is in type analysis: an all-0xFF `llvm.memset` over a range that is known-Pointer in the destination's Julia type should be typed from the destination (it is a poison/undef fill for the sret + `return_roots` ABI), not left untyped.

## Status of a SciML-free MWE

Reproduced standalone: the Julia-side half above — the 0xFF pointer-slot poisoning on a large inline field read out of a heap object. Enzyme, however, still types those synthetic memsets correctly, in combined mode, in split `autodiff_thunk`, and through `runtime_generic_augfwd`, for immutable and mutable containers and aggregates from 64 B to 304 B. The missing ingredient is the IR context of the real failure: the memset has to land in an Enzyme heap-promoted alloca (`!enzyme_fromstack` `malloc(304)`) whose element type Enzyme rebuilds only from the surviving non-poison stores. In the synthetic versions LLVM either forwards the sret buffer straight into the callee (eliding the poison) or Enzyme retains enough type information. Happy to keep reducing if that context can be forced more directly.

## Environment

```
Julia Version 1.12.6 (fails) / 1.11.9 (passes)
Enzyme v0.13.198
EnzymeCore v0.8.x
ModelingToolkit v11.38.0
SciMLBase v3.39.1
SciMLSensitivity v7.116.2
OrdinaryDiffEqTsit5 v2.1.2
DifferentiationInterface v0.7.20
ADTypes v1.22.3
SymbolicIndexingInterface v0.3.51
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.