Skip to content

Add GPUArraysCore extension: broadcast-based seed! for GPU dual arrays#816

Draft
ChrisRackauckas-Claude wants to merge 1 commit into
JuliaDiff:masterfrom
ChrisRackauckas-Claude:gpuarrayscore-seed-ext
Draft

Add GPUArraysCore extension: broadcast-based seed! for GPU dual arrays#816
ChrisRackauckas-Claude wants to merge 1 commit into
JuliaDiff:masterfrom
ChrisRackauckas-Claude:gpuarrayscore-seed-ext

Conversation

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Contributor

Summary

ForwardDiff 1.x rewrote the four seed! methods in src/apiutils.jl to write each
dual with a scalar setindex! loop over structural_eachindex:

for (i, idx) in zip(1:N, structural_eachindex(duals, x))
    duals[idx] = Dual{T,V,N}(x[idx], seeds[i])
end

On GPU arrays this triggers ERROR: Scalar indexing is disallowed, so every
ForwardDiff.jacobian / jacobian! call on a GPU array (e.g. CuArray) errors
in seed!. ForwardDiff 0.10 seeded with broadcast and worked on GPU arrays, so
this is a regression for GPU usage introduced in the 1.0 rewrite.

This was first hit downstream in SciML/ComplementaritySolve.jl#65, where the GPU
solver paths (NonlinearSolve's AutoForwardDiff jacobians) broke once the
environment resolved ForwardDiff ≥ 1.

Fix

Add a GPUArraysCore package extension (ForwardDiffGPUArraysCoreExt) that
overrides the four seed! methods for AbstractGPUArray{<:Dual} with
broadcast-based seeding, restoring the pre-1.0 behavior. GPU arrays are always
dense, one-based, and carry isbits element types, so the
structural_eachindex / unset-element handling of the generic methods is
unnecessary on this path and plain broadcast suffices.

This is a weak dependency, so it adds nothing for users who don't load a GPU
array package.

Scope

This covers the jacobian paths (vector and chunk mode), which is what seed!
feeds. ForwardDiff.gradient on GPU arrays has a second scalar-indexing site in
the gradient extraction path (extract_gradient_chunk!) that this PR does not
touch; fixing seed! is the necessary first step and resolves the jacobian
regression.

Tests

test/GPUArraysCoreTest.jl exercises all four seed! methods through
jacobian, jacobian!, and chunked configs using JLArrays, which emulates
GPU array semantics (including the scalar-indexing ban) on the CPU — so the
extension is covered in CI without a physical GPU.


Note

Opened as a draft by an agent on behalf of @ChrisRackauckas. Please ignore
until reviewed by @ChrisRackauckas.

🤖 Generated with Claude Code

The four `seed!` methods in src/apiutils.jl write each dual with a scalar
`setindex!` loop over `structural_eachindex`, which errors with "Scalar
indexing is disallowed" on GPU arrays. ForwardDiff 0.10 seeded with broadcast
and worked on GPU arrays, so jacobians on GPU arrays regressed in the 1.0
rewrite.

Add a GPUArraysCore package extension that overrides `seed!` for
`AbstractGPUArray{<:Dual}` with broadcast-based seeding, restoring the pre-1.0
behavior. GPU arrays are dense, one-based, and isbits-valued, so the
structural-index / unset-element handling of the generic methods is not needed
on this path.

Tested via JLArrays (which emulates the GPU scalar-indexing ban on the CPU) so
the extension is covered in CI without a physical GPU.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
@ChrisRackauckas

Copy link
Copy Markdown
Member

Because the broadcasts were purposefully removed due to performance things, we should at least recover the GPU functionality by handling AbstractGPUArray.

@codecov

codecov Bot commented Jun 13, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 90.88%. Comparing base (777420f) to head (18a692a).

Additional details and impacted files
@@            Coverage Diff             @@
##           master     #816      +/-   ##
==========================================
+ Coverage   90.74%   90.88%   +0.13%     
==========================================
  Files          11       12       +1     
  Lines        1070     1086      +16     
==========================================
+ Hits          971      987      +16     
  Misses         99       99              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

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