Skip to content

Fix batched ForwardWithPrimal primal slot width - #3488

Draft
ChrisRackauckas-Claude wants to merge 1 commit into
EnzymeAD:mainfrom
ChrisRackauckas-Claude:fix/batched-abi-primal-slot-width
Draft

ChrisRackauckas-Claude wants to merge 1 commit into
EnzymeAD:mainfrom
ChrisRackauckas-Claude:fix/batched-abi-primal-slot-width

Conversation

@ChrisRackauckas-Claude

@ChrisRackauckas-Claude ChrisRackauckas-Claude commented Aug 27, 2026 •

Copy link
Copy Markdown
Contributor

Ignore this PR until it has been reviewed by @ChrisRackauckas.

What changed and why

When batched ForwardWithPrimal returns a boxed primal alongside batched derivatives, the ABI wrapper has a batched derivative sret slot followed by a scalar primal sret slot. The wrapper currently applies the batch width to both slots, so it tries to extract the scalar primal pointer as an aggregate and segfaults in LLVMBuildExtractValue.

This keeps the existing scalar-width handling for Const returns and also assigns scalar width to the final primal slot. The regression test covers the distinct ForwardWithPrimal path and checks both tangent batches and the primal result.

Dependency

This follows #3486, which fixed the earlier undefined ABI-wrapper loop index and has now merged. This branch is rebased onto the resulting main commit and contains only the slot-width follow-up.

Failing before

I ran this SciML-free reproducer on commit 1187627c, after the undefined-index fix but before this slot-width fix:

using Enzyme

batched_uninferred_target(x) = (2 .* x,)
batched_uninferred_return(x) = Base.invokelatest(batched_uninferred_target, x)

x = [1.0, 2.0]
tangents = (ones(2), fill(3.0, 2))
return_activity = BatchDuplicated{Tuple{Vector{Float64}}, 2}
result = Enzyme.autodiff(
    Enzyme.set_runtime_activity(ForwardWithPrimal),
    Const(batched_uninferred_return),
    return_activity,
    BatchDuplicated(x, tangents),
)
@show result

It exited with status 139:

signal 11 (1): Segmentation fault
LLVMBuildExtractValue
extract_value! at LLVM/src/irbuilder.jl:299
create_abi_wrapper at Enzyme.jl/src/compiler.jl:3898

Passing after

I ran the same reproducer on this branch with Julia 1.12.7 and observed:

result = (var"1" = (var"1" = ([2.0, 2.0],), var"2" = ([6.0, 6.0],)), var"2" = ([2.0, 4.0],))

The targeted repository test group also passed:

$ julia +1.12 --startup-file=no --project=test test/runtests.jl typeunstable
Test Summary: | Pass  Total     Time
  Overall     |   39     39  1m05.8s
    SUCCESS

Pre-push checks:

$ ../git-runic --julia ~/.juliaup/bin/julia --project ../.runic-env -f 1187627c
runic did not modify any files

$ git diff --check
# no output

$ git diff --unified=0 1187627c -- | typos -
# no output

This repository does not expose a separate QA/Aqua/ExplicitImports/JET test target. I did not run the full suite, GPU tests, integration tests, or a docs build; this change does not touch documentation or public API.

Reviewer attention: the new condition is deliberately additive. Replacing the existing Const condition would regress its scalar-width case.

🤖 Generated with Claude Code
https://chatgpt.com/codex/tasks/01a03f84-af17-7a90-955c-10b5d980e4da

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude <noreply@anthropic.com>
Claude-Session: https://chatgpt.com/codex/tasks/01a03f84-af17-7a90-955c-10b5d980e4da

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants